# galilai-group/stable-worldmodel architecture

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

- Interactive diagram: https://gitdiagram.com/galilai-group/stable-worldmodel
- Repository: https://github.com/galilai-group/stable-worldmodel

## Overview

stable-worldmodel is a research platform for standardized world-model data collection, training, and model-predictive-control evaluation. The documented workflow configures a vectorized World over supported environments, records episodes through the dataset-format interface, loads clips for model training, then evaluates a policy that plans with a model and solver. Groups distinguish environments, orchestration, data, learning, and planning. The README documents additional dataset backends and reference baselines; those are represented without asserting unsampled implementation wiring.

Main entry points: [`stable_worldmodel/world/world.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/world.py) (World), [`stable_worldmodel/data/formats`](https://github.com/galilai-group/stable-worldmodel/tree/main/stable_worldmodel/data/formats) (Dataset formats), [`stable_worldmodel/world/env_pool.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/env_pool.py) (Environment pool), [`stable_worldmodel/data/format.py`](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/data/format.py) (Format registry)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_environment["Environments"]
  node_env_suites["Environment suites"]
end

subgraph group_orchestration["World orchestration"]
  node_world["World<br/>[world.py]"]
  node_env_pool["Environment pool<br/>[env_pool.py]"]
  node_wrappers["Environment wrappers<br/>[default.py]"]
end

subgraph group_data["Dataset pipeline"]
  node_format_registry["Format registry<br/>[format.py]"]
  node_dataset_readers["Dataset formats"]
  node_lance[("Lance backend<br/>[lance.py]")]
  node_dataset["Dataset abstractions<br/>[dataset.py]"]
end

subgraph group_learning["World models"]
  node_training_recipes["Training recipes"]
  node_world_models["World-model baselines"]
end

subgraph group_planning["Planning and evaluation"]
  node_policy["World-model policy<br/>[policy.py]"]
  node_solver["Planning solvers"]
  node_evaluator["Evaluation<br/>[evaluator.py]"]
  node_results["Episode results<br/>[world.py]"]
end

node_researcher(("Researcher"))

node_researcher -->|"configures"| node_world
node_world -->|"steps environments"| node_env_pool
node_env_pool -->|"processes observations"| node_wrappers
node_env_pool -->|"dispatches to"| node_env_suites
node_world -->|"records episodes"| node_format_registry
node_format_registry -->|"selects writer"| node_dataset_readers
node_format_registry -->|"selects reader"| node_dataset_readers
node_dataset_readers -->|"includes backend"| node_lance
node_researcher -->|"loads dataset"| node_dataset
node_dataset_readers -->|"provides samples"| node_dataset
node_dataset -->|"supplies clips"| node_training_recipes
node_training_recipes -->|"trains baselines"| node_world_models
node_researcher -->|"configures policy"| node_policy
node_policy -->|"requests plans"| node_solver
node_solver -->|"predicts with model"| node_world_models
node_world -->|"runs policy"| node_policy
node_world -->|"evaluates episodes"| node_evaluator
node_evaluator -->|"produces metrics"| node_results

click node_world "https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/world.py"
click node_env_pool "https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/env_pool.py"
click node_wrappers "https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/wrapper/default.py"
click node_env_suites "https://github.com/galilai-group/stable-worldmodel/tree/main/stable_worldmodel/envs"
click node_format_registry "https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/data/format.py"
click node_dataset_readers "https://github.com/galilai-group/stable-worldmodel/tree/main/stable_worldmodel/data/formats"
click node_lance "https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/data/formats/lance.py"
click node_dataset "https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/data/dataset.py"
click node_training_recipes "https://github.com/galilai-group/stable-worldmodel/tree/main/scripts/train"
click node_world_models "https://github.com/galilai-group/stable-worldmodel/tree/main/stable_worldmodel/wm"
click node_policy "https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/policy.py"
click node_solver "https://github.com/galilai-group/stable-worldmodel/tree/main/stable_worldmodel/planning/solver"
click node_evaluator "https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/planning/evaluator.py"
click node_results "https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/world.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_env_suites toneBlue
class node_world,node_env_pool,node_wrappers toneAmber
class node_format_registry,node_dataset_readers,node_lance,node_dataset toneMint
class node_training_recipes,node_world_models toneRose
class node_policy,node_solver,node_evaluator,node_results,node_researcher toneIndigo
```

## Components

### Environments

- [Environment suites](https://github.com/galilai-group/stable-worldmodel/tree/main/stable_worldmodel/envs): `stable_worldmodel/envs`

### World orchestration

- [World](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/world.py): `stable_worldmodel/world/world.py`
- [Environment pool](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/env_pool.py): `stable_worldmodel/world/env_pool.py`
- [Environment wrappers](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/wrapper/default.py): `stable_worldmodel/wrapper/default.py`

### Dataset pipeline

- [Format registry](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/data/format.py): `stable_worldmodel/data/format.py`
- [Dataset formats](https://github.com/galilai-group/stable-worldmodel/tree/main/stable_worldmodel/data/formats): `stable_worldmodel/data/formats`
- [Lance backend](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/data/formats/lance.py): `stable_worldmodel/data/formats/lance.py`
- [Dataset abstractions](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/data/dataset.py): `stable_worldmodel/data/dataset.py`

### World models

- [Training recipes](https://github.com/galilai-group/stable-worldmodel/tree/main/scripts/train): `scripts/train`
- [World-model baselines](https://github.com/galilai-group/stable-worldmodel/tree/main/stable_worldmodel/wm): `stable_worldmodel/wm`

### Planning and evaluation

- [World-model policy](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/policy.py): `stable_worldmodel/policy.py`
- [Planning solvers](https://github.com/galilai-group/stable-worldmodel/tree/main/stable_worldmodel/planning/solver): `stable_worldmodel/planning/solver`
- [Evaluation](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/planning/evaluator.py): `stable_worldmodel/planning/evaluator.py`
- [Episode results](https://github.com/galilai-group/stable-worldmodel/blob/main/stable_worldmodel/world/world.py): `stable_worldmodel/world/world.py`

### Other components

- **Researcher**

## Connections

- Researcher → World: configures
- World → Environment pool: steps environments
- Environment pool → Environment wrappers: processes observations
- Environment pool → Environment suites: dispatches to
- World → Format registry: records episodes
- Format registry → Dataset formats: selects writer
- Format registry → Dataset formats: selects reader
- Dataset formats → Lance backend: includes backend
- Researcher → Dataset abstractions: loads dataset
- Dataset formats → Dataset abstractions: provides samples
- Dataset abstractions → Training recipes: supplies clips
- Training recipes → World-model baselines: trains baselines
- Researcher → World-model policy: configures policy
- World-model policy → Planning solvers: requests plans
- Planning solvers → World-model baselines: predicts with model
- World → World-model policy: runs policy
- World → Evaluation: evaluates episodes
- Evaluation → Episode results: produces metrics