World Model RL: Teaching an Agent to Imagine

This project studies whether a learned neural world model can make a reinforcement-learning agent learn faster in a noisy GridWorld, and when imagined experience becomes harmful instead of helpful.

The core setting is an 8x8 stochastic GridWorld with 20% action noise. A neural world model is trained from random-policy transitions, frozen, and then used to generate imagined updates for Dyna-Q. The experiments compare:

• Q-learning (K=0): no model, no imagination.
• Dyna-Q (K=10) with a neural world model: one real update plus 10 imagined updates per real environment step.
• Value Iteration: oracle ceiling computed from the true transition model.
• Tabular / one-hot / improved local-view models: ablations for how imagined experience is generated.

The full write-up is in report/report.pdf.

Main Takeaways

1. A good world model improves sample efficiency.
   On the static map, Dyna-Q with a 10,000-transition world model reaches high success much earlier than Q-learning. It changes when the optimal policy is reached, not the final ceiling.

2. Imagined experience multiplies model error.
   A bad world model trained on only 200 transitions collapses the Dyna-Q agent to 0% success, worse than the no-model baseline. Since each real step is followed by 10 imagined updates, wrong Bellman targets dominate the correct real ones.

3. The static map is mostly a memorization problem.
   On one fixed, fully observable map, tabular Dyna-Q and Q-learning can already reach the Value-Iteration optimum. A learned model is useful for studying the method, but it cannot beat a perfect tabular memory on this task.

4. Representation matters under distribution shift.
   An improved world model adds a 3x3 local field of view. This gives the agent "eyes", but with only one static training map the local view is still mostly redundant with the absolute cell index. True generalization would require training across many maps and predicting relative motion rather than memorized absolute next states.

5. Dynamic obstacles finally make the local-view model meaningful.
   When walls move, the absolute cell index is no longer enough. The improved model can use the current local view to predict how nearby obstacles affect the next state, while tabular memory becomes stale and the one-hot model is phase-blind.

Environment

The project studies the standard model-based RL loop: the agent collects real experience, trains or queries a world model, and uses imagined transitions to accelerate value learning.

• Grid: 8x8, 64 states.
• Start: (0, 0).
• Goal: (7, 7).
• Actions: 0=up, 1=down, 2=left, 3=right.
• Noise: with probability 0.2, the intended action is replaced by a uniformly random direction.
• Walls: fixed obstacles in the static setting; moving periodic obstacles in the dynamic setting.
• Reward: +1.0 on reaching the goal, -0.01 otherwise.
• Episode limit: 100 steps.

The stochastic dynamics are intentional. They prevent the transition model from being a trivial lookup table and cap next-state prediction accuracy below 100%.

The static map below is the base environment used for the main experiments. Green marks the start, yellow marks the goal, and blocked cells create the navigation constraints that the policy must learn to route around.

World Models

One-Hot World Model

The baseline model receives:

one_hot(state)[64] + one_hot(action)[4] -> 68-dim input

It uses two small MLPs:

• Transition network: 68 -> 256 -> 256 -> 64
• Reward network: 68 -> 128 -> 64 -> 1

The transition network predicts a distribution over the 64 possible next cells. With 20% action noise, validation accuracy plateaus around the environment's irreducible stochastic ceiling rather than reaching 100%.

Improved Local-View World Model

The improved model augments the absolute state with a local observation:

one_hot(state)[64] + local_3x3_window[9] + one_hot(action)[4] -> 77-dim input

The 3x3 window marks nearby walls, free cells, boundaries, and the goal. The model predicts the next cell, next local view, and reward. Dropout over the absolute-state block and local-view block is used in the static-map experiment to test whether each modality can independently support prediction.

Algorithms

Dyna-Q uses the same Bellman update for real and imagined transitions:

Q(s,a) <- Q(s,a) + alpha * [r + gamma * max_a' Q(s',a') - Q(s,a)]

After every real environment step:

1. Apply one real Q-learning update.
2. Sample previously visited (state, action) pairs.
3. Query the frozen world model or tabular memory.
4. Apply K=10 imagined Q-updates.

Imagined pairs are sampled only from visited states, keeping planning inside the model's training support as much as possible.

Results And Visualizations

Static Environment

File	Description
results/figure1_visitation.png	Random-policy state visitation heatmap.
results/figure2_transition_accuracy.png	Transition-model validation accuracy and probe predictions.
results/figure3_learning_curves.png	Q-learning vs Dyna-Q vs Value Iteration learning curves.
results/figure4_value_heatmaps.png	Learned value functions compared with Value Iteration.
results/figure5_data_scaling.png	Model-quality failure analysis across 200 / 2,000 / 10,000 transitions.

Training Comparison

The learning curve shows the aggregate result, while the animation below shows the same story operationally: at matched training budgets, Dyna-Q reaches the goal earlier because each real transition is reused through imagined Bellman updates.

Animated Static-Map Rollouts

The GIFs in results/ visualize how policies improve during training and how different imagination sources behave on the same map.

Animation	Description
results/episode_rollout.gif	A greedy-policy rollout in the static GridWorld.
results/training_comparison.gif	Q-learning (K=0) vs Dyna-Q (K=10) at matched training budgets.
results/methods_comparison_static.gif	Tabular Dyna-Q, one-hot world model, and improved local-view world model side by side.

Improved Model And Generalization

File	Description
results/figure_improved_wm.png	Local-view model accuracy under state/window dropout.
results/figure_ablation_imagination.png	Same-map comparison of tabular, one-hot, and improved-model imagination.
results/newmap_layout.png	New unseen wall layout.
results/figure_newmap_generalization.png	Frozen-model transfer to a new map.

Dynamic Obstacles

The dynamic setting introduces a moving wall segment. This makes stale tabular memory and phase-blind one-hot prediction fail, while the improved local-view model has access to the current obstacle layout.

File	Description
results/dynamic_phases.png	Periodic wall configurations used by the dynamic environment.
results/dynamic_episode.gif	Rollout animation with moving obstacles and phase changes.
results/figure_dynamic_prediction.png	Prediction accuracy on dynamic-obstacle data.
results/figure_dynamic_learning_curves.png	Control performance in the dynamic environment.

Project Structure

World_Model_RL/
  gridworld.py                 Static and dynamic GridWorld environments
  utils.py                     State/action encoders and local-view utilities
  world_model.py               One-hot and improved world-model networks
  dyna_q.py                    Static Q-learning, Dyna-Q, and Value Iteration
  dyna_q_dynamic.py            Dynamic-environment Dyna-Q utilities
  plotting.py                  Figure and animation generation
  experiments.py               Main static experiment
  experiments_failure.py       Bad-model and data-scaling failure analysis
  experiments_improved_wm.py   Local-view world-model experiment
  experiments_ablation.py      Imagination-source ablation
  experiments_newmap.py        Frozen-model transfer to a new map
  experiments_methods_viz.py   Static side-by-side method animation
  experiments_dynamic.py       Dynamic-obstacle experiment
  results/                     Figures, GIFs, metrics, and trained weights
  report/                      Paper source and compiled report

Running

Install the required Python packages:

pip install torch numpy matplotlib

Run the main static experiment:

python experiments.py

Run the failure analysis:

python experiments_failure.py

Run the improved world-model and generalization experiments:

python experiments_improved_wm.py
python experiments_ablation.py
python experiments_newmap.py

Run the dynamic-obstacle experiment:

python experiments_dynamic.py

Most experiments write figures, GIFs, metrics, and model checkpoints into results/.

Context

This repository is the final project for the SJTU AI / Deep Learning and Reinforcement Learning course. It is best read as a controlled study of model-based RL: when a world model helps, when it fails, and how representation choice determines whether "imagination" is useful or destructive.