# mondo-robotics/pmt architecture

> Architecture diagram of the GitHub repository mondo-robotics/pmt, 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/mondo-robotics/pmt
- Repository: https://github.com/mondo-robotics/pmt

## Overview

PMT is a config-driven humanoid motion-tracking research system: users select a task, train or play a policy in simulation, and can evaluate flat tasks through an optional MuJoCo backend. A separate TCRS tool optimizes motion clips for terrain. The sampled code supports task derivation, motion-command behavior, policy components, and training storage; wiring of unsampled subsystems is included only where the README documents it.

Main entry points: [`pmt_tasks/tracking_env_cfg.py`](https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/tracking_env_cfg.py) (Tracking environments), [`TCRS/stair_mppi/mppi_foot_planner_smooth.py`](https://github.com/mondo-robotics/pmt/blob/main/TCRS/stair_mppi/mppi_foot_planner_smooth.py) (TCRS foot planner), [`motion_tracking_rl/runners/on_policy_runner.py`](https://github.com/mondo-robotics/pmt/blob/main/motion_tracking_rl/runners/on_policy_runner.py) (Training runners), [`pmt_tasks/builder.py`](https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/builder.py) (Task builder)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_tasks["Task configuration"]
  node_builder["Task builder<br/>[builder.py]"]
  node_taskcfg["Task configurations"]
  node_agentcfg["Agent configurations"]
  node_registry["Task registry<br/>[registry_gym.py]"]
end

subgraph group_simulation["Simulation and tracking"]
  node_environment["Tracking environments"]
  node_motion["Motion commands"]
  node_observations["Observations<br/>[observations.py]"]
  node_rewards["Rewards and termination<br/>[rewards.py]"]
  node_robot["G1 articulation<br/>[g1.py]"]
  node_terrain["Terrain utilities<br/>[terrain.py]"]
  node_simbackend["Isaac Lab"]
  node_mjlab["MuJoCo evaluation"]
end

subgraph group_learning["RL training"]
  node_train["Train entrypoint<br/>[train.py]"]
  node_play["Playback entrypoint<br/>[play.py]"]
  node_runner["Training runners"]
  node_distillrunner["Distillation runner"]
  node_algorithms["RL algorithms"]
  node_storage["Rollout storage<br/>[rollout_storage.py]"]
end

subgraph group_policies["Policy networks"]
  node_networks["Policy networks"]
  node_token_tracker["Motion token tracker<br/>[token_tracker.py]"]
end

subgraph group_optimizer["Terrain motion optimizer"]
  node_tcrs["TCRS foot planner"]
  node_tcrs_terrain["TCRS terrain model<br/>[terrain.py]"]
end

node_user(("Researcher"))
node_rawmotion["Raw motion clips"]
node_terrain_scene["Terrain scene"]
node_optimized["Optimized clips"]
node_external_sim(("MuJoCo / mjlab"))
node_checkpoint["Policy checkpoint"]

node_user -->|"starts training"| node_train
node_user -->|"starts playback"| node_play
node_train -->|"builds task"| node_builder
node_builder -->|"composes"| node_taskcfg
node_builder -->|"derives settings"| node_agentcfg
node_builder -->|"registers task"| node_registry
node_train -->|"dispatches runner"| node_runner
node_runner -->|"updates policy"| node_algorithms
node_runner -->|"stores rollouts"| node_storage
node_algorithms -->|"samples batches"| node_storage
node_algorithms -->|"optimizes"| node_networks
node_distillrunner -->|"trains student"| node_networks
node_taskcfg -->|"configures"| node_environment
node_environment -->|"runs in"| node_simbackend
node_environment -->|"uses commands"| node_motion
node_environment -->|"computes observations"| node_observations
node_environment -->|"computes rewards"| node_rewards
node_environment -->|"simulates"| node_robot
node_environment -->|"uses terrain"| node_terrain
node_motion -->|"loads clips"| node_rawmotion
node_play -->|"loads policy"| node_runner
node_play -->|"reads"| node_checkpoint
node_runner -->|"writes"| node_checkpoint
node_token_tracker -->|"provides architecture"| node_networks
node_user -->|"runs optimizer"| node_tcrs
node_tcrs -->|"optimizes motion"| node_rawmotion
node_tcrs -->|"queries terrain"| node_tcrs_terrain
node_tcrs -->|"uses scene"| node_terrain_scene
node_tcrs -->|"generates clips"| node_optimized
node_user -.->|"evaluates policy"| node_mjlab
node_mjlab -.->|"loads"| node_checkpoint
node_mjlab -.->|"runs on"| node_external_sim

click node_train "https://github.com/mondo-robotics/pmt/blob/main/scripts/train.py"
click node_play "https://github.com/mondo-robotics/pmt/blob/main/scripts/play.py"
click node_builder "https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/builder.py"
click node_taskcfg "https://github.com/mondo-robotics/pmt/tree/main/pmt_tasks/env_cfgs"
click node_agentcfg "https://github.com/mondo-robotics/pmt/tree/main/pmt_tasks/agent_cfgs"
click node_registry "https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/registry_gym.py"
click node_environment "https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/tracking_env_cfg.py"
click node_motion "https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/mdp/commands/multi_motion_command.py"
click node_observations "https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/mdp/observations.py"
click node_rewards "https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/mdp/rewards.py"
click node_robot "https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/robots/g1.py"
click node_terrain "https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/utils/terrain.py"
click node_simbackend "https://github.com/mondo-robotics/pmt/tree/main/pmt_tasks/backends"
click node_mjlab "https://github.com/mondo-robotics/pmt/blob/main/scripts/mjlab_eval_pmt_ckpt.py"
click node_runner "https://github.com/mondo-robotics/pmt/blob/main/motion_tracking_rl/runners/on_policy_runner.py"
click node_distillrunner "https://github.com/mondo-robotics/pmt/blob/main/motion_tracking_rl/runners/distillation_runner.py"
click node_algorithms "https://github.com/mondo-robotics/pmt/tree/main/motion_tracking_rl/algorithms"
click node_storage "https://github.com/mondo-robotics/pmt/blob/main/motion_tracking_rl/storage/rollout_storage.py"
click node_networks "https://github.com/mondo-robotics/pmt/tree/main/motion_tracking_rl/networks"
click node_token_tracker "https://github.com/mondo-robotics/pmt/blob/main/motion_tracking_rl/networks/perceptive_motion/token_tracker.py"
click node_tcrs "https://github.com/mondo-robotics/pmt/blob/main/TCRS/stair_mppi/mppi_foot_planner_smooth.py"
click node_tcrs_terrain "https://github.com/mondo-robotics/pmt/blob/main/TCRS/stair_mppi/terrain.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_builder,node_taskcfg,node_agentcfg,node_registry toneBlue
class node_environment,node_motion,node_observations,node_rewards,node_robot,node_terrain,node_simbackend,node_mjlab toneAmber
class node_train,node_play,node_runner,node_distillrunner,node_algorithms,node_storage toneMint
class node_networks,node_token_tracker toneRose
class node_tcrs,node_tcrs_terrain,node_user,node_rawmotion,node_terrain_scene,node_optimized,node_external_sim,node_checkpoint toneIndigo
```

## Components

### Task configuration

- [Task builder](https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/builder.py): `pmt_tasks/builder.py`
- [Task configurations](https://github.com/mondo-robotics/pmt/tree/main/pmt_tasks/env_cfgs): `pmt_tasks/env_cfgs`
- [Agent configurations](https://github.com/mondo-robotics/pmt/tree/main/pmt_tasks/agent_cfgs): `pmt_tasks/agent_cfgs`
- [Task registry](https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/registry_gym.py): `pmt_tasks/registry_gym.py`

### Simulation and tracking

- [Tracking environments](https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/tracking_env_cfg.py): `pmt_tasks/tracking_env_cfg.py`
- [Motion commands](https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/mdp/commands/multi_motion_command.py): `pmt_tasks/mdp/commands/multi_motion_command.py`
- [Observations](https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/mdp/observations.py): `pmt_tasks/mdp/observations.py`
- [Rewards and termination](https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/mdp/rewards.py): `pmt_tasks/mdp/rewards.py`
- [G1 articulation](https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/robots/g1.py): `pmt_tasks/robots/g1.py`
- [Terrain utilities](https://github.com/mondo-robotics/pmt/blob/main/pmt_tasks/utils/terrain.py): `pmt_tasks/utils/terrain.py`
- [Isaac Lab](https://github.com/mondo-robotics/pmt/tree/main/pmt_tasks/backends): `pmt_tasks/backends`
- [MuJoCo evaluation](https://github.com/mondo-robotics/pmt/blob/main/scripts/mjlab_eval_pmt_ckpt.py): `scripts/mjlab_eval_pmt_ckpt.py`

### RL training

- [Train entrypoint](https://github.com/mondo-robotics/pmt/blob/main/scripts/train.py): `scripts/train.py`
- [Playback entrypoint](https://github.com/mondo-robotics/pmt/blob/main/scripts/play.py): `scripts/play.py`
- [Training runners](https://github.com/mondo-robotics/pmt/blob/main/motion_tracking_rl/runners/on_policy_runner.py): `motion_tracking_rl/runners/on_policy_runner.py`
- [Distillation runner](https://github.com/mondo-robotics/pmt/blob/main/motion_tracking_rl/runners/distillation_runner.py): `motion_tracking_rl/runners/distillation_runner.py`
- [RL algorithms](https://github.com/mondo-robotics/pmt/tree/main/motion_tracking_rl/algorithms): `motion_tracking_rl/algorithms`
- [Rollout storage](https://github.com/mondo-robotics/pmt/blob/main/motion_tracking_rl/storage/rollout_storage.py): `motion_tracking_rl/storage/rollout_storage.py`

### Policy networks

- [Policy networks](https://github.com/mondo-robotics/pmt/tree/main/motion_tracking_rl/networks): `motion_tracking_rl/networks`
- [Motion token tracker](https://github.com/mondo-robotics/pmt/blob/main/motion_tracking_rl/networks/perceptive_motion/token_tracker.py): `motion_tracking_rl/networks/perceptive_motion/token_tracker.py`

### Terrain motion optimizer

- [TCRS foot planner](https://github.com/mondo-robotics/pmt/blob/main/TCRS/stair_mppi/mppi_foot_planner_smooth.py): `TCRS/stair_mppi/mppi_foot_planner_smooth.py`
- [TCRS terrain model](https://github.com/mondo-robotics/pmt/blob/main/TCRS/stair_mppi/terrain.py): `TCRS/stair_mppi/terrain.py`

### Other components

- **Researcher**
- **Raw motion clips**
- **Terrain scene**
- **Optimized clips**
- **MuJoCo / mjlab**
- **Policy checkpoint**

## Connections

- Researcher → Train entrypoint: starts training
- Researcher → Playback entrypoint: starts playback
- Train entrypoint → Task builder: builds task
- Task builder → Task configurations: composes
- Task builder → Agent configurations: derives settings
- Task builder → Task registry: registers task
- Train entrypoint → Training runners: dispatches runner
- Training runners → RL algorithms: updates policy
- Training runners → Rollout storage: stores rollouts
- RL algorithms → Rollout storage: samples batches
- RL algorithms → Policy networks: optimizes
- Distillation runner → Policy networks: trains student
- Task configurations → Tracking environments: configures
- Tracking environments → Isaac Lab: runs in
- Tracking environments → Motion commands: uses commands
- Tracking environments → Observations: computes observations
- Tracking environments → Rewards and termination: computes rewards
- Tracking environments → G1 articulation: simulates
- Tracking environments → Terrain utilities: uses terrain
- Motion commands → Raw motion clips: loads clips
- Playback entrypoint → Training runners: loads policy
- Playback entrypoint → Policy checkpoint: reads
- Training runners → Policy checkpoint: writes
- Motion token tracker → Policy networks: provides architecture
- Researcher → TCRS foot planner: runs optimizer
- TCRS foot planner → Raw motion clips: optimizes motion
- TCRS foot planner → TCRS terrain model: queries terrain
- TCRS foot planner → Terrain scene: uses scene
- TCRS foot planner → Optimized clips: generates clips
- Researcher → MuJoCo evaluation: evaluates policy
- MuJoCo evaluation → Policy checkpoint: loads
- MuJoCo evaluation → MuJoCo / mjlab: runs on