# sreeyukth/symbolic-world-model architecture

> Architecture diagram of the GitHub repository sreeyukth/symbolic-world-model, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-30). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/sreeyukth/symbolic-world-model
- Repository: https://github.com/sreeyukth/symbolic-world-model

## Overview

1. **What this repo is**
   - This is a Python research codebase for a **symbolic world model**: it treats a mathematical expression as the current world state, and a math operation as the transition action.
   - The core architectural question is whether a recurrent latent state can retain enough information to model procedural reasoning over symbolic math.

2. **Top-level pipeline**
   - The system flows in four stages: **generate trajectories → tokenize expressions → train recurrent world model → probe latent structure**.
   - The repo is organized so the symbolic domain is created outside the model, while the neural components only see tokenized expressions and operation labels.

3. **Data generation boundary**
   - `data/generator.py` is the source of synthetic training data.
   - It produces structured trajectories of symbolic expression transitions, likely using SymPy-based transformations from one expression to the next.
   - The generated dataset is stored as `data/trajectories.json` and is treated as an artifact rather than hand-authored data.

4. **Tokenization layer**
   - `data/tokenizer.py` is the interface between symbolic strings and neural inputs.
   - Expressions are represented in **prefix / Polish notation**, then mapped into a fixed token vocabulary.
   - This layer is important because it defines the model’s input boundary: symbolic syntax is normalized into discrete IDs before entering the encoder.

5. **Expression encoder**
   - `model/encoder.py` implements `ExprEncoder`.
   - Its job is to convert a tokenized expression into a compact latent vector via **embedding lookup + masked mean pooling**.
   - This is the first learned representation of the symbolic state; it abstracts away sequence length while preserving token content.

6. **Transition core / world-state memory**
   - `model/transition.py` contains the recurrent transition model, implemented as a **2-layer GRU**.
   - The transition input is the concatenation of the encoded expression latent and an operation embedding.
   - The GRU hidden state is the central “world state” (`h_t`): it carries procedural memory across steps and is the main recurrent subsystem in the architecture.

7. **Prediction heads**
   - `model/heads.py` defines two outputs from the recurrent state:
     - a **latent predictor head** that predicts the next expression’s latent embedding
     - an **operation predictor head** that predicts the next operation class
   - These are separate supervision targets, so the model learns both next-state structure and action classification.

8. **Full model composition**
   - `model/world_model.py` wraps the encoder, transition core, and heads into one end-to-end forward pass.
   - This is the main integration point where the symbolic input, action, recurrent state, and prediction outputs meet.
   - `model/test_model.py` provides a smoke test for tensor shapes and forward-pass correctness.

9. **Training subsystem**
   - `train.py` orchestrates dataset loading, batching, forward/backward passes, loss computation, and checkpointing.
   - The README indicates it is **MPS-aware**, meaning it is designed to run efficiently on Apple Silicon using PyTorch’s MPS backend.
   - Training outputs include `checkpoints/model.pt` and `training_curves.png`, making training both reproducible and inspectable.

10. **Latent-space analysis / probing**
   - `probe.py` is a post-training analysis tool that extracts hidden states and evaluates how much reasoning structure is encoded in them.
   - It produces PCA plots and linear probe metrics, stored under `probe_results/`.
   - This analysis subsystem is separate from training, which makes the repo not just a model implementation but also an experimental framework for representation studies.

11. **Configuration and reproducibility**
   - `config.py` centralizes hyperparameters, which keeps model size, training settings, and dataset-related choices in one place.
   - `requirements.txt` pins the Python dependencies needed for generation, training, and probing.
   - This separation suggests a clean distinction between experiment configuration and core model code.

12. **Key architectural boundary**
   - The most important boundary is between **symbolic mathematics** and **learned latent dynamics**.
   - On the left side are discrete expressions, operations, and generated trajectories; on the right side are embeddings, a GRU hidden state, and prediction heads.
   - The whole repo is built to study whether the recurrent hidden state organizes information about reasoning chains, operations, and step progression.

13. **Generated artifacts vs source code**
   - `data/trajectories.json`, `checkpoints/`, and `probe_results/` are experiment outputs, not source-of-truth logic.
   - They represent the three major runtime products of the system: the dataset, the trained weights, and the latent-space diagnostics.

Main entry points: [`model/world_model.py`](https://github.com/sreeyukth/symbolic-world-model/blob/main/model/world_model.py) (World Model), [`config.py`](https://github.com/sreeyukth/symbolic-world-model/blob/main/config.py) (Config), [`model/encoder.py`](https://github.com/sreeyukth/symbolic-world-model/blob/main/model/encoder.py) (Expr Encoder), [`model/transition.py`](https://github.com/sreeyukth/symbolic-world-model/blob/main/model/transition.py) (Transition)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_group_data["Data pipeline"]
  node_generator["Generator<br/>synthetic data<br/>[generator.py]"]
  node_tokenizer["Tokenizer<br/>symbol encoder<br/>[tokenizer.py]"]
end

subgraph group_group_model["Model core"]
  node_expr_encoder["Expr Encoder<br/>latent encoder<br/>[encoder.py]"]
  node_transition["Transition<br/>GRU memory<br/>[transition.py]"]
  node_heads["Heads<br/>[heads.py]"]
  node_world_model["World Model<br/>model wrapper<br/>[world_model.py]"]
end

subgraph group_group_runtime["Runtime"]
  node_config["Config<br/>hyperparameters<br/>[config.py]"]
  node_train["Train<br/>[train.py]"]
  node_probe["Probe<br/>analysis tool<br/>[probe.py]"]
  node_test_model["Model Test<br/>smoke test<br/>[test_model.py]"]
end

subgraph group_group_artifacts["Artifacts"]
  node_trajectories[("Trajectories<br/>generated dataset")]
  node_probe_results["Probe Results<br/>analysis output"]
  node_linear_report["Probe Report<br/>metric report"]
end

node_generator -->|"writes"| node_trajectories
node_trajectories -->|"feeds"| node_tokenizer
node_tokenizer -->|"encodes"| node_expr_encoder
node_expr_encoder -->|"state input"| node_transition
node_config -.->|"settings"| node_generator
node_config -.->|"settings"| node_train
node_config -.->|"settings"| node_probe
node_transition -->|"hidden state"| node_heads
node_heads -->|"outputs"| node_world_model
node_expr_encoder -->|"component"| node_world_model
node_transition -->|"component"| node_world_model
node_heads -->|"losses"| node_train
node_world_model -->|"trained by"| node_train
node_world_model -->|"hidden states"| node_probe
node_probe -->|"writes"| node_probe_results
node_probe_results -->|"contains"| node_linear_report
node_test_model -.->|"validates"| node_world_model

click node_config "https://github.com/sreeyukth/symbolic-world-model/blob/main/config.py"
click node_generator "https://github.com/sreeyukth/symbolic-world-model/blob/main/data/generator.py"
click node_tokenizer "https://github.com/sreeyukth/symbolic-world-model/blob/main/data/tokenizer.py"
click node_expr_encoder "https://github.com/sreeyukth/symbolic-world-model/blob/main/model/encoder.py"
click node_transition "https://github.com/sreeyukth/symbolic-world-model/blob/main/model/transition.py"
click node_heads "https://github.com/sreeyukth/symbolic-world-model/blob/main/model/heads.py"
click node_world_model "https://github.com/sreeyukth/symbolic-world-model/blob/main/model/world_model.py"
click node_train "https://github.com/sreeyukth/symbolic-world-model/blob/main/train.py"
click node_probe "https://github.com/sreeyukth/symbolic-world-model/blob/main/probe.py"
click node_test_model "https://github.com/sreeyukth/symbolic-world-model/blob/main/model/test_model.py"
click node_probe_results "https://github.com/sreeyukth/symbolic-world-model/tree/main/probe_results"
click node_linear_report "https://github.com/sreeyukth/symbolic-world-model/blob/main/probe_results/linear_probe_report.txt"

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_generator,node_tokenizer toneBlue
class node_expr_encoder,node_transition,node_heads,node_world_model toneAmber
class node_config,node_train,node_probe,node_test_model toneMint
class node_trajectories,node_probe_results,node_linear_report toneRose
```

## Components

### Data pipeline

Synthetic trajectory generation and tokenization boundary for symbolic expressions.

- [Generator](https://github.com/sreeyukth/symbolic-world-model/blob/main/data/generator.py): synthetic data · Builds symbolic trajectories of expression transitions. · `data/generator.py`
- [Tokenizer](https://github.com/sreeyukth/symbolic-world-model/blob/main/data/tokenizer.py): symbol encoder · Converts prefix expressions into fixed token IDs. · `data/tokenizer.py`

### Model core

Encoder, recurrent transition state, and prediction heads.

- [Expr Encoder](https://github.com/sreeyukth/symbolic-world-model/blob/main/model/encoder.py): latent encoder · Embeds token sequences and pools them into a compact expression vector. · `model/encoder.py`
- [Transition](https://github.com/sreeyukth/symbolic-world-model/blob/main/model/transition.py): GRU memory · Recurrent world-state core that updates hidden state from expression and action inputs. · `model/transition.py`
- [Heads](https://github.com/sreeyukth/symbolic-world-model/blob/main/model/heads.py): prediction heads · Predicts next expression latent and next operation class from the recurrent state. · `model/heads.py`
- [World Model](https://github.com/sreeyukth/symbolic-world-model/blob/main/model/world_model.py): model wrapper · End-to-end composition of encoder, transition core, and output heads. · `model/world_model.py`

### Runtime

Training and probing entrypoints plus shared configuration.

- [Config](https://github.com/sreeyukth/symbolic-world-model/blob/main/config.py): hyperparameters · Central experiment settings shared by generation, training, and probing. · `config.py`
- [Train](https://github.com/sreeyukth/symbolic-world-model/blob/main/train.py): training loop · Loads data, optimizes the model, and writes training artifacts. · `train.py`
- [Probe](https://github.com/sreeyukth/symbolic-world-model/blob/main/probe.py): analysis tool · Extracts hidden states and evaluates latent structure with probes and PCA. · `probe.py`
- [Model Test](https://github.com/sreeyukth/symbolic-world-model/blob/main/model/test_model.py): smoke test · Checks tensor shapes and forward-pass correctness. · `model/test_model.py`

### Artifacts

Generated outputs used for inspection rather than source logic.

- **Trajectories**: generated dataset · Synthetic symbolic transition data produced by the generator.
- [Probe Results](https://github.com/sreeyukth/symbolic-world-model/tree/main/probe_results): analysis output · Saved latent-space diagnostics and linear probe reports. · `probe_results`
- [Probe Report](https://github.com/sreeyukth/symbolic-world-model/blob/main/probe_results/linear_probe_report.txt): metric report · Text summary of linear probe evaluation results. · `probe_results/linear_probe_report.txt`

## Connections

- Generator → Trajectories: writes · Generates structured symbolic transition data.
- Trajectories → Tokenizer: feeds · Dataset expressions are normalized into token sequences.
- Tokenizer → Expr Encoder: encodes · Token IDs are embedded and pooled into latent vectors.
- Expr Encoder → Transition: state input · Expression latent becomes part of the recurrent transition input.
- Config → Generator: settings · Controls generation choices and dataset shape.
- Config → Train: settings · Provides training hyperparameters and backend choices.
- Config → Probe: settings · Shares analysis configuration with probing workflows.
- Transition → Heads: hidden state · Recurrent state is decoded by the prediction heads.
- Heads → World Model: outputs · Predictions are assembled into the full model forward pass.
- Expr Encoder → World Model: component · Encoder is wrapped by the end-to-end model.
- Transition → World Model: component · Transition core is wrapped by the end-to-end model.
- Heads → Train: losses · Training optimizes latent and operation prediction targets.
- World Model → Train: trained by · Train script drives forward/backward passes through the model.
- World Model → Probe: hidden states · Probe extracts recurrent representations for analysis.
- Probe → Probe Results: writes · Saves plots and probe metrics.
- Probe Results → Probe Report: contains · Report file stores linear probe evaluation summary.
- Model Test → World Model: validates · Smoke tests the integrated model interface.