RSSM: Dreaming Cheetah
A Recurrent State Space Model trained on Apple Silicon MPS visual observations to learn visual dynamics open-loop.
27.4x
GPU Acceleration Speedup
100,000
Dynamics Training Steps
5.02 M
Model Parameters
Apple MPS
Hardware Platform
Ground Truth
Reconstruction
Open-Loop Prediction
Context Phase
Dynamics Validation Metrics
-
Reconstruction MSE: 0.000640
During the first 10 steps (Context phase), the model uses real pixel observations to form posterior distributions, giving high-fidelity reconstructions. -
Prediction MSE: 0.000670
In the open-loop dreaming phase (steps 10 to 40), the model predicts future visual dynamics entirely using transition priors. The error remains low, showing high recurrent stability. -
Reward Prediction
The reward predictor accurately maps the latent states $(h_t, s_t)$ to environmental rewards, capturing the peaks and valleys of running speed even when dreaming.
RSSM Cell Information Flow
How information flows dynamically at each sequence step in the World Model.
Observation (x_t)
64x64x3 RGB image
➔
CNN Encoder
Converts image to 1024-dim
GRU Cell
h_t = GRU(h_t-1, s_t-1, a_t-1)
➔
Stochastic Post (q)
s_t ~ Posterior(h_t, embed_t)
➔
Stochastic Prior (p)
s_t ~ Prior(h_t) (Dreaming)
CNN Decoder
Reconstructs image from (h_t, s_t)
➔
Reward Predictor
Predicts task reward from (h_t, s_t)