[draft] add end-to-end RL training

[m1fisher]

Intent

Add the Dreamer 4 reinforcement learning training pipeline, enabling policy optimization in imagination. We introduce behavioral cloning + reward head finetuning (agent.py, train_agent.py) and imagination-based policy gradient training (train_imagination.py).

Context

train_agent.py: Finetunes the pretrained dynamics model by training Behavior Cloning and Reward prediction MLP heads on the dynamics model's hidden states. BC uses masked per-dimension supervision on expert data; the reward head uses symexp two-hot categorical predictions. Both heads use multi-token prediction (MTP) to forecast
L steps ahead. To facilitate realism in longer rollouts, includes a new experimental autoregressive rollout loss mode (optional) that unrolls the dynamics model for K steps feeding its own outputs back, supervised against ground-truth
latents.

train_imagination.py: Trains a policy and value function in imagined latent rollouts with frozen dynamics/reward model. Currently contains implementations of both Dreamer 3 REINFORCE (which mastered DMControl), and Dreamer 4 PMPO (which is less proven for continuous action-space control).

Status

End-to-end RL is not yet producing consistent reward improvement in Phase 3. Through extensive experimentation (18+ training runs across cartpole-swingup and walker-walk), we identified several potential challenges:

1. World model fidelity: The dynamics model, trained with single-step denoising, may drift during long rollouts. E.g., if playing cartpole in the interactive demo (both before and after dynamics fine-tuning), it is difficult to move the cart left past the halfway point in the screen. The new AR rollout
   finetuning loss is our attempted mitigation but is still experimental.
2. Weak per-step action sensitivity: In our imagination experiments, per-step reward varies by only std≈0.04 per step across random actions. Signal accumulates linearly (not compounding) over the horizon, giving ~0.6 std over 15 steps, a potentially too-weak signal.
3. Policy std collapse: For walker-walk task, REINFORCE consistently finds it easier to reduce the loss by collapsing std than by finding better actions. In our experiments, entropy regularization may either drown out the policy gradient (η=0.1) or fails to prevent collapse (η=3e-4).
4. Policy parameterization difficulty: We experimented with both TanhNormal and BoundedNormal distribution parameterizations. TanhNormal seemed to induce behavior collapse during Phase 3, but this may be incorrect cause attribution. More effort would be useful to find an ideal parameterization for the continuous policy, and then to use it consistenly across both training phases.

Validation

• Phase 2 BC+reward training verified on walker-walk (reward_corr=0.99 after task-specific finetuning), cartpole-swingup, cartpole-balance.
• Reward sensitivity to actions measured across horizons 1-32 with both fixed and random action protocols.
• Policy drift analysis comparing early vs late checkpoints: actions change (0.03-0.18 raw diff) but returns improve inconsistently (~50/50 on cartpole, briefly positive then degrading on walker).
• Value head convergence verified (V tracks R within 200 warmup steps)
• Multiple loss functions tested: PMPO, PPO, REINFORCE with various normalization schemes.
• AR rollout finetuning reduces multi-step prediction error 4-5x (MSE 0.014 → 0.003) for K=4 steps, but does not seem to fix the interactive demo's fidelity (i.e. the cartpole issue mentioned above).

A detailed report of our efforts and findings thus far is attached here

[nicklashansen]

This looks really cool, thanks for bringing it to my attention! Let me know if you need my input on anything :-)