Teach a PlayStation agent.
Understand every step.

1. What we achieved—and what remains unsolved

We built a real-game reinforcement-learning system for the first five levels of Crash Bandicoot: Warped, a shared model that can act in all five levels, demonstration-based training, evaluation and films, cloud infrastructure, and two distinct distributed-training experiment tracks.

The strongest completed real-game systems result is 3.32× higher end-to-end training throughput with four GPUs, compared with one, at matching global experience and optimizer-update budgets. Two GPUs achieved approximately 1.86×. Both results repeated closely across two seeds.

This does not establish improved playing ability. All six models in that scaling comparison scored 0/20 evaluation completions. Separate demonstration-trained models complete Orient Express from the demonstrated starting state, but small changes in starting timing still break that success. The longer overnight campaign had not produced an additional level completion at this snapshot.

The project therefore teaches three separate questions:

A successful answer to one does not imply a successful answer to the others.

2. The journey through the project

We first extended PS-Env beyond one platforming level. The first five levels involve platforming, swimming, tiger riding, a chase, and jet-ski control. That required level-specific actions, memory interpretation, starting states and completion checks—not merely five names in a menu.

We added named presets, local/cloud execution, checkpoints, model loading, recorded gameplay and a hosted console. We then explored one shared visual model with separate action outputs for each level, alongside independent level-specific training.

Poor returns and repeated failures led us to inspect reward components, build scripted searches for candidate demonstrations, replay those candidates, and train policies to imitate their actions. We compared demonstration initialization, continued demonstration replay, learning from scratch, and practice from route sections.

A separate transformer project explored distributed systems in a more conventional neural-network workload: single-device training, DDP, FSDP, DeepSpeed ZeRO, NCCL communication, Megatron tensor parallelism, memory limits, tuning, profiling and recovery. We then returned those systems lessons to real PS-Env training using synchronous gradient sharing.

Finally, we expanded the real-game comparison to one, two and four GPUs with fixed update counts, and launched longer reward experiments and additional demonstration searches overnight.

The UI and presentation are different deliverables. At this snapshot the live experiments page contained the new six-run scaling comparison; the Machenta slide deck still contained the older 1.64× pilot result. A live dashboard updating does not automatically update a static presentation.

3. Vocabulary: the units that matter

The policy receives four stacked 84×84 grayscale observations. Stacking helps it infer movement from changes between images. The stack contains successive sampled observations, not necessarily four adjacent game frames.

A decision is not necessarily a fresh greedy network choice: exploration can select or repeat an action. The frame-skip setting remains eight frames per decision, with level-specific button-tap behavior where required.

Example: 200,000 decisions correspond to about 1.6 million emulated action frames, excluding reset and evaluation overhead. With four workers, those are 200,000 decisions in total, not 200,000 per worker.

Episode return and training loss are also different. A decreasing DQN loss means its current sampled predictions fit its targets more closely. It does not establish that the targets encode desirable behavior, or that the policy can finish a level.

4. How a single PS-Env learner works

Figure 1. Two connected loops: playing creates data; learning uses that data to change the policy.

The model is a convolutional DQN. It predicts a value for each available action: how much future discounted reward taking that action is expected to produce. During greedy evaluation it chooses the highest-valued action. During training an epsilon-greedy exploration mechanism sometimes chooses exploratory actions instead.

Double DQN separates choosing a bootstrap action from evaluating that action using the target network. A delayed target network helps stabilize the moving learning target. In simplified one-step notation:

chosen_next_action = argmax Q_online(next_observation, action)
target = reward + gamma * Q_target(next_observation, chosen_next_action)

For a true terminal transition, there is no future-value term. Our implementation also supports three-step returns: it combines several consecutive rewards and bootstraps after that sequence. Episode and level boundaries must not be crossed accidentally. Artificial time limits need careful treatment because truncating an attempt is different from the game reaching a terminal state.

The implementation uses Adam, Huber loss, gradient clipping and target-network synchronization. These ingredients stabilize optimization; none guarantees successful exploration or an appropriate reward objective.

The actual tensor shapes in our code

B is the number of observations processed together, A is the number of discrete actions, and dimensions below are batch × channels × height × width. One grayscale image has one channel; stacking four images produces four channels. The convolutional network has no language input and is not a transformer.

The convolutions have no padding. For example, the first spatial dimension is floor((84 - 8) / 4) + 1 = 20. ReLU keeps its input shape. “512 units” describes a hidden vector, not 512 actions or 512 game states. The two single-level dueling branches have separate 3136→512 weight matrices.

The five-level SharedDQN reuses the same convolutional shapes, then has one shared 3136→512 layer. Each level has a small 512→A advantage head and a 512→1 value head. The caller supplies the level identifier to route the head; it is not concatenated as a language input. A forward call for one level produces that level's [B, A] output, not a single concatenation of all five action spaces.

These counts were computed from instantiated code, include biases, and count one policy, excluding its delayed target copy, optimizer state, gradients, replay and activations. Historical checkpoints can use different flags: inspect their saved config rather than assuming every Toad Village model uses 27 actions. The shared model is smaller than a single-level dueling model because it shares the large hidden layer rather than maintaining two separate 3136→512 branches.

In DQN the target network has the same shape as the online policy, but is updated by copying weights rather than by its own optimizer. Double DQN concerns target construction; dueling concerns the output branches. They are independent choices. With 1/2/4 GPUs in our PS-Env benchmark, every GPU receives the complete policy and target network; adding GPUs does not add layers or split this model into pieces.

from agent.model import NatureDQN
import torch
net = NatureDQN(num_actions=7, dueling=True)
x = torch.zeros(2, 4, 84, 84, dtype=torch.uint8)
print(net(x).shape)                          # torch.Size([2, 7])
print(sum(p.numel() for p in net.parameters()))  # 3294376

From the original project report · What PS-Env Is

Historical record · September 2026 reconstruction and single-GPU experiments. These results are separate from the later shared-model scaling experiments in this guide.

PS-Env is an OpenEnv-style reinforcement-learning environment for PlayStation 1 games. At the bottom sits the game itself, a Crash Bandicoot: Warped disc image and a PS1 BIOS, neither of which ever enters the repository. The SwanStation libretro core runs them headlessly; it is downloaded as a binary and loaded at runtime with dlopen. A small C++ host exposes exactly the four things we need from the core: press buttons, grab the frame, read memory, and save or restore state. pybind11 lifts that host into Python, where a FastAPI server offers reset() and step(action) over a WebSocket JSON protocol. The agent, a Nature DQN, sits at the far end of that socket and only ever sees stacked frames and rewards.

The Loop

reset() restores a save state, so every episode begins on the same frame of Toad Village with the same four lives. step(a) turns one of 13 discrete actions into a button bitmask, holds it for eight emulator frames (frame skip 8; the recovered design used 2, and the training story below explains the change), and returns the last of those frames as an 84 by 84 grayscale image together with a reward and a done flag. The actor keeps the last four frames, so the network's input is 4 by 84 by 84.

What the Agent Sees

the emulator's frame

t − 3

t − 2

t − 1

t

What the Agent Can Do

Thirteen discrete actions, each a PS1 button bitmask whose bit index is the libretro joypad id (Cross is bit 0, Square bit 1, the d-pad bits 4 to 7). Cross jumps, Square spins, and Up is forward, because Toad Village runs away from the camera.

• 0NOOP0x0000
• 1UP0x0010
• 2DOWN0x0020
• 3LEFT0x0040
• 4RIGHT0x0080
• 5JUMP0x0001
• 6SPIN0x0002
• 7JUMP_LEFT0x0041
• 8JUMP_RIGHT0x0081
• 9JUMP_UP0x0011
• 10SPIN_LEFT0x0042
• 11SPIN_RIGHT0x0082
• 12SPIN_UP0x0012

Rewards Live in RAM

Nothing the reward needs is on screen in a usable form. Fruit collected, boxes broken and lives left live somewhere in the console's 2 MB of RAM, at addresses nobody documented in 1998. The team found six of them with a memory scanner (search for a known value, change the game state, search again, until one address survives), and the reward function reads them as single bytes every step with no image processing at all. The 2026 reproduction added Crash's position (CRASH_Z, a 32-bit word) for dense progress shaping and for the level-exit check, since the recovered "78 = level done" value was never observed on the real disc.

Signal	Reward	Read from
holding Up (forward)	+0.1	the action mask
pressing Down (backward)	−0.02	the action mask
every step	−0.0002	always
Wumpa fruit collected	+0.5 each	MAIN_WUMPA 0x000AD22D, BONUS_WUMPA 0x000AD239
box broken	+1.0 each	MAIN_BOXES 0x000AD255, BONUS_BOXES 0x000AD241
life lost	−25	LIVES 0x000AD231 decreased
game over (episode ends)	−5,000	a life lost while LIVES was already 0
level complete (episode ends)	+5,000	WARP_STATE 0x0006F5E9 goes from 32 to 52
forward progress	+0.01 per unit	new maximum of CRASH_Z travelled since reset (164.5K units for the whole level)

Two more charges exist only in the trainer, and the training story explains why they were needed: a stall cut (an episode ends after a run of steps without new progress) that costs a lost life, and a cost of 1 on every step that makes no new progress.

The Training Architecture

The talk's architecture diagram shows the distributed design: actors on an epsilon ladder from 1.0 down to 0.01, each with its own emulator, push transitions to Redis; a GPU learner samples batches of 32, trains the policy net (three convolutions, a 512-unit layer, 13 outputs), hard-copies it to the target net and publishes weights; TensorBoard reads its logs. The reproduction below ran the single-process trainer (agent.train) against one emulator server instead, because the emulator, not the GPU, was the bottleneck: about 50 steps per second at frame skip 8 on a g5.xlarge, with the A10G mostly idle.

5. Sharing across levels versus sharing across GPUs

Figure 2. This entire five-head network is copied onto every GPU in a distributed run.

These are independent design choices.

Across levels: our shared model has one visual feature extractor and a shared hidden representation, followed by five level-specific action heads. An action suitable for swimming need not have the same meaning as an action suitable for tiger riding. Balanced replay gives all five levels a contribution to learning.

Across GPUs: synchronous data parallelism places a complete copy of that entire shared model on every GPU. Different workers gather different experience, but synchronize the model updates.

Therefore, four GPUs do not mean four heads, and five levels do not require five GPUs. Every distributed replica contains all five heads. Each worker assigns its next episode to a level with relatively little collected experience.

There is also parallel experimentation: four GPUs can run four unrelated trials with different rewards or methods. Those models never average gradients. That is how the overnight gameplay ablations are scheduled. It is different from the one-/two-/four-GPU scaling trials, which run one distributed experiment at a time.

6. Exactly how our distributed update works

Figure 3. The batch stays at 80 globally. The arrows entering NCCL carry gradients, not raw gameplay.

Our PS-Env implementation uses torchrun to launch one process per GPU. PyTorch initializes the process group; CUDA workers communicate through NCCL. CPU validation uses Gloo.

Each rank owns:

• A complete online network, target network and optimizer.
• A local native emulator and independent exploration randomness.
• Replay buffers partitioned by level.

Initial model parameters are broadcast from rank zero. During training, each rank gathers transitions locally. At the configured update boundary, all ranks check whether their replay partitions are ready. They then sample equal contributions from the five levels.

For the overnight comparison, the global batch is always 80:

Each rank calculates its local loss and gradients. The implementation flattens the gradients into one tensor, sums it across ranks with all_reduce, divides by the number of ranks, copies the averaged gradients back, clips them, and performs Adam's update.

If the four gradients for one parameter were 2, 4, 6 and 8, the averaged gradient would be 5. Adam then uses that gradient together with its optimizer state; the parameter change is not necessarily just minus five times the learning rate.

Equal local sample counts matter. Averaging worker means is equivalent to averaging all samples only when workers have appropriate equal weights. Unequal batches require weighted aggregation. We validate exact batch divisibility instead of silently dropping samples.

What travels between GPUs? Primarily the gradients, plus synchronization/control information and occasional statistics. Raw replay experience stays local. This is not a central experience database, and replay is not shared simply because the model weights are synchronized.

This implementation uses explicit gradient averaging. It follows the synchronous data-parallel principle but does not wrap the PS-Env model in PyTorch DistributedDataParallel. The production DDP wrapper offers mechanisms such as gradient bucketing and communication/computation overlap that our straightforward flattened implementation does not claim to reproduce. See PyTorch DDP and distributed communication.

7. How we checked correctness

We checked three different things, each answering a different question:

1. Reference-gradient agreement: on identical global data, averaged gradients matched a reference learner within approximately 1.27×10⁻⁷ maximum absolute difference.
2. Replica agreement: all distributed replicas ended with identical policy hashes. This detects divergence between workers; it does not prove the objective is correct.
3. Real-game plumbing: short runs exercised actual emulators, updates, checkpoints, evaluations and films. Short smoke-test episodes are not evidence of mastery.

The latest two-GPU smoke collected 800 global decisions and performed exactly 11 scheduled updates, with identical replica weights. A separate five-level demonstration/reward smoke completed 2,000 decisions, 40 updates and 20 evaluation conditions.

Testing the inner algorithm alone was insufficient: a later command-line validation bug prevented all six search jobs from starting even though the underlying search function had passed a direct smoke. We fixed the CLI and added tests with and without a prefix. Test the actual entry point that automation launches.

8. Reading the scaling result correctly

The improved comparison ran 200,000 global decisions per trial, a global batch of 80 and a global replay capacity of 40,000. Updates started at decision 50,000, inclusive, and occurred every 20 global decisions:

number of updates = (200,000 − 50,000) / 20 + 1 = 7,501

If any replay partition was unready when this fixed schedule began, the trial failed rather than silently changing the update count. All six trials passed.

Speedup is throughput divided by the one-GPU throughput. Scaling efficiency is speedup divided by GPU count. It describes how much of ideal linear scaling we achieved; it is not GPU utilization.

The workload took approximately 55, 29 and 16½ minutes respectively, before subsequent evaluation. These are end-to-end training measurements including experience collection and training-loop overhead, not isolated neural-network kernel timings. The one-/two-/four-GPU runs used the same four-GPU host with different numbers of active training devices.

The result is encouraging because it repeats across seeds at matching update counts. It is limited because two seeds are few, trajectories still differ, and fixed-start evaluations do not establish broad generalization. All six scored 0/20 completions.

The earlier pilot measured 51.54 versus 84.33 decisions/s, or 1.64×. However, it performed 20,801 versus 18,764 updates. Different replay readiness changed the amount of learning work. That pilot motivated the explicit fixed-update schedule. Do not present the change from 1.64× to 1.86× as an isolated optimization gain: the benchmark protocol also changed.

9. The separate transformer experiments

Figure 4. Both arrangements use four GPUs, but they answer different experimental questions.

The transformer track trained a small synthetic next-token prediction task. A token is a sequence element; tokens/s is not interchangeable with game decisions/s. The task was deliberately simple, not a claim of language-model quality.

FSDP and ZeRO are different implementations with related sharding goals. NCCL is the communication library they may use, not another sharding strategy. See PyTorch FSDP, DeepSpeed's ZeRO tutorial and PyTorch's FSDP/activation-memory discussion.

We used genuine Megatron Core tensor parallelism, with pipeline size one. Its model and kernels differed from TinyGPT, so its throughput is not an apples-to-apples framework speed comparison.

The small two-node pilot: more GPUs were slower

Fourteen GPU trials completed on two separate A10G instances using private TCP networking. For the same small FP32 TinyGPT workload:

These measurements are consistent with communication overhead outweighing useful parallel work for a small model. They do not prove a universal framework ranking. FSDP had the lowest measured allocation peak in this pilot. DeepSpeed's default communication buckets were large relative to the tiny model, so its memory peak was not a tuned comparison.

A separate 16 MiB NCCL reduction test verified its numerical output and measured about 26.03 ms mean latency, with 0.644 GB/s effective algorithm bandwidth. That is not physical-link bandwidth. A small Megatron run completed at 21,847 tokens/s, but used a different architecture and only one repetition.

The expanded campaign

The retained final phase summaries contain 139 completed trials, 18 OOM outcomes and 3 other failures:

Thus the smaller Megatron pilot succeeded, but the expanded Megatron configurations did not fit. Those statements are compatible and should remain separate.

Actual accelerator work here was on A10G, including a four-GPU g5.12xlarge host. Requested H100 capacity did not turn into an H100 run. We do not claim hands-on H100, H200 or B200 results from these experiments.

GPU reference: what we used, what we did not

A GPU name describes an accelerator. An AWS instance name describes the whole rented machine: GPUs, CPUs, host memory and networking. The two names are not interchangeable.

These are specific product variants. For example, H100 variants do not all have identical memory specifications. The H100/H200/B200 figures above refer to NVIDIA's HGX SXM comparison, not every cloud offering bearing those names. NVIDIA lists their respective device-memory bandwidths as 3.35, 4.8 and up to 8 TB/s. These are manufacturer specifications, not rates measured by our code. Sources: AWS G5/A10G specifications and NVIDIA HGX hardware specifications.

The machines behind the GPU counts

A vCPU is an AWS scheduling unit, not a claim of one dedicated physical core. Instance resources are documented in AWS's accelerated-instance specification table. The one-GPU PS-Env scaling baseline used one active GPU inside the four-GPU host. It was not a comparison against the smaller g5.xlarge's CPU and memory resources.

The four-GPU host has four separate device-memory spaces. Sharding changes how software places data across them; it does not turn them into one ordinary local allocation.

Host RAM holds CPU-side data, including our emulator state and replay arrays. GPU memory (VRAM) holds device tensors and neural-network workspaces. Extra host RAM does not automatically fix CUDA out-of-memory errors. Explicit offloading can move selected state to the CPU, but introduces data movement; our PS-Env learner does not claim an offloading implementation.

GB and GiB differ. GB normally denotes billions of bytes; GiB denotes powers of two. AWS markets A10G as 24 GB, while its technical instance table reports roughly 22 GiB per device. Unit conventions, reservations and runtime availability matter. Use the actual runtime memory report and measured allocation headroom when sizing a job, rather than treating the product label as guaranteed allocatable space.

Capacity, bandwidth and compute answer different questions

HBM means high-bandwidth memory. Its bandwidth figure describes access to memory attached to a GPU; it is not internet speed or NCCL collective throughput. NVLink/NVSwitch describe inter-GPU connectivity in applicable systems. The cited NVIDIA HGX systems use these links; using NCCL on our A10G instances does not mean we measured an HGX/NVSwitch fabric. Our early two-node pilot communicated over private TCP networking.

Why A10G was a useful choice here

We could provision it and verify the actual software stack, whereas the H100 capacity attempts did not yield a runnable experiment. It supplied enough GPU memory for the PS-Env network and enough devices to study one-/two-/four-worker synchronization. The four-GPU machine also provided CPU and host-memory resources for emulators and replay.

For the larger transformer configurations, A10G's memory boundary was useful evidence: it let us observe when replication failed, when sharding helped, and when activation checkpointing substantially increased the tested batch size. That does not mean A10G is universally the cheapest or fastest choice.

A larger accelerator would make sense as a new experiment when a representative workload is demonstrably memory-, compute-, or communication-limited. It would not automatically improve the reward function, extend an incomplete demonstration, or teach recovery from unfamiliar game states. We have no measured H100/H200/B200 speedup to substitute into our A10G results.

A sensible next hardware comparison would keep the model and learning protocol fixed, measure end-to-end throughput and quality, and record CPU resources, interconnect topology, GPU memory peaks, precision and software versions. GPU price alone would be insufficient: the relevant practical quantity is the cost and elapsed time to reach the desired evaluated result.

10. Why we deliberately caused out-of-memory errors

Figure 5. This is a conceptual picture of model-state storage. Activations and temporary buffers also consume memory.

GPU memory holds several things: model parameters, gradients, optimizer state, intermediate activations, communication buffers and temporary allocations. Sharding model state reduces some of these categories. It does not automatically eliminate activation memory.

In a controlled capacity probe we increase batch size until the workload no longer fits. A recorded OOM is an expected boundary observation, provided the runner captures it and proceeds safely. An unexpected OOM in a normal scheduled run is still a failure. We must not relabel every failure as deliberate.

For the tested larger BF16 transformer at sequence length 2,048, the retained capacity results were:

These are sampled boundaries for one model/configuration, not exact mathematical maxima or promises for other workloads. Capacity probes were short; fitting does not establish long-run convergence. Activation checkpointing's large improvement shows why identifying the category consuming memory matters.

When reading memory statistics, distinguish allocated tensor memory, allocator-reserved memory and process/device usage reported by nvidia-smi. They measure different things. Likewise, high GPU utilization does not by itself prove high useful throughput, and low utilization does not by itself identify the bottleneck.

Shapes of the separate transformer benchmark

The transformer benchmarks train bench/model.py::TinyGPT on a synthetic next-token task, not gameplay images. Here T is sequence length, D hidden width, H attention heads and V vocabulary size. These attention heads are not our five game-level output heads.

These are the expanded campaign presets in bench/campaign.py; smoke tests and the early Megatron pilot used smaller overrides. Read each result's configuration before comparing. The conceptual attention matrix is [B,H,T,T], but fused scaled-dot-product attention need not materialize that entire matrix. Doubling sequence length therefore does not justify assuming every measured memory category quadruples.

DDP replicates this architecture. FSDP/ZeRO change storage and communication of model state, without replacing the logical layers. Megatron tensor parallelism partitions computation inside transformer layers; it uses a separate implementation and requires its own equivalence checks. None of these synthetic transformer runs turns TinyGPT into the PS-Env policy.

11. Demonstrations and the gameplay learning experiments

Figure 6. A successful scripted route is an input to learning; it is not itself a learned-policy win.

Our demonstrations were produced by scripted action-sequence search, not by claiming a human expert had played them. Search uses emulator save states to try alternatives, rejects undesirable branches and keeps candidate paths according to heuristics such as survival or spatial coverage. A beam of 12 retains up to 12 candidate branches at each search layer; it does not mean the search tested only 12 sequences overall.

This privileged save-state search is different from a deployed policy acting once through the game. Its output is a candidate action trace. We bind that trace to a starting-state hash and replay it through the actual environment. A warp transition alone is insufficient to claim completion.

Four names that describe different parts of training

These are overlapping descriptions, not four mutually exclusive neural-network architectures. The same image-to-action network can be behavior-cloned, trained with reinforcement learning, and practiced using a reverse curriculum.

BC is a subset of imitation learning. “Demonstration-trained” is a useful provenance label, but name the objective too: BC only, BC followed by RL, or RL with continued demonstration replay. Reverse curriculum can be combined with any suitable learning objective; it does not require copying teacher actions. Our section-practice implementation is related to reverse curriculum but a local section goal is not always the true level exit.

References: behavioral cloning documentation, imitation-learning methods, and Salimans and Chen's demonstration-based reverse-start method. Implementation: tools/learning_trial.py, agent/curriculum.py, and tools/train_ride_sections.py.

# Our BC objective: demonstrated actions are class labels.
action_scores = policy(demo_images)       # [B, A]
loss = torch.nn.functional.cross_entropy(action_scores, demo_actions)  # scalar; labels [B]
optimizer.zero_grad()
loss.backward()
optimizer.step()

The scores are used as classification logits during BC and as Q estimates during subsequent DQN learning. BC alone does not calibrate them as expected future rewards.

Comparison 1: what each term commits us to

The important distinction is data source versus method family versus learning objective versus practice schedule. A label such as “demonstration-trained model” leaves several implementation choices unspecified. The following tables make those choices explicit; “depends on the method” is a real distinction, not a missing definition.

Comparison 2: what happens during training and gameplay

Comparison 3: benefits, limits and convincing evidence

The “retained snapshot” row refers to the evidence dated in this guide. It is not a live claim about the newer level 1 campaign. That campaign uses the historical 13-action ordinary DQN and full-start training, without reconstructing its earlier curriculum or success buffer.

The concrete recipes we compared in code

These rows are actual experimental combinations. The four broad terms above are not four experiment arms.

In tools/learning_trial.py, a normal per-task replay batch for a task with enough demonstration data uses 16 new-replay samples + 16 demo samples. The combined batch pays the TD objective; flagged demo samples also pay the large-margin action objective. BC pretraining instead uses 32 demonstration examples per task and cross-entropy. These are separate losses and training phases. The configured demo margin is 0.8; it is not an 80% demo sampling share.

For the section arm in this tool, the curriculum uses a five-result window, an 80% success threshold and a 50% full-start share. Historical Toad Village recipes have different settings; do not transfer these numbers to every run labelled “curriculum.” A current frontier can advance on section completion in this arm, so its success signal differs from an exit-based curriculum.

Worked example: an obstacle near the end of a level

Imagine a verified route has 100 decisions, and decision 90 is just before a difficult jump. These numbers are illustrative, not a new measured PS-Env result.

A minimal sketch of the difference is:

# Behavioral cloning: observations and teacher actions define a supervised loss.
scores = policy(demo_images)                 # [B, A]
bc_loss = cross_entropy(scores, teacher_actions)
update_weights(bc_loss)

# Reverse curriculum: choose where an RL episode begins.
start = curriculum.pick()                  # None = full start; index = snapshot
observation = env.reset(start=start)
trajectory = play_episode(policy, observation)
update_from_rl_replay(trajectory)
curriculum.report(start, trajectory.success)

This is pseudocode: play_episode, update_weights, trajectory.success and update_from_rl_replay stand for the real implementations. It demonstrates why BC and reverse curriculum can coexist: one determines a loss, the other determines a reset. For section practice, success must be explicitly identified as section success rather than silently treated as level completion.

How to describe a model precisely in the presentation

The audited demonstration set used in the current experiments contains:

We convert replayed observations and actions into imitation data. Behavior cloning trains the network to reproduce those action choices. Continued demonstration replay mixes demonstration transitions with newly collected experience and adds a demonstration action-margin objective. Section practice starts from verified intermediate snapshots; a section success is not a full-start win.

We compared scratch, demonstration initialization, initialization plus continued replay, and replay plus section practice. In an early snapshot all 11 finished demonstration-based trials retained their pre-RL checkpoint under our completion/milestone selector. This means no later checkpoint beat that selection score; ties do not prove every later policy was identical or worse in every respect.

The failure under tiny timing shifts is an important generalization result. Learning the demonstrated sequence of situations is easier than learning robust recovery when the game deviates from that sequence.

From the original project report · Charts

Historical record · September 2026 reconstruction and single-GPU experiments. These results are separate from the later shared-model scaling experiments in this guide.

Five charts rendered with matplotlib from the trainer's per-episode log lines across the 27 training logs, by tools/plot_training.py (the originals live in docs/charts/). Buckets are by the step at which an episode ended.

12. Reward design: what the results actually tell us

The early level-one recipe included a particular direction/progress assumption and large death penalties. Applying it unchanged to swimming or curved routes could produce inappropriate guidance. Later presets disabled direction-specific shaping and introduced capped exploration bonuses. Exploration helped provide feedback, but exploring new cells does not necessarily mean progressing toward an exit.

One clear mismatch under the previous reward setup was Under Pressure: a policy reaching 8/8 reference landmarks earned about -3.5, while a policy reaching only 1/8 and dying earned about +6.2. Makin Waves could earn more than 200 with little reference-route progress because exploration contributed heavily.

The overnight reward package reduced exploration, reduced the per-step cost, strengthened life/game-over penalties, and added +5 for reaching each ordered verified landmark once per episode. Completion detection remained unchanged. Demonstration rewards were recomputed through the same calculator, with every resulting observation checked against the stored data.

This changes several coefficients together. It tests a reward package, not the isolated effect of one coefficient. Within each new paired comparison the other settings match; comparisons against older campaigns also involve changes to learning rate, exploration and starting-state variation.

At this snapshot, greedy checkpoint evaluation showed:

A route reaching 8/8 or 17/17 is at the end of our reference, not necessarily the level. That metric then saturates and cannot distinguish progress beyond the recorded route. Likewise, a level without calibrated landmarks should show “not calibrated,” not a fabricated zero-percent completion metric.

More experience can help, but it can also reinforce a poor objective or fail to discover a narrow successful trajectory. We have evidence that reward ranking is imperfect and robustness is weak. We have not proved that reward is the sole cause, or that the current revision solves it.

13. Infrastructure is part of the experiment

The stack includes Terraform-managed workers, a pinned Python/CUDA environment, native emulator builds, private game assets, restricted SSH access, systemd campaign services, checkpointing, a private S3 artifact destination, a separate collector and the hosted console.

On the GPU workers, presigned upload URLs grant access to specific artifact objects rather than placing broad AWS credentials inside training jobs. The console collects metrics and selected models independently of the operator's laptop. Secrets and private assets are not exposed through the metrics API.

We publish checkpoints by atomic rename so readers do not mistake a partially written file for a usable model. Campaign summaries distinguish queued, running, completed, failed, interrupted and unrun work. Missing evaluation is not zero completions.

A checkpoint containing network and optimizer state still may not be an exact restart. Our PS-Env distributed checkpoint does not preserve replay contents and active emulator state. Loading it starts from learned parameters with new experience collection; it is a warm start. Exact recovery requires identifying and restoring every relevant source of state, including randomness and sampling position.

The overnight controller has a finite runtime, per-job limits, periodic uploads and a separate host shutdown timer. Instance-initiated shutdown terminates the GPU instance. At the snapshot it was still running, with hard termination scheduled for 09:01 Qatar. That is a schedule, not evidence that termination has already occurred.

14. Failures and the lessons they gave us

The three expanded recovery failures were DDP with four GPUs and ZeRO-3 with two/four GPUs. Seven other recovery trials completed. It would be inaccurate to summarize this as universally successful checkpoint recovery.

15. How to evaluate the next experiment

Start with one question and the evidence that could change your mind.

For systems scaling, control model, global batch, total decisions/tokens, update count, precision, replay capacity, checkpoint frequency and timing scope. Repeat seeds or measurements, reverse run order where useful, and preserve failures. Report speedup and efficiency alongside quality.

For learning quality, compare against both scratch and demonstration-only baselines. Keep rewards fixed when comparing training methods, or keep training settings fixed when comparing rewards. Evaluate unseen starting conditions and review films. A deterministic repeat from the same state is not a new independent robustness trial.

For cost, distinguish active-GPU efficiency from the cloud bill. These one-/two-/four-GPU PS-Env measurements all ran on the same four-GPU instance: using one GPU did not reduce that host's hourly price. A 3.32× speedup is not automatically a 3.32× bill reduction. End-to-end cost includes idle time, setup, evaluation, storage and artifact transfer. This document is not an audited AWS invoice.

For memory, identify whether parameters, optimizer state, activations or temporary buffers are limiting the workload before choosing sharding, checkpointing or a larger accelerator.

16. Exercises using our own project

1. Calculate global batch. Four GPUs each sample four transitions per level across five levels. How many samples contribute to one update? Answer: 80.
2. Calculate update count. Start at 50,000 inclusive, stop at 200,000 inclusive, update every 20. Answer: 7,501. If you omit “inclusive,” you miss one update.
3. Calculate efficiency. Four GPUs give 3.32× speedup. Answer: 3.32/4=83% scaling efficiency, not 83% GPU utilization.
4. Spot a misleading claim. “Four GPUs got 0/20, one GPU got 0/20, so the models are equally good.” Unsupported: this test only failed to observe a completion in either; other behavior and untested conditions may differ.
5. Inspect synchronization. Read average_gradients and find where summing becomes averaging. Predict what forgetting the division would do to gradient magnitude.
6. Inspect restart semantics. List what core.save preserves and what the emulator/replay loop keeps only in memory. Explain why loading weights is not an exact continuation.
7. Design a reward ablation. Change only the exploration coefficient while holding data, starts, seeds and budgets fixed. Decide beforehand which gameplay metrics would count as improvement.
8. Design a robustness check. Choose starting perturbations excluded from training. Explain why replaying the exact winning start twenty times is inadequate.

Do not begin by launching every framework at once. First reproduce a small correct single-device loop, then identical-data gradient checks, then two-device execution, then a controlled scaling measurement. Memory-sharding and tensor-parallel experiments become easier to understand once that baseline is solid.

17. Code and evidence map

Paths in this table are relative to the corresponding repository. Generated run files are generally ignored by Git; retain collected evidence separately.

In the separate distributed-training-benchmarks repository, begin with README.md, evidence/RESULTS.md, bench/, and presentation/evidence/expanded-*.json. The phase summaries support the expanded outcome counts and the per-trial memory files support the tested capacity boundaries. The PS-Env six-run record is overnight-distributed.json in the collected learning artifacts.

The source of project-specific numbers is our retained experiment evidence, not the framework documentation. The linked primary documentation explains the techniques; it does not validate our particular results.

A defensible description of the work: “I built and evaluated a multi-level reinforcement-learning system, implemented synchronous data-parallel training with PyTorch/NCCL, measured 3.32× end-to-end throughput on four A10G GPUs at controlled update counts, and compared replicated, sharded and tensor-parallel training in a separate transformer benchmark. I also investigated memory limits, recovery failures, reward misalignment and demonstration generalization. Faster execution is established; robust multi-level gameplay remains unfinished.”

18. Where open robotics models fit

Research extension checked on 14 September 2026. None of the external models in this chapter was trained or deployed by our PS-Env experiments. These are examples for understanding the field, not a ranking of every available release.

A vision-language model (VLM) combines visual input with language. It might identify objects, answer a question about an image, or describe a goal. Understanding an instruction is different from producing correctly scaled robot controls.

A vision-language-action model (VLA) conditions action prediction on visual information and language, often with robot state as an additional input. Its output may be discrete action tokens or a continuous sequence of movements. A VLA is a kind of policy. A robot policy does not have to be a VLA: our DQN has no language input, and geometric demonstration-conditioned policies need not contain an LLM.

Model roles are distinct from training methods. A VLA may use imitation training, while distributed training determines how its parameter updates are computed across devices.

“Diffusion” here usually describes generating actions, not generating pictures. Action chunking means predicting several future actions together. The controller can execute part of the chunk, observe again, and replan. Longer chunks can amortize inference cost but may delay reaction to unexpected events. This is a control-design tradeoff, not the same as increasing the training batch.

Representative open projects

These sizes describe the cited releases, not a claim about every newer checkpoint. LeRobot is a robotics software/data ecosystem, rather than the name of a single policy. Open X-Embodiment denotes a robot-dataset collection used by projects such as OpenVLA and Octo; it is not a model or an optimizer.

The OpenVLA licensing distinction is documented in its official README. For other projects, use their model cards, repository licenses and upstream model terms rather than assuming the word “open” resolves every component.

Why we cannot just load one into Crash Bandicoot

A robot model may expect multiple RGB cameras, joint positions and continuous end-effector actions. PS-Env provides game images and discrete controller actions. The visual distribution, action meaning, timing and training objective differ.

Adapting a representation is plausible; directly sending robot-action outputs to a game controller is not a meaningful integration. We would need a defined observation adapter, an action head or action mapping, task data, and a held-out evaluation. A pretrained model's success on robot manipulation does not establish competence in our game.

19. In-context learning: change the examples, not the weights

Fine-tuning uses new task data to update model parameters. In-context learning (ICL) changes behavior by providing examples in the model's input while keeping its parameters fixed during that adaptation. In robotics, the examples may be demonstrations containing observations and actions, rather than text question-and-answer pairs.

A model can require substantial pretraining to acquire this ability. “No updates for the new task” does not mean “no training ever.”

In notation, fine-tuning changes the parameters: theta → theta_task, then predicts policy(observation; theta_task). In-context adaptation instead predicts policy(observation, demonstrations; theta) with the same theta.

For a model to benefit, it must have an appropriate context interface and the learned ability or algorithmic structure to use it. Appending arbitrary demonstrations to an ordinary DQN does not create ICL. Similarly, an instruction-following VLA does not automatically support few-shot trajectory context just because it understands text.

One way to train that ability is to construct task episodes: provide some demonstrations as context and ask the model to predict actions in another trajectory from the same task. Test on held-out tasks afterward. The context examples and prediction targets must be separated, and evaluation must not feed the answers it is supposed to measure back into the context.

What changes for distributed training?

Distributed pretraining and fixed-weight in-context adaptation can coexist:

1. During pretraining, workers receive different context/query task episodes and average gradients to train a common model.
2. At deployment, the trained model receives a new context and predicts actions without an optimizer update.
3. Serving multiple inference requests or robot streams in parallel is not the same operation as distributed training.

Demonstration context consumes input-processing time and memory. Longer context can increase attention or graph-processing cost. Reusing cached context representations may help when the architecture supports it, but we must measure latency and prediction quality. A larger context is not an unlimited free substitute for training.

20. Connecting this to Edward Johns and Vitalis Vosylius

Edward Johns's lab describes research in robot learning, including imitation and efficient adaptation. For this guide, the strongest connections are demonstrated capabilities in the papers below. I did not establish a dedicated DDP/FSDP/NCCL-method publication by him in the material reviewed. Do not turn “scalable robot learning” into an unsupported claim that he developed a particular distributed-training framework. These academic sources also do not establish Machenta's private model architecture. Lab research index, biography.

A concrete training detail from Instant Policy

The paper's Appendix E reports 2.5 million optimization steps, approximately five days on a single RTX 3080 Ti, followed by a 50k-step learning-rate cooldown, using float16 and compilation. Its scaling study compares 69M, 117M and 178M model variants and training durations. These are the authors' results, not our runs. Parallel generation of pseudo-demonstrations supplies training data; that is distinct from reporting a multi-GPU gradient-synchronization benchmark. Published paper, Sections 4–5 and Appendix E.

The useful comparison is therefore between three kinds of scaling: generating more relevant data, training a larger or longer-trained model, and executing a fixed training workload faster. We measured the third directly; our gameplay results show why the first two and the structure of the task still matter.

What to emphasize in the conversation with Ed

A relevant explanation would be: “We measured a repeatable systems speedup, but the stronger behavior came from demonstrations, and small timing shifts exposed weak recovery. That made us separate throughput from adaptation quality and ask whether demonstration-conditioned representations or recovery-focused data would be more useful than simply increasing steps.”

Useful questions grounded in these publications include:

• When should we prefer retrieval/alignment, fine-tuning or in-context imitation for a new task?
• Which demonstration variations teach recovery rather than just repeat the nominal behavior?
• How should we measure total human time to a usable policy?
• In a shared policy, how do we detect one task harming another?
• When is generating better task/context episodes more valuable than accelerating gradient updates?

These are proposed discussion questions, not claims about Machenta's internal technology or Ed's personal preferences.

21. Small code examples and a practical next experiment

Example A: context can change an output without an optimizer

This executable example uses nearest-example attention over simple vectors. It isolates conditioning; it is not a learned robot policy, a VLA, or an implementation of Instant Policy. Its scalar outputs are illustrative values, not commands for hardware.

import torch


def action_from_context(query, demo_observations, demo_actions):
    distances = torch.cdist(query, demo_observations).square()
    attention = torch.softmax(-distances / 0.1, dim=-1)
    return attention @ demo_actions


observation = torch.tensor([[1.0, 0.0]])
examples = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
context_a = torch.tensor([[0.2], [-0.2]])
context_b = torch.tensor([[0.8], [-0.2]])

with torch.inference_mode():
    print(action_from_context(observation, examples, context_a))  # ~0.2
    print(action_from_context(observation, examples, context_b))  # ~0.8

The observation and computation are unchanged; the demonstration input changes. A learned context-conditioned policy would learn a much richer interpretation of observations, task relationships and action geometry. This toy does not demonstrate that capability.

The runnable version, with output checks, is tools/learning_examples/robotics_context.py. It was executed locally. No external robotics checkpoint was downloaded to run it.

Example B: training an imitation model with DDP

This is a training-loop sketch, not a standalone launch script or a VLA-specific API. Assume a CUDA process group, one local device per rank, a policy on that device, equal valid sample counts per rank, batch tensors on that device, and a loader that partitions examples across workers. The policy here directly predicts normalized continuous actions.

import torch
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.nn.functional as F

# policy and train_loader are application-specific; setup precedes this snippet.
ddp_policy = DDP(policy, device_ids=[local_rank])
optimizer = torch.optim.AdamW(ddp_policy.parameters(), lr=1e-5)

for batch in train_loader:
    optimizer.zero_grad(set_to_none=True)
    predicted_actions = ddp_policy(batch["observation"], batch["context"])
    loss = F.mse_loss(predicted_actions, batch["target_actions"])
    loss.backward()  # DDP synchronizes gradients during backward.
    optimizer.step()

For discrete actions, a classification loss may be appropriate; diffusion or flow policies use their own noise/velocity prediction objectives. Mixed embodiments also need action-space handling, masks and normalization. This sketch shows where gradient sharing fits, not a universal robotics loss function. PyTorch documents the synchronization behavior in DDP.

At adaptation time an ICL-capable model instead follows the shape action = policy(observation, new_context) inside inference mode, with no backward() or optimizer step. eval() alone does not disable gradient tracking; it selects evaluation behavior for applicable layers.

Example C: an action output still needs the correct interpretation

OpenVLA's original interface includes predict_action and a dataset-specific unnorm_key. The following illustrates the inference call after loading the compatible processor/model and constructing its image-plus-instruction inputs:

# Interface pattern from the original OpenVLA README, not executed here.
# "bridge_orig" is meaningful for the matching dataset/robot convention.
with torch.inference_mode():
    action = vla.predict_action(
        **inputs, unnorm_key="bridge_orig", do_sample=False
    )

The normalization key is not cosmetic. Output coordinates, units, action horizon and robot conventions must match the intended environment. This is why a robot checkpoint does not plug directly into PS-Env. Use the release's official setup and inference example; this fragment intentionally does not claim a complete installation or execution.

A next experiment that connects the ideas

A useful proposed comparison would keep evaluation tasks fixed and compare demonstration-only BC, BC plus recovery examples, and a genuinely context-conditioned policy. Train the latter on separated context/query trajectories across many task variations; evaluate it on withheld variations without adapting weights.

Measure completion, failure recovery, number of demonstrations, adaptation time, inference latency and total operator time. Only afterward compare one-/two-/four-GPU pretraining throughput at matching global work. Otherwise we risk mixing a change in learning method with a change in execution speed again.

For PS-Env, a first step is building a richer collection of task variations and recovery examples. One successful fixed-start trace and two partial routes are insufficient evidence that a general in-context policy will learn robust adaptation. The diagrams and code explain a direction we could test; they do not upgrade the current DQN into a robotics foundation model.