Sync local repo state

.gitignore

@@ -0,0 +1,16 @@
+.pytest_cache/
+.ruff_cache/
+.venv/
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+*.so
+*.egg-info/
+.DS_Store
+
+artifacts/
+checkpoints/
+metrics/
+data/
+logs/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 alpatrykos
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,87 @@
----
-license: mit
----
+# MiniDreamer
+
+MiniDreamer is a PlaNet-style world model project for `MiniGrid-FourRooms-v0`. It learns a recurrent latent dynamics model from partial RGB observations, predicts reward and episode termination, and uses discrete CEM planning in latent space.
+
+The repository contains:
+
+- MiniGrid RGB environment wrappers and bootstrap trajectory collection
+- Episode-aware replay buffer with reproducible train/val/test splits
+- CNN encoder, Gaussian RSSM, reward/done heads, optional decoder
+- Discrete CEM planner with termination-aware return scoring
+- PPO baseline entrypoint with a MiniGrid-compatible CNN feature extractor
+- Evaluation code, configs, scripts, tests, and project documentation
+
+A complete baseline training run has been executed. A summary is recorded in [results.md](/Users/patryktargosinski/minidreamer/results.md), while the frozen baseline artifacts remain gitignored under `artifacts/world_model/`.
+
+## Layout
+
+```text
+configs/
+docs/
+notebooks/
+scripts/
+src/
+tests/
+```
+
+Core code lives under `src/minidreamer/`, with CLI entrypoints at `src/train_world_model.py` and `src/evaluate.py`.
+
+## Setup
+
+Use Python 3.11 or 3.12. The project metadata is defined in [pyproject.toml](/Users/patryktargosinski/minidreamer/pyproject.toml).
+
+```bash
+python3.11 -m venv .venv
+source .venv/bin/activate
+pip install -e ".[dev]"
+```
+
+## Main Commands
+
+Bootstrap replay collection:
+
+```bash
+./scripts/collect_random.sh
+```
+
+World-model pipeline:
+
+```bash
+./scripts/train_world_model.sh
+```
+
+By default, the script writes new experiments to `artifacts/world_model_experiment/`. To choose a different experiment directory without touching the frozen baseline, set `MINIDREAMER_OUTPUT_DIR`:
+
+```bash
+MINIDREAMER_OUTPUT_DIR=artifacts/world_model_restricted_actions ./scripts/train_world_model.sh
+```
+
+Resume an interrupted world-model run from a checkpoint:
+
+```bash
+python3.11 src/train_world_model.py \
+  --config configs/fourrooms_world_model.yaml \
+  --output-dir artifacts/world_model \
+  --replay-dir artifacts/world_model/replay \
+  --resume-checkpoint artifacts/world_model/checkpoints/world_model_env_steps_90021.pt
+```
+
+Planner evaluation from a checkpoint:
+
+```bash
+./scripts/eval_planner.sh /path/to/checkpoint.pt /path/to/replay
+```
+
+PPO baseline:
+
+```bash
+./scripts/train_ppo.sh
+```
+
+## Notes
+
+- The latest completed run summary is in [results.md](/Users/patryktargosinski/minidreamer/results.md).
+- The baseline run in `artifacts/world_model/` is intentionally frozen as the reference artifact.
+- New world-model experiments should write to separate directories under `artifacts/`.
+- The trainer refuses to overwrite an existing run directory unless you resume with `--resume-checkpoint` or explicitly pass `--allow-overwrite-existing-output`.
+- Metrics, replay snapshots, and checkpoints are intentionally gitignored.

configs/fourrooms_ppo.yaml

@@ -0,0 +1,31 @@
+project:
+  name: minidreamer-fourrooms-ppo
+  seed: 0
+
+env:
+  id: MiniGrid-FourRooms-v0
+  rgb_partial_obs: true
+  image_only: true
+  resize: [64, 64]
+  normalize_obs: true
+
+ppo:
+  total_timesteps: 100000
+  num_envs: 4
+  learning_rate: 0.0003
+  n_steps: 256
+  batch_size: 256
+  n_epochs: 4
+  gamma: 0.99
+  gae_lambda: 0.95
+  clip_range: 0.2
+  ent_coef: 0.01
+  vf_coef: 0.5
+  features_dim: 256
+  device: auto
+
+evaluation:
+  episodes: 100
+  seeds: [0, 1, 2]
+  eval_every_env_steps: 10000
+

configs/fourrooms_world_model.yaml

@@ -0,0 +1,68 @@
+project:
+  name: minidreamer-fourrooms
+  seed: 0
+
+env:
+  id: MiniGrid-FourRooms-v0
+  rgb_partial_obs: true
+  image_only: true
+  resize: [64, 64]
+  normalize_obs: true
+  action_space: full
+
+replay:
+  capacity_episodes: 5000
+  sequence_length: 32
+  batch_size: 32
+  train_fraction: 0.8
+  val_fraction: 0.1
+  test_fraction: 0.1
+  split_key: episode_id
+
+collection:
+  bootstrap_env_steps: 5000
+  bootstrap_success_threshold: 20
+  bootstrap_fallback_policy: restricted_random_3_actions
+  bootstrap_env_step_cap: 20000
+  collect_steps_per_iteration: 1000
+  gradient_updates_per_iteration: 1000
+  train_collect_ratio: 1.0
+  random_action_fraction_after_planner: 0.3
+
+model:
+  embedding_dim: 256
+  deter_dim: 256
+  stoch_dim: 32
+  hidden_dim: 256
+  use_decoder: true
+  min_std: 0.1
+
+training:
+  optimizer: adam
+  lr: 0.0003
+  grad_clip_norm: 100.0
+  train_steps: 100000
+  beta_kl: 1.0
+  free_nats: 1.0
+  beta_recon: 1.0
+  beta_reward: 1.0
+  beta_done: 1.0
+  device: null
+
+planner:
+  type: discrete_cem
+  horizon: 8
+  candidates: 256
+  elites: 32
+  iterations: 4
+  discount: 0.99
+  use_done_mask: true
+
+evaluation:
+  episodes: 100
+  seeds: [0, 1, 2]
+  eval_every_env_steps: 10000
+
+comparison:
+  env_steps: [10000, 25000, 50000, 100000]
+

docs/spec_clarifications.md

@@ -0,0 +1,26 @@
+# Spec Clarifications
+
+This document records the implementation choices the spec left open enough that code needed an explicit default.
+
+## Implemented choices
+
+1. The Python package lives under `src/minidreamer/`, with thin CLI entrypoints at `src/train_world_model.py` and `src/evaluate.py`. That keeps imports stable while still matching the spec's requested top-level scripts.
+2. The encoder and decoder use `padding=1` in all `4x4, stride=2` convolutions so `64x64` inputs shrink cleanly to `4x4` and decode symmetrically back to `64x64`.
+3. CEM planning uses the prior mean during latent imagination instead of sampling the stochastic latent. That reduces planner variance and makes candidate ranking deterministic given the current model parameters.
+4. Replay sampling pads only at the tail of an in-episode chunk and applies a transition mask so padded steps do not contribute to loss terms.
+5. Bootstrap and online training both default to `train_collect_ratio = 1.0`, so the initial bootstrap replay produces one gradient update per collected environment step unless an explicit `gradient_updates_per_iteration` override is set.
+6. Evaluation computes one-step held-out metrics over full episodes and open-loop rollout metrics for horizons `1`, `5`, and `10` using actual held-out action sequences.
+
+## Remaining optional extensions
+
+These are intentionally not implemented in v1 because the spec marked them as later improvements or ablations:
+
+- KL balancing beyond free nats
+- uncertainty penalties or ensembles in the planner
+- actor-critic imagination learning
+- richer sparse-reward heads beyond scalar reward regression
+
+## Budget semantics
+
+Collection currently finishes complete episodes rather than cutting trajectories mid-episode. That keeps replay episodes semantically clean for recurrent training, but it means a run can land slightly above a requested step target when the last episode crosses the threshold. If exact step-matched checkpoints become mandatory for reporting, the next refinement should snapshot metrics at the first checkpoint at or above each budget and report the realized environment step count alongside the nominal target.
+

minidreamer_project_spec.md

@@ -0,0 +1,911 @@
+# MiniDreamer: A PlaNet-style World Model for Pixel-based MiniGrid Planning
+
+## 1. Project name
+
+**MiniDreamer: A PlaNet-style world model for pixel-based MiniGrid planning**
+
+## 2. Goal
+
+Build a small research-grade world model agent that learns a latent dynamics model from pixel observations in `MiniGrid-FourRooms-v0`, then uses the learned model for short-horizon planning.
+
+The project should answer:
+
+> Can a compact latent dynamics model trained from partial RGB observations support useful planning in a sparse-reward gridworld?
+
+MiniGrid is a good first target because it is designed as a simple, fast, customizable RL benchmark suite, with discrete actions and goal-oriented environments. FourRooms specifically asks an agent to navigate a four-room maze to reach a green goal; the registered environment is `MiniGrid-FourRooms-v0`.
+
+Reference: [MiniGrid documentation](https://minigrid.farama.org/index.html)
+
+---
+
+## 3. Scope
+
+### In scope
+
+- Pixel-based observations using:
+  - `RGBImgPartialObsWrapper`
+  - `ImgObsWrapper`
+- Environment:
+  - MVP: `MiniGrid-FourRooms-v0`
+  - Extension: `MiniGrid-DoorKey-*`, `MiniGrid-LockedRoom-v0`, or `MiniGrid-Dynamic-Obstacles-*`
+- Learned latent dynamics model
+- Reward and termination prediction
+- MPC/CEM-style planning in latent space
+- PPO baseline for comparison
+- Ablations:
+  - with image reconstruction vs. without
+  - latent size
+  - planning horizon
+  - stochastic vs. deterministic latent state
+
+### Out of scope for v1
+
+- Full DreamerV3 reproduction
+- Actor-critic learning inside imagination
+- Minecraft/Crafter-scale environments
+- Large video-generation-style world models
+- Language-conditioned BabyAI tasks
+
+---
+
+## 4. Technical background
+
+This should be a **PlaNet-lite** implementation. PlaNet learns a latent dynamics model from images and chooses actions through online planning in latent space, using both deterministic and stochastic transition components.
+
+Reference: [PlaNet: Learning Latent Dynamics for Planning from Pixels](https://arxiv.org/abs/1811.04551)
+
+DreamerV3 is the more mature descendant: it learns a world model and improves behavior by imagining future scenarios, but its full implementation is too much for a first project. Use DreamerV3 as conceptual inspiration, not as the implementation target.
+
+Reference: [DreamerV3](https://arxiv.org/abs/2301.04104)
+
+MiniGrid’s default observation is a compact symbolic encoding, not raw pixels. For this project, use `RGBImgPartialObsWrapper` to obtain RGB pixel observations, then `ImgObsWrapper` to keep only the image tensor.
+
+Reference: [MiniGrid wrappers](https://minigrid.farama.org/api/wrapper/)
+
+---
+
+## 5. Core hypothesis
+
+### Main hypothesis
+
+A learned recurrent latent dynamics model can support better sample efficiency than a model-free baseline in early training, even if final performance is lower or less stable.
+
+### Secondary hypotheses
+
+1. **Reconstruction loss may help representation learning early**, but may hurt planning if it forces the latent state to model irrelevant pixels.
+2. **Short planning horizons will work better than long ones**, because model error compounds over imagined rollouts.
+3. **Stochastic latent states should outperform deterministic-only states** under partial observability.
+4. **Reward prediction quality matters more than pixel reconstruction quality** for planning performance.
+
+---
+
+## 6. Environment specification
+
+### MVP environment
+
+```python
+env_id = "MiniGrid-FourRooms-v0"
+```
+
+FourRooms has a discrete action space of size 7. The meaningful actions for this task are mostly `left`, `right`, and `forward`; `pickup`, `drop`, `toggle`, and `done` are listed but unused in FourRooms.
+
+Reference: [MiniGrid FourRooms environment](https://minigrid.farama.org/environments/minigrid/FourRoomsEnv/)
+
+### Wrappers
+
+```python
+import gymnasium as gym
+from minigrid.wrappers import RGBImgPartialObsWrapper, ImgObsWrapper
+
+env = gym.make("MiniGrid-FourRooms-v0")
+env = RGBImgPartialObsWrapper(env)
+env = ImgObsWrapper(env)
+```
+
+### Observation
+
+Use RGB partial observation.
+
+Normalize to:
+
+```python
+obs = obs.astype("float32") / 255.0
+```
+
+Recommended internal size:
+
+```text
+64 x 64 x 3
+```
+
+Resize if needed.
+
+### Action space
+
+Use the full 7-action space for baseline compatibility.
+
+Optional ablation:
+
+```text
+restricted_action_space = {left, right, forward}
+```
+
+If bootstrap data collection yields too few successful episodes, it is acceptable to use the restricted 3-action random policy during bootstrap only. If that fallback is used, record it explicitly in the report and still count all bootstrap environment steps toward the world-model sample-efficiency budget.
+
+### Reward
+
+FourRooms gives a success reward based on step count and zero reward for failure; the documented reward is:
+
+```text
+1 - 0.9 * step_count / max_steps
+```
+
+For model training, predict:
+
+```text
+reward_t ∈ R
+done_t ∈ {0, 1}
+```
+
+### Episode end semantics
+
+Use the Gymnasium step API explicitly:
+
+```python
+obs_next, reward, terminated, truncated, info = env.step(action)
+```
+
+Store all three episode-end signals:
+
+```text
+terminated_t = goal reached / environment terminal condition
+truncated_t = episode stopped by time limit
+done_t = terminated_t or truncated_t
+```
+
+Training and planning semantics:
+
+- `reward_t` is the scalar emitted after taking `action_t`
+- `done_t` is the episode-end target used by the world model and planner
+- `success_t = 1` iff `terminated_t == 1` and `reward_t > 0`
+
+Do not collapse `terminated` and `truncated` in logged metrics; report them separately when debugging failures.
+
+---
+
+## 7. Model architecture
+
+### 7.1 Overview
+
+```text
+obs_t ──► CNN encoder ──► embedding e_t
+                         │
+action_{t-1} ────────────┤
+                         ▼
+                  RSSM latent model
+                         │
+              ┌──────────┼──────────┐
+              ▼          ▼          ▼
+        reward head   done head   decoder? optional
+```
+
+### 7.2 Latent state
+
+Use a simplified **RSSM-style** state:
+
+```text
+h_t = deterministic recurrent state
+z_t = stochastic latent state
+s_t = concat(h_t, z_t)
+```
+
+Recommended dimensions:
+
+```yaml
+embedding_dim: 256
+deterministic_state_dim: 256
+stochastic_state_dim: 32
+num_stochastic_classes: null  # use Gaussian latent for v1
+```
+
+### 7.3 Encoder
+
+CNN encoder:
+
+```text
+Input: 64x64x3 RGB
+Conv 4x4 stride 2, channels 32
+Conv 4x4 stride 2, channels 64
+Conv 4x4 stride 2, channels 128
+Conv 4x4 stride 2, channels 256
+Flatten
+Linear -> embedding_dim
+```
+
+Use `LayerNorm` or `BatchNorm` only if training is unstable. Start simple.
+
+### 7.4 Recurrent dynamics
+
+Prior:
+
+```text
+p(z_t | h_t)
+```
+
+Posterior:
+
+```text
+q(z_t | h_t, e_t)
+```
+
+Transition:
+
+```text
+h_t = GRU(h_{t-1}, concat(z_{t-1}, one_hot(action_{t-1})))
+```
+
+Gaussian latent:
+
+```text
+prior_mean, prior_std = prior_net(h_t)
+post_mean, post_std = posterior_net(h_t, e_t)
+z_t ~ Normal(post_mean, post_std)
+```
+
+Use reparameterization.
+
+Clamp std:
+
+```python
+std = softplus(raw_std) + 0.1
+```
+
+### 7.4.1 Sequence indexing and initialization
+
+Use transition tuples with explicit alignment:
+
+```text
+(obs_t, action_t, reward_t, terminated_t, truncated_t, done_t, obs_{t+1})
+```
+
+Training-time convention:
+
+1. Infer posterior `s_t` from `obs_t` and the previous recurrent state.
+2. Apply `action_t` in the transition model to form the prior for step `t+1`.
+3. Condition on `obs_{t+1}` to infer posterior `s_{t+1}`.
+4. Predict `reward_t`, `done_t`, and optional reconstruction of `obs_{t+1}` from `s_{t+1}`.
+
+Sequence-start convention:
+
+- initialize `h_0` and `z_0` to zeros
+- use a zero action embedding in place of `action_{-1}`
+- never sample a training chunk that crosses an episode boundary
+- if padding is needed for batching, pad only within an episode and apply a loss mask so padded steps do not contribute to any loss
+
+### 7.5 Prediction heads
+
+Reward head:
+
+```text
+MLP(s_{t+1}) -> scalar reward_t
+```
+
+Done head:
+
+```text
+MLP(s_{t+1}) -> Bernoulli logit for done_t
+```
+
+Optional decoder:
+
+```text
+MLP/ConvTranspose(s_{t+1}) -> reconstructed RGB observation obs_{t+1}
+```
+
+For v1, include the decoder behind a config flag:
+
+```yaml
+use_decoder: true | false
+```
+
+---
+
+## 8. Loss function
+
+Total loss:
+
+```text
+L = β_reward * L_reward + β_done * L_done + β_kl * max(L_kl, free_nats) + β_recon * L_recon
+```
+
+### Reward loss
+
+```text
+L_reward = MSE(pred_reward_t, reward_t)
+```
+
+Alternative for sparse rewards:
+
+```text
+two_head_reward:
+  reward_occurrence: BCE
+  reward_value_given_success: MSE
+```
+
+Do not start there unless MSE fails.
+
+### Done loss
+
+```text
+L_done = BCEWithLogits(done_logit_t, done_t)
+```
+
+### KL loss
+
+```text
+L_kl = KL(q(z_t | h_t, e_t) || p(z_t | h_t))
+```
+
+Use free-nats from the start in v1. KL balancing is a later stabilization improvement.
+
+Initial config:
+
+```yaml
+beta_reward: 1.0
+beta_done: 1.0
+beta_kl: 1.0
+free_nats: 1.0
+```
+
+### Reconstruction loss
+
+If decoder is enabled:
+
+```text
+L_recon = MSE(reconstructed_obs_{t+1}, obs_{t+1})
+```
+
+If decoder is disabled, set:
+
+```yaml
+beta_recon: 0.0
+```
+
+Initial config:
+
+```yaml
+beta_recon: 1.0
+```
+
+Ablate with:
+
+```yaml
+beta_recon: 0.0
+```
+
+---
+
+## 9. Data collection
+
+### Phase 1: bootstrap dataset
+
+Collect:
+
+```yaml
+bootstrap_env_steps: 5000
+policy: random_full_action_space
+fallback_if_successes_too_low: random_restricted_action_space
+min_success_episodes_before_planning: 20
+bootstrap_env_step_cap: 20000
+```
+
+Use environment-step budgets, not fixed episode counts, so the early comparison points against PPO remain meaningful.
+
+Bootstrap rule:
+
+- start with full-action random data collection
+- if the replay buffer contains fewer than `20` successful episodes after `5000` steps, continue collecting with restricted-action random `{left, right, forward}`
+- do not make sample-efficiency claims for planner performance before the replay buffer contains at least `20` successful episodes
+
+Store transitions:
+
+```text
+obs_t
+action_t
+reward_t
+terminated_t
+truncated_t
+done_t
+obs_{t+1}
+episode_id
+step_id
+```
+
+### Phase 2: mixed dataset
+
+After bootstrap, alternate fixed collection/training rounds:
+
+```yaml
+collect_steps_per_iteration: 1000
+gradient_updates_per_iteration: 1000
+```
+
+Default collection mix:
+
+```text
+70% planner policy
+30% random exploration
+```
+
+This prevents early model errors from collapsing exploration. `train_collect_ratio = 1.0` means one gradient update per newly collected environment step.
+
+### Replay format and splits
+
+Use chunked trajectory sampling:
+
+```yaml
+sequence_length: 32
+batch_size: 32
+```
+
+Sample contiguous sequences, not independent transitions, because the recurrent model needs temporal structure.
+
+Maintain episode-level dataset splits:
+
+```yaml
+train_fraction: 0.8
+val_fraction: 0.1
+test_fraction: 0.1
+split_key: episode_id
+```
+
+Assign splits by episode id, not by transition, and keep the split fixed as more episodes are collected. Use train episodes for optimization, val episodes for model selection and debugging, and held-out test episodes only for final world-model reporting.
+
+---
+
+## 10. Planner
+
+### MVP planner: CEM over discrete actions
+
+At each real environment step:
+
+1. Encode current observation.
+2. Update posterior latent state.
+3. Sample candidate action sequences.
+4. Roll them forward using the learned prior model.
+5. Score imagined returns.
+6. Execute the first action from the best sequence.
+7. Replan at the next step.
+
+### Planning config
+
+```yaml
+planning_horizon: 8
+num_candidates: 256
+num_elites: 32
+cem_iterations: 4
+discount: 0.99
+```
+
+Because actions are discrete, maintain categorical probabilities over actions per timestep:
+
+```text
+π_h ∈ R[horizon, num_actions]
+```
+
+CEM loop:
+
+```text
+sample action sequences from π
+roll out latent dynamics
+score by termination-aware discounted reward
+select top-k elites
+update π toward elite action frequencies
+```
+
+### Score function
+
+```text
+done_prob_t = sigmoid(pred_done_logit_t)
+alive_0 = 1
+alive_{t+1} = alive_t * (1 - done_prob_t)
+score = Σ_t γ^t * alive_t * predicted_reward_t
+```
+
+This makes the planner termination-aware and prevents CEM from exploiting trajectories that unrealistically continue after predicted episode end.
+
+Optional later refinement:
+
+```text
+score += λ * predicted_success_probability
+score -= μ * predicted_done_probability_before_reward
+```
+
+Keep v1 simple: use only discounted reward with the alive-mask above.
+
+---
+
+## 11. Baselines
+
+### Baseline 1: random policy
+
+Measure:
+
+```text
+success_rate
+mean_return
+mean_episode_length
+```
+
+### Baseline 2: PPO
+
+Use Stable-Baselines3 PPO with a custom CNN feature extractor. MiniGrid’s own training tutorial shows PPO training with Stable-Baselines3 and notes that a custom feature extractor is needed because the default CNN architecture does not directly support MiniGrid’s observation space.
+
+Reference: [MiniGrid training tutorial](https://minigrid.farama.org/content/training/)
+
+Compare PPO and world model under equal environment steps:
+
+```yaml
+env_steps:
+  - 10_000
+  - 25_000
+  - 50_000
+  - 100_000
+```
+
+### Comparison protocol
+
+Use the same wrapped observation setup for PPO and the world-model agent unless an experiment explicitly states otherwise.
+
+Budget accounting:
+
+- count every environment interaction used for world-model data collection toward the world-model budget, including bootstrap random data and later planner-driven data
+- do not count offline gradient updates toward environment-step budgets
+- do not count evaluation episodes toward training budgets
+
+Checkpointing protocol at budget `B`:
+
+1. PPO is trained for exactly `B` environment steps, then evaluated.
+2. The world-model pipeline is allowed to collect exactly `B` environment steps in total, updating the model online according to Section 9, then evaluated.
+3. Report the latest checkpoint at budget `B`, not the best checkpoint seen so far.
+4. Use the same seeds and the same number of evaluation episodes for both methods.
+
+---
+
+## 12. Evaluation metrics
+
+### Agent metrics
+
+```yaml
+success_rate: fraction of episodes reaching goal
+mean_return: average episodic return
+median_return: robust episodic return
+mean_episode_length: average steps per episode
+env_steps_to_80_percent_success: sample efficiency metric
+```
+
+### World model metrics
+
+```yaml
+reward_mse: reward prediction error
+done_bce: termination prediction error
+kl_loss: posterior-prior divergence
+reconstruction_mse: if decoder enabled
+open_loop_reward_error_h1_h5_h10: reward rollout error over horizons
+open_loop_done_accuracy_h1_h5_h10: done prediction accuracy over horizons
+```
+
+World-model metrics must be reported on held-out validation or test episodes, never on the training replay used for optimization. For final tables and plots, prefer the held-out test split.
+
+### Planning metrics
+
+```yaml
+imagined_return_vs_real_return_correlation
+planner_action_entropy
+model_rollout_horizon_sensitivity
+```
+
+---
+
+## 13. Experiment matrix
+
+### MVP experiments
+
+| ID | Model | Decoder | Latent | Horizon | Actions | Purpose |
+|---:|---|---:|---:|---:|---|---|
+| E0 | Random | n/a | n/a | n/a | 7 | lower bound |
+| E1 | PPO | n/a | n/a | n/a | 7 | model-free baseline |
+| E2 | World model | yes | 32 | 8 | 7 | default model |
+| E3 | World model | no | 32 | 8 | 7 | test reconstruction value |
+| E4 | World model | yes | 16 | 8 | 7 | latent bottleneck |
+| E5 | World model | yes | 64 | 8 | 7 | larger latent |
+| E6 | World model | yes | 32 | 4 | 7 | short planning |
+| E7 | World model | yes | 32 | 16 | 7 | long planning |
+
+### Extension experiments
+
+| ID | Environment | Purpose |
+|---:|---|---|
+| X1 | DoorKey | test object interaction |
+| X2 | LockedRoom | test longer-horizon partial observability |
+| X3 | DynamicObstacles | test non-stationary local hazards |
+| X4 | FourRooms larger/custom | test layout generalization |
+
+---
+
+## 14. Expected deliverables
+
+### Code deliverables
+
+```text
+minidreamer/
+  configs/
+    fourrooms_world_model.yaml
+    fourrooms_ppo.yaml
+  src/
+    envs/
+      make_env.py
+    data/
+      replay_buffer.py
+      collect_random.py
+    models/
+      encoder.py
+      rssm.py
+      heads.py
+      decoder.py
+      world_model.py
+    planning/
+      cem.py
+      evaluate_planner.py
+    baselines/
+      train_ppo.py
+    train_world_model.py
+    evaluate.py
+  scripts/
+    collect_random.sh
+    train_world_model.sh
+    eval_planner.sh
+    train_ppo.sh
+  notebooks/
+    rollout_debug.ipynb
+    results_analysis.ipynb
+  tests/
+    test_env.py
+    test_replay_buffer.py
+    test_rssm_shapes.py
+    test_cem_planner.py
+```
+
+### Research deliverables
+
+```text
+README.md
+results.md
+plots/
+  learning_curves.png
+  success_rate_vs_env_steps.png
+  model_error_vs_rollout_horizon.png
+  reconstruction_examples.png
+  imagined_vs_real_rollouts.png
+```
+
+### Minimum publishable artifact
+
+A short report with:
+
+1. Problem statement
+2. Method
+3. Environment setup
+4. Baselines
+5. Main learning curves
+6. Ablations
+7. Failure cases
+8. Next steps
+
+---
+
+## 15. Milestones
+
+### Milestone 1: environment and data pipeline
+
+Acceptance criteria:
+
+- Can create `MiniGrid-FourRooms-v0`
+- Can wrap it into RGB-only observation mode
+- Can collect and save bootstrap trajectories with explicit `terminated`, `truncated`, and `done` flags
+- Can reload trajectory chunks for sequence training
+- Train/val/test episode splits are reproducible
+- Shape tests pass
+
+### Milestone 2: world model training
+
+Acceptance criteria:
+
+- RSSM forward pass works on sequence batches
+- Reward, done, KL, and optional reconstruction losses train without NaNs
+- One-step reward/done prediction beats trivial constant predictor on held-out validation episodes
+- Model checkpointing works
+
+### Milestone 3: open-loop model evaluation
+
+Acceptance criteria:
+
+- Can visualize real vs. reconstructed observations
+- Can roll model forward for 1, 5, and 10 imagined steps
+- Can report reward and done prediction error by horizon on held-out episodes
+
+### Milestone 4: CEM planner
+
+Acceptance criteria:
+
+- Planner can choose actions from latent state
+- Planner runs online in the environment
+- Planner scoring uses termination-aware reward masking
+- Success rate beats random policy over 100 evaluation episodes
+
+### Milestone 5: baseline comparison
+
+Acceptance criteria:
+
+- PPO baseline runs on same observation setup
+- World model and PPO are compared at fixed environment-step budgets, with bootstrap data counted toward the world-model budget
+- At least 3 seeds per method
+
+### Milestone 6: ablations and report
+
+Acceptance criteria:
+
+- Run decoder/no-decoder ablation
+- Run at least 2 planning horizons
+- Run at least 2 latent sizes
+- Produce final plots and written analysis
+
+---
+
+## 16. Suggested default config
+
+```yaml
+project:
+  name: minidreamer-fourrooms
+  seed: 0
+
+env:
+  id: MiniGrid-FourRooms-v0
+  rgb_partial_obs: true
+  image_only: true
+  resize: [64, 64]
+  normalize_obs: true
+  action_space: full
+
+replay:
+  capacity_episodes: 5000
+  sequence_length: 32
+  batch_size: 32
+  train_fraction: 0.8
+  val_fraction: 0.1
+  test_fraction: 0.1
+  split_key: episode_id
+
+collection:
+  bootstrap_env_steps: 5000
+  bootstrap_success_threshold: 20
+  bootstrap_fallback_policy: restricted_random_3_actions
+  bootstrap_env_step_cap: 20000
+  collect_steps_per_iteration: 1000
+  train_collect_ratio: 1.0
+  random_action_fraction_after_planner: 0.3
+
+model:
+  embedding_dim: 256
+  deter_dim: 256
+  stoch_dim: 32
+  hidden_dim: 256
+  use_decoder: true
+  min_std: 0.1
+
+training:
+  optimizer: adam
+  lr: 0.0003
+  grad_clip_norm: 100.0
+  train_steps: 100000
+  beta_kl: 1.0
+  free_nats: 1.0
+  beta_recon: 1.0
+  beta_reward: 1.0
+  beta_done: 1.0
+
+planner:
+  type: discrete_cem
+  horizon: 8
+  candidates: 256
+  elites: 32
+  iterations: 4
+  discount: 0.99
+  use_done_mask: true
+
+evaluation:
+  episodes: 100
+  seeds: [0, 1, 2]
+  eval_every_env_steps: 10000
+```
+
+---
+
+## 17. Acceptance criteria for v1
+
+The project is successful if:
+
+1. The world model trains stably for at least 3 seeds.
+2. The CEM planner beats random policy on `MiniGrid-FourRooms-v0`.
+3. The report shows learning curves for world-model planner, PPO, and random.
+4. The report includes at least one ablation showing whether reconstruction helps.
+5. The repo can reproduce the main result from a clean config file and fixed environment-step comparison protocol.
+
+A strong v1 result would be:
+
+```text
+world_model_success_rate > random_success_rate
+```
+
+An excellent v1 result would be:
+
+```text
+world_model reaches useful success rate with fewer env steps than PPO
+```
+
+Do not require the world model to beat PPO asymptotically. For a first project, sample efficiency and interpretable failure analysis are more important.
+
+---
+
+## 18. Main risks
+
+### Risk 1: sparse reward makes reward model hard to train
+
+Mitigations:
+
+- start with random + heuristic exploratory data
+- add goal-reaching trajectories if random rarely succeeds
+- optionally train first on `MiniGrid-Empty-*`
+- use success classification in addition to reward regression
+
+### Risk 2: model learns visual reconstruction but not controllable dynamics
+
+Mitigations:
+
+- ablate decoder
+- track reward/done prediction separately from reconstruction
+- evaluate imagined-vs-real return correlation
+
+### Risk 3: CEM planner exploits model errors
+
+Mitigations:
+
+- short planning horizon
+- replan every step
+- penalize high predicted uncertainty if using ensemble later
+- mix random exploration into data collection
+
+### Risk 4: partial observability hurts Markov assumptions
+
+Mitigations:
+
+- keep recurrent state
+- compare deterministic-only vs. stochastic recurrent latent
+- optionally add frame stacking as a diagnostic, not as the main solution
+
+---
+
+## 19. Recommended implementation order
+
+1. Implement environment wrapper and trajectory collection.
+2. Implement replay buffer with sequence sampling.
+3. Implement encoder + RSSM + reward/done heads.
+4. Train model on random trajectories.
+5. Add decoder and visualization.
+6. Add open-loop rollout diagnostics.
+7. Implement discrete CEM planner.
+8. Evaluate against random.
+9. Add PPO baseline.
+10. Run ablations.
+
+The key design constraint: **do not optimize for final score first**. Optimize for observability: diagnostics, rollout plots, prediction errors, and failure cases. That will make this a research project instead of just another RL training script.

notebooks/results_analysis.ipynb

@@ -0,0 +1,25 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# Results Analysis\n",
+        "\n",
+        "Use this notebook to load `metrics/*.jsonl`, aggregate seed-level runs, and generate the plots referenced in `results.md`."
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 5
+}

notebooks/rollout_debug.ipynb

@@ -0,0 +1,26 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# Rollout Debug\n",
+        "\n",
+        "Use this notebook to inspect replay sequences, reconstructions, and imagined-vs-real reward rollouts after training has been run."
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 5
+}
+

plots/.gitkeep

@@ -0,0 +1 @@
+

pyproject.toml

@@ -0,0 +1,46 @@
+[build-system]
+requires = ["setuptools>=69", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "minidreamer"
+version = "0.1.0"
+description = "PlaNet-style world model planning for pixel-based MiniGrid"
+readme = "README.md"
+requires-python = ">=3.11,<3.13"
+license = { file = "LICENSE" }
+authors = [{ name = "OpenAI Codex" }]
+dependencies = [
+  "gymnasium>=0.29,<1.1",
+  "minigrid>=2.3.1",
+  "numpy>=1.26,<3",
+  "pillow>=10.0",
+  "pyyaml>=6.0",
+  "torch>=2.3,<3",
+  "tqdm>=4.66",
+  "matplotlib>=3.8",
+  "stable-baselines3>=2.3.0",
+]
+
+[project.optional-dependencies]
+dev = [
+  "pytest>=8.2",
+]
+
+[project.scripts]
+minidreamer-collect-random = "minidreamer.data.collect_random:main"
+minidreamer-train-world-model = "train_world_model:main"
+minidreamer-evaluate = "evaluate:main"
+minidreamer-train-ppo = "minidreamer.baselines.train_ppo:main"
+
+[tool.setuptools]
+package-dir = { "" = "src" }
+py-modules = ["train_world_model", "evaluate"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["minidreamer*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["src"]

results.md

@@ -0,0 +1,127 @@
+# Results
+
+## Status
+
+A baseline world-model training run completed on `2026-04-21` on Apple Silicon using the `mps` backend. This run is now frozen as the reference artifact under `artifacts/world_model/`. Final world-model summaries now come from `artifacts/world_model/metrics/run_summary.json`, `artifacts/world_model/metrics/final_eval_latest.json`, and `artifacts/world_model/metrics/planner_eval_latest_clean.json`, the PPO baseline summary comes from `artifacts/ppo/metrics/run_summary.json`, and generated figures live under `plots/`.
+
+## Problem Statement
+
+Train a PlaNet-style world model for `MiniGrid-FourRooms-v0` from partial RGB observations, then evaluate a latent-space discrete CEM planner against a random baseline.
+
+## Method
+
+- CNN encoder + Gaussian RSSM world model with reward, done, and reconstruction heads.
+- Discrete CEM planning in latent space.
+- Replay-buffer training with episode-aware train/val/test splits.
+- Config: `configs/fourrooms_world_model.yaml`.
+
+## Environment Setup
+
+- Device: `mps`
+- Target environment steps: `100000`
+- Final realized environment steps: `100004`
+- Total gradient updates completed: `10000`
+- Final replay summary:
+  - Episodes: `1022`
+  - Success episodes: `42`
+  - Train/val/test episodes: `819 / 116 / 87`
+
+## Baselines
+
+Persisted comparison metrics were recorded at `97124` env steps:
+
+- Random baseline success rate: `0.0`
+- Random baseline mean return: `0.0`
+- Random baseline mean episode length: `100.0`
+
+Completed PPO baseline after `100000` training env steps:
+
+- PPO success rate: `0.14`
+- PPO mean return: `0.1336`
+- PPO median return: `0.0`
+- PPO return std: `0.3313`
+- PPO mean episode length: `86.71`
+
+## Main Metrics
+
+Persisted evaluation metrics at `97124` env steps:
+
+- Planner success rate: `0.10`
+- Planner mean return: `0.0568`
+- Planner median return: `0.0`
+- Planner mean episode length: `94.8`
+- Planner action entropy: `0.6569`
+
+Held-out world-model metrics at the same checkpoint:
+
+- Reward MSE: `1.98e-6`
+- Done BCE: `0.1808`
+- KL loss: `0.8996`
+- Reconstruction MSE: `0.0109`
+
+Open-loop rollout quality:
+
+- Done accuracy @1/@5/@10: `0.9880 / 0.9960 / 0.9973`
+- Reward error @1/@5/@10: `1.98e-6 / 1.83e-6 / 1.69e-6`
+
+Canonical clean final planner evaluation at `world_model_latest.pt`:
+
+- Evaluation budget: `100` planner episodes and `100` random episodes
+- Planner success rate: `0.04`
+- Planner mean return: `0.03415`
+- Planner median return: `0.0`
+- Planner mean episode length: `96.65`
+- Planner action entropy: `0.9142`
+- Final random baseline success rate: `0.03`
+- Final random baseline mean return: `0.02073`
+- Final random baseline mean episode length: `98.03`
+
+Small-sample final world-model diagnostics at `world_model_latest.pt`:
+
+- Evaluation budget: `5` held-out world-model episodes
+- Final reward MSE: `4.95e-6`
+- Final done BCE: `0.1513`
+- Final KL loss: `0.8833`
+- Final reconstruction MSE: `0.0096`
+- Final done accuracy @1/@5/@10: `0.9880 / 0.9971 / 0.9978`
+- Final reward error @1/@5/@10: `4.98e-6 / 5.09e-6 / 5.12e-6`
+
+## Comparison
+
+At roughly matched data budgets, the PPO baseline produced stronger direct control performance than the world-model planner:
+
+- PPO at `100000` env steps: `0.14` success rate, `0.1336` mean return, `86.71` mean episode length over `100` evaluation episodes.
+- Final clean planner eval at `100004` env steps: `0.04` success rate, `0.03415` mean return, `96.65` mean episode length over `100` evaluation episodes.
+- Final clean random eval at `100004` env steps: `0.03` success rate, `0.02073` mean return, `98.03` mean episode length over `100` evaluation episodes.
+
+The predictive model stayed numerically strong through the end of training, but that did not translate into a robust planner win on FourRooms at this training budget. After removing evaluation-time action noise, the final planner is slightly better than random, but still clearly below PPO.
+
+## Ablations
+
+No ablation runs have been recorded yet.
+
+## Failure Cases / Operational Notes
+
+- The initial long run stopped near `90021` env steps when the local machine hit severe disk pressure.
+- Training was resumed successfully from `artifacts/world_model/checkpoints/world_model_env_steps_90021.pt` after adding checkpoint-resume support to the trainer.
+- Only one scheduled planner evaluation row is persisted in `eval_metrics.jsonl`. The resumed segment completed without crossing another configured evaluation boundary before the final save, so the end-of-run planner metrics are stored separately in `final_eval_latest.json` and `planner_eval_latest_clean.json`.
+- The scheduled planner row in `eval_metrics.jsonl` and the planner section inside `final_eval_latest.json` were recorded before the evaluation-noise fix. They are kept for provenance, but `planner_eval_latest_clean.json` is the canonical final planner result.
+
+## Visualizations
+
+- Generated from `artifacts/world_model/metrics/train_metrics.jsonl`, `artifacts/world_model/metrics/eval_metrics.jsonl`, `artifacts/world_model/metrics/final_eval_latest.json`, and `artifacts/world_model/metrics/planner_eval_latest_clean.json` using `scripts/generate_results_plots.py`.
+- The success-rate chart uses the clean final planner evaluation for the last point.
+- `plots/learning_curves.png`
+- `plots/success_rate_vs_env_steps.png`
+- `plots/model_error_vs_rollout_horizon.png`
+
+## Artifact Locations
+
+- Metrics: `artifacts/world_model/metrics/`
+- Replay: `artifacts/world_model/replay/`
+- Checkpoints: `artifacts/world_model/checkpoints/`
+- Final checkpoint: `artifacts/world_model/checkpoints/world_model_latest.pt`
+- Canonical clean final planner eval: `artifacts/world_model/metrics/planner_eval_latest_clean.json`
+- PPO metrics: `artifacts/ppo/metrics/`
+- PPO checkpoint: `artifacts/ppo/checkpoints/ppo_latest.zip`
+- Generated plots: `plots/`

scripts/collect_random.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+export PYTHONPATH="${PYTHONPATH:-}:$(pwd)/src"
+python3.11 src/minidreamer/data/collect_random.py \
+  --config configs/fourrooms_world_model.yaml \
+  --output-dir artifacts/bootstrap_replay \
+  "$@"
+

scripts/eval_planner.sh

@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [ "$#" -lt 2 ]; then
+  echo "Usage: $0 CHECKPOINT REPLAY_DIR [extra evaluate args]" >&2
+  exit 1
+fi
+
+checkpoint="$1"
+replay_dir="$2"
+shift 2
+
+export PYTHONPATH="${PYTHONPATH:-}:$(pwd)/src"
+python3.11 src/evaluate.py \
+  planner \
+  --config configs/fourrooms_world_model.yaml \
+  --checkpoint "${checkpoint}" \
+  "$@"
+
+python3.11 src/evaluate.py \
+  world-model \
+  --config configs/fourrooms_world_model.yaml \
+  --checkpoint "${checkpoint}" \
+  --replay-dir "${replay_dir}" \
+  "$@"
+

scripts/generate_results_plots.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+
+
+ROOT = Path(__file__).resolve().parents[1]
+METRICS_DIR = ROOT / "artifacts" / "world_model" / "metrics"
+PLOTS_DIR = ROOT / "plots"
+
+
+def load_jsonl(path: Path) -> list[dict]:
+    rows = []
+    if not path.exists():
+        return rows
+    for line in path.read_text(encoding="utf-8").splitlines():
+        line = line.strip()
+        if line:
+            rows.append(json.loads(line))
+    return rows
+
+
+def load_json(path: Path) -> dict | None:
+    if not path.exists():
+        return None
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def rolling_mean(values: list[float], window: int) -> list[float]:
+    if not values:
+        return []
+    out: list[float] = []
+    running_sum = 0.0
+    for idx, value in enumerate(values):
+        running_sum += value
+        if idx >= window:
+            running_sum -= values[idx - window]
+        out.append(running_sum / min(idx + 1, window))
+    return out
+
+
+def generate_learning_curves(train_rows: list[dict]) -> None:
+    if not train_rows:
+        return
+    steps = list(range(1, len(train_rows) + 1))
+    window = min(250, len(train_rows))
+    loss = rolling_mean([row["loss"] for row in train_rows], window)
+    kl = rolling_mean([row["kl_loss"] for row in train_rows], window)
+    done = rolling_mean([row["done_loss"] for row in train_rows], window)
+    recon = rolling_mean([row["recon_loss"] for row in train_rows], window)
+
+    fig, ax = plt.subplots(figsize=(10, 6))
+    ax.plot(steps, loss, label="total loss", linewidth=2.0)
+    ax.plot(steps, kl, label="kl loss", linewidth=1.5)
+    ax.plot(steps, done, label="done loss", linewidth=1.5)
+    ax.plot(steps, recon, label="recon loss", linewidth=1.5)
+    ax.set_title("World Model Training Curves")
+    ax.set_xlabel("Gradient update")
+    ax.set_ylabel("Smoothed loss")
+    ax.legend()
+    ax.grid(alpha=0.25)
+    fig.tight_layout()
+    fig.savefig(PLOTS_DIR / "learning_curves.png", dpi=200)
+    plt.close(fig)
+
+
+def generate_success_plot(
+    eval_rows: list[dict],
+    final_eval: dict | None,
+    clean_planner_eval: dict | None,
+) -> None:
+    steps: list[int] = []
+    planner_success: list[float] = []
+    random_success: list[float] = []
+
+    for row in eval_rows:
+        steps.append(int(row["env_steps"]))
+        planner_success.append(float(row["planner/success_rate"]))
+        random_success.append(float(row["random/success_rate"]))
+
+    if clean_planner_eval is not None:
+        steps.append(int(clean_planner_eval["metadata"]["env_steps"]))
+        planner_success.append(float(clean_planner_eval["planner_clean"]["success_rate"]))
+        random_success.append(float(clean_planner_eval["random"]["success_rate"]))
+    elif final_eval is not None:
+        steps.append(int(final_eval["metadata"]["env_steps"]))
+        planner_success.append(float(final_eval["planner"]["success_rate"]))
+        random_success.append(float(final_eval["random"]["success_rate"]))
+
+    if not steps:
+        return
+
+    paired = sorted(zip(steps, planner_success, random_success), key=lambda item: item[0])
+    steps = [item[0] for item in paired]
+    planner_success = [item[1] for item in paired]
+    random_success = [item[2] for item in paired]
+
+    fig, ax = plt.subplots(figsize=(9, 5))
+    ax.plot(steps, planner_success, marker="o", linewidth=2.0, label="planner success rate")
+    ax.plot(steps, random_success, marker="o", linewidth=2.0, label="random success rate")
+    ax.set_title("Success Rate vs Environment Steps")
+    ax.set_xlabel("Environment steps")
+    ax.set_ylabel("Success rate")
+    ax.set_ylim(-0.02, 1.02)
+    ax.legend()
+    ax.grid(alpha=0.25)
+    fig.tight_layout()
+    fig.savefig(PLOTS_DIR / "success_rate_vs_env_steps.png", dpi=200)
+    plt.close(fig)
+
+
+def generate_rollout_error_plot(final_eval: dict | None) -> None:
+    if final_eval is None:
+        return
+    horizons = [1, 5, 10]
+    reward_errors = [
+        float(final_eval["world_model"][f"open_loop_reward_error_h{h}"])
+        for h in horizons
+    ]
+    done_accuracy = [
+        float(final_eval["world_model"][f"open_loop_done_accuracy_h{h}"])
+        for h in horizons
+    ]
+
+    fig, ax1 = plt.subplots(figsize=(9, 5))
+    ax1.bar([str(h) for h in horizons], reward_errors, color="#2b6cb0", alpha=0.8)
+    ax1.set_xlabel("Open-loop horizon")
+    ax1.set_ylabel("Reward MSE", color="#2b6cb0")
+    ax1.tick_params(axis="y", labelcolor="#2b6cb0")
+    ax1.set_title("Model Error vs Rollout Horizon")
+    ax1.grid(alpha=0.2, axis="y")
+
+    ax2 = ax1.twinx()
+    ax2.plot([str(h) for h in horizons], done_accuracy, color="#c05621", marker="o", linewidth=2.0)
+    ax2.set_ylabel("Done accuracy", color="#c05621")
+    ax2.tick_params(axis="y", labelcolor="#c05621")
+    ax2.set_ylim(0.95, 1.001)
+
+    fig.tight_layout()
+    fig.savefig(PLOTS_DIR / "model_error_vs_rollout_horizon.png", dpi=200)
+    plt.close(fig)
+
+
+def main() -> None:
+    PLOTS_DIR.mkdir(parents=True, exist_ok=True)
+    train_rows = load_jsonl(METRICS_DIR / "train_metrics.jsonl")
+    eval_rows = load_jsonl(METRICS_DIR / "eval_metrics.jsonl")
+    final_eval = load_json(METRICS_DIR / "final_eval_latest.json")
+    clean_planner_eval = load_json(METRICS_DIR / "planner_eval_latest_clean.json")
+
+    generate_learning_curves(train_rows)
+    generate_success_plot(eval_rows, final_eval, clean_planner_eval)
+    generate_rollout_error_plot(final_eval)
+
+
+if __name__ == "__main__":
+    main()

scripts/train_ppo.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+export PYTHONPATH="${PYTHONPATH:-}:$(pwd)/src"
+python3.11 src/minidreamer/baselines/train_ppo.py \
+  --config configs/fourrooms_ppo.yaml \
+  --output-dir artifacts/ppo \
+  "$@"
+

scripts/train_world_model.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+output_dir="${MINIDREAMER_OUTPUT_DIR:-artifacts/world_model_experiment}"
+
+export PYTHONPATH="${PYTHONPATH:-}:$(pwd)/src"
+python3.11 src/train_world_model.py \
+  --config configs/fourrooms_world_model.yaml \
+  --output-dir "${output_dir}" \
+  "$@"

src/evaluate.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from minidreamer.config import load_config
+from minidreamer.data.replay_buffer import ReplayBuffer
+from minidreamer.evaluation import evaluate_random_policy, evaluate_world_model
+from minidreamer.envs.make_env import make_env_from_config
+from minidreamer.planning.evaluate_planner import evaluate_planner
+from minidreamer.serialization import load_world_model_checkpoint
+
+
+def build_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Evaluate MiniDreamer components.")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    random_parser = subparsers.add_parser("random", help="Evaluate a random policy.")
+    random_parser.add_argument("--config", type=Path, required=True)
+
+    planner_parser = subparsers.add_parser("planner", help="Evaluate a trained planner.")
+    planner_parser.add_argument("--config", type=Path, required=True)
+    planner_parser.add_argument("--checkpoint", type=Path, required=True)
+    planner_parser.add_argument(
+        "--random-action-fraction",
+        type=float,
+        default=0.0,
+        help="Optional evaluation-time action noise. Defaults to 0.0 for a clean planner evaluation.",
+    )
+
+    world_model_parser = subparsers.add_parser("world-model", help="Evaluate held-out world model metrics.")
+    world_model_parser.add_argument("--config", type=Path, required=True)
+    world_model_parser.add_argument("--checkpoint", type=Path, required=True)
+    world_model_parser.add_argument("--replay-dir", type=Path, required=True)
+    world_model_parser.add_argument("--split", type=str, default="val", choices=["train", "val", "test"])
+    return parser
+
+
+def main() -> None:
+    parser = build_arg_parser()
+    args = parser.parse_args()
+    config = load_config(args.config)
+
+    if args.command == "random":
+        print(evaluate_random_policy(config))
+        return
+
+    env = make_env_from_config(config, seed=config.get("project", {}).get("seed", 0))
+    action_dim = env.action_space.n
+    env.close()
+    model, _, metadata = load_world_model_checkpoint(args.checkpoint, action_dim=action_dim, map_location="cpu")
+
+    if args.command == "planner":
+        print({
+            "metadata": metadata,
+            **evaluate_planner(
+                config,
+                model,
+                random_action_fraction=args.random_action_fraction,
+            ),
+        })
+        return
+
+    replay = ReplayBuffer.load(args.replay_dir)
+    print({"metadata": metadata, **evaluate_world_model(config, model, replay, split=args.split)})
+
+
+if __name__ == "__main__":
+    main()

src/minidreamer/__init__.py

@@ -0,0 +1,6 @@
+"""MiniDreamer package."""
+
+__all__ = [
+    "config",
+]
+

src/minidreamer/baselines/__init__.py

@@ -0,0 +1,2 @@
+"""Baseline agents."""
+

src/minidreamer/baselines/train_ppo.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import gymnasium as gym
+import torch
+from torch import nn
+
+from minidreamer.config import ensure_run_dirs, load_config
+from minidreamer.envs.make_env import make_env_from_config
+from minidreamer.utils.common import seed_everything
+
+try:
+    from stable_baselines3 import PPO
+    from stable_baselines3.common.evaluation import evaluate_policy
+    from stable_baselines3.common.monitor import Monitor
+    from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
+    from stable_baselines3.common.vec_env import DummyVecEnv
+except ImportError as exc:  # pragma: no cover - exercised only when dependency is missing.
+    PPO = None
+    IMPORT_ERROR = exc
+else:
+    IMPORT_ERROR = None
+
+
+class MiniGridCNNExtractor(BaseFeaturesExtractor):
+    def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 256) -> None:
+        super().__init__(observation_space, features_dim)
+        if len(observation_space.shape) != 3:
+            raise ValueError(f"Expected 3D image observations, got {observation_space.shape}.")
+        self.channel_first = observation_space.shape[0] in (1, 3)
+        channels = observation_space.shape[0] if self.channel_first else observation_space.shape[2]
+        self.cnn = nn.Sequential(
+            nn.Conv2d(channels, 32, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Flatten(),
+        )
+        with torch.no_grad():
+            sample = torch.as_tensor(observation_space.sample()[None]).float()
+            if not self.channel_first:
+                sample = sample.permute(0, 3, 1, 2)
+            flattened_dim = self.cnn(sample).shape[1]
+        self.linear = nn.Sequential(nn.Linear(flattened_dim, features_dim), nn.ReLU())
+
+    def forward(self, observations: torch.Tensor) -> torch.Tensor:
+        if observations.dim() == 4 and observations.shape[1] not in (1, 3):
+            observations = observations.permute(0, 3, 1, 2)
+        return self.linear(self.cnn(observations.float()))
+
+
+def build_env(config: dict, seed: int, rank: int):
+    def _make():
+        env = make_env_from_config(config, seed=seed + rank)
+        return Monitor(env)
+
+    return _make
+
+
+def train_ppo(config: dict, output_dir: str | Path) -> dict[str, float]:
+    if PPO is None:
+        raise ImportError(
+            "stable-baselines3 is required for PPO training."
+        ) from IMPORT_ERROR
+    ppo_cfg = config["ppo"]
+    seed = config.get("project", {}).get("seed", 0)
+    seed_everything(seed)
+    run_dirs = ensure_run_dirs(output_dir)
+
+    env_fns = [build_env(config, seed, rank) for rank in range(ppo_cfg.get("num_envs", 4))]
+    vec_env = DummyVecEnv(env_fns)
+    policy_kwargs = {
+        "features_extractor_class": MiniGridCNNExtractor,
+        "features_extractor_kwargs": {"features_dim": ppo_cfg.get("features_dim", 256)},
+    }
+    model = PPO(
+        "CnnPolicy",
+        vec_env,
+        policy_kwargs=policy_kwargs,
+        learning_rate=ppo_cfg.get("learning_rate", 3e-4),
+        n_steps=ppo_cfg.get("n_steps", 256),
+        batch_size=ppo_cfg.get("batch_size", 256),
+        n_epochs=ppo_cfg.get("n_epochs", 4),
+        gamma=ppo_cfg.get("gamma", 0.99),
+        gae_lambda=ppo_cfg.get("gae_lambda", 0.95),
+        clip_range=ppo_cfg.get("clip_range", 0.2),
+        ent_coef=ppo_cfg.get("ent_coef", 0.01),
+        vf_coef=ppo_cfg.get("vf_coef", 0.5),
+        seed=seed,
+        device=ppo_cfg.get("device", "auto"),
+        verbose=1,
+    )
+    model.learn(total_timesteps=ppo_cfg["total_timesteps"])
+    model.save(Path(run_dirs["checkpoints"]) / "ppo_latest")
+    mean_reward, std_reward = evaluate_policy(
+        model,
+        vec_env,
+        n_eval_episodes=config["evaluation"]["episodes"],
+        deterministic=True,
+    )
+    vec_env.close()
+    return {"mean_reward": float(mean_reward), "std_reward": float(std_reward)}
+
+
+def build_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Train a PPO baseline on MiniGrid pixels.")
+    parser.add_argument("--config", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    return parser
+
+
+def main() -> None:
+    parser = build_arg_parser()
+    args = parser.parse_args()
+    config = load_config(args.config)
+    summary = train_ppo(config, args.output_dir)
+    print(summary)
+
+
+if __name__ == "__main__":
+    main()

src/minidreamer/config.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import copy
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+ConfigDict = dict[str, Any]
+
+
+def load_config(path: str | Path) -> ConfigDict:
+    path = Path(path)
+    with path.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    if not isinstance(config, dict):
+        raise ValueError(f"Config at {path} must be a mapping.")
+    return config
+
+
+def merge_dicts(base: ConfigDict, overrides: ConfigDict) -> ConfigDict:
+    merged = copy.deepcopy(base)
+    for key, value in overrides.items():
+        if isinstance(value, dict) and isinstance(merged.get(key), dict):
+            merged[key] = merge_dicts(merged[key], value)
+        else:
+            merged[key] = value
+    return merged
+
+
+def save_config(config: ConfigDict, path: str | Path) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as handle:
+        yaml.safe_dump(config, handle, sort_keys=False)
+
+
+def ensure_run_dirs(base_dir: str | Path) -> dict[str, Path]:
+    base = Path(base_dir)
+    paths = {
+        "base": base,
+        "checkpoints": base / "checkpoints",
+        "metrics": base / "metrics",
+        "plots": base / "plots",
+        "replay": base / "replay",
+    }
+    for path in paths.values():
+        path.mkdir(parents=True, exist_ok=True)
+    return paths
+
+
+def deep_get(config: ConfigDict, *keys: str, default: Any = None) -> Any:
+    current: Any = config
+    for key in keys:
+        if not isinstance(current, dict) or key not in current:
+            return default
+        current = current[key]
+    return current
+

src/minidreamer/envs/__init__.py

@@ -0,0 +1,2 @@
+"""Environment helpers."""
+

src/minidreamer/envs/make_env.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable
+
+import gymnasium as gym
+import numpy as np
+import torch
+from gymnasium import spaces
+from minigrid.wrappers import ImgObsWrapper, RGBImgPartialObsWrapper
+from PIL import Image
+
+
+@dataclass(frozen=True)
+class EnvSpec:
+    env_id: str
+    resize: tuple[int, int] = (64, 64)
+    normalize_obs: bool = True
+    rgb_partial_obs: bool = True
+    image_only: bool = True
+
+
+class ResizeNormalizeObservation(gym.ObservationWrapper):
+    def __init__(
+        self,
+        env: gym.Env,
+        resize: tuple[int, int] | None = (64, 64),
+        normalize: bool = True,
+    ) -> None:
+        super().__init__(env)
+        self.resize = resize
+        self.normalize = normalize
+        base_space = env.observation_space
+        if not isinstance(base_space, spaces.Box):
+            raise TypeError("MiniDreamer expects a Box observation space after wrappers.")
+        channels = base_space.shape[-1]
+        if resize is None:
+            height, width = base_space.shape[:2]
+        else:
+            height, width = resize
+        low, high = (0.0, 1.0) if normalize else (0, 255)
+        dtype = np.float32 if normalize else np.uint8
+        self.observation_space = spaces.Box(
+            low=low,
+            high=high,
+            shape=(height, width, channels),
+            dtype=dtype,
+        )
+
+    def observation(self, observation: np.ndarray) -> np.ndarray:
+        obs = observation
+        if self.resize is not None and tuple(obs.shape[:2]) != self.resize:
+            pil_image = Image.fromarray(obs.astype(np.uint8))
+            pil_image = pil_image.resize((self.resize[1], self.resize[0]), Image.Resampling.BILINEAR)
+            obs = np.asarray(pil_image)
+        if self.normalize:
+            return obs.astype(np.float32) / 255.0
+        return obs.astype(np.uint8)
+
+
+def make_env(
+    env_id: str = "MiniGrid-FourRooms-v0",
+    seed: int | None = None,
+    resize: tuple[int, int] = (64, 64),
+    normalize_obs: bool = True,
+    rgb_partial_obs: bool = True,
+    image_only: bool = True,
+    render_mode: str | None = None,
+) -> gym.Env:
+    env = gym.make(env_id, render_mode=render_mode)
+    env = gym.wrappers.RecordEpisodeStatistics(env)
+    if rgb_partial_obs:
+        env = RGBImgPartialObsWrapper(env)
+    if image_only:
+        env = ImgObsWrapper(env)
+    env = ResizeNormalizeObservation(env, resize=resize, normalize=normalize_obs)
+    if seed is not None:
+        env.reset(seed=seed)
+        env.action_space.seed(seed)
+    return env
+
+
+def make_env_from_config(config: dict, seed: int | None = None) -> gym.Env:
+    env_cfg = config["env"]
+    return make_env(
+        env_id=env_cfg["id"],
+        seed=seed,
+        resize=tuple(env_cfg.get("resize", (64, 64))),
+        normalize_obs=env_cfg.get("normalize_obs", True),
+        rgb_partial_obs=env_cfg.get("rgb_partial_obs", True),
+        image_only=env_cfg.get("image_only", True),
+    )
+
+
+def observation_to_tensor(observation: np.ndarray, device: torch.device | None = None) -> torch.Tensor:
+    if observation.ndim != 3:
+        raise ValueError(f"Expected HWC observation, got shape {observation.shape}.")
+    tensor = torch.from_numpy(observation).permute(2, 0, 1).float()
+    return tensor.to(device) if device is not None else tensor
+
+
+def batch_observations_to_tensor(
+    observations: np.ndarray,
+    device: torch.device | None = None,
+) -> torch.Tensor:
+    if observations.ndim != 5:
+        raise ValueError(f"Expected BT HWC observations, got shape {observations.shape}.")
+    tensor = torch.from_numpy(observations).permute(0, 1, 4, 2, 3).float()
+    return tensor.to(device) if device is not None else tensor
+
+
+def action_subset(action_space_n: int, names: Iterable[str] | None = None) -> list[int]:
+    if names is None:
+        return list(range(action_space_n))
+    lookup = {
+        "left": 0,
+        "right": 1,
+        "forward": 2,
+        "pickup": 3,
+        "drop": 4,
+        "toggle": 5,
+        "done": 6,
+    }
+    return [lookup[name] for name in names]
+

src/minidreamer/evaluation.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from minidreamer.data.replay_buffer import Episode, ReplayBuffer
+from minidreamer.envs.make_env import make_env_from_config
+from minidreamer.models.world_model import WorldModel
+
+
+def evaluate_random_policy(config: dict, episodes: int | None = None, seed: int | None = None) -> dict[str, float]:
+    eval_cfg = config["evaluation"]
+    episodes = episodes or eval_cfg["episodes"]
+    seed = config.get("project", {}).get("seed", 0) if seed is None else seed
+    env = make_env_from_config(config, seed=seed)
+    rng = np.random.default_rng(seed)
+    returns = []
+    lengths = []
+    successes = []
+    for episode_idx in range(episodes):
+        obs, _ = env.reset(seed=seed + episode_idx)
+        total_return = 0.0
+        terminated = False
+        truncated = False
+        length = 0
+        while not (terminated or truncated):
+            action = int(rng.integers(0, env.action_space.n))
+            obs, reward, terminated, truncated, _ = env.step(action)
+            total_return += float(reward)
+            length += 1
+        returns.append(total_return)
+        lengths.append(length)
+        successes.append(float(terminated and total_return > 0.0))
+    env.close()
+    returns_array = np.asarray(returns, dtype=np.float32)
+    lengths_array = np.asarray(lengths, dtype=np.float32)
+    successes_array = np.asarray(successes, dtype=np.float32)
+    return {
+        "success_rate": float(successes_array.mean()),
+        "mean_return": float(returns_array.mean()),
+        "median_return": float(np.median(returns_array)),
+        "mean_episode_length": float(lengths_array.mean()),
+    }
+
+
+def _episode_to_batch(episode: Episode, device: torch.device) -> dict[str, torch.Tensor]:
+    batch = {
+        "obs": torch.from_numpy(episode.obs[None]).permute(0, 1, 4, 2, 3).float().to(device),
+        "actions": torch.from_numpy(episode.actions[None]).long().to(device),
+        "rewards": torch.from_numpy(episode.rewards[None]).float().to(device),
+        "terminated": torch.from_numpy(episode.terminated[None]).float().to(device),
+        "truncated": torch.from_numpy(episode.truncated[None]).float().to(device),
+        "done": torch.from_numpy(episode.done[None]).float().to(device),
+        "mask": torch.ones((1, episode.length), dtype=torch.float32, device=device),
+    }
+    return batch
+
+
+def _sequence_state(model: WorldModel, episode: Episode, start_idx: int) -> Any:
+    state = model.posterior_step(model.initial_state(1), None, episode.obs[0], sample=False)
+    for idx in range(start_idx):
+        state = model.posterior_step(state, int(episode.actions[idx]), episode.obs[idx + 1], sample=False)
+    return state
+
+
+def _discounted_return(rewards: np.ndarray, done: np.ndarray, discount: float) -> float:
+    total = 0.0
+    alive = 1.0
+    for step, reward in enumerate(rewards):
+        total += (discount**step) * alive * float(reward)
+        alive *= 1.0 - float(done[step])
+    return total
+
+
+def evaluate_world_model(
+    config: dict,
+    model: WorldModel,
+    replay: ReplayBuffer,
+    split: str = "val",
+    max_episodes: int | None = None,
+) -> dict[str, float]:
+    device = model.device
+    model.eval()
+    metrics: dict[str, list[float]] = defaultdict(list)
+    horizons = [1, 5, 10]
+    discount = float(config["planner"]["discount"])
+    episodes = replay.episode_ids(split)
+    if max_episodes is not None:
+        episodes = episodes[:max_episodes]
+
+    with torch.no_grad():
+        for episode_id in episodes:
+            episode = replay.episodes[episode_id]
+            batch = _episode_to_batch(episode, device)
+            outputs = model.observe_sequence(batch["obs"], batch["actions"], sample=False)
+            reward_mse = F.mse_loss(outputs.reward_pred, batch["rewards"], reduction="none").mean()
+            done_bce = F.binary_cross_entropy_with_logits(outputs.done_logits, batch["done"], reduction="none").mean()
+            kl = model.rssm.kl_divergence(
+                outputs.post_mean,
+                outputs.post_std,
+                outputs.prior_mean,
+                outputs.prior_std,
+            ).mean()
+            metrics["reward_mse"].append(float(reward_mse.cpu()))
+            metrics["done_bce"].append(float(done_bce.cpu()))
+            metrics["kl_loss"].append(float(kl.cpu()))
+            if outputs.reconstructions is not None:
+                recon_mse = F.mse_loss(outputs.reconstructions, batch["obs"][:, 1:], reduction="none").mean()
+                metrics["reconstruction_mse"].append(float(recon_mse.cpu()))
+
+            for horizon in horizons:
+                if episode.length < horizon:
+                    continue
+                reward_errors = []
+                done_correct = []
+                imagined_returns = []
+                real_returns = []
+                for start_idx in range(episode.length - horizon + 1):
+                    state = _sequence_state(model, episode, start_idx)
+                    actions = torch.from_numpy(episode.actions[start_idx : start_idx + horizon]).long().to(device)
+                    rollout = model.score_action_sequences(
+                        state,
+                        actions.unsqueeze(0),
+                        discount=discount,
+                        use_done_mask=True,
+                    )
+                    reward_pred = rollout["reward_pred"].squeeze(0).cpu().numpy()
+                    done_prob = rollout["done_prob"].squeeze(0).cpu().numpy()
+                    done_pred = (done_prob >= 0.5).astype(np.float32)
+                    real_rewards = episode.rewards[start_idx : start_idx + horizon]
+                    real_done = episode.done[start_idx : start_idx + horizon]
+                    reward_errors.append(np.mean((reward_pred - real_rewards) ** 2))
+                    done_correct.append(np.mean(done_pred == real_done))
+                    imagined_returns.append(float(rollout["scores"].squeeze(0).cpu()))
+                    real_returns.append(_discounted_return(real_rewards, real_done, discount))
+                metrics[f"open_loop_reward_error_h{horizon}"].append(float(np.mean(reward_errors)))
+                metrics[f"open_loop_done_accuracy_h{horizon}"].append(float(np.mean(done_correct)))
+                if len(imagined_returns) > 1 and np.std(real_returns) > 0.0 and np.std(imagined_returns) > 0.0:
+                    correlation = np.corrcoef(imagined_returns, real_returns)[0, 1]
+                    metrics["imagined_return_vs_real_return_correlation"].append(float(correlation))
+
+    return {name: float(np.mean(values)) for name, values in metrics.items() if values}
+

src/minidreamer/models/__init__.py

@@ -0,0 +1,2 @@
+"""Model components for MiniDreamer."""
+

src/minidreamer/models/decoder.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from torch import nn
+
+
+class ConvDecoder(nn.Module):
+    def __init__(self, feature_dim: int, out_channels: int = 3) -> None:
+        super().__init__()
+        self.projection = nn.Sequential(
+            nn.Linear(feature_dim, 256 * 4 * 4),
+            nn.ReLU(),
+        )
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(32, out_channels, kernel_size=4, stride=2, padding=1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, features):
+        hidden = self.projection(features).view(-1, 256, 4, 4)
+        return self.decoder(hidden)
+

src/minidreamer/models/encoder.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class ConvEncoder(nn.Module):
+    def __init__(self, in_channels: int = 3, embedding_dim: int = 256) -> None:
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, 32, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(),
+        )
+        self.projection = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(256 * 4 * 4, embedding_dim),
+            nn.ReLU(),
+        )
+
+    def forward(self, obs: torch.Tensor) -> torch.Tensor:
+        hidden = self.conv(obs)
+        return self.projection(hidden)
+

src/minidreamer/models/heads.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from torch import nn
+
+
+class MLPHead(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, hidden_dim: int = 256) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden_dim),
+            nn.ELU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ELU(),
+            nn.Linear(hidden_dim, out_dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class RewardHead(MLPHead):
+    def __init__(self, in_dim: int, hidden_dim: int = 256) -> None:
+        super().__init__(in_dim=in_dim, out_dim=1, hidden_dim=hidden_dim)
+
+
+class DoneHead(MLPHead):
+    def __init__(self, in_dim: int, hidden_dim: int = 256) -> None:
+        super().__init__(in_dim=in_dim, out_dim=1, hidden_dim=hidden_dim)
+

src/minidreamer/models/rssm.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+@dataclass
+class RSSMState:
+    deter: torch.Tensor
+    stoch: torch.Tensor
+    mean: torch.Tensor
+    std: torch.Tensor
+
+    def features(self) -> torch.Tensor:
+        return torch.cat([self.deter, self.stoch], dim=-1)
+
+    def detach(self) -> "RSSMState":
+        return RSSMState(
+            deter=self.deter.detach(),
+            stoch=self.stoch.detach(),
+            mean=self.mean.detach(),
+            std=self.std.detach(),
+        )
+
+    def repeat(self, count: int) -> "RSSMState":
+        return RSSMState(
+            deter=self.deter.repeat(count, 1),
+            stoch=self.stoch.repeat(count, 1),
+            mean=self.mean.repeat(count, 1),
+            std=self.std.repeat(count, 1),
+        )
+
+
+class RSSM(nn.Module):
+    def __init__(
+        self,
+        action_dim: int,
+        embedding_dim: int = 256,
+        deter_dim: int = 256,
+        stoch_dim: int = 32,
+        hidden_dim: int = 256,
+        min_std: float = 0.1,
+    ) -> None:
+        super().__init__()
+        self.action_dim = action_dim
+        self.embedding_dim = embedding_dim
+        self.deter_dim = deter_dim
+        self.stoch_dim = stoch_dim
+        self.hidden_dim = hidden_dim
+        self.min_std = min_std
+
+        self.input_net = nn.Sequential(
+            nn.Linear(stoch_dim + action_dim, hidden_dim),
+            nn.ELU(),
+        )
+        self.gru = nn.GRUCell(hidden_dim, deter_dim)
+        self.prior_net = nn.Sequential(
+            nn.Linear(deter_dim, hidden_dim),
+            nn.ELU(),
+            nn.Linear(hidden_dim, 2 * stoch_dim),
+        )
+        self.posterior_net = nn.Sequential(
+            nn.Linear(deter_dim + embedding_dim, hidden_dim),
+            nn.ELU(),
+            nn.Linear(hidden_dim, 2 * stoch_dim),
+        )
+
+    def initial(self, batch_size: int, device: torch.device) -> RSSMState:
+        zeros_deter = torch.zeros(batch_size, self.deter_dim, device=device)
+        zeros_stoch = torch.zeros(batch_size, self.stoch_dim, device=device)
+        return RSSMState(
+            deter=zeros_deter,
+            stoch=zeros_stoch,
+            mean=zeros_stoch,
+            std=torch.ones_like(zeros_stoch),
+        )
+
+    def _stats(self, tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        mean, raw_std = torch.chunk(tensor, 2, dim=-1)
+        std = F.softplus(raw_std) + self.min_std
+        return mean, std
+
+    def _action_one_hot(self, action: torch.Tensor) -> torch.Tensor:
+        action = action.long().view(-1)
+        return F.one_hot(action, num_classes=self.action_dim).float()
+
+    def _next_deter(self, prev_state: RSSMState, action: torch.Tensor) -> torch.Tensor:
+        action_one_hot = self._action_one_hot(action)
+        gru_input = self.input_net(torch.cat([prev_state.stoch, action_one_hot], dim=-1))
+        return self.gru(gru_input, prev_state.deter)
+
+    def prior(self, deter: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return self._stats(self.prior_net(deter))
+
+    def posterior(self, deter: torch.Tensor, embed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return self._stats(self.posterior_net(torch.cat([deter, embed], dim=-1)))
+
+    @staticmethod
+    def sample(mean: torch.Tensor, std: torch.Tensor, sample: bool = True) -> torch.Tensor:
+        if sample:
+            return mean + torch.randn_like(std) * std
+        return mean
+
+    def observe(
+        self,
+        prev_state: RSSMState,
+        prev_action: torch.Tensor | None,
+        embed: torch.Tensor,
+        sample: bool = True,
+    ) -> tuple[RSSMState, tuple[torch.Tensor, torch.Tensor]]:
+        if prev_action is None:
+            deter = prev_state.deter
+        else:
+            deter = self._next_deter(prev_state, prev_action)
+        prior_mean, prior_std = self.prior(deter)
+        post_mean, post_std = self.posterior(deter, embed)
+        stoch = self.sample(post_mean, post_std, sample=sample)
+        state = RSSMState(deter=deter, stoch=stoch, mean=post_mean, std=post_std)
+        return state, (prior_mean, prior_std)
+
+    def imagine(
+        self,
+        prev_state: RSSMState,
+        action: torch.Tensor,
+        sample: bool = False,
+    ) -> tuple[RSSMState, tuple[torch.Tensor, torch.Tensor]]:
+        deter = self._next_deter(prev_state, action)
+        mean, std = self.prior(deter)
+        stoch = self.sample(mean, std, sample=sample)
+        state = RSSMState(deter=deter, stoch=stoch, mean=mean, std=std)
+        return state, (mean, std)
+
+    @staticmethod
+    def kl_divergence(
+        post_mean: torch.Tensor,
+        post_std: torch.Tensor,
+        prior_mean: torch.Tensor,
+        prior_std: torch.Tensor,
+    ) -> torch.Tensor:
+        log_var_ratio = 2.0 * (torch.log(prior_std) - torch.log(post_std))
+        var_ratio = (post_std / prior_std) ** 2
+        mean_term = ((post_mean - prior_mean) / prior_std) ** 2
+        return 0.5 * torch.sum(var_ratio + mean_term + log_var_ratio - 1.0, dim=-1)
+

src/minidreamer/models/world_model.py

@@ -0,0 +1,267 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from minidreamer.models.decoder import ConvDecoder
+from minidreamer.models.encoder import ConvEncoder
+from minidreamer.models.heads import DoneHead, RewardHead
+from minidreamer.models.rssm import RSSM, RSSMState
+from minidreamer.utils.common import masked_mean
+
+
+@dataclass
+class WorldModelOutputs:
+    states: list[RSSMState]
+    prior_mean: torch.Tensor
+    prior_std: torch.Tensor
+    post_mean: torch.Tensor
+    post_std: torch.Tensor
+    reward_pred: torch.Tensor
+    done_logits: torch.Tensor
+    reconstructions: torch.Tensor | None
+
+
+class WorldModel(nn.Module):
+    def __init__(
+        self,
+        action_dim: int,
+        embedding_dim: int = 256,
+        deter_dim: int = 256,
+        stoch_dim: int = 32,
+        hidden_dim: int = 256,
+        use_decoder: bool = True,
+        min_std: float = 0.1,
+        obs_channels: int = 3,
+    ) -> None:
+        super().__init__()
+        self.action_dim = action_dim
+        self.embedding_dim = embedding_dim
+        self.deter_dim = deter_dim
+        self.stoch_dim = stoch_dim
+        self.hidden_dim = hidden_dim
+        self.use_decoder = use_decoder
+        self.obs_channels = obs_channels
+
+        self.encoder = ConvEncoder(in_channels=obs_channels, embedding_dim=embedding_dim)
+        self.rssm = RSSM(
+            action_dim=action_dim,
+            embedding_dim=embedding_dim,
+            deter_dim=deter_dim,
+            stoch_dim=stoch_dim,
+            hidden_dim=hidden_dim,
+            min_std=min_std,
+        )
+        feature_dim = deter_dim + stoch_dim
+        self.reward_head = RewardHead(feature_dim, hidden_dim=hidden_dim)
+        self.done_head = DoneHead(feature_dim, hidden_dim=hidden_dim)
+        self.decoder = ConvDecoder(feature_dim, out_channels=obs_channels) if use_decoder else None
+
+    @classmethod
+    def from_config(cls, config: dict, action_dim: int, obs_channels: int = 3) -> "WorldModel":
+        model_cfg = config["model"]
+        return cls(
+            action_dim=action_dim,
+            embedding_dim=model_cfg["embedding_dim"],
+            deter_dim=model_cfg["deter_dim"],
+            stoch_dim=model_cfg["stoch_dim"],
+            hidden_dim=model_cfg["hidden_dim"],
+            use_decoder=model_cfg.get("use_decoder", True),
+            min_std=model_cfg.get("min_std", 0.1),
+            obs_channels=obs_channels,
+        )
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def initial_state(self, batch_size: int) -> RSSMState:
+        return self.rssm.initial(batch_size=batch_size, device=self.device)
+
+    def encode(self, obs: torch.Tensor) -> torch.Tensor:
+        if obs.dim() < 4:
+            raise ValueError(f"Expected at least 4 dims for observations, got {obs.shape}.")
+        leading_shape = obs.shape[:-3]
+        flat_obs = obs.reshape(-1, *obs.shape[-3:])
+        flat_embeddings = self.encoder(flat_obs)
+        return flat_embeddings.reshape(*leading_shape, -1)
+
+    def observe_sequence(
+        self,
+        obs: torch.Tensor,
+        actions: torch.Tensor,
+        sample: bool = True,
+    ) -> WorldModelOutputs:
+        if obs.dim() != 5:
+            raise ValueError(f"Expected obs shape [B, T+1, C, H, W], got {obs.shape}.")
+        if actions.dim() != 2:
+            raise ValueError(f"Expected actions shape [B, T], got {actions.shape}.")
+        batch_size, time_steps = actions.shape
+        embeddings = self.encode(obs)
+        state, _ = self.rssm.observe(self.initial_state(batch_size), None, embeddings[:, 0], sample=sample)
+
+        states = [state]
+        prior_means = []
+        prior_stds = []
+        post_means = []
+        post_stds = []
+        rewards = []
+        done_logits = []
+        reconstructions = []
+
+        for t in range(time_steps):
+            next_state, (prior_mean, prior_std) = self.rssm.observe(
+                state,
+                actions[:, t],
+                embeddings[:, t + 1],
+                sample=sample,
+            )
+            features = next_state.features()
+            prior_means.append(prior_mean)
+            prior_stds.append(prior_std)
+            post_means.append(next_state.mean)
+            post_stds.append(next_state.std)
+            rewards.append(self.reward_head(features).squeeze(-1))
+            done_logits.append(self.done_head(features).squeeze(-1))
+            if self.decoder is not None:
+                reconstructions.append(self.decoder(features))
+            states.append(next_state)
+            state = next_state
+
+        recon_tensor = torch.stack(reconstructions, dim=1) if reconstructions else None
+        return WorldModelOutputs(
+            states=states,
+            prior_mean=torch.stack(prior_means, dim=1),
+            prior_std=torch.stack(prior_stds, dim=1),
+            post_mean=torch.stack(post_means, dim=1),
+            post_std=torch.stack(post_stds, dim=1),
+            reward_pred=torch.stack(rewards, dim=1),
+            done_logits=torch.stack(done_logits, dim=1),
+            reconstructions=recon_tensor,
+        )
+
+    def compute_losses(self, batch: dict[str, torch.Tensor], config: dict[str, Any]) -> dict[str, torch.Tensor]:
+        outputs = self.observe_sequence(batch["obs"], batch["actions"], sample=True)
+        training_cfg = config["training"]
+        rewards = batch["rewards"]
+        done = batch["done"]
+        mask = batch["mask"]
+
+        reward_loss = masked_mean(F.mse_loss(outputs.reward_pred, rewards, reduction="none"), mask)
+        done_loss = masked_mean(
+            F.binary_cross_entropy_with_logits(outputs.done_logits, done, reduction="none"),
+            mask,
+        )
+        kl_per_step = self.rssm.kl_divergence(
+            outputs.post_mean,
+            outputs.post_std,
+            outputs.prior_mean,
+            outputs.prior_std,
+        )
+        free_nats = torch.full_like(kl_per_step, float(training_cfg.get("free_nats", 1.0)))
+        kl_loss = masked_mean(torch.maximum(kl_per_step, free_nats), mask)
+
+        if outputs.reconstructions is not None and training_cfg.get("beta_recon", 0.0) > 0.0:
+            recon_target = batch["obs"][:, 1:]
+            recon_error = F.mse_loss(outputs.reconstructions, recon_target, reduction="none").mean(dim=(2, 3, 4))
+            recon_loss = masked_mean(recon_error, mask)
+        else:
+            recon_loss = torch.zeros((), device=self.device)
+
+        total_loss = (
+            float(training_cfg.get("beta_reward", 1.0)) * reward_loss
+            + float(training_cfg.get("beta_done", 1.0)) * done_loss
+            + float(training_cfg.get("beta_kl", 1.0)) * kl_loss
+            + float(training_cfg.get("beta_recon", 0.0)) * recon_loss
+        )
+        return {
+            "loss": total_loss,
+            "reward_loss": reward_loss.detach(),
+            "done_loss": done_loss.detach(),
+            "kl_loss": kl_loss.detach(),
+            "recon_loss": recon_loss.detach(),
+            "reward_mse": F.mse_loss(outputs.reward_pred, rewards, reduction="none").mul(mask).sum() / mask.sum().clamp_min(1.0),
+            "done_bce": F.binary_cross_entropy_with_logits(outputs.done_logits, done, reduction="none").mul(mask).sum() / mask.sum().clamp_min(1.0),
+        }
+
+    def posterior_step(
+        self,
+        prev_state: RSSMState,
+        prev_action: int | torch.Tensor | None,
+        observation: np.ndarray | torch.Tensor,
+        sample: bool = False,
+    ) -> RSSMState:
+        obs_tensor = self._prepare_single_observation(observation)
+        embed = self.encode(obs_tensor)
+        if prev_action is None:
+            action_tensor = None
+        else:
+            action_tensor = torch.as_tensor(prev_action, device=self.device).view(1)
+        state, _ = self.rssm.observe(prev_state, action_tensor, embed, sample=sample)
+        return state
+
+    def imagine_rollout(
+        self,
+        start_state: RSSMState,
+        action_sequences: torch.Tensor,
+        sample: bool = False,
+    ) -> dict[str, torch.Tensor]:
+        if action_sequences.dim() == 1:
+            action_sequences = action_sequences.unsqueeze(0)
+        batch_size, horizon = action_sequences.shape
+        state = start_state if start_state.deter.shape[0] == batch_size else start_state.repeat(batch_size)
+        rewards = []
+        done_logits = []
+        states = []
+        for t in range(horizon):
+            state, _ = self.rssm.imagine(state, action_sequences[:, t], sample=sample)
+            features = state.features()
+            rewards.append(self.reward_head(features).squeeze(-1))
+            done_logits.append(self.done_head(features).squeeze(-1))
+            states.append(state)
+        return {
+            "states": states,
+            "reward_pred": torch.stack(rewards, dim=1),
+            "done_logits": torch.stack(done_logits, dim=1),
+        }
+
+    def score_action_sequences(
+        self,
+        start_state: RSSMState,
+        action_sequences: torch.Tensor,
+        discount: float = 0.99,
+        use_done_mask: bool = True,
+    ) -> dict[str, torch.Tensor]:
+        rollout = self.imagine_rollout(start_state, action_sequences, sample=False)
+        reward_pred = rollout["reward_pred"]
+        done_prob = torch.sigmoid(rollout["done_logits"])
+        alive = torch.ones(reward_pred.shape[0], device=self.device)
+        scores = torch.zeros_like(alive)
+        for t in range(reward_pred.shape[1]):
+            scores = scores + (discount**t) * alive * reward_pred[:, t]
+            if use_done_mask:
+                alive = alive * (1.0 - done_prob[:, t])
+        rollout["scores"] = scores
+        rollout["done_prob"] = done_prob
+        return rollout
+
+    def _prepare_single_observation(self, observation: np.ndarray | torch.Tensor) -> torch.Tensor:
+        if torch.is_tensor(observation):
+            obs_tensor = observation.to(self.device).float()
+        else:
+            obs_tensor = torch.as_tensor(observation, device=self.device).float()
+        if obs_tensor.dim() == 3 and obs_tensor.shape[-1] in (1, 3):
+            obs_tensor = obs_tensor.permute(2, 0, 1).unsqueeze(0)
+        elif obs_tensor.dim() == 3 and obs_tensor.shape[0] in (1, 3):
+            obs_tensor = obs_tensor.unsqueeze(0)
+        elif obs_tensor.dim() == 4:
+            pass
+        else:
+            raise ValueError(f"Unsupported observation shape {tuple(obs_tensor.shape)}.")
+        return obs_tensor
+

src/minidreamer/planning/__init__.py

@@ -0,0 +1,2 @@
+"""Planning utilities."""
+

src/minidreamer/planning/cem.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+from minidreamer.models.rssm import RSSMState
+from minidreamer.models.world_model import WorldModel
+
+
+@dataclass
+class PlannerOutput:
+    action: int
+    sequence: list[int]
+    score: float
+    policy: torch.Tensor
+    entropy: float
+
+
+class DiscreteCEMPlanner:
+    def __init__(
+        self,
+        world_model: WorldModel,
+        action_dim: int,
+        horizon: int = 8,
+        candidates: int = 256,
+        elites: int = 32,
+        iterations: int = 4,
+        discount: float = 0.99,
+        use_done_mask: bool = True,
+        smoothing: float = 1e-3,
+    ) -> None:
+        self.world_model = world_model
+        self.action_dim = action_dim
+        self.horizon = horizon
+        self.candidates = candidates
+        self.elites = min(elites, candidates)
+        self.iterations = iterations
+        self.discount = discount
+        self.use_done_mask = use_done_mask
+        self.smoothing = smoothing
+
+    @classmethod
+    def from_config(cls, world_model: WorldModel, action_dim: int, config: dict) -> "DiscreteCEMPlanner":
+        planner_cfg = config["planner"]
+        return cls(
+            world_model=world_model,
+            action_dim=action_dim,
+            horizon=planner_cfg["horizon"],
+            candidates=planner_cfg["candidates"],
+            elites=planner_cfg["elites"],
+            iterations=planner_cfg["iterations"],
+            discount=planner_cfg["discount"],
+            use_done_mask=planner_cfg.get("use_done_mask", True),
+        )
+
+    def _sample_sequences(self, probs: torch.Tensor) -> torch.Tensor:
+        flat = probs.unsqueeze(0).expand(self.candidates, -1, -1).reshape(-1, self.action_dim)
+        sampled = torch.multinomial(flat, num_samples=1, replacement=True)
+        return sampled.view(self.candidates, self.horizon)
+
+    def plan(self, state: RSSMState) -> PlannerOutput:
+        device = self.world_model.device
+        probs = torch.full(
+            (self.horizon, self.action_dim),
+            fill_value=1.0 / self.action_dim,
+            device=device,
+        )
+        best_sequence = None
+        best_score = torch.tensor(float("-inf"), device=device)
+
+        for _ in range(self.iterations):
+            action_sequences = self._sample_sequences(probs)
+            scores = self.world_model.score_action_sequences(
+                state,
+                action_sequences,
+                discount=self.discount,
+                use_done_mask=self.use_done_mask,
+            )["scores"]
+            elite_indices = torch.topk(scores, k=self.elites, largest=True).indices
+            elites = action_sequences[elite_indices]
+            elite_freq = F.one_hot(elites, num_classes=self.action_dim).float().mean(dim=0)
+            probs = elite_freq + self.smoothing
+            probs = probs / probs.sum(dim=-1, keepdim=True)
+
+            iteration_best_idx = scores.argmax()
+            iteration_best_score = scores[iteration_best_idx]
+            if iteration_best_score > best_score:
+                best_score = iteration_best_score
+                best_sequence = action_sequences[iteration_best_idx]
+
+        if best_sequence is None:
+            raise RuntimeError("CEM planner failed to sample any action sequence.")
+        entropy = float((-(probs * probs.clamp_min(1e-8).log()).sum(dim=-1)).mean().detach().cpu())
+        return PlannerOutput(
+            action=int(best_sequence[0].item()),
+            sequence=[int(action.item()) for action in best_sequence],
+            score=float(best_score.detach().cpu()),
+            policy=probs.detach().cpu(),
+            entropy=entropy,
+        )
+

src/minidreamer/planning/evaluate_planner.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from minidreamer.envs.make_env import make_env_from_config
+from minidreamer.planning.cem import DiscreteCEMPlanner
+
+
+@dataclass
+class PlannerEpisode:
+    success: bool
+    total_return: float
+    length: int
+    terminated: bool
+    truncated: bool
+    planner_entropy: float
+
+
+def run_planner_episode(
+    env,
+    world_model,
+    planner: DiscreteCEMPlanner,
+    rng: np.random.Generator,
+    seed: int | None = None,
+    random_action_fraction: float = 0.0,
+) -> PlannerEpisode:
+    obs, _ = env.reset(seed=seed)
+    world_model.eval()
+    with torch.no_grad():
+        state = world_model.posterior_step(world_model.initial_state(1), None, obs, sample=False)
+        total_return = 0.0
+        length = 0
+        terminated = False
+        truncated = False
+        entropies: list[float] = []
+
+        while not (terminated or truncated):
+            if rng.random() < random_action_fraction:
+                action = int(env.action_space.sample())
+            else:
+                plan = planner.plan(state)
+                action = plan.action
+                entropies.append(plan.entropy)
+            obs, reward, terminated, truncated, _ = env.step(action)
+            total_return += float(reward)
+            length += 1
+            if not (terminated or truncated):
+                state = world_model.posterior_step(state, action, obs, sample=False)
+
+    return PlannerEpisode(
+        success=bool(terminated and total_return > 0.0),
+        total_return=total_return,
+        length=length,
+        terminated=bool(terminated),
+        truncated=bool(truncated),
+        planner_entropy=float(np.mean(entropies)) if entropies else float("nan"),
+    )
+
+
+def evaluate_planner(
+    config: dict,
+    world_model,
+    episodes: int | None = None,
+    seed: int | None = None,
+    random_action_fraction: float = 0.0,
+) -> dict[str, float]:
+    eval_cfg = config["evaluation"]
+    episodes = episodes or eval_cfg["episodes"]
+    seed = config.get("project", {}).get("seed", 0) if seed is None else seed
+    env = make_env_from_config(config, seed=seed)
+    planner = DiscreteCEMPlanner.from_config(world_model, env.action_space.n, config)
+    rng = np.random.default_rng(seed)
+    results = [
+        run_planner_episode(
+            env,
+            world_model,
+            planner,
+            rng,
+            seed=seed + episode_idx,
+            random_action_fraction=random_action_fraction,
+        )
+        for episode_idx in range(episodes)
+    ]
+    env.close()
+
+    returns = np.asarray([result.total_return for result in results], dtype=np.float32)
+    lengths = np.asarray([result.length for result in results], dtype=np.float32)
+    successes = np.asarray([result.success for result in results], dtype=np.float32)
+    entropies = np.asarray([result.planner_entropy for result in results], dtype=np.float32)
+    return {
+        "success_rate": float(successes.mean()),
+        "mean_return": float(returns.mean()),
+        "median_return": float(np.median(returns)),
+        "mean_episode_length": float(lengths.mean()),
+        "planner_action_entropy": float(np.nanmean(entropies)),
+    }

src/minidreamer/serialization.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from minidreamer.models.world_model import WorldModel
+
+
+def save_world_model_checkpoint(
+    path: str | Path,
+    model: WorldModel,
+    config: dict[str, Any],
+    optimizer: torch.optim.Optimizer | None = None,
+    metadata: dict[str, Any] | None = None,
+) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    payload = {
+        "model_state": model.state_dict(),
+        "config": config,
+        "metadata": metadata or {},
+    }
+    if optimizer is not None:
+        payload["optimizer_state"] = optimizer.state_dict()
+    torch.save(payload, path)
+
+
+def load_world_model_checkpoint(
+    path: str | Path,
+    action_dim: int,
+    map_location: str | torch.device | None = None,
+) -> tuple[WorldModel, dict[str, Any], dict[str, Any]]:
+    payload = torch.load(path, map_location=map_location, weights_only=False)
+    config = payload["config"]
+    model = WorldModel.from_config(config, action_dim=action_dim)
+    model.load_state_dict(payload["model_state"])
+    return model, config, payload.get("metadata", {})
+

src/minidreamer/utils/__init__.py

@@ -0,0 +1,2 @@
+"""Utility helpers for MiniDreamer."""
+

src/minidreamer/utils/common.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import json
+import random
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def get_device(device: str | None = None) -> torch.device:
+    if device is not None:
+        return torch.device(device)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def to_numpy(value: Any) -> np.ndarray:
+    if isinstance(value, np.ndarray):
+        return value
+    if torch.is_tensor(value):
+        return value.detach().cpu().numpy()
+    return np.asarray(value)
+
+
+def masked_mean(values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    mask = mask.to(values.dtype)
+    denom = torch.clamp(mask.sum(), min=1.0)
+    return (values * mask).sum() / denom
+
+
+def write_json(path: str | Path, payload: dict[str, Any]) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2, sort_keys=True)
+
+
+def write_jsonl(path: str | Path, rows: list[dict[str, Any]]) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as handle:
+        for row in rows:
+            handle.write(json.dumps(row, sort_keys=True))
+            handle.write("\n")
+

src/train_world_model.py

@@ -0,0 +1,334 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from torch.nn.utils import clip_grad_norm_
+from tqdm import trange
+
+from minidreamer.config import ensure_run_dirs, load_config, merge_dicts, save_config
+from minidreamer.data.collect_random import collect_bootstrap_dataset
+from minidreamer.data.replay_buffer import ReplayBuffer
+from minidreamer.evaluation import evaluate_random_policy, evaluate_world_model
+from minidreamer.envs.make_env import make_env_from_config
+from minidreamer.models.world_model import WorldModel
+from minidreamer.planning.cem import DiscreteCEMPlanner
+from minidreamer.planning.evaluate_planner import evaluate_planner
+from minidreamer.serialization import save_world_model_checkpoint
+from minidreamer.utils.common import get_device, seed_everything, write_json, write_jsonl
+
+
+def train_world_model_updates(
+    model: WorldModel,
+    replay: ReplayBuffer,
+    optimizer: torch.optim.Optimizer,
+    config: dict[str, Any],
+    num_updates: int,
+    device: torch.device,
+) -> list[dict[str, float]]:
+    if num_updates <= 0:
+        return []
+    model.train()
+    logs: list[dict[str, float]] = []
+    progress = trange(num_updates, desc="world-model-updates", leave=False)
+    for _ in progress:
+        batch = ReplayBuffer.batch_to_torch(replay.sample_sequences(split="train"), device=device)
+        losses = model.compute_losses(batch, config)
+        optimizer.zero_grad(set_to_none=True)
+        losses["loss"].backward()
+        clip_grad_norm_(model.parameters(), float(config["training"].get("grad_clip_norm", 100.0)))
+        optimizer.step()
+        log_row = {
+            "loss": float(losses["loss"].detach().cpu()),
+            "reward_loss": float(losses["reward_loss"].cpu()),
+            "done_loss": float(losses["done_loss"].cpu()),
+            "kl_loss": float(losses["kl_loss"].cpu()),
+            "recon_loss": float(losses["recon_loss"].cpu()),
+        }
+        logs.append(log_row)
+        progress.set_postfix({key: f"{value:.3f}" for key, value in log_row.items()})
+    return logs
+
+
+def optimizer_to_device(optimizer: torch.optim.Optimizer, device: torch.device) -> None:
+    for state in optimizer.state.values():
+        for key, value in state.items():
+            if torch.is_tensor(value):
+                state[key] = value.to(device)
+
+
+def load_training_state(
+    checkpoint_path: str | Path,
+    config: dict[str, Any],
+    action_dim: int,
+    device: torch.device,
+) -> tuple[dict[str, Any], WorldModel, torch.optim.Optimizer, dict[str, Any]]:
+    payload = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    resolved_config = merge_dicts(payload["config"], config)
+    model = WorldModel.from_config(resolved_config, action_dim=action_dim).to(device)
+    model.load_state_dict(payload["model_state"])
+    optimizer = torch.optim.Adam(model.parameters(), lr=float(resolved_config["training"]["lr"]))
+    optimizer_state = payload.get("optimizer_state")
+    if optimizer_state is not None:
+        optimizer.load_state_dict(optimizer_state)
+        optimizer_to_device(optimizer, device)
+    return resolved_config, model, optimizer, payload.get("metadata", {})
+
+
+def find_existing_run_artifacts(base_dir: str | Path) -> list[Path]:
+    base = Path(base_dir)
+    if not base.exists():
+        return []
+
+    artifact_files = [
+        base / "metrics" / "run_summary.json",
+        base / "metrics" / "train_metrics.jsonl",
+        base / "metrics" / "eval_metrics.jsonl",
+        base / "checkpoints" / "world_model_latest.pt",
+        base / "replay" / "metadata.json",
+    ]
+    found = [path for path in artifact_files if path.exists()]
+    if found:
+        return found
+
+    for subdir_name in ("checkpoints", "metrics", "replay"):
+        subdir = base / subdir_name
+        if subdir.exists():
+            for child in subdir.iterdir():
+                found.append(child)
+                break
+    return found
+
+
+def collect_planner_steps(
+    env,
+    replay: ReplayBuffer,
+    model: WorldModel,
+    planner: DiscreteCEMPlanner,
+    num_steps: int,
+    random_action_fraction: float,
+    rng: np.random.Generator,
+) -> dict[str, int]:
+    collected_steps = 0
+    episodes = 0
+    success_episodes = 0
+    model.eval()
+    while collected_steps < num_steps:
+        obs, _ = env.reset()
+        observations = [obs]
+        actions: list[int] = []
+        rewards: list[float] = []
+        terminated_flags: list[float] = []
+        truncated_flags: list[float] = []
+        done_flags: list[float] = []
+        terminated = False
+        truncated = False
+
+        with torch.no_grad():
+            state = model.posterior_step(model.initial_state(1), None, obs, sample=False)
+            while not (terminated or truncated):
+                if rng.random() < random_action_fraction:
+                    action = int(env.action_space.sample())
+                else:
+                    action = planner.plan(state).action
+                obs, reward, terminated, truncated, _ = env.step(action)
+                actions.append(action)
+                rewards.append(float(reward))
+                terminated_flags.append(float(terminated))
+                truncated_flags.append(float(truncated))
+                done_flags.append(float(terminated or truncated))
+                observations.append(obs)
+                collected_steps += 1
+                if terminated or truncated:
+                    break
+                state = model.posterior_step(state, action, obs, sample=False)
+
+        replay.add_episode(
+            obs=np.asarray(observations, dtype=np.float32),
+            actions=np.asarray(actions, dtype=np.int64),
+            rewards=np.asarray(rewards, dtype=np.float32),
+            terminated=np.asarray(terminated_flags, dtype=np.float32),
+            truncated=np.asarray(truncated_flags, dtype=np.float32),
+            done=np.asarray(done_flags, dtype=np.float32),
+        )
+        episodes += 1
+        success_episodes += int(bool(terminated and np.sum(rewards) > 0.0))
+    return {
+        "env_steps": collected_steps,
+        "episodes": episodes,
+        "success_episodes": success_episodes,
+    }
+
+
+def run_training(
+    config: dict[str, Any],
+    output_dir: str | Path,
+    replay_dir: str | Path | None = None,
+    resume_checkpoint: str | Path | None = None,
+    allow_overwrite_existing_output: bool = False,
+) -> dict[str, Any]:
+    seed = config.get("project", {}).get("seed", 0)
+    seed_everything(seed)
+    existing_artifacts = find_existing_run_artifacts(output_dir)
+    if existing_artifacts and resume_checkpoint is None and not allow_overwrite_existing_output:
+        preview = ", ".join(str(path) for path in existing_artifacts[:3])
+        raise FileExistsError(
+            f"Refusing to overwrite existing run directory '{output_dir}'. "
+            f"Found existing artifacts: {preview}. "
+            "Choose a new --output-dir, resume with --resume-checkpoint, "
+            "or pass --allow-overwrite-existing-output to overwrite intentionally."
+        )
+    run_dirs = ensure_run_dirs(output_dir)
+    device = get_device(config.get("training", {}).get("device"))
+
+    env = make_env_from_config(config, seed=seed)
+    action_dim = env.action_space.n
+    env.close()
+
+    if replay_dir is not None and Path(replay_dir).exists():
+        replay = ReplayBuffer.load(replay_dir)
+        collection_summary = {"replay_loaded": replay.summary()}
+    else:
+        replay, collection_summary = collect_bootstrap_dataset(config, output_dir=run_dirs["replay"], seed=seed)
+
+    resume_metadata: dict[str, Any] = {}
+    if resume_checkpoint is not None:
+        config, model, optimizer, resume_metadata = load_training_state(
+            checkpoint_path=resume_checkpoint,
+            config=config,
+            action_dim=action_dim,
+            device=device,
+        )
+    else:
+        model = WorldModel.from_config(config, action_dim=action_dim).to(device)
+        optimizer = torch.optim.Adam(model.parameters(), lr=float(config["training"]["lr"]))
+
+    save_config(config, run_dirs["base"] / "resolved_config.yaml")
+    training_logs: list[dict[str, float]] = []
+    evaluation_logs: list[dict[str, float]] = []
+
+    train_collect_ratio = float(config["collection"].get("train_collect_ratio", 1.0))
+    total_updates_budget = int(config["training"]["train_steps"])
+    if resume_checkpoint is not None:
+        updates_done = int(resume_metadata.get("updates_done", 0))
+        checkpoint_env_steps = int(resume_metadata.get("env_steps", 0))
+        if replay.env_steps > checkpoint_env_steps and updates_done < total_updates_budget:
+            collect_steps_per_iteration = max(1, int(config["collection"].get("collect_steps_per_iteration", 1)))
+            per_iteration_updates = int(
+                config["collection"].get(
+                    "gradient_updates_per_iteration",
+                    round(collect_steps_per_iteration * train_collect_ratio),
+                )
+            )
+            missed_iterations = max(0, round((replay.env_steps - checkpoint_env_steps) / collect_steps_per_iteration))
+            catch_up_updates = min(total_updates_budget - updates_done, per_iteration_updates * missed_iterations)
+            catch_up_logs = train_world_model_updates(model, replay, optimizer, config, catch_up_updates, device)
+            training_logs.extend(catch_up_logs)
+            updates_done += len(catch_up_logs)
+    else:
+        initial_updates = min(total_updates_budget, max(1, int(round(replay.env_steps * train_collect_ratio))))
+        training_logs.extend(train_world_model_updates(model, replay, optimizer, config, initial_updates, device))
+        updates_done = len(training_logs)
+
+    comparison_budgets = config.get("comparison", {}).get("env_steps", [replay.env_steps])
+    target_env_steps = int(max(comparison_budgets))
+    rng = np.random.default_rng(seed)
+    env = make_env_from_config(config, seed=seed)
+    planner = DiscreteCEMPlanner.from_config(model, env.action_space.n, config)
+    eval_every_steps = int(config["evaluation"].get("eval_every_env_steps", target_env_steps))
+    next_eval_step = replay.env_steps
+
+    while replay.env_steps < target_env_steps and updates_done < total_updates_budget:
+        collect_steps = min(
+            int(config["collection"]["collect_steps_per_iteration"]),
+            target_env_steps - replay.env_steps,
+        )
+        collection_row = collect_planner_steps(
+            env,
+            replay,
+            model,
+            planner,
+            num_steps=collect_steps,
+            random_action_fraction=float(config["collection"].get("random_action_fraction_after_planner", 0.0)),
+            rng=rng,
+        )
+        updates = int(config["collection"].get("gradient_updates_per_iteration", round(collection_row["env_steps"] * train_collect_ratio)))
+        updates = min(updates, total_updates_budget - updates_done)
+        training_logs.extend(train_world_model_updates(model, replay, optimizer, config, updates, device))
+        updates_done = len(training_logs)
+        replay.save(run_dirs["replay"])
+
+        if replay.env_steps >= next_eval_step:
+            world_model_metrics = evaluate_world_model(config, model, replay, split="val", max_episodes=10)
+            planner_metrics = evaluate_planner(config, model, episodes=min(10, config["evaluation"]["episodes"]), seed=seed)
+            random_metrics = evaluate_random_policy(config, episodes=min(10, config["evaluation"]["episodes"]), seed=seed)
+            eval_row = {
+                "env_steps": replay.env_steps,
+                "updates_done": updates_done,
+                **{f"world_model/{key}": value for key, value in world_model_metrics.items()},
+                **{f"planner/{key}": value for key, value in planner_metrics.items()},
+                **{f"random/{key}": value for key, value in random_metrics.items()},
+            }
+            evaluation_logs.append(eval_row)
+            next_eval_step += eval_every_steps
+            save_world_model_checkpoint(
+                run_dirs["checkpoints"] / f"world_model_env_steps_{replay.env_steps}.pt",
+                model,
+                config,
+                optimizer=optimizer,
+                metadata={"env_steps": replay.env_steps, "updates_done": updates_done},
+            )
+
+    env.close()
+    save_world_model_checkpoint(
+        run_dirs["checkpoints"] / "world_model_latest.pt",
+        model,
+        config,
+        optimizer=optimizer,
+        metadata={"env_steps": replay.env_steps, "updates_done": updates_done},
+    )
+    write_json(run_dirs["metrics"] / "collection_summary.json", collection_summary)
+    write_jsonl(run_dirs["metrics"] / "train_metrics.jsonl", training_logs)
+    write_jsonl(run_dirs["metrics"] / "eval_metrics.jsonl", evaluation_logs)
+    summary = {
+        "replay": replay.summary(),
+        "updates_done": updates_done,
+        "device": str(device),
+    }
+    write_json(run_dirs["metrics"] / "run_summary.json", summary)
+    return summary
+
+
+def build_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Train the MiniDreamer world model.")
+    parser.add_argument("--config", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument("--replay-dir", type=Path, default=None, help="Optional existing replay directory.")
+    parser.add_argument("--resume-checkpoint", type=Path, default=None, help="Optional checkpoint to resume from.")
+    parser.add_argument(
+        "--allow-overwrite-existing-output",
+        action="store_true",
+        help="Allow overwriting an existing run directory when not resuming.",
+    )
+    return parser
+
+
+def main() -> None:
+    parser = build_arg_parser()
+    args = parser.parse_args()
+    config = load_config(args.config)
+    summary = run_training(
+        config,
+        args.output_dir,
+        replay_dir=args.replay_dir,
+        resume_checkpoint=args.resume_checkpoint,
+        allow_overwrite_existing_output=args.allow_overwrite_existing_output,
+    )
+    print(summary)
+
+
+if __name__ == "__main__":
+    main()

tests/test_cem_planner.py

@@ -0,0 +1,31 @@
+import torch
+
+from minidreamer.planning.cem import DiscreteCEMPlanner
+
+
+class DummyWorldModel:
+    def __init__(self):
+        self.device = torch.device("cpu")
+
+    def score_action_sequences(self, state, action_sequences, discount=0.99, use_done_mask=True):
+        target = torch.tensor([1, 2, 0, 1], device=action_sequences.device)
+        scores = -(action_sequences != target).float().sum(dim=-1)
+        return {"scores": scores}
+
+
+def test_discrete_cem_planner_finds_high_scoring_sequence():
+    torch.manual_seed(0)
+    planner = DiscreteCEMPlanner(
+        world_model=DummyWorldModel(),
+        action_dim=3,
+        horizon=4,
+        candidates=512,
+        elites=64,
+        iterations=5,
+        discount=1.0,
+        use_done_mask=False,
+    )
+    output = planner.plan(state=object())
+    assert output.action == 1
+    assert output.sequence == [1, 2, 0, 1]
+

tests/test_env.py

@@ -0,0 +1,19 @@
+import numpy as np
+
+from minidreamer.envs.make_env import make_env
+
+
+def test_make_env_returns_normalized_rgb_observation():
+    env = make_env(seed=0)
+    obs, _ = env.reset()
+    assert obs.shape == (64, 64, 3)
+    assert obs.dtype == np.float32
+    assert 0.0 <= float(obs.min()) <= float(obs.max()) <= 1.0
+
+    next_obs, reward, terminated, truncated, _ = env.step(env.action_space.sample())
+    assert next_obs.shape == (64, 64, 3)
+    assert isinstance(float(reward), float)
+    assert isinstance(bool(terminated), bool)
+    assert isinstance(bool(truncated), bool)
+    env.close()
+

tests/test_replay_buffer.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from minidreamer.data.replay_buffer import ReplayBuffer
+
+
+def make_episode(length: int, reward: float = 0.0):
+    obs = np.random.rand(length + 1, 64, 64, 3).astype(np.float32)
+    actions = np.arange(length, dtype=np.int64) % 7
+    rewards = np.full(length, reward, dtype=np.float32)
+    terminated = np.zeros(length, dtype=np.float32)
+    truncated = np.zeros(length, dtype=np.float32)
+    done = np.zeros(length, dtype=np.float32)
+    terminated[-1] = 1.0
+    done[-1] = 1.0
+    return obs, actions, rewards, terminated, truncated, done
+
+
+def test_replay_buffer_sampling_and_padding(tmp_path: Path):
+    buffer = ReplayBuffer(capacity_episodes=10, sequence_length=8, batch_size=4)
+    for episode_id, length in enumerate((3, 5, 9)):
+        obs, actions, rewards, terminated, truncated, done = make_episode(length, reward=float(episode_id))
+        buffer.add_episode(obs, actions, rewards, terminated, truncated, done, episode_id=episode_id)
+
+    available_split = next(split for split in ("train", "val", "test") if buffer.episode_ids(split))
+    batch = buffer.sample_sequences(split=available_split, batch_size=2, rng=np.random.default_rng(0))
+    assert batch["obs"].shape == (2, 9, 64, 64, 3)
+    assert batch["actions"].shape == (2, 8)
+    assert batch["mask"].shape == (2, 8)
+    assert np.all(batch["mask"].sum(axis=1) >= 1)
+
+    save_dir = tmp_path / "replay"
+    buffer.save(save_dir)
+    loaded = ReplayBuffer.load(save_dir)
+    assert loaded.summary()["episodes"] == buffer.summary()["episodes"]
+    assert loaded.summary()["env_steps"] == buffer.summary()["env_steps"]
+
+
+def test_replay_buffer_torch_batch_shapes():
+    buffer = ReplayBuffer(capacity_episodes=4, sequence_length=4, batch_size=2)
+    obs, actions, rewards, terminated, truncated, done = make_episode(5, reward=1.0)
+    buffer.add_episode(obs, actions, rewards, terminated, truncated, done)
+    available_split = next(split for split in ("train", "val", "test") if buffer.episode_ids(split))
+    batch = buffer.sample_sequences(split=available_split, batch_size=2, rng=np.random.default_rng(1))
+    tensor_batch = ReplayBuffer.batch_to_torch(batch)
+    assert tensor_batch["obs"].shape == (2, 5, 3, 64, 64)
+    assert tensor_batch["actions"].dtype == torch.int64
+

tests/test_rssm_shapes.py

@@ -0,0 +1,43 @@
+import torch
+
+from minidreamer.models.world_model import WorldModel
+
+
+def test_world_model_sequence_shapes_and_loss():
+    torch.manual_seed(0)
+    model = WorldModel(
+        action_dim=7,
+        embedding_dim=128,
+        deter_dim=128,
+        stoch_dim=16,
+        hidden_dim=128,
+        use_decoder=True,
+    )
+    obs = torch.rand(4, 33, 3, 64, 64)
+    actions = torch.randint(0, 7, (4, 32))
+    outputs = model.observe_sequence(obs, actions, sample=False)
+    assert outputs.reward_pred.shape == (4, 32)
+    assert outputs.done_logits.shape == (4, 32)
+    assert outputs.prior_mean.shape == (4, 32, 16)
+    assert outputs.reconstructions is not None
+    assert outputs.reconstructions.shape == (4, 32, 3, 64, 64)
+
+    batch = {
+        "obs": obs,
+        "actions": actions,
+        "rewards": torch.zeros(4, 32),
+        "done": torch.zeros(4, 32),
+        "mask": torch.ones(4, 32),
+    }
+    config = {
+        "training": {
+            "beta_reward": 1.0,
+            "beta_done": 1.0,
+            "beta_kl": 1.0,
+            "beta_recon": 1.0,
+            "free_nats": 1.0,
+        }
+    }
+    losses = model.compute_losses(batch, config)
+    assert torch.isfinite(losses["loss"])
+