Latent Dynamics / Rollout kernel

[superposition]

Description

Efficient latent imagination/rollout for model-based planning.

Operation

# Rollout dynamics model in latent space
# z0: (B, D) initial latent
# actions: (B, H, action_dim) action sequence
# dynamics: learned transition function
latent_trajectory = latent_rollout(z0, actions, dynamics_fn, horizon=H)
# Returns: (B, H+1, D) latent states

Implementation

• Fused multi-step rollout (avoid kernel launch overhead)
• Parallel rollout for multiple action sequences (beam search)
• Support for stochastic dynamics (sampling)
• Gradient through rollout for differentiable planning

Tests

• Matches sequential dynamics application
• Correct trajectory length (H+1 states)
• Gradient flows through full rollout
• Works with deterministic and stochastic dynamics
• Parallel rollout produces independent trajectories
• Memory efficient for long horizons

Used In

Dreamer, TD-MPC, IRIS, model-predictive control, imagination-based planning