Train pixel games with vectorized updates

Saved run results

Dependency cache ready: /Users/sanbu/Documents/Codex/2026-10-20/revie/work/chapter-lab-runs/jax-en/workspaces/jax/.cache/online-experiments/hands-on-modern-rl-experiment07-jax-games-d3ef1adbd087.ready
Repository: /Users/sanbu/Documents/Codex/2026-10-20/revie/work/chapter-lab-runs/jax-en/workspaces/jax
Experiment runtime: /Users/sanbu/Documents/Codex/2026-10-20/revie/work/chapter-lab-runs/jax-en/workspaces/jax/modelscope-space/hands-on-modern-rl-experiment07-jax-games

Available tasks:
  asterix              Asterix MinAtar · Asterix-MinAtar
  breakout             Breakout MinAtar · Breakout-MinAtar
  freeway              Freeway MinAtar · Freeway-MinAtar
  space-invaders       Space Invaders MinAtar · SpaceInvaders-MinAtar

Selected: Breakout MinAtar
Algorithm: JAX REINFORCE
Epoch schedule: 2 epochs × 50 steps = 100 total steps
Every epoch is evaluated and saved as a separately selectable model.

Device: CPU
Starting the same training generator used by the live Studio...

/Users/sanbu/Documents/Codex/2026-10-20/revie/work/studio-venv/lib/python3.12/site-packages/jax/_src/ops/scatter.py:105: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=int32 to dtype=bool with jax_numpy_dtype_promotion=standard. In future JAX releases this will result in an error.
  warnings.warn(
/Users/sanbu/Documents/Codex/2026-10-20/revie/work/studio-venv/lib/python3.12/site-packages/jax/_src/ops/scatter.py:105: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=float32 to dtype=bool with jax_numpy_dtype_promotion=standard. In future JAX releases this will result in an error.
  warnings.warn(

JIT compiling policy update · observation=(10, 10, 4) actions=3

JAX epoch=1/2 episode=50 loss=-0.010270 eval_return=1.600 replay_return=2.000
SAVE model=policy.npz replay=learned-policy.gif frames=18

JAX epoch=2/2 episode=100 loss=-0.010691 eval_return=0.000 replay_return=0.000
SAVE model=policy.npz replay=learned-policy.gif frames=6

Saved 2 independently selectable JAX policies and semantic replays

Final phase: complete
Final score: 0.0
Saved 2 selectable epoch models in this run.
Model index: /Users/sanbu/Documents/Codex/2026-10-20/revie/work/chapter-lab-runs/jax-en/workspaces/jax/modelscope-space/hands-on-modern-rl-experiment07-jax-games/artifacts/notebook-models.json

Run the inference cell below to choose an epoch model and visualize its exact learned policy.

VBox(children=(Dropdown(description='Epoch model:', layout=Layout(width='95%'), options=(('Epoch 2/2 · 100 ste…

The newest epoch is selected by default. Choose another saved model, then click Visualize selected policy.

{
  "score": 0.0,
  "steps": 100,
  "epochs": 2,
  "checkpoints": 2,
  "preview": "learned-policy.gif"
}

{
  "score": 0.0,
  "steps": 100,
  "epochs": 2,
  "checkpoints": 2,
  "preview": "learned-policy.gif"
}