# nju-rl/meta-dt architecture

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

- Interactive diagram: https://gitdiagram.com/nju-rl/meta-dt
- Repository: https://github.com/nju-rl/meta-dt

## Overview

This repository implements Meta-DT, an offline meta-reinforcement-learning research system. The documented workflow is to train SAC agents across tasks, turn their experience into offline datasets, train a context encoder/world model, then train and evaluate Meta Decision Transformers in few-shot or zero-shot settings. Environment implementations supply task-specific interactions; context models encode experience and predict reward or next state, while trajectory models provide sequence-modeling components. The README also describes pretrained world-model and dataset inputs. Training entry points are not all present in the supplied tree, so some orchestration links are based on README descriptions rather than inspected source.

Main entry points: [`train_context.py`](https://github.com/nju-rl/meta-dt/blob/main/train_context.py) (Context training), [`data_collection/sac.py`](https://github.com/nju-rl/meta-dt/blob/main/data_collection/sac.py) (SAC trainer), [`meta_dt/trainer.py`](https://github.com/nju-rl/meta-dt/blob/main/meta_dt/trainer.py) (Meta-DT trainer), [`meta_dt/evaluation.py`](https://github.com/nju-rl/meta-dt/blob/main/meta_dt/evaluation.py) (Meta-DT evaluation)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_collection["Task Data Collection"]
  node_sac["SAC trainer<br/>[sac.py]"]
  node_collector_model["SAC networks<br/>[model.py]"]
  node_replay["Replay memory<br/>[replay_memory.py]"]
  node_collection_eval["Collection evaluation<br/>[evaluation.py]"]
  node_dataset_build["Dataset generation<br/>[get_datasets.py]"]
  node_offline_data["Offline datasets<br/>[dataset.py]"]
end

subgraph group_contexts["Context World Model"]
  node_context_train["Context training<br/>[train_context.py]"]
  node_context_data["Context dataset<br/>[dataset.py]"]
  node_context_encoder["Context encoder<br/>[model.py]"]
  node_reward_decoder["Reward decoder<br/>[model.py]"]
  node_state_decoder["State decoder<br/>[model.py]"]
end

subgraph group_policy["Meta Decision Transformer"]
  node_meta_dataset["Meta-DT dataset<br/>[dataset.py]"]
  node_meta_trainer["Meta-DT trainer<br/>[trainer.py]"]
  node_meta_model["Meta-DT model<br/>[model.py]"]
  node_trajectory_transformer["Trajectory transformer<br/>[trajectory_gpt2.py]"]
  node_meta_eval["Meta-DT evaluation<br/>[evaluation.py]"]
end

subgraph group_environment["Task Environments"]
  node_task_envs["Task environments<br/>[envs.py]"]
  node_mujoco_tasks["MuJoCo tasks"]
end

node_researcher(("Researcher"))
node_external_mujoco["MuJoCo runtime"]

node_researcher -->|"starts training"| node_sac
node_sac -->|"uses networks"| node_collector_model
node_sac -->|"stores experience"| node_replay
node_sac -->|"interacts with"| node_task_envs
node_task_envs -->|"selects tasks"| node_mujoco_tasks
node_mujoco_tasks -.->|"simulates in"| node_external_mujoco
node_researcher -->|"requests datasets"| node_dataset_build
node_dataset_build -->|"uses trajectories"| node_replay
node_dataset_build -->|"produces datasets"| node_offline_data
node_researcher -->|"starts training"| node_context_train
node_context_train -->|"reads experience"| node_offline_data
node_context_train -->|"prepares batches"| node_context_data
node_context_train -->|"trains encoder"| node_context_encoder
node_context_train -->|"trains decoder"| node_reward_decoder
node_context_train -->|"trains decoder"| node_state_decoder
node_researcher -->|"starts training"| node_meta_trainer
node_meta_trainer -->|"uses batches"| node_meta_dataset
node_meta_trainer -->|"trains model"| node_meta_model
node_meta_model -->|"uses transformer"| node_trajectory_transformer
node_meta_trainer -->|"trains from"| node_offline_data
node_meta_eval -->|"evaluates model"| node_meta_model
node_meta_eval -->|"runs tasks"| node_task_envs
node_meta_eval -.->|"conditions on context"| node_context_encoder

click node_sac "https://github.com/nju-rl/meta-dt/blob/main/data_collection/sac.py"
click node_collector_model "https://github.com/nju-rl/meta-dt/blob/main/data_collection/model.py"
click node_replay "https://github.com/nju-rl/meta-dt/blob/main/data_collection/replay_memory.py"
click node_collection_eval "https://github.com/nju-rl/meta-dt/blob/main/data_collection/evaluation.py"
click node_dataset_build "https://github.com/nju-rl/meta-dt/blob/main/get_datasets.py"
click node_offline_data "https://github.com/nju-rl/meta-dt/blob/main/decision_transformer/dataset.py"
click node_context_train "https://github.com/nju-rl/meta-dt/blob/main/train_context.py"
click node_context_data "https://github.com/nju-rl/meta-dt/blob/main/context/dataset.py"
click node_context_encoder "https://github.com/nju-rl/meta-dt/blob/main/context/model.py"
click node_reward_decoder "https://github.com/nju-rl/meta-dt/blob/main/context/model.py"
click node_state_decoder "https://github.com/nju-rl/meta-dt/blob/main/context/model.py"
click node_meta_dataset "https://github.com/nju-rl/meta-dt/blob/main/meta_dt/dataset.py"
click node_meta_trainer "https://github.com/nju-rl/meta-dt/blob/main/meta_dt/trainer.py"
click node_meta_model "https://github.com/nju-rl/meta-dt/blob/main/meta_dt/model.py"
click node_trajectory_transformer "https://github.com/nju-rl/meta-dt/blob/main/meta_dt/trajectory_gpt2.py"
click node_meta_eval "https://github.com/nju-rl/meta-dt/blob/main/meta_dt/evaluation.py"
click node_task_envs "https://github.com/nju-rl/meta-dt/blob/main/src/envs.py"
click node_mujoco_tasks "https://github.com/nju-rl/meta-dt/tree/main/src/tp_envs/environments/mujoco"

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_sac,node_collector_model,node_replay,node_collection_eval,node_dataset_build,node_offline_data toneBlue
class node_context_train,node_context_data,node_context_encoder,node_reward_decoder,node_state_decoder toneAmber
class node_meta_dataset,node_meta_trainer,node_meta_model,node_trajectory_transformer,node_meta_eval toneMint
class node_task_envs,node_mujoco_tasks toneRose
class node_researcher,node_external_mujoco toneIndigo
```

## Components

### Task Data Collection

- [SAC trainer](https://github.com/nju-rl/meta-dt/blob/main/data_collection/sac.py): `data_collection/sac.py`
- [SAC networks](https://github.com/nju-rl/meta-dt/blob/main/data_collection/model.py): `data_collection/model.py`
- [Replay memory](https://github.com/nju-rl/meta-dt/blob/main/data_collection/replay_memory.py): `data_collection/replay_memory.py`
- [Collection evaluation](https://github.com/nju-rl/meta-dt/blob/main/data_collection/evaluation.py): `data_collection/evaluation.py`
- [Dataset generation](https://github.com/nju-rl/meta-dt/blob/main/get_datasets.py): `get_datasets.py`
- [Offline datasets](https://github.com/nju-rl/meta-dt/blob/main/decision_transformer/dataset.py): `decision_transformer/dataset.py`

### Context World Model

- [Context training](https://github.com/nju-rl/meta-dt/blob/main/train_context.py): `train_context.py`
- [Context dataset](https://github.com/nju-rl/meta-dt/blob/main/context/dataset.py): `context/dataset.py`
- [Context encoder](https://github.com/nju-rl/meta-dt/blob/main/context/model.py): `context/model.py`
- [Reward decoder](https://github.com/nju-rl/meta-dt/blob/main/context/model.py): `context/model.py`
- [State decoder](https://github.com/nju-rl/meta-dt/blob/main/context/model.py): `context/model.py`

### Meta Decision Transformer

- [Meta-DT dataset](https://github.com/nju-rl/meta-dt/blob/main/meta_dt/dataset.py): `meta_dt/dataset.py`
- [Meta-DT trainer](https://github.com/nju-rl/meta-dt/blob/main/meta_dt/trainer.py): `meta_dt/trainer.py`
- [Meta-DT model](https://github.com/nju-rl/meta-dt/blob/main/meta_dt/model.py): `meta_dt/model.py`
- [Trajectory transformer](https://github.com/nju-rl/meta-dt/blob/main/meta_dt/trajectory_gpt2.py): `meta_dt/trajectory_gpt2.py`
- [Meta-DT evaluation](https://github.com/nju-rl/meta-dt/blob/main/meta_dt/evaluation.py): `meta_dt/evaluation.py`

### Task Environments

- [Task environments](https://github.com/nju-rl/meta-dt/blob/main/src/envs.py): `src/envs.py`
- [MuJoCo tasks](https://github.com/nju-rl/meta-dt/tree/main/src/tp_envs/environments/mujoco): `src/tp_envs/environments/mujoco`

### Other components

- **Researcher**
- **MuJoCo runtime**

## Connections

- Researcher → SAC trainer: starts training
- SAC trainer → SAC networks: uses networks
- SAC trainer → Replay memory: stores experience
- SAC trainer → Task environments: interacts with
- Task environments → MuJoCo tasks: selects tasks
- MuJoCo tasks → MuJoCo runtime: simulates in
- Researcher → Dataset generation: requests datasets
- Dataset generation → Replay memory: uses trajectories
- Dataset generation → Offline datasets: produces datasets
- Researcher → Context training: starts training
- Context training → Offline datasets: reads experience
- Context training → Context dataset: prepares batches
- Context training → Context encoder: trains encoder
- Context training → Reward decoder: trains decoder
- Context training → State decoder: trains decoder
- Researcher → Meta-DT trainer: starts training
- Meta-DT trainer → Meta-DT dataset: uses batches
- Meta-DT trainer → Meta-DT model: trains model
- Meta-DT model → Trajectory transformer: uses transformer
- Meta-DT trainer → Offline datasets: trains from
- Meta-DT evaluation → Meta-DT model: evaluates model
- Meta-DT evaluation → Task environments: runs tasks
- Meta-DT evaluation → Context encoder: conditions on context