This simulator is best viewed on a desktop or laptop — interactive grids and timelines are easier to explore on a larger screen.

ARROW Simulator

ARROW: Augmented Replay for RObust World models

A model-based continual-reinforcement-learning method that extends DreamerV3 with an augmented replay memory: a recency-biased FIFO store paired with a long-term, distribution-matching buffer. The simulator links the paper's method, code (Code/ARROW_and_DV3), and experiment configs through the --arrow-replay-ratio and swap_sched design choices. Sections follow the code path in train.py: setup (curriculum, replay buffer), then the per-epoch loop (collect → evaluate → train → advance).

1 ARROW

ARROW is not a separate policy architecture. It is a DreamerV3-style world-model agent whose replay distribution is redesigned for continual learning under a fixed memory budget.

AGENT  Actor + Critic

The actor selects actions and the critic estimates value from latent world-model states. Both are trained on imagined rollouts rather than directly on raw observations. Code: ac.py.

MODEL  World Model

A recurrent state-space model learns latent dynamics, image reconstruction, reward, and continuation predictions from an internal state (z, h). Code: wm.py and rssm.py.

MEMORY  Augmented Replay

The contribution is the dual replay distribution: a FIFO buffer for recent experience and a LTDM reservoir that approximates the global history. Code: replay.py.

219observations total replay budget, matched against DreamerV3 and TES-SAC
512timesteps per stored spliced rollout
6 tasksper curriculum in Atari and CoinRun
No task IDsshared model and memory across the whole sequence

Contribution framing: ARROW is a working recipe for distribution-aware replay in a world-model agent, not a new neural architecture or a new prioritized sampling rule.

2 Architecture

ARROW keeps the DreamerV3-style architecture and changes the replay system that trains it. The agent is organized as a world model, an actor-critic controller, and an augmented replay memory.

World model: RSSM latent dynamics

h is a deterministic GRU state that carries temporal context.
z is a stochastic latent state with 32 discrete variables and 32 classes.

The model updates ht+1 = GRU(ht, zt, at), infers posterior latents from observations, and uses prior latents for open-loop imagination.

Encoder, decoder, and prediction heads

Observations are 64x64 RGB frames encoded by a CNN into latent states. The world model reconstructs images and predicts reward and continuation signals, trained with KL balancing to stabilize latent transitions.

Actor-critic controller

The actor and critic are MLPs with 2 layers and 512 hidden features. They consume concat(z, h), not raw pixels, and support 18 Atari actions or 15 CoinRun actions from the configs.

Replay interface

The architecture receives training data as 512-step spliced rollouts. ARROW's FIFO+LTDM replay changes the distribution of world-model updates while leaving the RSSM and actor-critic architecture matched to DreamerV3.

REAL  Environment rollout

Code: generate_trajectories. The actor, or a random policy at epoch 0, interacts with the environment; observations, actions, rewards, continuations, and resets are stored.

obs → encode → (z,h) → actor → action → real env → reward, next obs

DREAM  Imagined rollout

Code: dream_rollout. The actor-critic is trained from short imagined trajectories generated by the RSSM, with no additional environment interaction.

(z,h) → actor → action → WM predicts reward + next (z,h) → repeat ×16

3 Curriculum setup — swap_sched

At startup, config.get_env_schedule() builds SequentialEnvironments. The configs instantiate two task suites and three curriculum settings: default order, reversed order, and two-cycle training. All use the same total budget over 6 tasks; swap_sched controls task duration.

Cycle 1 · epochs 0–539
Epoch 0 / 539
Play animates epochs 0-539; dragging the slider pauses playback.

The agent receives pixels, actions, rewards, continuation flags, and reset flags, but no explicit task-change label or ID. Evaluation is performed every 10th epoch on all 6 tasks in the selected suite.

4 Replay buffer setup

At startup, config.get_replay_buffer() constructs the augmented store. The total replay budget is fixed at 2¹⁹ = 524,288 observations for all methods. ARROW changes how this budget is allocated between short-term and long-term storage, and uses matching minibatch sampling weights.

FIFO 50%
LTDM 50%
512
FIFO slots
512
LTDM slots
0.50
FIFO sample weight
0.50
LTDM sample weight
262,144
FIFO observations
262,144
LTDM observations
524,288
Total (= 2¹⁹, fixed)
512
Steps per slot

In code, _arrow_fifo_ltdm_capacity_ns() splits 2 × data_n_max = 1024 trajectory slots, while _arrow_fifo_ltdm_sampling_weights() assigns the corresponding read probabilities. DreamerV3 and TES-SAC use the same total capacity as one FIFO buffer.

5 One training epoch (train.py)

Mirrors the main loop in train.py: collect, optionally evaluate, train the world model, train the actor-critic, then envs.step(). Config values: n_sync=4, gen_seq_len=4096, steps_per_batch=1000, and ac_train_steps=800. Atari uses env_repeat=4; CoinRun uses env_repeat=1.

① Active task
② Collect
③ replay.add
④ Evaluate
⑤ Train WM
⑥ Train AC
⑦ envs.step()

6 Writing experience to both buffers

After generate_trajectories, replay.add() writes each epoch's 32 new 512-step spliced rollouts. The same incoming chunks are offered to both buffers: FIFO always writes them by overwriting the oldest entries, while LTDM accepts entries through reservoir sampling.

newly written chunk in the current epoch FIFO slot · number = task 1–6 LTDM slot · number = task 1–6 hold Task colors to visualize the per-task composition

FIFO buffer (512 slots)

A short-term store with deterministic recency bias: new chunks replace the oldest chunks, concentrating capacity on the current task.

Task mix in FIFO right now
empty — add an epoch

LTDM buffer (512 slots)

A long-term distribution-matching store: each chunk receives a random key and replaces the lowest-key slot only if its key is larger.

Task mix in LTDM right now
empty — add an epoch
Epoch 0 / 539

Task experience distribution — slots per task (1–6)

Column height is the number of buffer slots from each task. FIFO concentrates on recent tasks; LTDM approaches a broader sample of the historical stream as the curriculum progresses.

The grids are scaled for display (512 real slots become 256 cells at 50/50). All 32 chunks are offered to both buffers: fewer LTDM highlights simply mean fewer reservoir replacements in that epoch.

7 Sampling a minibatch — one buffer per update

Each world-model gradient update calls replay.minibatch(). ARROW first selects one replay buffer according to the configured weights, then samples a uniform minibatch from that buffer.

1
2
3
Press "Next step →" to begin
Three stages define one replay minibatch.

Repeated draws approximate the configured sampling weights

The buffer choice is stochastic at each update, but the empirical draw frequencies converge toward the ratio selected in section 4.

0
FIFO draws
0
LTDM draws

Within the selected buffer, sampling is uniform. This is not Prioritized Experience Replay: LTDM's random key determines retention, not training priority.

8 Why this reduces forgetting

The recency problem

A single FIFO buffer eventually becomes dominated by the current task. The world model is then optimized mainly on recent observations, so its latent dynamics for older tasks degrade, and the actor trains inside an imagined environment that no longer represents the full curriculum.

ARROW's mechanism

The LTDM buffer retains a distribution-matched sample of previous experience under the same total budget. World-model updates therefore continue to include older tasks, stabilizing the latent simulator that the actor-critic uses for imagined training.

The ratio ablation traces a stability-plasticity axis. On Atari, allocating more capacity to LTDM generally reduces forgetting but weakens forward transfer; on CoinRun, where tasks share structure, differences between ratios are often within seed variability. AR50 is therefore a robust default, not a tuned optimum.