This simulator is best viewed on a desktop or laptop — interactive grids and timelines are easier to explore on a larger screen.
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).
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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
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
swap_schedAt 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.
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.
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.
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.
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.
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.
A short-term store with deterministic recency bias: new chunks replace the oldest chunks, concentrating capacity on the current task.
A long-term distribution-matching store: each chunk receives a random key and replaces the lowest-key slot only if its key is larger.
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.
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.
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.
Within the selected buffer, sampling is uniform. This is not Prioritized Experience Replay: LTDM's random key determines retention, not training priority.
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.
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.