## Mathematics

### The data decides what can be learned

The state is the scene's 16 planar numbers (pusher and three pucks, position and velocity); the action is the pusher's commanded velocity, held for 0.1 s. Transitions come from a collection policy, and the choice of policy decides what the model sees. Uniformly random actions would rarely bring the pusher to a puck, so the model would learn little about contact, which is the only hard part of this scene (Lesson 21.4 showed the same coverage effect for the cart-pole). The policy here is the expert of Level 22 plus Gaussian noise ($\sigma = 0.3$), with uniformly random actions in 20% of steps. Over 1000 episodes (120 000 transitions) the pusher touches a puck in 40.9% of steps.

### The model and its rollout

An MLP with three hidden layers of 256 SiLU units predicts the normalized change of state, $\Delta s = (f_\theta(s, a) - s)$, scaled by the training data's statistics, and is trained for 20 000 steps on single transitions:

$$\mathcal L(\theta) = \mathbb E\, \big\| (f_\theta(s_t, a_t) - s_{t+1}) \oslash \sigma_s \big\|^2 .$$

An $H$-step prediction feeds each output back in, under the logged actions:

$$\hat s_{t+k+1} = f_\theta(\hat s_{t+k}, a_{t+k}),\qquad \hat s_t = s_t .$$

### Error against horizon

Output (2, 3): from random states of 100 held-out episodes, the mean distance between predicted and true positions:

| Model | quantity | 1 step | 5 | 10 | 20 | 50 |
|---|---|---|---|---|---|---|
| nothing moves ($\hat s = s_t$) | pusher | 12.98 | 44.88 | 78.57 | 123.24 | 188.80 |
| | pushed puck | 3.48 | 17.51 | 34.63 | 68.93 | 143.32 |
| linear, least squares | pusher | 0.37 | 1.60 | 3.17 | 6.03 | 12.64 |
| | pushed puck | 2.69 | 14.84 | 30.80 | 61.45 | 129.71 |
| MLP, one-step objective | pusher | 1.18 | 4.17 | 7.57 | 13.11 | 27.56 |
| | pushed puck | 0.99 | 3.83 | 8.13 | 20.59 | 76.66 |

(mm; one step is 0.1 s.) Three readings:
- **The baselines calibrate the numbers.** "Nothing moves" says how far things move; a model is useful only where it beats that. The MLP's 77 mm at 50 steps is half of the 143 mm the puck actually travels in 5 s: better than not predicting, far from good.
- **A simpler model wins where the physics is simple.** The free pusher is close to a linear system (a velocity servo), and the linear model predicts it two to three times better than the MLP at every horizon (3.2 times at one step, 2.2 times at fifty). On the pushed puck, which only moves through contact, the linear model is barely better than not predicting at all. Contact is the nonlinear part.
- **The error grows faster than linearly at long horizons for the puck** (8 mm at 10 steps, 21 at 20, 77 at 50), and more slowly than linearly for the pusher (7.6, 13.1 and 27.6 mm). A slope steeper than 1 on log-log axes means errors are compounding, not just accumulating.

### How one-step accuracy is lost

Output (4) follows one held-out episode (the green puck toward the upper zone), predicted for 6 s from its first state under its logged actions:

| time | pusher, true | pusher, predicted | puck, true | puck, predicted |
|---|---|---|---|---|
| 0 s | (−0.213, 0.121) | (−0.213, 0.121) | (−0.088, −0.001) | (−0.088, −0.001) |
| 2 s | (−0.088, 0.001) | (−0.096, −0.004) | (−0.044, 0.006) | (−0.053, −0.001) |
| 4 s | (−0.014, −0.026) | (−0.007, −0.030) | (0.043, −0.035) | (0.037, −0.004) |
| 6 s | (0.140, −0.007) | (0.173, −0.015) | (0.182, 0.007) | (0.044, 0.004) |

(m.) For 4 s the predicted pusher stays within a centimetre of the real one. Somewhere after that, the predicted pusher's few millimetres of error put it beside the predicted puck rather than behind it, and the predicted contact ends. The real pusher keeps pushing; the real puck travels another 14 cm while the predicted one stays at x = 0.044 m. A per-step error of about a millimetre became a qualitative error (contact or no contact), and a qualitative error does not shrink (inference from this trajectory; Lesson 23.5 measures the consequence for control).

> [!recommendation] Evaluating a learned dynamics model
> Report open-loop error at several horizons, on held-out data from the distribution the model will be used on, next to "nothing moves" and a simple model of your own. Split the error by regime (contact and free motion here). Look at trajectories, not only averages: the failure that matters is often discrete.

## Implementation

```io
INPUT: tabletop.xml; the standard datasets and model (mjcourse.worldmodel.standard)
PROCESS: collect 1000 + 100 episodes; fit three models; open-loop errors at five horizons; one trajectory
OUTPUT: tables; runs/l23_2/horizon.png; runs/l23_2/dynamics.json

def rollout(self, s0, actions):
    s = np.atleast_2d(s0).astype(float)
    out = [s]
    for t in range(actions.shape[1]):
        s = s + np.c_[s, actions[:, t], np.ones(len(s))] @ self.w
        out.append(s)
    return np.stack(out, 1)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 3.8), constrained_layout=True)
colours = {"nothing moves": "#7a7a7a", "linear": "#E69F00", "MLP, one-step objective": "#0072B2"}
for name, err in table.items():
    ax1.plot(HORIZONS, [1000 * err[h]["target"] for h in HORIZONS], marker="o", color=colours[name], label=name)
ax1.set_xscale("log")
ax1.set_yscale("log")
ax1.set_xlabel("prediction horizon (steps of 0.1 s)")
ax1.set_ylabel("error of the pushed puck (mm)")
ax1.grid(alpha=0.3, which="both")
ax1.legend(frameon=False)
true, pred, i = np.array(trajectory["true"]), np.array(trajectory["pred"]), 4 + 4 * trajectory["target"]
ax2.plot(true[:, 1], true[:, 0], color="#0072B2", label="pusher, MuJoCo")
ax2.plot(pred[:, 1], pred[:, 0], "--", color="#0072B2", label="pusher, learned model")
ax2.plot(true[:, i + 1], true[:, i], color="#D55E00", label="puck, MuJoCo")
ax2.plot(pred[:, i + 1], pred[:, i], "--", color="#D55E00", label="puck, learned model")
ax2.invert_xaxis()
ax2.set_aspect("equal")
ax2.set_xlabel("y (m)")
ax2.set_ylabel("x (m)")
ax2.legend(frameon=False, fontsize=8)
ax2.set_title("one held-out episode, 6 s open loop", fontsize=10)
fig.savefig(OUT / "horizon.png", dpi=150)
plt.close(fig)

def b64(t) -> str:
    return base64.b64encode(np.ascontiguousarray(t.detach().numpy(), dtype="<f4").tobytes()).decode()

layers = [m for m in mlp.net if isinstance(m, torch.nn.Linear)]
blob = {"layers": [{"out": layer.out_features, "in": layer.in_features, "w": b64(layer.weight), "b": b64(layer.bias)}
                   for layer in layers],
        "activation": "silu", **{k: getattr(mlp, k).numpy().round(8).tolist() for k in ("s_mean", "s_std", "d_mean", "d_std")},
        "state": list(tt.STATE_NAMES), "dt": tt.CONTROL_DT, "max_speed": tt.MAX_SPEED,
        "source": "examples/l23_2_learned_dynamics.py; mjcourse.worldmodel.standard model 'one_step' (1000 guided episodes, "
                  "seed 0, 20 000 steps)"}
(OUT / "dynamics.json").write_text(json.dumps(blob))

The model is trained by `mjcourse.worldmodel.dynamics.fit`:

```python file=src/mjcourse/worldmodel/dynamics.py
from __future__ import annotations

import time
from dataclasses import dataclass
from itertools import pairwise

import numpy as np
import torch
from torch import nn

STATE, ACTION = 16, 2


@dataclass
class Normalizer:
    s_mean: np.ndarray
    s_std: np.ndarray
    d_mean: np.ndarray
    d_std: np.ndarray


class DynamicsModel(nn.Module):
    def __init__(self, norm: Normalizer, hidden: int = 256, layers: int = 3):
        super().__init__()
        sizes = [STATE + ACTION] + [hidden] * layers
        mods = []
        for a, b in pairwise(sizes):
            mods += [nn.Linear(a, b), nn.SiLU()]
        self.net = nn.Sequential(*mods, nn.Linear(hidden, STATE))
        for name, value in (("s_mean", norm.s_mean), ("s_std", norm.s_std), ("d_mean", norm.d_mean), ("d_std", norm.d_std)):
            self.register_buffer(name, torch.as_tensor(value, dtype=torch.float32))

    def forward(self, s: torch.Tensor, a: torch.Tensor) -> torch.Tensor:
        x = torch.cat([(s - self.s_mean) / self.s_std, a], dim=-1)
        return s + self.net(x) * self.d_std + self.d_mean

    @torch.no_grad()
    def rollout(self, s0: np.ndarray, actions: np.ndarray) -> np.ndarray:
        s = torch.as_tensor(np.atleast_2d(s0), dtype=torch.float32)
        a = torch.as_tensor(actions, dtype=torch.float32)
        out = [s]
        for t in range(a.shape[1]):
            s = self(s, a[:, t])
            out.append(s)
        return torch.stack(out, 1).numpy()


def normalizer(s: np.ndarray, s2: np.ndarray) -> Normalizer:
    d = s2 - s
    return Normalizer(s.mean(0), s.std(0) + 1e-6, d.mean(0), d.std(0) + 1e-6)


def fit(episodes: list[dict], steps: int = 20_000, batch: int = 256, lr: float = 1e-3, seed: int = 0,
        objective: str = "one_step", k: int = 5, hidden: int = 256, log=None) -> tuple[DynamicsModel, list[float]]:
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    s = np.concatenate([e["states"][:-1] for e in episodes])
    s2 = np.concatenate([e["states"][1:] for e in episodes])
    model = DynamicsModel(normalizer(s, s2), hidden=hidden)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, steps)
    window = 1 if objective == "one_step" else k
    # Flat arrays: states of every episode back to back; a window may start at t if t + window stays in the episode.
    flat_s = torch.as_tensor(np.concatenate([e["states"] for e in episodes]))
    flat_a = torch.as_tensor(np.concatenate([np.concatenate([e["actions"], e["actions"][-1:]]) for e in episodes]))
    offsets = np.cumsum([0] + [len(e["states"]) for e in episodes])
    starts = torch.as_tensor(np.concatenate([o + np.arange(len(e["actions"]) - window + 1) for o, e in zip(offsets, episodes)]))
    losses, start = [], time.perf_counter()
    for step in range(steps):
        idx = starts[torch.as_tensor(rng.integers(0, len(starts), batch))]
        s_t = flat_s[idx]
        loss = 0.0
        for h in range(window):
            s_t = model(s_t, flat_a[idx + h])
            loss = loss + (((s_t - flat_s[idx + h + 1]) / model.s_std) ** 2).mean()
        loss = loss / window
        opt.zero_grad()
        loss.backward()
        opt.step()
        sched.step()
        losses.append(float(loss.detach()))
        if log and (step + 1) % 5000 == 0:
            log(f"step {step + 1}: loss {np.mean(losses[-5000:]):.5f} ({time.perf_counter() - start:.0f} s)")
    return model, losses


def horizon_errors(model: DynamicsModel, episodes: list[dict], horizons=(1, 5, 10, 20, 50), starts_per_episode: int = 4,
                   seed: int = 0) -> dict:
    from mjcourse.tabletop import OBJECTS

    rng = np.random.default_rng(seed)
    out = {h: {"pusher": [], "target": [], "worst_puck": []} for h in horizons}
    hmax = max(horizons)
    for e in episodes:
        target = OBJECTS.index(e["task"][0])
        for t0 in rng.integers(0, len(e["actions"]) - hmax + 1, starts_per_episode):
            pred = model.rollout(e["states"][t0][None], e["actions"][None, t0:t0 + hmax])[0]
            for h in horizons:
                true, p = e["states"][t0 + h], pred[h]
                puck = [np.linalg.norm(p[4 + 4 * i:6 + 4 * i] - true[4 + 4 * i:6 + 4 * i]) for i in range(3)]
                out[h]["pusher"].append(np.linalg.norm(p[:2] - true[:2]))
                out[h]["target"].append(puck[target])
                out[h]["worst_puck"].append(max(puck))
    return {h: {k: float(np.mean(v)) for k, v in d.items()} for h, d in out.items()}