How facebookresearch/td_jepa works
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 Training workspace · evaluations/dmc.py DMC evaluation · ogbench.py OGBench evaluation · envs/dmc.py DMC environments
Components 26
Training and evaluation
- Training workspace
train.py - Evaluation runners
metamotivo/evaluations - DMC evaluation
metamotivo/evaluations/dmc.py - OGBench evaluation
metamotivo/evaluations/ogbench.py
RL algorithms
- FB agent
metamotivo/agents/fb/agent.py - FB model
metamotivo/agents/fb/model.py - SF agent
metamotivo/agents/sf/agent.py - SF model
metamotivo/agents/sf/model.py - TD-JEPA agent
metamotivo/agents/td_jepa/agent.py - TD-JEPA model
metamotivo/agents/td_jepa/model.py - Other baselines
metamotivo/agents
Offline data
- Trajectory loading
metamotivo/data_loading/base.py - Replay buffers
metamotivo/buffers - DMC data adapter
metamotivo/data_loading/dmc.py - OGBench data adapter
metamotivo/data_loading/ogbench.py
Environments
- DMC environments
metamotivo/envs/dmc.py - DMC task suite
metamotivo/envs/dmc_tasks/dmc.py - OGBench environments
metamotivo/envs/ogbench.py - Env wrappers
metamotivo/envs/utils/wrappers.py - Policy rollout
metamotivo/envs/utils/rollout.py
Shared model layers
- Config base
metamotivo/base.py - Model persistence
metamotivo/base_model.py - Neural networks
metamotivo/nn_models.py - Observation normalizers
metamotivo/normalizers.py - Gym space conversion
metamotivo/envs/utils/gym_spaces.py
Other components
- Experimenter
Connections 39
- Experimenter → Training workspace: configures runs
README.md - Training workspace → DMC data adapter: selects data
train.py - Training workspace → OGBench data adapter: selects data
train.py - Training workspace → DMC environments: builds environment
train.py - Training workspace → OGBench environments: builds environment
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
metamotivo/agents/__init__.py - Training workspace → Evaluation runners: runs evaluations
train.py - Evaluation runners → DMC evaluation: dispatches
metamotivo/evaluations/dmc.py - Evaluation runners → OGBench evaluation: dispatches
metamotivo/evaluations/ogbench.py - Trajectory loading → Replay buffers: fills transitions
metamotivo/data_loading/base.py - DMC data adapter → Trajectory loading: uses loader
metamotivo/data_loading/dmc.py - OGBench data adapter → Trajectory loading: uses loader
metamotivo/data_loading/ogbench.py - DMC data adapter → OGBench environments: imports environment
- DMC environments → DMC task suite: builds tasks
metamotivo/envs/dmc.py - DMC environments → Env wrappers: applies wrappers
metamotivo/envs/dmc.py - OGBench environments → Env wrappers: applies wrappers
metamotivo/envs/ogbench.py - DMC evaluation → DMC environments: builds tasks
metamotivo/evaluations/dmc.py - DMC evaluation → DMC task suite: uses task suite
metamotivo/evaluations/dmc.py - DMC evaluation → Policy rollout: rolls out policy
metamotivo/evaluations/dmc.py - OGBench evaluation → OGBench environments: builds environment
metamotivo/evaluations/ogbench.py - OGBench evaluation → Policy rollout: rolls out policy
metamotivo/evaluations/ogbench.py - FB agent → FB model: builds model
metamotivo/agents/fb/agent.py - SF agent → SF model: builds model
metamotivo/agents/sf/agent.py - TD-JEPA agent → TD-JEPA model: builds model
metamotivo/agents/td_jepa/agent.py - FB agent → Neural networks: uses networks
metamotivo/agents/fb/agent.py - SF agent → Neural networks: uses networks
metamotivo/agents/sf/agent.py - TD-JEPA agent → Neural networks: uses networks
metamotivo/agents/td_jepa/agent.py - FB model → Model persistence: extends model base
metamotivo/agents/fb/model.py - SF model → Model persistence: extends model base
metamotivo/agents/sf/model.py - TD-JEPA model → Model persistence: extends model base
metamotivo/agents/td_jepa/model.py - FB model → Observation normalizers: uses normalizers
metamotivo/agents/fb/model.py - SF model → Observation normalizers: uses normalizers
metamotivo/agents/sf/model.py - TD-JEPA model → Observation normalizers: uses normalizers
metamotivo/agents/td_jepa/model.py - Model persistence → Gym space conversion: serializes spaces
metamotivo/base_model.py - DMC evaluation → Training workspace: reports results
train.py - OGBench evaluation → Training workspace: reports results
train.py