# eloialonso/iris architecture

> Architecture diagram of the GitHub repository eloialonso/iris, 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/eloialonso/iris
- Repository: https://github.com/eloialonso/iris

## Overview

IRIS trains a sample-efficient Atari agent by collecting real environment episodes, learning a discrete visual tokenizer and autoregressive world model, then optimizing an actor-critic over imagined trajectories. The central runtime is Trainer: it coordinates collection, dataset-backed component training, evaluation and logging. The README also documents checkpoint-based evaluation/resumption, reconstruction generation, and live, imagined, or recorded-episode playback. Several of those entry points are represented at their exact tree paths; sampled excerpts do not establish every internal call, so uncertain wiring is omitted.

Main entry points: [`src/models/world_model.py`](https://github.com/eloialonso/iris/blob/main/src/models/world_model.py) (Transformer world model), [`src/models/tokenizer/tokenizer.py`](https://github.com/eloialonso/iris/blob/main/src/models/tokenizer/tokenizer.py) (Discrete tokenizer), [`src/trainer.py`](https://github.com/eloialonso/iris/blob/main/src/trainer.py) (Trainer), [`src/collector.py`](https://github.com/eloialonso/iris/blob/main/src/collector.py) (Episode collector)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_runtime["Training runtime"]
  node_config["Run configuration<br/>[trainer.yaml]"]
  node_main["Training entry<br/>[main.py]"]
  node_trainer["Trainer<br/>[trainer.py]"]
end

subgraph group_collection["Experience collection"]
  node_collector["Episode collector<br/>[collector.py]"]
  node_envs["Environment runners"]
  node_wrappers["Atari preprocessing<br/>[wrappers.py]"]
  node_episode["Episode records<br/>[episode.py]"]
  node_dataset[("Episode dataset<br/>[dataset.py]")]
end

subgraph group_learning["World-model learning"]
  node_agent["Agent<br/>[agent.py]"]
  node_tokenizer["Discrete tokenizer<br/>[tokenizer.py]"]
  node_tokenizer_nets["Encoder decoder<br/>[nets.py]"]
  node_world_model["Transformer world model<br/>[world_model.py]"]
  node_transformer["Autoregressive transformer<br/>[transformer.py]"]
  node_actor_critic["Actor critic<br/>[actor_critic.py]"]
  node_wm_env["Imagined environment<br/>[world_model_env.py]"]
end

subgraph group_interaction["Evaluation and playback"]
  node_eval_script["Run evaluation<br/>[eval.py]"]
  node_reconstruction["Reconstructions"]
  node_play["Playback entry<br/>[play.py]"]
  node_game["Interactive game<br/>[game.py]"]
  node_replay["Episode replay"]
end

node_user(("Researcher"))
node_atari["Atari environment"]
node_wandb[("Weights &amp; Biases")]

node_user -->|"starts training"| node_main
node_config -->|"configures run"| node_main
node_main -->|"creates trainer"| node_trainer
node_trainer -->|"collects episodes"| node_collector
node_trainer -->|"trains agent"| node_agent
node_trainer -->|"logs metrics"| node_wandb
node_collector -->|"steps environments"| node_envs
node_envs -.->|"preprocesses frames"| node_wrappers
node_wrappers -->|"wraps environment"| node_atari
node_collector -->|"stores episodes"| node_dataset
node_collector -->|"records transitions"| node_episode
node_agent -->|"trains tokenizer"| node_tokenizer
node_tokenizer -->|"encodes and decodes"| node_tokenizer_nets
node_agent -->|"trains dynamics"| node_world_model
node_world_model -->|"encodes observations"| node_tokenizer
node_world_model -->|"models sequences"| node_transformer
node_agent -->|"optimizes policy"| node_actor_critic
node_actor_critic -->|"imagines with"| node_tokenizer
node_actor_critic -->|"imagines with"| node_world_model
node_wm_env -->|"encodes and decodes"| node_tokenizer
node_wm_env -->|"predicts imagined steps"| node_world_model
node_user -->|"evaluates run"| node_eval_script
node_user -->|"starts playback"| node_play
node_play -->|"launches interface"| node_game
node_game -.->|"plays imagined world"| node_wm_env
node_game -.->|"replays episodes"| node_replay
node_replay -->|"loads saved episodes"| node_episode
node_user -->|"generates reconstructions"| node_reconstruction

click node_config "https://github.com/eloialonso/iris/blob/main/config/trainer.yaml"
click node_main "https://github.com/eloialonso/iris/blob/main/src/main.py"
click node_trainer "https://github.com/eloialonso/iris/blob/main/src/trainer.py"
click node_agent "https://github.com/eloialonso/iris/blob/main/src/agent.py"
click node_collector "https://github.com/eloialonso/iris/blob/main/src/collector.py"
click node_envs "https://github.com/eloialonso/iris/tree/main/src/envs"
click node_wrappers "https://github.com/eloialonso/iris/blob/main/src/envs/wrappers.py"
click node_episode "https://github.com/eloialonso/iris/blob/main/src/episode.py"
click node_dataset "https://github.com/eloialonso/iris/blob/main/src/dataset.py"
click node_tokenizer "https://github.com/eloialonso/iris/blob/main/src/models/tokenizer/tokenizer.py"
click node_tokenizer_nets "https://github.com/eloialonso/iris/blob/main/src/models/tokenizer/nets.py"
click node_world_model "https://github.com/eloialonso/iris/blob/main/src/models/world_model.py"
click node_transformer "https://github.com/eloialonso/iris/blob/main/src/models/transformer.py"
click node_actor_critic "https://github.com/eloialonso/iris/blob/main/src/models/actor_critic.py"
click node_wm_env "https://github.com/eloialonso/iris/blob/main/src/envs/world_model_env.py"
click node_eval_script "https://github.com/eloialonso/iris/blob/main/scripts/eval.py"
click node_reconstruction "https://github.com/eloialonso/iris/blob/main/src/make_reconstructions.py"
click node_play "https://github.com/eloialonso/iris/blob/main/src/play.py"
click node_game "https://github.com/eloialonso/iris/blob/main/src/game/game.py"
click node_replay "https://github.com/eloialonso/iris/blob/main/src/game/episode_replay_env.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_config,node_main,node_trainer toneBlue
class node_collector,node_envs,node_wrappers,node_episode,node_dataset,node_wandb toneAmber
class node_agent,node_tokenizer,node_tokenizer_nets,node_world_model,node_transformer,node_actor_critic,node_wm_env toneMint
class node_eval_script,node_reconstruction,node_play,node_game,node_replay toneRose
class node_user,node_atari toneIndigo
```

## Components

### Training runtime

- [Run configuration](https://github.com/eloialonso/iris/blob/main/config/trainer.yaml): `config/trainer.yaml`
- [Training entry](https://github.com/eloialonso/iris/blob/main/src/main.py): `src/main.py`
- [Trainer](https://github.com/eloialonso/iris/blob/main/src/trainer.py): `src/trainer.py`

### Experience collection

- [Episode collector](https://github.com/eloialonso/iris/blob/main/src/collector.py): `src/collector.py`
- [Environment runners](https://github.com/eloialonso/iris/tree/main/src/envs): `src/envs`
- [Atari preprocessing](https://github.com/eloialonso/iris/blob/main/src/envs/wrappers.py): `src/envs/wrappers.py`
- [Episode records](https://github.com/eloialonso/iris/blob/main/src/episode.py): `src/episode.py`
- [Episode dataset](https://github.com/eloialonso/iris/blob/main/src/dataset.py): `src/dataset.py`

### World-model learning

- [Agent](https://github.com/eloialonso/iris/blob/main/src/agent.py): `src/agent.py`
- [Discrete tokenizer](https://github.com/eloialonso/iris/blob/main/src/models/tokenizer/tokenizer.py): `src/models/tokenizer/tokenizer.py`
- [Encoder decoder](https://github.com/eloialonso/iris/blob/main/src/models/tokenizer/nets.py): `src/models/tokenizer/nets.py`
- [Transformer world model](https://github.com/eloialonso/iris/blob/main/src/models/world_model.py): `src/models/world_model.py`
- [Autoregressive transformer](https://github.com/eloialonso/iris/blob/main/src/models/transformer.py): `src/models/transformer.py`
- [Actor critic](https://github.com/eloialonso/iris/blob/main/src/models/actor_critic.py): `src/models/actor_critic.py`
- [Imagined environment](https://github.com/eloialonso/iris/blob/main/src/envs/world_model_env.py): `src/envs/world_model_env.py`

### Evaluation and playback

- [Run evaluation](https://github.com/eloialonso/iris/blob/main/scripts/eval.py): `scripts/eval.py`
- [Reconstructions](https://github.com/eloialonso/iris/blob/main/src/make_reconstructions.py): `src/make_reconstructions.py`
- [Playback entry](https://github.com/eloialonso/iris/blob/main/src/play.py): `src/play.py`
- [Interactive game](https://github.com/eloialonso/iris/blob/main/src/game/game.py): `src/game/game.py`
- [Episode replay](https://github.com/eloialonso/iris/blob/main/src/game/episode_replay_env.py): `src/game/episode_replay_env.py`

### Other components

- **Researcher**
- **Atari environment**
- **Weights & Biases**

## Connections

- Researcher → Training entry: starts training
- Run configuration → Training entry: configures run
- Training entry → Trainer: creates trainer
- Trainer → Episode collector: collects episodes
- Trainer → Agent: trains agent
- Trainer → Weights & Biases: logs metrics
- Episode collector → Environment runners: steps environments
- Environment runners → Atari preprocessing: preprocesses frames
- Atari preprocessing → Atari environment: wraps environment
- Episode collector → Episode dataset: stores episodes
- Episode collector → Episode records: records transitions
- Agent → Discrete tokenizer: trains tokenizer
- Discrete tokenizer → Encoder decoder: encodes and decodes
- Agent → Transformer world model: trains dynamics
- Transformer world model → Discrete tokenizer: encodes observations
- Transformer world model → Autoregressive transformer: models sequences
- Agent → Actor critic: optimizes policy
- Actor critic → Discrete tokenizer: imagines with
- Actor critic → Transformer world model: imagines with
- Imagined environment → Discrete tokenizer: encodes and decodes
- Imagined environment → Transformer world model: predicts imagined steps
- Researcher → Run evaluation: evaluates run
- Researcher → Playback entry: starts playback
- Playback entry → Interactive game: launches interface
- Interactive game → Imagined environment: plays imagined world
- Interactive game → Episode replay: replays episodes
- Episode replay → Episode records: loads saved episodes
- Researcher → Reconstructions: generates reconstructions