# jrobine/sgf architecture

> Architecture diagram of the GitHub repository jrobine/sgf, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/jrobine/sgf
- Repository: https://github.com/jrobine/sgf

## Overview

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: [`src/main.py`](https://github.com/jrobine/sgf/blob/main/src/main.py) (Training entry), [`src/trainer.py`](https://github.com/jrobine/sgf/blob/main/src/trainer.py) (Training coordinator), [`src/agent.py`](https://github.com/jrobine/sgf/blob/main/src/agent.py) (Agent trainer), [`src/wm.py`](https://github.com/jrobine/sgf/blob/main/src/wm.py) (World model)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_runtime["Training runtime"]
  node_entry["Training entry<br/>[main.py]"]
  node_config["Run configuration<br/>[default.yaml]"]
  node_coordinator["Training coordinator<br/>[trainer.py]"]
end

subgraph group_environment["Environment and data"]
  node_atari_factory["Atari setup<br/>[envs.py]"]
  node_atari["Atari environment"]
  node_replay[("Replay buffer<br/>[replay.py]")]
end

subgraph group_learning["Model and policy learning"]
  node_agent["Action-stack agent<br/>[agent.py]"]
  node_world_model["World model<br/>[wm.py]"]
  node_wm_trainer["World-model trainer<br/>[wm.py]"]
  node_agent_trainer["Agent trainer<br/>[agent.py]"]
  node_policy["Actor-critic policy<br/>[ac.py]"]
  node_policy_trainer["Policy trainer<br/>[ac.py]"]
end

subgraph group_support["Shared support"]
  node_nets["Network components<br/>[nets.py]"]
  node_utils["Training utilities<br/>[utils.py]"]
end

node_researcher(("Researcher"))
node_wandb["Weights &amp; Biases"]

node_researcher -->|"starts run"| node_entry
node_entry -->|"loads settings"| node_config
node_entry -->|"initializes, logs"| node_wandb
node_entry -->|"creates environment"| node_atari_factory
node_atari_factory -->|"wraps Atari"| node_atari
node_entry -->|"constructs"| node_policy
node_entry -->|"constructs"| node_agent
node_entry -->|"constructs"| node_world_model
node_entry -->|"starts training"| node_coordinator
node_coordinator -->|"steps, resets"| node_atari
node_coordinator -->|"selects actions"| node_agent
node_coordinator -->|"encodes observations"| node_world_model
node_coordinator -->|"records transitions"| node_replay
node_coordinator -->|"trains model"| node_wm_trainer
node_coordinator -->|"trains agent"| node_agent_trainer
node_wm_trainer -->|"samples transitions"| node_replay
node_wm_trainer -->|"optimizes, evaluates"| node_world_model
node_agent_trainer -->|"samples starting states"| node_replay
node_agent_trainer -->|"encodes, imagines"| node_world_model
node_agent_trainer -->|"trains policy"| node_policy_trainer
node_agent_trainer -->|"creates evaluation env"| node_atari_factory
node_agent -->|"requests actions"| node_policy
node_policy_trainer -->|"updates actor, critic"| node_policy
node_world_model -->|"builds networks"| node_nets
node_policy -->|"builds networks"| node_nets
node_policy_trainer -->|"uses optimizers"| node_nets
node_coordinator -->|"aggregates metrics"| node_utils
node_agent_trainer -->|"collects episodes"| node_utils
node_entry -->|"saves artifacts"| node_wandb

click node_entry "https://github.com/jrobine/sgf/blob/main/src/main.py"
click node_config "https://github.com/jrobine/sgf/blob/main/configs/default.yaml"
click node_coordinator "https://github.com/jrobine/sgf/blob/main/src/trainer.py"
click node_atari_factory "https://github.com/jrobine/sgf/blob/main/src/envs.py"
click node_replay "https://github.com/jrobine/sgf/blob/main/src/replay.py"
click node_agent "https://github.com/jrobine/sgf/blob/main/src/agent.py"
click node_world_model "https://github.com/jrobine/sgf/blob/main/src/wm.py"
click node_wm_trainer "https://github.com/jrobine/sgf/blob/main/src/wm.py"
click node_agent_trainer "https://github.com/jrobine/sgf/blob/main/src/agent.py"
click node_policy "https://github.com/jrobine/sgf/blob/main/src/ac.py"
click node_policy_trainer "https://github.com/jrobine/sgf/blob/main/src/ac.py"
click node_nets "https://github.com/jrobine/sgf/blob/main/src/nets.py"
click node_utils "https://github.com/jrobine/sgf/blob/main/src/utils.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_entry,node_config,node_coordinator toneBlue
class node_atari_factory,node_atari,node_replay toneAmber
class node_agent,node_world_model,node_wm_trainer,node_agent_trainer,node_policy,node_policy_trainer toneMint
class node_nets,node_utils toneRose
class node_researcher,node_wandb toneIndigo
```

## Components

### Training runtime

- [Training entry](https://github.com/jrobine/sgf/blob/main/src/main.py): `src/main.py`
- [Run configuration](https://github.com/jrobine/sgf/blob/main/configs/default.yaml): `configs/default.yaml`
- [Training coordinator](https://github.com/jrobine/sgf/blob/main/src/trainer.py): `src/trainer.py`

### Environment and data

- [Atari setup](https://github.com/jrobine/sgf/blob/main/src/envs.py): `src/envs.py`
- **Atari environment**
- [Replay buffer](https://github.com/jrobine/sgf/blob/main/src/replay.py): `src/replay.py`

### Model and policy learning

- [Action-stack agent](https://github.com/jrobine/sgf/blob/main/src/agent.py): `src/agent.py`
- [World model](https://github.com/jrobine/sgf/blob/main/src/wm.py): `src/wm.py`
- [World-model trainer](https://github.com/jrobine/sgf/blob/main/src/wm.py): `src/wm.py`
- [Agent trainer](https://github.com/jrobine/sgf/blob/main/src/agent.py): `src/agent.py`
- [Actor-critic policy](https://github.com/jrobine/sgf/blob/main/src/ac.py): `src/ac.py`
- [Policy trainer](https://github.com/jrobine/sgf/blob/main/src/ac.py): `src/ac.py`

### Shared support

- [Network components](https://github.com/jrobine/sgf/blob/main/src/nets.py): `src/nets.py`
- [Training utilities](https://github.com/jrobine/sgf/blob/main/src/utils.py): `src/utils.py`

### Other components

- **Researcher**
- **Weights & Biases**

## Connections

- 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