rssmlite

A lightweight, well-documented, PyTorch-only library for learning RSSM-based (Recurrent State-Space Model) world models — in the DreamerV2/V3 tradition — on Gymnasium environments. Designed to run end-to-end on a free Google Colab T4 GPU with zero external datasets.

Status: core library complete and released at v0.1.0. RSSM, ReplayBuffer, RSSMAgent, four environment configs, scripts/train.py CLI, checkpoint save/resume, and full documentation are all implemented and tested (29 tests passing).

World model component: production-quality. Reconstruction MSE ~0.004 on CartPole, near-perfect pole-angle tracking confirmed on a real Colab T4.

Actor-critic component: research-grade. Correct implementation of DreamerV3-style imagination-based AC with EMA critic target, but does not reliably solve CartPole within 200k steps on a small model. This is a known limitation of imagination-based RL at small model scale — documented honestly in configs/cartpole.yaml. The P2 roadmap (TransformerDynamics ablation) is the next planned improvement.

Install

pip install rssmlite[envs]        # + gymnasium
pip install rssmlite[envs,viz]    # + matplotlib for plots

From source (development):

git clone https://github.com/Mattral/rssmlite.git
cd rssmlite
pip install -e ".[dev,envs,viz]"
pytest tests/   # 21 passed

Quick start

import gymnasium as gym
from rssmlite import RSSMAgent

env = gym.make("CartPole-v1")
agent = RSSMAgent.from_config("configs/cartpole.yaml", env=env)
agent.train(env, steps=200_000, checkpoint_dir="./ckpts")

# Inspect what the model learned
agent.imagine_rollout(steps=15)
agent.visualize_latent_space()

Or via the CLI:

python scripts/train.py --config configs/cartpole.yaml
python scripts/train.py --config configs/cartpole.yaml --resume ckpts/checkpoint_50000.pt

Design philosophy

• PyTorch only. No JAX, no TensorFlow.
• Readable over clever. Every class should be understandable in under 10 minutes top-to-bottom. If it's longer, it gets split.
• Config-driven experiments, code-driven architecture. Hyperparameters and environment choice live in YAML; the model code doesn't change per run.
• Self-contained. No external dataset downloads — the environment is the data source.

Available environments

Documentation

• docs/architecture.md — RSSM math, design rationale, file map
• docs/getting_started.md — install, first training run, evaluation
• docs/api_reference.md — full public API
• docs/colab_guide.md — Colab setup, Drive checkpointing, disconnect recovery

Known simplifications (documented, not hidden)

• Continuous actions (Pendulum) have no tanh squashing — actions can technically exceed the env's valid range; Gym clips silently. Fine for current results; to fix before reporting serious Pendulum numbers.
• train.steps is a total budget across the whole run including pre-resume time. Resuming from step 80 000 with steps: 200000 trains until 200 000.
• Device management is explicit: call .to(device) on agent.rssm, agent.actor, agent.critic manually. See docs/colab_guide.md.

License

MIT — see LICENSE.