Recurrent State-Space Models › Code Implementation
World Models HandbookAll chapters Recurrent State-Space Models › Code ImplementationOne section · Sections in this chapter7 / 10 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
The simulator below uses an explicit-Euler spring-damper transition with step size , plus discrete process noise. It is the discrete model being studied, not the exact sampled transition of a continuous-time stochastic oscillator. Observation measures the position in state ; action affects . We use consistent arbitrary position and time units, so velocity has position-per-time units. The learned models receive positions and actions, not the true state, simulator coefficients, or individual noise realizations.
import numpy as np
import torch
import torch.nn as nn
rng = np.random.default_rng(7)
dt = 0.1
A = np.array([[1.0, dt], [-4.0 * dt, 1.0 - 0.5 * dt]]) # spring-damper dynamics
B = np.array([[0.0], [dt]]) # actions push on velocity
Q = 0.01 * np.eye(2) # process noise
obs_std = 0.25 # sensor noise on position
def simulate_batch(n_seqs, T, rng):
"""Spring-damper with force actions; only noisy positions are observed."""
states = np.zeros((n_seqs, T + 1, 2))
obs = np.zeros((n_seqs, T, 1))
actions = rng.normal(scale=0.5, size=(n_seqs, T, 1))
states[:, 0] = rng.normal(scale=[1.0, 0.5], size=(n_seqs, 2))
for t in range(T):
noise = rng.multivariate_normal(np.zeros(2), Q, size=n_seqs)
states[:, t + 1] = (
states[:, t] @ A.T + np.outer(actions[:, t], B.ravel()) + noise
)
obs[:, t] = states[:, t, 0:1] + rng.normal(
scale=obs_std, size=(n_seqs, 1)
)
return states, obs, actionsNow we generate a training set and a held-out validation set, and pick one demonstration sequence for plotting later.
states, obs, actions = simulate_batch(480, 50, rng)
train_states, val_states = states[:400], states[400:]
train_obs, val_obs = obs[:400], obs[400:]
train_act, val_act = actions[:400], actions[400:]
demo_idx = 3
t_grid = np.arange(50)
env_demo_pos = states[demo_idx, :50, 0]
env_demo_vel = states[demo_idx, :50, 1]
env_demo_obs = obs[demo_idx, :, 0]sequences: train=400, val=80, length T=50 true state: (position, velocity); observation: noisy position only position range: [-3.23, +3.23] velocity range: [-6.25, +6.08] sensor noise std: 0.25
The model receives the noisy position observations shown here and their recorded actions. The true state curves are for our inspection. Try estimating the direction of motion from a short noisy history: a single position generally does not determine velocity, while earlier positions and applied forces provide additional information. Useful predictive memory need not contain a coordinate explicitly named velocity; equivalent learned features can carry motion information.
The model
The RSSM below mirrors the equations of the architecture section exactly: a GRU cell for the deterministic path, two small heads for the prior and posterior Gaussians, and a decoder. Note that the posterior head is the only component that receives the observation as direct input; the decoder touches the observation only indirectly, through z_t sampled from the posterior.
class RSSM(nn.Module):
"""PlaNet-style recurrent state-space model.
h_t = GRU(h_{t-1}, [z_{t-1}, a_{t-1}]) deterministic memory
z_t ~ N(mu, diag(sigma**2)) from prior(h_t) or posterior(h_t, o_t)
o_t is decoded from [h_t, z_t]
"""
def __init__(self, obs_dim=1, act_dim=1, h_dim=48, z_dim=8, hidden=64):
super().__init__()
self.h_dim, self.z_dim = h_dim, z_dim
self.cell = nn.GRUCell(z_dim + act_dim, h_dim)
self.prior_net = nn.Sequential(
nn.Linear(h_dim, hidden), nn.ELU(), nn.Linear(hidden, 2 * z_dim)
)
self.posterior_net = nn.Sequential(
nn.Linear(h_dim + obs_dim, hidden),
nn.ELU(),
nn.Linear(hidden, 2 * z_dim),
)
self.decoder = nn.Sequential(
nn.Linear(h_dim + z_dim, hidden),
nn.ELU(),
nn.Linear(hidden, obs_dim),
)
@staticmethod
def gauss_head(params):
mu, log_std = params.chunk(2, dim=-1)
return mu, log_std.clamp(-5.0, 2.0)The gauss_head splits each head into a mean and log standard deviation. Exponentiation makes the scale positive; clamping the log scale to bounds it to approximately . The mean is not bounded, and this clamp does not prevent posterior collapse or every numerical failure. The GRU uses before the current observation is encoded, matching the timing convention above.
The training objective
The KL helper implements the analytic conditional diagonal-Gaussian expression. The loop updates , computes both heads, samples the posterior, decodes, and accumulates squared reconstruction error plus weighted clipped KL. This is a single-trajectory estimate of a modified variational objective, not the literal unmodified ELBO displayed earlier.
The free-nat floor is 1.0, applied after summing latent dimensions but separately for each sequence and timestep. The KL coefficient is 0.1. With our scalar squared-error reduction, its fixed-Gaussian interpretation uses decoder variance before free-nat clipping. This is a chosen observation-model scale, not an estimate of the simulator's sensor variance . We log raw KL as well as the clipped loss term. Neither their scalar ratio nor convergence near the threshold establishes that the prior has learned all predictable information.
def gaussian_kl(mu_q, ls_q, mu_p, ls_p):
"""Diagonal Gaussian KL; ls values are log standard deviations, not variances."""
return (
ls_p
- ls_q
+ (torch.exp(2 * ls_q) + (mu_q - mu_p) ** 2) / (2 * torch.exp(2 * ls_p))
- 0.5
).sum(dim=-1)
def train_rssm(
model,
obs,
act,
steps=1500,
batch=24,
lr=1e-3,
free_nats=1.0,
kl_weight=0.1,
seed=0,
):
torch.manual_seed(seed)
opt = torch.optim.Adam(model.parameters(), lr=lr)
n_seq, T, _ = obs.shape
obs_t = torch.tensor(obs, dtype=torch.float32)
act_t = torch.tensor(act, dtype=torch.float32)
history = {"step": [], "recon": [], "kl": [], "raw_kl": []}
for step in range(steps):
idx = torch.randint(0, n_seq, (batch,))
o_b, a_b = obs_t[idx], act_t[idx]
h = torch.zeros(batch, model.h_dim)
z = torch.zeros(batch, model.z_dim)
a_prev = torch.zeros(batch, act_t.shape[-1])
recon_losses, kl_losses, raw_kls = [], [], []
for t in range(T):
h = model.cell(torch.cat([z, a_prev], dim=-1), h)
mu_p, ls_p = model.gauss_head(model.prior_net(h))
mu_q, ls_q = model.gauss_head(
model.posterior_net(torch.cat([h, o_b[:, t]], dim=-1))
)
z = mu_q + torch.randn_like(mu_q) * torch.exp(ls_q)
pred = model.decoder(torch.cat([h, z], dim=-1))
recon_losses.append(((pred - o_b[:, t]) ** 2).sum(dim=-1))
kl_t = gaussian_kl(mu_q, ls_q, mu_p, ls_p)
raw_kls.append(kl_t)
kl_losses.append(torch.clamp(kl_t, min=free_nats))
a_prev = a_b[:, t]
recon = torch.stack(recon_losses).mean()
kl = torch.stack(kl_losses).mean()
loss = (
torch.stack(recon_losses).sum(0)
+ kl_weight * torch.stack(kl_losses).sum(0)
).mean()
opt.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), 10.0)
opt.step()
if step % 10 == 0 or step == steps - 1:
history["step"].append(step)
history["recon"].append(recon.item())
history["kl"].append(kl.item())
history["raw_kl"].append(torch.stack(raw_kls).mean().item())
return historyThe clamp is applied before the timestep sum. Strictly below the threshold, this KL contribution has zero derivative; other loss paths still train the model. At the threshold, the implementation chooses a subgradient. The loss sums 50 timesteps and then averages the minibatch. Because every sequence has the same length here, averaging over time instead would change the objective by a constant scale; variable-length weighting would need an explicit choice.
torch.manual_seed(0) # seed weight initialization before construction
rssm = RSSM()
history = train_rssm(rssm, train_obs, train_act, steps=1500, seed=0)final reconstruction MSE per step: 0.0175 sensor noise variance: 0.0625 final KL per step: 1.101 nats (free-nat floor = 1.0) final raw KL per step: 1.017 nats
Read raw and clipped KL separately. A clipped value exactly at 1 would only imply that every contributing raw value was at or below 1; it would not prove optimal prediction or identify a noise floor. Posterior reconstruction can fall below the sensor variance because its encoder sees the same noisy observation used as its target. That is allowed by the objective, not an accidental input leak. By contrast, a history-only prediction of a fresh observation has expected squared error at least the independent sensor variance. Clean-position evaluation below asks whether the posterior output denoises, and the Kalman reference shows how far it is from the known-model optimum.
In the seeded reproduction, reconstruction fell to 0.0175 per step. The logged clipped KL began near 1, rose, and finished at 1.101 nats; raw KL finished at 1.017 nats. The final values are last-minibatch diagnostics, not held-out averages. The initial increase is another reason not to interpret every training trace as a prior steadily catching up to a fixed posterior.
A deterministic control model
The baseline GRU receives and predicts with a squared-error point decoder. Its hidden memory could encode motion information or distribution parameters, but this particular output head returns only a point. It differs from the mixture-density RNN in World Models. We leave its small training budget explicit rather than describing the comparison as an isolated test of stochasticity.
class DetSSM(nn.Module):
"""Deterministic RNN world model: h_t = f(h_{t-1}, o_{t-1}, a_{t-1})."""
def __init__(self, obs_dim=1, act_dim=1, h_dim=48, hidden=64):
super().__init__()
self.h_dim = h_dim
self.cell = nn.GRUCell(obs_dim + act_dim, h_dim)
self.decoder = nn.Sequential(
nn.Linear(h_dim, hidden), nn.ELU(), nn.Linear(hidden, obs_dim)
)
def train_det(model, obs, act, steps=800, batch=24, lr=1e-3, seed=1):
torch.manual_seed(seed)
opt = torch.optim.Adam(model.parameters(), lr=lr)
n_seq, T, _ = obs.shape
obs_t = torch.tensor(obs, dtype=torch.float32)
act_t = torch.tensor(act, dtype=torch.float32)
final_loss = None
for step in range(steps):
idx = torch.randint(0, n_seq, (batch,))
o_b, a_b = obs_t[idx], act_t[idx]
h = torch.zeros(batch, model.h_dim)
mse = torch.zeros(batch)
for t in range(1, T):
h = model.cell(torch.cat([o_b[:, t - 1], a_b[:, t - 1]], dim=-1), h)
pred = model.decoder(h)
mse = mse + ((pred - o_b[:, t]) ** 2).sum(dim=-1)
loss = mse.mean() / (T - 1)
opt.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), 10.0)
opt.step()
final_loss = loss.item()
return final_losstorch.manual_seed(
1
) # initialize independently of earlier training's RNG consumption
det_model = DetSSM()
det_final_loss = train_det(det_model, train_obs, train_act, steps=800, seed=1)deterministic model final teacher-forced MSE: 0.0997 (teacher-forced: the real previous observation is fed at every step)
Teacher forcing feeds the real previous observation during baseline training; recursive evaluation feeds the prediction instead. This change can expose compounding errors, but does not guarantee that every horizon worsens. The RSSM also changes operating regime when it replaces future posterior updates with prior draws. It supplies an ensemble while this baseline supplies no interval. That difference in output format does not establish that the ensemble is well calibrated.
Filtering versus imagination
Now the two evaluation regimes. The filter runs the posterior at every step: the model always gets to look. The imagination burns in the belief on the first BURN_IN real observations, then rolls out with the prior only, under the true future actions, sampling N_SAMPLES independent futures. The deterministic baseline rolls out by feeding its own predictions back in place of observations, the classic deterministic world-model rollout.
@torch.no_grad()
def rssm_filter(model, obs_seq, act_seq, seed=0, n_samples=64):
"""Mean decoder output over posterior paths; first path supplies probe features."""
if len(obs_seq) == 0 or len(obs_seq) != len(act_seq) or n_samples < 1:
raise ValueError(
"Need aligned nonempty sequences and positive n_samples."
)
g = torch.Generator().manual_seed(seed)
T = obs_seq.shape[0]
h = torch.zeros(n_samples, model.h_dim)
z = torch.zeros(n_samples, model.z_dim)
a_prev = torch.zeros(n_samples, act_seq.shape[-1])
preds, feats, kls = [], [], []
for t in range(T):
h = model.cell(torch.cat([z, a_prev], dim=-1), h)
mu_p, ls_p = model.gauss_head(model.prior_net(h))
mu_q, ls_q = model.gauss_head(
model.posterior_net(
torch.cat([h, obs_seq[t : t + 1].expand(n_samples, -1)], dim=-1)
)
)
z = mu_q + torch.randn(mu_q.shape, generator=g) * torch.exp(ls_q)
preds.append(model.decoder(torch.cat([h, z], dim=-1)).mean(0).numpy())
feats.append(torch.cat([h, z], dim=-1).numpy()[0])
kls.append(gaussian_kl(mu_q, ls_q, mu_p, ls_p).mean().item())
a_prev = act_seq[t : t + 1].expand(n_samples, -1)
return np.array(preds), np.array(feats), np.array(kls)
@torch.no_grad()
def rssm_imagine(model, obs_seq, act_seq, burn_in=10, n_samples=256, seed=0):
"""Independent posterior paths, followed by prior-only decoder-mean paths."""
if (
not 1 <= burn_in < len(obs_seq)
or len(obs_seq) != len(act_seq)
or n_samples < 1
):
raise ValueError(
"Need aligned sequences, 1 <= burn_in < T, and positive n_samples."
)
g = torch.Generator().manual_seed(seed)
T = obs_seq.shape[0]
h = torch.zeros(n_samples, model.h_dim)
z = torch.zeros(n_samples, model.z_dim)
a_prev = torch.zeros(n_samples, act_seq.shape[-1])
for t in range(burn_in):
h = model.cell(torch.cat([z, a_prev], dim=-1), h)
mu_q, ls_q = model.gauss_head(
model.posterior_net(
torch.cat([h, obs_seq[t : t + 1].expand(n_samples, -1)], dim=-1)
)
)
z = mu_q + torch.randn(mu_q.shape, generator=g) * torch.exp(ls_q)
a_prev = act_seq[t : t + 1].expand(n_samples, -1)
imag = []
for t in range(burn_in, T):
h = model.cell(torch.cat([z, a_prev], dim=-1), h)
mu_p, ls_p = model.gauss_head(model.prior_net(h))
z = mu_p + torch.randn(mu_p.shape, generator=g) * torch.exp(ls_p)
imag.append(model.decoder(torch.cat([h, z], dim=-1)).numpy())
a_prev = act_seq[t : t + 1].repeat(n_samples, 1)
return np.stack(imag, axis=0) # (T - burn_in, n_samples, obs_dim)
@torch.no_grad()
def det_imagine(model, obs_seq, act_seq, burn_in=10):
"""Burn in on real observations, then feed the model its own predictions."""
if not 1 <= burn_in < len(obs_seq) or len(obs_seq) != len(act_seq):
raise ValueError("Need aligned sequences and 1 <= burn_in < T.")
T = obs_seq.shape[0]
h = torch.zeros(1, model.h_dim)
for t in range(1, burn_in):
h = model.cell(
torch.cat([obs_seq[t - 1 : t], act_seq[t - 1 : t]], dim=-1), h
)
o_prev = obs_seq[burn_in - 1 : burn_in] # last real observation
imag = []
for t in range(burn_in - 1, T - 1):
h = model.cell(torch.cat([o_prev, act_seq[t : t + 1]], dim=-1), h)
o_prev = model.decoder(h)
imag.append(o_prev.numpy())
return np.concatenate(imag, axis=0) # (T - burn_in, obs_dim)The two RSSM passes reuse the transition and decoder, while changing which distribution supplies the latent when observations are withheld. Future filtering and imagination therefore have different information, even for a perfectly specified model. Here each ensemble member follows its own sampled posterior history during burn-in before continuing with independent prior noise. This retains uncertainty across installed latent states instead of replicating one sampled endpoint. We decode conditional observation means without adding observation-likelihood noise: the bands below describe this decoder-mean ensemble scored against clean position, not the model's full noisy-observation predictive distribution. Its suitability as a clean-state interval is an empirical question, since clean positions were not training targets.
The reference filter uses the exact simulator parameters. At a measurement, it updates a Gaussian mean and covariance with and measurement variance . Between measurements, the mean follows , and covariance follows . The first term propagates remaining uncertainty; the second adds process noise. Starting from the posterior after observation 9 and omitting further measurement updates gives the exact open-loop reference for positions 10 through 49. Its covariance includes uncertainty after burn-in, not only newly added noise.
The memoryless reference deliberately ignores history and actions. It uses the known marginal position variance under the simulator's random action distribution and returns . Its expected error is , illustrating why beating raw-observation echo alone does not prove temporal inference. Both analytical references have privileged model knowledge and are labeled accordingly.
def kalman_reference(obs_seq, act_seq, stop_observing=None):
"""Exact known-model state moments, before or after the observation cutoff."""
T = len(obs_seq)
cutoff = T if stop_observing is None else stop_observing
if len(act_seq) != T or not 0 <= cutoff <= T:
raise ValueError("Need aligned sequences and a cutoff in [0, T].")
m, P = np.zeros(2), np.diag([1.0, 0.25])
C = np.array([1.0, 0.0])
means, covariances = [], []
for t in range(T):
if t < cutoff:
K = P @ C / (C @ P @ C + obs_std**2)
m = m + K * (float(obs_seq[t, 0]) - C @ m)
J = np.eye(2) - np.outer(K, C)
P = J @ P @ J.T + obs_std**2 * np.outer(K, K)
means.append(m.copy())
covariances.append(P.copy())
m = A @ m + B[:, 0] * float(act_seq[t, 0])
P = A @ P @ A.T + Q
return np.array(means), np.array(covariances)
def point_and_sample_rmse(samples, targets):
"""Inputs (sequence, horizon, sample) and (sequence, horizon)."""
if (
samples.ndim != 3
or samples.shape[:2] != targets.shape
or samples.shape[2] < 1
):
raise ValueError(
"Expected aligned sequence/horizon axes and a nonempty sample axis."
)
mean = samples.mean(axis=2)
point = np.sqrt(((mean - targets) ** 2).mean(axis=0))
sample = np.sqrt(((samples - targets[..., None]) ** 2).mean(axis=(0, 2)))
return point, sample
def interval_metrics(lower, upper, targets, alpha=0.05):
"""Per-horizon coverage, width and central interval score (smaller is better)."""
if lower.shape != upper.shape or lower.shape != targets.shape:
raise ValueError(
"Interval bounds and targets must have matching shapes."
)
if not 0 < alpha < 1 or np.any(lower > upper):
raise ValueError("Require ordered bounds and alpha in (0, 1).")
width = upper - lower
score = width + (2 / alpha) * (
np.maximum(lower - targets, 0) + np.maximum(targets - upper, 0)
)
return (
((lower <= targets) & (targets <= upper)).mean(0),
width.mean(0),
score.mean(0),
)
def wilson_interval(proportion, n, z=1.96):
"""Approximate binomial 95% confidence interval for each marginal coverage."""
if n <= 0 or np.any((proportion < 0) | (proportion > 1)):
raise ValueError("Need positive sample size and proportions in [0, 1].")
center = (proportion + z * z / (2 * n)) / (1 + z * z / n)
half = (
z
* np.sqrt(proportion * (1 - proportion) / n + z * z / (4 * n * n))
/ (1 + z * z / n)
)
return center - half, center + halfWe evaluate all 80 held-out sequences. Each RSSM imagination uses 256 particles, independently sampled through burn-in as well as the future. Posterior point estimates average 64 paths; probe features below use the first path, and their labels are never used to train the RSSM. A central sample interval uses the 2.5th and 97.5th percentiles, without assuming Gaussian decoded outputs. Its finite-particle endpoints still have Monte Carlo error.
At each horizon, coverage has 80 independent truth sequences, not 80 times 256 independent outcomes. Wilson intervals describe uncertainty in that marginal proportion; they are not simultaneous confidence bands across horizons, whose outcomes are correlated. We also report interval width and the central interval score: width plus times any distance by which truth misses the interval, where . This penalizes both unnecessary width and missed outcomes. All computations happen before plotting.
BURN_IN, N_SAMPLES, N_EVAL = 10, 256, len(val_obs)
T = val_obs.shape[1]
H = T - BURN_IN
filter_preds, feat_list, kl_list = [], [], []
imag_list, det_list, true_list = [], [], []
oracle_filters, oracle_means, oracle_covs = [], [], []
for i in range(N_EVAL):
o_seq = torch.tensor(val_obs[i], dtype=torch.float32)
a_seq = torch.tensor(val_act[i], dtype=torch.float32)
f_pred, feats, kls = rssm_filter(rssm, o_seq, a_seq, seed=1000 + i)
imag_list.append(
rssm_imagine(rssm, o_seq, a_seq, BURN_IN, N_SAMPLES, seed=2000 + i)
)
det_list.append(det_imagine(det_model, o_seq, a_seq, BURN_IN))
filter_preds.append(f_pred)
feat_list.append(feats)
kl_list.append(kls)
true_list.append(val_states[i, :T, 0])
oracle_filters.append(kalman_reference(val_obs[i], val_act[i])[0][:, 0])
om, op = kalman_reference(val_obs[i], val_act[i], stop_observing=BURN_IN)
oracle_means.append(om[BURN_IN:, 0])
oracle_covs.append(op[BURN_IN:, 0, 0])
true_pos = np.stack(true_list) # (N_EVAL, T)
filt_arr = np.stack(filter_preds) # (N_EVAL, T, 1)
imag = np.stack(imag_list) # (N_EVAL, H, S, 1)
det_arr = np.stack(det_list) # (N_EVAL, H, 1)
val_feats = np.stack(feat_list) # (N_EVAL, T, h_dim + z_dim)
val_kl = np.stack(kl_list) # (N_EVAL, T)
true_future = true_pos[:, BURN_IN:] # (N_EVAL, H)
imag_mean = imag[..., 0].mean(axis=2) # (N_EVAL, H)
imag_std = imag[..., 0].std(axis=2)
imag_rmse, imag_sample_rmse = point_and_sample_rmse(imag[..., 0], true_future)
det_rmse = np.sqrt(((det_arr[..., 0] - true_future) ** 2).mean(axis=0))
filter_err = np.sqrt(((filt_arr[..., 0] - true_pos) ** 2).mean(axis=0))
filter_rmse = np.sqrt(((filt_arr[..., 0] - true_pos) ** 2).mean())
echo_rmse = np.sqrt(((val_obs[:N_EVAL, :, 0] - true_pos) ** 2).mean())
imag_lower, imag_upper = np.quantile(imag[..., 0], [0.025, 0.975], axis=2)
coverage, band_width, interval_score = interval_metrics(
imag_lower, imag_upper, true_future
)
coverage_low, coverage_high = wilson_interval(coverage, N_EVAL)
oracle_mean, oracle_var = np.stack(oracle_means), np.stack(oracle_covs)
oracle_rmse = np.sqrt(((oracle_mean - true_future) ** 2).mean(0))
oracle_expected_rmse = np.sqrt(oracle_var.mean(0))
oracle_filter_rmse = np.sqrt(
((np.stack(oracle_filters) - true_pos) ** 2).mean()
)
oracle_lower = oracle_mean - 1.959963984540054 * np.sqrt(oracle_var)
oracle_upper = oracle_mean + 1.959963984540054 * np.sqrt(oracle_var)
oracle_coverage, oracle_width, oracle_score = interval_metrics(
oracle_lower, oracle_upper, true_future
)
marginal_P = np.diag([1.0, 0.25])
shrink_gains = []
for t in range(T):
shrink_gains.append(marginal_P[0, 0] / (marginal_P[0, 0] + obs_std**2))
marginal_P = A @ marginal_P @ A.T + Q + 0.5**2 * (B @ B.T)
shrink_pred = val_obs[:N_EVAL, :, 0] * np.array(shrink_gains)
shrink_rmse = np.sqrt(((shrink_pred - true_pos) ** 2).mean())
horizons = np.arange(1, H + 1)echo-the-observation RMSE: 0.249
RSSM filtering RMSE (vs true pos): 0.210
known-model memoryless RMSE: 0.231
known-model Kalman filtering RMSE: 0.146
RSSM mean RMSE: h=1 0.198 h=10 0.304 h=40 0.547
deterministic RMSE: h=1 0.190 h=10 0.300 h=40 0.498
Kalman mean RMSE: h=1 0.189 h=10 0.280 h=40 0.488
RSSM sample-path RMSE at h=40: 0.927 (different estimand)
h=1: RSSM interval coverage 100.0% (Wilson interval 95.4%–100.0%), width 1.622, interval score 1.622
Kalman interval coverage 95.0%, width 0.742, interval score 0.930
h=10: RSSM interval coverage 100.0% (Wilson interval 95.4%–100.0%), width 2.108, interval score 2.108
Kalman interval coverage 97.5%, width 1.161, interval score 1.211
h=40: RSSM interval coverage 98.8% (Wilson interval 93.3%–99.8%), width 2.893, interval score 2.901
Kalman interval coverage 93.8%, width 1.912, interval score 2.048
mean prior-posterior KL on val: 1.02 nats/stepRead the filter errors against all three references, not echo alone. A memoryless estimate can beat echo through shrinkage toward a marginal mean without consulting any earlier observations. The Kalman filter uses the supplied simulator parameters, which the neural models do not receive, and is an estimation oracle rather than a fair data-only competitor. The two neural point forecasts are now scored identically: first average RSSM samples, then compute squared error across held-out sequences. Individual-sample RMSE is reported separately because it also includes ensemble variance. Coverage and interval width assess a different output property; neither turns point error into a decomposition of model bias and physical randomness. Raw held-out KL is another diagnostic, not a count of useful facts the model has learned.
In the seeded CPU reproduction (PyTorch 2.13.0, NumPy 2.4.2, two CPU threads), posterior-mean position RMSE was 0.210, versus 0.249 for echo, 0.231 for known-model memoryless shrinkage, and 0.146 for Kalman filtering. This RSSM does denoise on these sequences, but falls short of the known-model reference. At horizon 40, ensemble-mean RMSE was 0.547, versus 0.498 for the point baseline and 0.488 for the Kalman mean. These are one-seed results, not a statistically established ranking across training runs. The Monte Carlo ensemble mean also has finite-particle estimation error.
The sample-path RMSE at horizon 40 was 0.927. This is not a contradiction: for the same target, mean squared sample error equals squared ensemble-mean error plus ensemble variance. At a fixed ensemble mean, greater ensemble variance raises the sample-path metric, even when that variance correctly represents possible futures. That is why the point-error plot uses the ensemble mean and the interval diagnostics are reported separately.
The RSSM intervals covered 100% of held-out positions at horizons 1 and 10, and 98.8% at horizon 40. Their mean widths were 1.622, 2.108, and 2.893 position units, compared with Kalman widths of 0.742, 1.161, and 1.912. The central interval scores were also worse than the oracle's at those horizons. This evidence favors a diagnosis of unnecessarily broad decoder-mean intervals here, not a claim of successful calibration. At horizon 40 the coverage estimate's Wilson interval was approximately 93.3%–99.8%; 80 outcomes still leave substantial tail uncertainty.
One sequence, two regimes
The next figure shows a single held-out sequence, which makes the regime split vivid. Before the vertical line the model filters with real observations; after it, imagination runs on the prior alone.
demo = 0
t_all = np.arange(T)
t_imag = np.arange(BURN_IN, T)
demo_true = true_pos[demo]
demo_obs = val_obs[demo, :, 0]
demo_filter = filt_arr[demo, :, 0]
demo_imag_mean = imag_mean[demo]
demo_imag_lower, demo_imag_upper = imag_lower[demo], imag_upper[demo]
demo_det = det_arr[demo, :, 0]
The posterior trace continues using observations after the cutoff only as an offline comparison; the imagined paths do not. The interval can fluctuate in width with both the inferred dynamics and finite Monte Carlo sampling. Inspect where truth leaves it and compare that with the aggregate coverage and width below. The point baseline has no uncertainty head, so its line alone supplies no interval. Conversely, a visibly broad RSSM band does not certify reliable uncertainty.
Error versus horizon
Aggregating over all evaluated sequences shows the horizon-dependent errors quantitatively.
Both neural point forecasts and the exact Kalman mean incur error against stochastic futures. Even the oracle cannot predict the realized process noise. Its covariance recursion supplies the irreducible conditional mean-squared-error reference; finite held-out RMSE also fluctuates with the sampled outcomes. An excess over the oracle in expectation can reflect model or inference defects, while a crossing on a small test set is not proof that a learned model beats the conditional-mean optimum. The posterior curve uses observations throughout and is shown only for scale. It need not be flat, and none of these curves is required to rise monotonically.
What the belief actually contains
Finally, we test linear recoverability of physical-state information. The RSSM is frozen. A ridge readout sees training features and true-state labels, and is evaluated on held-out sequences. Those labels train the probe, not the world model. Success supports a relationship between the learned representation and velocity under this data distribution; it does not establish how the RSSM uses that information or whether its full belief is sufficient.
probe_feats, probe_targets = [], []
for i in range(0, len(train_obs), 4):
o_seq = torch.tensor(train_obs[i], dtype=torch.float32)
a_seq = torch.tensor(train_act[i], dtype=torch.float32)
_, feats, _ = rssm_filter(rssm, o_seq, a_seq, seed=i, n_samples=64)
probe_feats.append(feats)
probe_targets.append(train_states[i, :T, :])
Xp = np.concatenate(probe_feats)
yp = np.concatenate(probe_targets)
Xv = val_feats.reshape(-1, val_feats.shape[-1])
yv = np.stack([val_states[i, :T, :] for i in range(N_EVAL)]).reshape(-1, 2)
def ridge_r2(X_tr, y_tr, X_te, y_te, lam=1e-2):
x_mean, y_mean = X_tr.mean(0), y_tr.mean(0)
Xc, yc = X_tr - x_mean, y_tr - y_mean
gram = Xc.T @ Xc + lam * np.eye(X_tr.shape[1])
w = np.linalg.solve(gram, Xc.T @ yc)
pred = (X_te - x_mean) @ w + y_mean
ss_res = ((y_te - pred) ** 2).sum(axis=0)
ss_tot = ((y_te - y_te.mean(axis=0)) ** 2).sum(axis=0)
if np.any(ss_tot == 0):
raise ValueError("R-squared requires nonconstant held-out targets.")
return 1.0 - ss_res / ss_tot
r2 = ridge_r2(Xp, yp, Xv, yv)
## Fit PCA on training features; project held-out features without using labels.
Xall = val_feats.reshape(-1, val_feats.shape[-1])
Xc = Xall - Xp.mean(axis=0)
_, _, Vt = np.linalg.svd(Xp - Xp.mean(axis=0), full_matrices=False)
belief_pc = Xc @ Vt[:2].T
train_pc = (Xp - Xp.mean(axis=0)) @ Vt[:2].T
pca_r2 = ridge_r2(train_pc, yp, belief_pc, yv)
vel_color = np.stack([val_states[i, :T, 1] for i in range(N_EVAL)]).reshape(-1)linear probe R^2, belief -> position: 0.94 linear probe R^2, belief -> velocity: 0.93 two-PC probe R^2, projection -> velocity: 0.05 (velocity labels train only the separate probe, not the RSSM)
The full-feature probe and the two-dimensional PCA answer different questions. The probe tests whether a supervised linear readout can recover state information on held-out sequences. PCA retains directions of high feature variance without using velocity labels. Mixed colors in its first two components therefore do not contradict strong recoverability from all 56 features; informative directions may lie elsewhere. A weak linear probe would not rule out nonlinear information either. The held-out split reduces direct trajectory memorization, but additional seeds, input-only or untrained-feature controls, and tests under new action distributions would be needed for stronger conclusions.
The useful result is more limited than autonomous discovery of a named physical variable: a predictive model can learn features from noisy positions and recorded actions from which a separate readout recovers velocity-related information. It was not given true velocities, the spring constant, or damping during world-model training. Good probe scores do not by themselves establish a calibrated posterior, good long-horizon forecasts, or causal use of the recovered feature. The experiment keeps those questions separate.
The reproduced held-out full-feature probe scored about 0.94 for position and 0.93 for velocity. Restricting the probe to the first two training-fitted principal components reduced velocity to about 0.05. The mixed colors in this projection are therefore a useful warning: the directions with most representation variance are not necessarily those carrying the state information of interest. Neither number establishes Markov sufficiency.