# facebookresearch/td_jepa architecture

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

- Interactive diagram: https://gitdiagram.com/facebookresearch/td_jepa
- Repository: https://github.com/facebookresearch/td_jepa

## Overview

TD-JEPA is a research framework for training and comparing zero-shot reinforcement-learning agents across DeepMind Control/ExORL and OGBench, using state or RGB observations. The principal workflow is configured by the training workspace: build an environment and load offline trajectories, train a selected agent from replay data, then evaluate it on tasks and log results. The graph preserves the documented algorithm families and evaluation domains. Some baseline implementations and runtime connections are unsampled; those subsystems are represented at their supported directory paths, without guessing internal calls.

Main entry points: [`train.py`](https://github.com/facebookresearch/td_jepa/blob/main/train.py) (Training workspace), [`metamotivo/evaluations/dmc.py`](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/evaluations/dmc.py) (DMC evaluation), [`metamotivo/evaluations/ogbench.py`](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/evaluations/ogbench.py) (OGBench evaluation), [`metamotivo/envs/dmc.py`](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/envs/dmc.py) (DMC environments)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_orchestration["Training and evaluation"]
  node_workspace["Training workspace<br/>[train.py]"]
  node_evaluation["Evaluation runners"]
  node_dmc_eval["DMC evaluation<br/>[dmc.py]"]
  node_ogbench_eval["OGBench evaluation<br/>[ogbench.py]"]
end

subgraph group_algorithms["RL algorithms"]
  node_fb_agent["FB agent<br/>[agent.py]"]
  node_fb_model["FB model<br/>[model.py]"]
  node_sf_agent["SF agent<br/>[agent.py]"]
  node_sf_model["SF model<br/>[model.py]"]
  node_tdjepa_agent["TD-JEPA agent<br/>[agent.py]"]
  node_tdjepa_model["TD-JEPA model<br/>[model.py]"]
  node_other_agents["Other baselines"]
end

subgraph group_data["Offline data"]
  node_loader["Trajectory loading<br/>[base.py]"]
  node_dataset_buffers[("Replay buffers")]
  node_dmc_loader["DMC data adapter<br/>[dmc.py]"]
  node_ogbench_loader["OGBench data adapter<br/>[ogbench.py]"]
end

subgraph group_environments["Environments"]
  node_dmc_env["DMC environments<br/>[dmc.py]"]
  node_dmc_tasks["DMC task suite<br/>[dmc.py]"]
  node_ogbench_env["OGBench environments<br/>[ogbench.py]"]
  node_env_utils["Env wrappers<br/>[wrappers.py]"]
  node_rollout["Policy rollout<br/>[rollout.py]"]
end

subgraph group_models["Shared model layers"]
  node_base_config["Config base<br/>[base.py]"]
  node_base_model["Model persistence<br/>[base_model.py]"]
  node_networks["Neural networks<br/>[nn_models.py]"]
  node_normalizers["Observation normalizers<br/>[normalizers.py]"]
  node_spaces["Gym space conversion<br/>[gym_spaces.py]"]
end

node_researcher(("Experimenter"))

node_researcher -->|"configures runs"| node_workspace
node_workspace -->|"selects data"| node_dmc_loader
node_workspace -->|"selects data"| node_ogbench_loader
node_workspace -->|"builds environment"| node_dmc_env
node_workspace -->|"builds environment"| node_ogbench_env
node_workspace -.->|"builds or loads"| node_fb_agent
node_workspace -.->|"builds or loads"| node_sf_agent
node_workspace -.->|"builds or loads"| node_tdjepa_agent
node_workspace -->|"selects baseline"| node_other_agents
node_workspace -->|"runs evaluations"| node_evaluation
node_evaluation -.->|"dispatches"| node_dmc_eval
node_evaluation -.->|"dispatches"| node_ogbench_eval
node_loader -->|"fills transitions"| node_dataset_buffers
node_dmc_loader -->|"uses loader"| node_loader
node_ogbench_loader -->|"uses loader"| node_loader
node_dmc_loader -->|"imports environment"| node_ogbench_env
node_dmc_env -->|"builds tasks"| node_dmc_tasks
node_dmc_env -->|"applies wrappers"| node_env_utils
node_ogbench_env -->|"applies wrappers"| node_env_utils
node_dmc_eval -->|"builds tasks"| node_dmc_env
node_dmc_eval -->|"uses task suite"| node_dmc_tasks
node_dmc_eval -->|"rolls out policy"| node_rollout
node_ogbench_eval -->|"builds environment"| node_ogbench_env
node_ogbench_eval -->|"rolls out policy"| node_rollout
node_fb_agent -->|"builds model"| node_fb_model
node_sf_agent -->|"builds model"| node_sf_model
node_tdjepa_agent -->|"builds model"| node_tdjepa_model
node_fb_agent -->|"uses networks"| node_networks
node_sf_agent -->|"uses networks"| node_networks
node_tdjepa_agent -->|"uses networks"| node_networks
node_fb_model -->|"extends model base"| node_base_model
node_sf_model -->|"extends model base"| node_base_model
node_tdjepa_model -->|"extends model base"| node_base_model
node_fb_model -->|"uses normalizers"| node_normalizers
node_sf_model -->|"uses normalizers"| node_normalizers
node_tdjepa_model -->|"uses normalizers"| node_normalizers
node_base_model -->|"serializes spaces"| node_spaces
node_dmc_eval -.->|"reports results"| node_workspace
node_ogbench_eval -.->|"reports results"| node_workspace

click node_workspace "https://github.com/facebookresearch/td_jepa/blob/main/train.py"
click node_evaluation "https://github.com/facebookresearch/td_jepa/tree/main/metamotivo/evaluations"
click node_dmc_eval "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/evaluations/dmc.py"
click node_ogbench_eval "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/evaluations/ogbench.py"
click node_fb_agent "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/agents/fb/agent.py"
click node_fb_model "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/agents/fb/model.py"
click node_sf_agent "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/agents/sf/agent.py"
click node_sf_model "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/agents/sf/model.py"
click node_tdjepa_agent "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/agents/td_jepa/agent.py"
click node_tdjepa_model "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/agents/td_jepa/model.py"
click node_other_agents "https://github.com/facebookresearch/td_jepa/tree/main/metamotivo/agents"
click node_loader "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/data_loading/base.py"
click node_dataset_buffers "https://github.com/facebookresearch/td_jepa/tree/main/metamotivo/buffers"
click node_dmc_loader "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/data_loading/dmc.py"
click node_ogbench_loader "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/data_loading/ogbench.py"
click node_dmc_env "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/envs/dmc.py"
click node_dmc_tasks "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/envs/dmc_tasks/dmc.py"
click node_ogbench_env "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/envs/ogbench.py"
click node_env_utils "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/envs/utils/wrappers.py"
click node_rollout "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/envs/utils/rollout.py"
click node_base_config "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/base.py"
click node_base_model "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/base_model.py"
click node_networks "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/nn_models.py"
click node_normalizers "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/normalizers.py"
click node_spaces "https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/envs/utils/gym_spaces.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_workspace,node_evaluation,node_dmc_eval,node_ogbench_eval toneBlue
class node_fb_agent,node_fb_model,node_sf_agent,node_sf_model,node_tdjepa_agent,node_tdjepa_model,node_other_agents toneAmber
class node_loader,node_dataset_buffers,node_dmc_loader,node_ogbench_loader toneMint
class node_dmc_env,node_dmc_tasks,node_ogbench_env,node_env_utils,node_rollout toneRose
class node_base_config,node_base_model,node_networks,node_normalizers,node_spaces,node_researcher toneIndigo
```

## Components

### Training and evaluation

- [Training workspace](https://github.com/facebookresearch/td_jepa/blob/main/train.py): `train.py`
- [Evaluation runners](https://github.com/facebookresearch/td_jepa/tree/main/metamotivo/evaluations): `metamotivo/evaluations`
- [DMC evaluation](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/evaluations/dmc.py): `metamotivo/evaluations/dmc.py`
- [OGBench evaluation](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/evaluations/ogbench.py): `metamotivo/evaluations/ogbench.py`

### RL algorithms

- [FB agent](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/agents/fb/agent.py): `metamotivo/agents/fb/agent.py`
- [FB model](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/agents/fb/model.py): `metamotivo/agents/fb/model.py`
- [SF agent](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/agents/sf/agent.py): `metamotivo/agents/sf/agent.py`
- [SF model](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/agents/sf/model.py): `metamotivo/agents/sf/model.py`
- [TD-JEPA agent](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/agents/td_jepa/agent.py): `metamotivo/agents/td_jepa/agent.py`
- [TD-JEPA model](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/agents/td_jepa/model.py): `metamotivo/agents/td_jepa/model.py`
- [Other baselines](https://github.com/facebookresearch/td_jepa/tree/main/metamotivo/agents): `metamotivo/agents`

### Offline data

- [Trajectory loading](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/data_loading/base.py): `metamotivo/data_loading/base.py`
- [Replay buffers](https://github.com/facebookresearch/td_jepa/tree/main/metamotivo/buffers): `metamotivo/buffers`
- [DMC data adapter](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/data_loading/dmc.py): `metamotivo/data_loading/dmc.py`
- [OGBench data adapter](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/data_loading/ogbench.py): `metamotivo/data_loading/ogbench.py`

### Environments

- [DMC environments](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/envs/dmc.py): `metamotivo/envs/dmc.py`
- [DMC task suite](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/envs/dmc_tasks/dmc.py): `metamotivo/envs/dmc_tasks/dmc.py`
- [OGBench environments](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/envs/ogbench.py): `metamotivo/envs/ogbench.py`
- [Env wrappers](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/envs/utils/wrappers.py): `metamotivo/envs/utils/wrappers.py`
- [Policy rollout](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/envs/utils/rollout.py): `metamotivo/envs/utils/rollout.py`

### Shared model layers

- [Config base](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/base.py): `metamotivo/base.py`
- [Model persistence](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/base_model.py): `metamotivo/base_model.py`
- [Neural networks](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/nn_models.py): `metamotivo/nn_models.py`
- [Observation normalizers](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/normalizers.py): `metamotivo/normalizers.py`
- [Gym space conversion](https://github.com/facebookresearch/td_jepa/blob/main/metamotivo/envs/utils/gym_spaces.py): `metamotivo/envs/utils/gym_spaces.py`

### Other components

- **Experimenter**

## Connections

- Experimenter → Training workspace: configures runs (evidence: `README.md`)
- Training workspace → DMC data adapter: selects data (evidence: `train.py`)
- Training workspace → OGBench data adapter: selects data (evidence: `train.py`)
- Training workspace → DMC environments: builds environment (evidence: `train.py`)
- Training workspace → OGBench environments: builds environment (evidence: `train.py`)
- Training workspace → FB agent: builds or loads
- Training workspace → SF agent: builds or loads
- Training workspace → TD-JEPA agent: builds or loads
- Training workspace → Other baselines: selects baseline (evidence: `metamotivo/agents/__init__.py`)
- Training workspace → Evaluation runners: runs evaluations (evidence: `train.py`)
- Evaluation runners → DMC evaluation: dispatches (evidence: `metamotivo/evaluations/dmc.py`)
- Evaluation runners → OGBench evaluation: dispatches (evidence: `metamotivo/evaluations/ogbench.py`)
- Trajectory loading → Replay buffers: fills transitions (evidence: `metamotivo/data_loading/base.py`)
- DMC data adapter → Trajectory loading: uses loader (evidence: `metamotivo/data_loading/dmc.py`)
- OGBench data adapter → Trajectory loading: uses loader (evidence: `metamotivo/data_loading/ogbench.py`)
- DMC data adapter → OGBench environments: imports environment
- DMC environments → DMC task suite: builds tasks (evidence: `metamotivo/envs/dmc.py`)
- DMC environments → Env wrappers: applies wrappers (evidence: `metamotivo/envs/dmc.py`)
- OGBench environments → Env wrappers: applies wrappers (evidence: `metamotivo/envs/ogbench.py`)
- DMC evaluation → DMC environments: builds tasks (evidence: `metamotivo/evaluations/dmc.py`)
- DMC evaluation → DMC task suite: uses task suite (evidence: `metamotivo/evaluations/dmc.py`)
- DMC evaluation → Policy rollout: rolls out policy (evidence: `metamotivo/evaluations/dmc.py`)
- OGBench evaluation → OGBench environments: builds environment (evidence: `metamotivo/evaluations/ogbench.py`)
- OGBench evaluation → Policy rollout: rolls out policy (evidence: `metamotivo/evaluations/ogbench.py`)
- FB agent → FB model: builds model (evidence: `metamotivo/agents/fb/agent.py`)
- SF agent → SF model: builds model (evidence: `metamotivo/agents/sf/agent.py`)
- TD-JEPA agent → TD-JEPA model: builds model (evidence: `metamotivo/agents/td_jepa/agent.py`)
- FB agent → Neural networks: uses networks (evidence: `metamotivo/agents/fb/agent.py`)
- SF agent → Neural networks: uses networks (evidence: `metamotivo/agents/sf/agent.py`)
- TD-JEPA agent → Neural networks: uses networks (evidence: `metamotivo/agents/td_jepa/agent.py`)
- FB model → Model persistence: extends model base (evidence: `metamotivo/agents/fb/model.py`)
- SF model → Model persistence: extends model base (evidence: `metamotivo/agents/sf/model.py`)
- TD-JEPA model → Model persistence: extends model base (evidence: `metamotivo/agents/td_jepa/model.py`)
- FB model → Observation normalizers: uses normalizers (evidence: `metamotivo/agents/fb/model.py`)
- SF model → Observation normalizers: uses normalizers (evidence: `metamotivo/agents/sf/model.py`)
- TD-JEPA model → Observation normalizers: uses normalizers (evidence: `metamotivo/agents/td_jepa/model.py`)
- Model persistence → Gym space conversion: serializes spaces (evidence: `metamotivo/base_model.py`)
- DMC evaluation → Training workspace: reports results (evidence: `train.py`)
- OGBench evaluation → Training workspace: reports results (evidence: `train.py`)