How eloialonso/iris works
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: world_model.py Transformer world model · tokenizer.py Discrete tokenizer · trainer.py Trainer · collector.py Episode collector
Components 23
Training runtime
- Run configuration
config/trainer.yaml - Training entry
src/main.py - Trainer
src/trainer.py
Experience collection
- Episode collector
src/collector.py - Environment runners
src/envs - Atari preprocessing
src/envs/wrappers.py - Episode records
src/episode.py - Episode dataset
src/dataset.py
World-model learning
- Agent
src/agent.py - Discrete tokenizer
src/models/tokenizer/tokenizer.py - Encoder decoder
src/models/tokenizer/nets.py - Transformer world model
src/models/world_model.py - Autoregressive transformer
src/models/transformer.py - Actor critic
src/models/actor_critic.py - Imagined environment
src/envs/world_model_env.py
Evaluation and playback
- Run evaluation
scripts/eval.py - Reconstructions
src/make_reconstructions.py - Playback entry
src/play.py - Interactive game
src/game/game.py - Episode replay
src/game/episode_replay_env.py
Other components
- Researcher
- Atari environment
- Weights & Biases
Connections 28
- 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