How jrobine/sgf works
SGF trains an Atari agent using a learned world model: the trainer collects environment transitions into replay, updates the model, then trains the policy on imagined trajectories. It also supports agent/world-model evaluation and W&B metrics/model artifacts. The sampled excerpts establish the main orchestration and training calls; unsampled portions of the trainer and model code limit detail about some scheduling and evaluation internals.
Main entry points: main.py Training entry · trainer.py Training coordinator · agent.py Agent trainer · wm.py World model
Components 16
Training runtime
- Training entry
src/main.py - Run configuration
configs/default.yaml - Training coordinator
src/trainer.py
Environment and data
- Atari setup
src/envs.py - Atari environment
- Replay buffer
src/replay.py
Model and policy learning
- Action-stack agent
src/agent.py - World model
src/wm.py - World-model trainer
src/wm.py - Agent trainer
src/agent.py - Actor-critic policy
src/ac.py - Policy trainer
src/ac.py
Shared support
- Network components
src/nets.py - Training utilities
src/utils.py
Other components
- Researcher
- Weights & Biases
Connections 29
- Researcher → Training entry: starts run
- Training entry → Run configuration: loads settings
- Training entry → Weights & Biases: initializes, logs
- Training entry → Atari setup: creates environment
- Atari setup → Atari environment: wraps Atari
- Training entry → Actor-critic policy: constructs
- Training entry → Action-stack agent: constructs
- Training entry → World model: constructs
- Training entry → Training coordinator: starts training
- Training coordinator → Atari environment: steps, resets
- Training coordinator → Action-stack agent: selects actions
- Training coordinator → World model: encodes observations
- Training coordinator → Replay buffer: records transitions
- Training coordinator → World-model trainer: trains model
- Training coordinator → Agent trainer: trains agent
- World-model trainer → Replay buffer: samples transitions
- World-model trainer → World model: optimizes, evaluates
- Agent trainer → Replay buffer: samples starting states
- Agent trainer → World model: encodes, imagines
- Agent trainer → Policy trainer: trains policy
- Agent trainer → Atari setup: creates evaluation env
- Action-stack agent → Actor-critic policy: requests actions
- Policy trainer → Actor-critic policy: updates actor, critic
- World model → Network components: builds networks
- Actor-critic policy → Network components: builds networks
- Policy trainer → Network components: uses optimizers
- Training coordinator → Training utilities: aggregates metrics
- Agent trainer → Training utilities: collects episodes
- Training entry → Weights & Biases: saves artifacts