Learned simulator

Overview

A learned simulator is a network that maps the recent state of a system, plus any known external inputs, to the state one time step later. A forecast of any length is produced by calling this model repeatedly and feeding each prediction back as the newest input 2212.12794 app. §3.1. Some variants change the output. TrafficSim predicts a plan for every actor over a planning window, then applies only the first part of that plan before it predicts again 2101.06557 §3.3.

The shared problem is error build-up. A model that is checked on single steps can still drift when it is run for many steps, because its own small mistakes become its next inputs. The papers handle this by training through rollouts. GraphCast trains on rollouts of up to 12 steps 2212.12794 app. §4.2. TrafficSim runs its policy in closed loop for the whole 12 s simulation during training 2101.06557 §3.4. In the TrafficSim ablation, training through the rollout cut the collision rate from 5.92% to 1.28% at the same planning length 2101.06557 tab. 2.

Three cards cover this concept. The first is a weather forecaster. The second is a traffic simulator for self-driving tests. The third is a survey of world models, which describes the same kind of one-step model for planning in reinforcement learning, not for forecasting 2411.14499 §3.1.

Variants

Mesh graph network with residual output (GraphCast)

The grid is encoded onto a multi-mesh, a refined icosahedron whose edges from all refinement levels are kept. Coarse edges give long-range links, fine edges give local ones. One message-passing step encodes grid to mesh, 16 unshared processor layers run on the mesh, and one step decodes back to the grid. The output is a normalised change added to the current state 2212.12794 app. §3.3 2212.12794 app. §3.4 2212.12794 app. §3.5 2212.12794 app. §3.6. Card: GraphCast: Learning skillful medium-range global weather forecasting.

Reported setup: 0.25° grid of 721 × 1440 points, 6-hour step, up to 40 steps (10 days), 36.7 million parameters 2212.12794 app. §2.2 2212.12794 app. §2.3 2212.12794 §GraphCast.

Joint multi-actor traffic simulator with unrolled training (TrafficSim)

Here the one-step model is replaced by a policy for all actors in a scene. At each tick the model encodes a rasterised map once, and for each actor encodes its last 3 s of boxes with a recurrent network. A graph-based interaction module then links the actors. A single scene-level latent sample is decoded into a joint plan for every actor, so the actors' plans fit together 2101.06557 §3.1 2101.06557 §3.2 2101.06557 app. A. Each sample gives one whole-scene plan, and several samples from the same start give several plausible scenes 2101.06557 §3.2. The simulator applies the first κ steps of the plan and then plans again 2101.06557 §3.3. On the GTX 1080 Ti the authors measured 0.83 s, 0.45 s and 0.24 s for 2 Hz, 1 Hz and 0.5 Hz inference, with collision rates of 0.50%, 0.85% and 0.96% 2101.06557 tab. 4.

Training runs the policy in closed loop for 12 s. The first 8 s have expert labels and use an imitation term, a variational bound on the future given the ground truth. Beyond that point the imitation term is dropped. A collision penalty is applied to samples from the prior 2101.06557 §3.4 2101.06557 eq. 3. The collision term lowered the collision rate from 0.60% to 0.50% in the ablation 2101.06557 tab. 2. At simulation time, rejection sampling or gradient steps on the latent can reduce collisions further without retraining. In the reported runs they lowered the collision rate to 0.33% and 0.12%, and raised the rule-violation rate from 2.77% to about 3.0% 2101.06557 tab. 5 2101.06557 app. B. Card: TrafficSim: Learning to Simulate Realistic Multi-Agent Behaviors.

Author-reported results on the ATG4D dataset with 15 samples per scenario 2101.06557 tab. 1: scenario collision rate 0.50%, against 1.19% for the rule-based IDM and 2.90% for the next learned baseline, ILVM. The minimum ADE over samples is 0.57 m, the lowest in the table. The rule-violation rate is 2.77%, against 0.25% for IDM 2101.06557 tab. 1. A planner trained on about 15k synthetic scenes from this simulator reaches a collision rate of 10.73% on real test scenarios, against 10.56% when trained on real data 2101.06557 tab. 3. These results come from one dataset and are not independently checked.

One-step transition model for planning (model-based RL)

The survey treats a learned world model as the transition function of a Markov decision process whose reward function is known. It fits that function by supervised learning on real transitions, with a squared one-step error 2411.14499 §3.1 2411.14499 eq. 1 or a probabilistic one 2411.14499 eq. 2. A planner then picks the action sequence with the highest predicted return over a horizon τ 2411.14499 eq. 3. The survey names Monte Carlo sampling of action sequences, probabilistic ensembles with trajectory sampling, Monte Carlo tree search over learned values, and latent-space trajectory optimisation as ways to search 2411.14499 §3.1. The survey gives no architecture or training schedule for the one-step model. It lists cloud-side roles for such models as data engines, reinforcement learning environments and policy evaluators, without reporting results for them 2411.14499 §5.5. Card: Understanding World or Predicting Future? A Comprehensive Survey of World Models.

When to use which

The variants serve different goals. GraphCast forecasts a system from a long trajectory archive. TrafficSim simulates other agents that react to each other and must avoid collisions. The planning variant chooses actions. Only the first two report results, and their numbers come from different tasks and protocols, so this table is about constraints, not about which variant is better.

Claims are author-reported. The GraphCast headline is that the model beats HRES-fc0 on 90.3% of 1380 targets at 10 days [arxiv:2212.12794 §Forecast verification results]. The Pangu-Weather comparison uses a different initialisation and target protocol, so its 99.2% is not comparable to the headline 2212.12794 app. §6. The survey reports no experiments of its own; its benchmark table names evaluation suites without scores 2411.14499 tab. 8.

Implementation sketch

The rollout loop is the common core. The GraphCast training objective is a weighted, rollout-based squared error. The survey's planning loop uses the same one-step model but scores candidate action sequences by predicted return instead of forecasting a fixed horizon 2411.14499 eq. 3. The card's pseudocode is its own rendering of each paper's description, not released code.

# Roll the one-step model forward  [arxiv:2212.12794 app. §3.1]
# Input:  prev_state (t-1 step), state (t), forcings_for(step)
# Output: list of forecast states

def forecast(prev_state, state, num_steps):
    forecasts = []
    for step in range(num_steps):
        next_state = one_step(prev_state, state, forcings_for(step))
        forecasts.append(next_state)
        # Feed the prediction back in as the newest input
        prev_state, state = state, next_state
    return forecasts

# Autoregressive weighted-MSE training loss  [arxiv:2212.12794 app. §4.2]
# Input:  batch of trajectories, rollout length (ramped up to 12)
# Output: scalar loss; gradients flow through every rollout step

def training_loss(batch, rollout_steps):
    prev_state, state = batch.initial_states
    total = 0.0
    for step in range(rollout_steps):
        pred = one_step(prev_state, state, batch.forcings[step])
        target = batch.targets[step]
        total = total + weighted_squared_error(pred, target)
        prev_state, state = state, pred
    return total / rollout_steps

Variants differ in the one-step model (here a mesh graph network with residual output) and in the rollout schedule. The paper ramps the rollout from 2 to 12 steps over 11,000 updates at a fixed learning rate of 3e-7, after two phases with one step 2212.12794 app. §4.5.

The traffic simulator uses the same outer loop but with three differences. Each tick draws one scene latent shared by all actors. Only the first κ steps of the joint plan are kept. The loss mixes an imitation term, which is dropped after the labelled part, with a collision term on prior samples 2101.06557 §3.3 2101.06557 §3.4. The paper's printed formulas for the mixing weight and the collision loss do not produce the behaviour the text describes, so an implementer must choose the intended form 2101.06557 app. A. Its optimiser, learning rate and batch size are not stated 2101.06557 §4.

Contradictions

The three cards report no numbers for the same task, so there is no direct numerical disagreement between them. Inside the TrafficSim paper, the text calls the rule-violation rate similar to IDM's 2101.06557 §4, while the table gives 2.77% against 0.25% for IDM 2101.06557 tab. 1. Use the table values. The survey notes that the term "world model" is itself contested 2411.14499 §1, so read the planning variant as a different scope, not as a rival result. The protocol differences within the GraphCast comparisons are listed under "When to use which".

Papers

• GraphCast: Learning skillful medium-range global weather forecasting: a mesh graph network learned one-step weather model, rolled out autoregressively and trained on ERA5.
• TrafficSim: Learning to Simulate Realistic Multi-Agent Behaviors: a joint multi-actor traffic policy with a scene latent, trained by unrolling through the simulator with an imitation and a collision loss.
• Understanding World or Predicting Future? A Comprehensive Survey of World Models: a survey that frames a learned one-step transition model for planning in model-based reinforcement learning. Reports no experiments.