How nju-rl/meta-dt works
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 Context training · sac.py SAC trainer · trainer.py Meta-DT trainer · evaluation.py Meta-DT evaluation
Components 20
Task Data Collection
- SAC trainer
data_collection/sac.py - SAC networks
data_collection/model.py - Replay memory
data_collection/replay_memory.py - Collection evaluation
data_collection/evaluation.py - Dataset generation
get_datasets.py - Offline datasets
decision_transformer/dataset.py
Context World Model
- Context training
train_context.py - Context dataset
context/dataset.py - Context encoder
context/model.py - Reward decoder
context/model.py - State decoder
context/model.py
Meta Decision Transformer
- Meta-DT dataset
meta_dt/dataset.py - Meta-DT trainer
meta_dt/trainer.py - Meta-DT model
meta_dt/model.py - Trajectory transformer
meta_dt/trajectory_gpt2.py - Meta-DT evaluation
meta_dt/evaluation.py
Task Environments
- Task environments
src/envs.py - MuJoCo tasks
src/tp_envs/environments/mujoco
Other components
- Researcher
- MuJoCo runtime
Connections 23
- 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