# Training Code Structure

This page describes the Python training codebase under `gear_sonic/`, covering directory layout, the training pipeline, configuration system, key modules, and evaluation scripts.

---

## Directory Layout

```
gear_sonic/
├── train_agent_trl.py          # Main training entry point
├── eval_agent_trl.py           # Single-checkpoint evaluation
├── eval_exp.py                 # Checkpoint monitor (continuous eval)
├── config/                     # Hydra configuration hierarchy
│   ├── base.yaml               # Global defaults (seed, num_envs, paths)
│   ├── base_eval.yaml          # Eval-specific global defaults
│   ├── eval_exp.yaml           # Checkpoint monitor config
│   ├── base/                   # Hydra plumbing (output dirs, resolvers)
│   ├── algo/                   # PPO hyperparameters
│   ├── actor_critic/           # Actor-critic architecture configs
│   │   ├── encoders/           # Per-encoder MLP configs (g1, smpl, teleop)
│   │   ├── decoders/           # Decoder MLP configs (g1_kin, g1_dyn)
│   │   ├── critics/            # Critic backbone configs
│   │   ├── quantizers/         # FSQ quantizer config
│   │   └── universal_token/    # Assembled encoder+decoder+quantizer presets
│   ├── aux_losses/             # Auxiliary loss definitions
│   ├── callbacks/              # Training callback configs
│   ├── exp/                    # Experiment presets (compose all pieces)
│   ├── manager_env/            # Environment MDP component configs
│   ├── opt/                    # Logging options (wandb)
│   └── trainer/                # Trainer class selection
├── envs/                       # IsaacLab environment wrappers
│   ├── manager_env/
│   │   ├── modular_tracking_env_cfg.py   # Scene, sensors, robot articulation
│   │   ├── robots/             # Per-robot configs (g1.py, h2.py)
│   │   └── mdp/                # MDP components (see below)
│   ├── wrapper/
│   │   └── manager_env_wrapper.py  # RL-facing env wrapper
│   └── env_utils/              # Joint ordering utilities
├── trl/                        # Training modules (PPO, actor-critic, losses)
│   ├── trainer/
│   │   ├── ppo_trainer.py          # Base PPO trainer
│   │   └── ppo_trainer_aux_loss.py # PPO + auxiliary losses (SONIC)
│   ├── modules/
│   │   ├── actor_critic_modules.py     # Actor, Critic classes
│   │   ├── universal_token_modules.py  # UniversalTokenModule (SONIC ATM)
│   │   ├── base_module.py              # Shared MLP building blocks
│   │   └── data_utils.py              # Batch/data helpers
│   ├── losses/
│   │   └── token_losses.py     # Reconstruction & latent auxiliary losses
│   ├── callbacks/              # Runtime callbacks
│   │   ├── im_eval_callback.py     # Imitation evaluation metrics
│   │   ├── im_resample_callback.py # Adaptive motion resampling
│   │   ├── model_save_callback.py  # Checkpoint saving
│   │   ├── wandb_callback.py       # W&B logging
│   │   └── read_eval_callback.py   # Read eval results from disk
│   └── utils/                  # Math, rotation, scheduling utilities
├── utils/                      # Shared utilities
│   ├── motion_lib/             # Motion library loading (PKL format)
│   ├── mujoco_sim/             # MuJoCo sim-to-sim bridge
│   └── teleop/                 # VR teleoperation helpers
├── data/                       # Robot models, URDF/USD assets
├── data_process/               # Motion data conversion scripts
└── scripts/                    # MuJoCo sim loop, misc tools
```

---

## Training Pipeline

Running `python gear_sonic/train_agent_trl.py +exp=manager/universal_token/all_modes/sonic_release` executes the following steps:

### 1. Configuration Loading

The entry point uses `@hydra.main(config_path="config", config_name="base")`. The `+exp=...` argument selects an experiment preset that composes all sub-configs:

```
base.yaml                          # Global defaults
  └── +exp=manager/universal_token/all_modes/sonic_release
        ├── /algo: ppo_im_phc      # PPO hyperparameters
        ├── /actor_critic: universal_token/all_mlp_v1
        │     ├── encoders/g1_mf_mlp, smpl_mlp, teleop_mlp
        │     ├── decoders/g1_kin_mf_mlp, g1_dyn_mlp
        │     ├── quantizers/fsq
        │     └── critics/mlp
        ├── /manager_env: base_env  # Environment config
        │     ├── observations/{tokenizer, policy, critic}
        │     ├── rewards/tracking/base_5point_local_feet_acc
        │     ├── terminations/tracking/base_adaptive_strict_ori_foot_xyz
        │     └── events/tracking/level0_4
        ├── /aux_losses: universal_token/g1_recon_and_all_latent
        ├── /trainer: trl_ppo_aux
        └── /callbacks: model_save, wandb, read_eval, im_resample
```

### 2. Simulator and Accelerator Init

After config resolution, the script:
1. Parses TRL `PPOConfig` / `ScriptArguments` / `ModelConfig` from the config dict.
2. Creates a HuggingFace `Accelerator` for multi-GPU support (DDP).
3. Launches the IsaacLab `AppLauncher` to start the Isaac Sim runtime.
4. Saves `config.yaml` and `meta.yaml` to the experiment directory.

### 3. Environment Creation

`create_manager_env()` instantiates the IsaacLab `ManagerBasedRLEnv` from the composed environment config, then wraps it with `ManagerEnvWrapper`:

```
ManagerBasedRLEnv (IsaacLab)
  └── ManagerEnvWrapper
        ├── Observation spaces (policy, critic, tokenizer groups)
        ├── Motion command manager (motion_lib)
        ├── Action transform module (optional, for pretrained ATM)
        └── Keyboard / visualization hooks
```

### 4. Policy and Value Model Creation

The actor and critic are instantiated from the algo config. For SONIC training, the actor backbone is `UniversalTokenModule`:

```python
# Simplified from train_agent_trl.py
policy = custom_instantiate(config.algo.config.actor, env_config=env.config, ...)
value_model = custom_instantiate(config.algo.config.critic, env_config=env.config, ...)
```

The `Actor` wraps `UniversalTokenModule` as its backbone and adds a diagonal Gaussian distribution for exploration. The `Critic` wraps a separate MLP backbone.

### 5. PPO Training Loop

The `TRLAuxLossPPOTrainer.train()` method runs the main loop:

```
for iteration in range(num_learning_iterations):
    # 1. Rollout: collect num_steps_per_env transitions
    for step in range(num_steps_per_env):
        actions = policy.rollout(obs_dict)
        obs_dict, rewards, dones, infos = env.step(actions)
        store(obs, actions, rewards, values, log_probs)

    # 2. GAE: compute advantages and returns
    advantages = generalized_advantage_estimation(rewards, values, dones)

    # 3. PPO update: num_ppo_epochs over mini-batches
    for epoch in range(num_ppo_epochs):
        for mini_batch in shuffle_and_split(rollout_data):
            policy_loss = clipped_surrogate_objective(...)
            value_loss  = clipped_value_loss(...)
            aux_loss    = sum(coef_i * aux_loss_i)  # encoder reconstruction, etc.
            total_loss  = policy_loss + value_loss_coef * value_loss
                        + aux_loss_scale * aux_loss
            optimizer.step(total_loss)

    # 4. Post-update: sync running stats, adaptive sampling, callbacks
    update_scheduled_params(...)     # learning rate, domain randomization
    callbacks.on_step_end(...)       # checkpointing, evaluation, logging
```

---

## Configuration System

The configuration system uses [Hydra](https://hydra.cc/) with config groups and composition.

### Hierarchy

| Level | Path | Purpose |
|---|---|---|
| **Global** | `config/base.yaml` | Seed, num_envs, paths, wandb toggle |
| **Algorithm** | `config/algo/ppo_im_phc.yaml` | PPO hyperparameters, learning rates, epochs |
| **Actor-Critic** | `config/actor_critic/` | Network architecture (encoders, decoders, critic) |
| **Environment** | `config/manager_env/` | Observations, rewards, terminations, events |
| **Auxiliary Losses** | `config/aux_losses/` | Reconstruction and latent alignment losses |
| **Trainer** | `config/trainer/` | Trainer class selection (PPO or PPO+AuxLoss) |
| **Callbacks** | `config/callbacks/` | Checkpointing, evaluation, W&B logging |
| **Experiment** | `config/exp/` | Preset that composes all the above |

### Experiment Presets

Experiment configs live under `config/exp/` and use the `@package _global_` directive to set values at the root level. They compose all component configs via `defaults`:

```yaml
# config/exp/manager/universal_token/all_modes/sonic_release.yaml
defaults:
  - /algo: ppo_im_phc
  - /manager_env: base_env
  - override /actor_critic: universal_token/all_mlp_v1
  - override /manager_env/observations/tokenizer: unitoken_all_noz
  - override /manager_env/observations/policy: local_dir_hist
  - override /manager_env/rewards: tracking/base_5point_local_feet_acc
  - override /manager_env/terminations: tracking/base_adaptive_strict_ori_foot_xyz
  - override /manager_env/events: tracking/level0_4
  # ...
```

### Key Config Parameters

| Parameter | Default | Description |
|---|---|---|
| `num_envs` | 4096 | Number of parallel simulation environments |
| `algo.config.num_learning_iterations` | 100000 | Total training iterations |
| `algo.config.num_steps_per_env` | 32 | Rollout horizon per iteration |
| `algo.config.num_learning_epochs` | 5 | PPO epochs per iteration |
| `algo.config.num_mini_batches` | 4 | Mini-batches per PPO epoch |
| `algo.config.actor_learning_rate` | 2e-5 | Actor learning rate |
| `algo.config.critic_learning_rate` | 1e-3 | Critic learning rate |
| `algo.config.clip_param` | 0.2 | PPO clipping parameter |
| `algo.config.init_noise_std` | 0.05 | Initial exploration noise std |
| `algo.config.save_interval` | 500 | Checkpoint save frequency (iterations) |

---

## Universal Token Module

The `UniversalTokenModule` implements SONIC's action transform module (ATM) -- the core architecture that maps diverse motion inputs into a shared token space.

### Architecture

```
                  ┌─────────────┐
  G1 obs    ───►  │  G1 Encoder │──┐
                  └─────────────┘  │
                  ┌─────────────┐  │    ┌─────────┐     ┌─────────────┐
  Teleop obs───►  │Teleop Encdr │──┼──► │   FSQ   │──►  │ G1 Dynamic  │──► joint actions
                  └─────────────┘  │    │Quantizer│     │   Decoder   │
                  ┌─────────────┐  │    └─────────┘     └─────────────┘
  SMPL obs  ───►  │ SMPL Encoder│──┘          │
                  └─────────────┘             │         ┌─────────────┐
                                              └───────► │G1 Kinematic │──► (aux loss only)
                                                        │   Decoder   │
                                                        └─────────────┘
```

**Encoders** map different observation modalities into a shared latent space. Each encoder is an MLP that takes modality-specific tokenizer observations and outputs a fixed-size latent vector. During training, one encoder is sampled per environment according to `encoder_sample_probs`.

**FSQ Quantizer** discretizes the continuous latent into a finite set of tokens using Finite Scalar Quantization. Each latent dimension is independently quantized to one of `fsq_level_list` discrete levels. This produces a compact, discrete token representation.

**Decoders** reconstruct outputs from the quantized tokens plus proprioception:
- **G1 Dynamic Decoder** (`g1_dyn`): Produces joint-space actions fed to the actuators. This is the only decoder used at deployment time.
- **G1 Kinematic Decoder** (`g1_kin`): Reconstructs future motion frames from tokens. Used only during training to compute reconstruction auxiliary losses.

### Latent Residual Mode

For downstream tasks (e.g., object manipulation), an external policy can inject corrections into the token space without retraining the base ATM:

| Mode | Behavior |
|---|---|
| `post_quantization` (default) | Residual added after FSQ quantization |
| `pre_quantization` | Residual added before FSQ; the sum gets quantized |
| `pre_quantization_replace` | Latent is replaced entirely by the residual |

### Encoder Sampling

During training, each environment is randomly assigned an encoder per episode according to `encoder_sample_probs`. The `encoder_index` observation tells the module which encoder produced the current token. At deployment, only one encoder is active (selected by the observation configuration).

---

## Environment Structure

The training environment is built on IsaacLab's `ManagerBasedRLEnv` and uses a modular MDP design where each component is configured independently via YAML.

### MDP Components

All MDP components live in `gear_sonic/envs/manager_env/mdp/`:

| Module | Config path | Description |
|---|---|---|
| `observations.py` | `config/manager_env/observations/` | Observation terms for policy, critic, and tokenizer groups |
| `actions.py` | `config/manager_env/actions/` | Joint position action space |
| `rewards.py` | `config/manager_env/rewards/` | Reward terms (tracking, regularization) |
| `terminations.py` | `config/manager_env/terminations/` | Episode termination conditions |
| `events.py` | `config/manager_env/events/` | Domain randomization events |
| `commands.py` | `config/manager_env/commands/` | Motion command generation (motion library) |
| `curriculum.py` | `config/manager_env/curriculum/` | Curriculum schedules |
| `terrain.py` | (inline) | Terrain generation |
| `recorders.py` | `config/manager_env/recorders/` | Video recording |

### Observation Groups

Observations are split into groups, each with its own config file:

| Group | Purpose | Example terms |
|---|---|---|
| **policy** | Direct input to the policy MLP | joint_pos, joint_vel, base_ang_vel, gravity_dir, last_actions |
| **critic** | Privileged observations for the value function | All policy obs + base_lin_vel, body_pos, body_ori |
| **tokenizer** | Input to the UniversalTokenModule encoders | Multi-future joint commands, SMPL joints, VR targets, anchor orientations |

### Reward Terms

Reward configs compose individual terms from `config/manager_env/rewards/terms/`. Key tracking rewards:

| Term | Description |
|---|---|
| `tracking_relative_body_pos` | Track reference body positions (5-point: root, wrists, feet) |
| `tracking_relative_body_ori` | Track reference body orientations |
| `tracking_anchor_pos` | Track root anchor position |
| `tracking_anchor_ori` | Track root anchor orientation |
| `tracking_body_linvel` | Track reference body linear velocities |
| `tracking_body_angvel` | Track reference body angular velocities |
| `action_rate_l2` | Penalize action jerk |
| `feet_acc` | Penalize foot acceleration (smoothness) |

### ManagerEnvWrapper

`ManagerEnvWrapper` bridges the IsaacLab environment with the RL training loop. It handles:
- Flattening observation dicts for the policy
- Applying the optional pretrained action transform module
- Motion replay mode
- Debug visualization and keyboard controls

---

## Evaluation Scripts

### eval_agent_trl.py -- Single Checkpoint

Loads a single checkpoint and runs evaluation in Isaac Sim. Automatically reads the training `config.yaml` from the checkpoint directory to reconstruct the full configuration.

```bash
# Interactive visualization
python gear_sonic/eval_agent_trl.py +checkpoint=path/to/model.pt +headless=False ++num_envs=1

# Headless with video rendering
python gear_sonic/eval_agent_trl.py +checkpoint=path/to/model.pt +headless=True \
    ++num_envs=16 +run_once=True \
    ++manager_env.config.save_rendering_dir=path/to/output \
    ++manager_env.config.render_results=True \
    +manager_env/recorders=render
```

Key features:
- Merges training config with eval overrides (`eval_overrides` in config)
- Removes train-only events and terminations automatically
- Supports `+run_once=True` to exit after all environments complete one episode
- Handles `+metrics_file` to render worst-performing motions from a prior eval

### eval_exp.py -- Checkpoint Monitor

`CheckpointEvaluator` continuously monitors an experiment directory for new checkpoints and evaluates them sequentially. It runs as a companion process alongside training.

```bash
python gear_sonic/eval_exp.py ++experiment_dir=path/to/experiment
```

For each new checkpoint, it:
1. Runs metrics evaluation (launches `eval_agent_trl.py` via subprocess)
2. Runs video rendering for the hardest motions
3. Logs results and videos to W&B (resuming the training run)
4. Marks each checkpoint as evaluated to avoid redundant work

Configuration (`config/eval_exp.yaml`):

| Parameter | Description |
|---|---|
| `experiment_dir` | Path to the training experiment directory |
| `scan_interval` | Seconds between checkpoint scans (default: 60) |
| `num_eval_envs` | Number of environments for metric evaluation |
| `num_render_videos` | Number of videos to render per checkpoint |
| `eval_frequency` | Only evaluate every N-th checkpoint (default: all) |
| `single_pass` | Evaluate pending checkpoints once and exit |

---

## Key Classes Reference

| Class | Module | Description |
|---|---|---|
| `Actor` | `trl/modules/actor_critic_modules.py` | Policy network: backbone + diagonal Gaussian. Maintains observation buffer for temporal models. |
| `Critic` | `trl/modules/actor_critic_modules.py` | Value function network: backbone + scalar output. Supports running mean/std normalization. |
| `UniversalTokenModule` | `trl/modules/universal_token_modules.py` | SONIC ATM: multi-encoder, FSQ quantizer, multi-decoder. Computes auxiliary reconstruction losses. |
| `TRLPPOTrainer` | `trl/trainer/ppo_trainer.py` | Base PPO trainer adapted from HuggingFace TRL. Handles rollout collection, GAE, and gradient updates. |
| `TRLAuxLossPPOTrainer` | `trl/trainer/ppo_trainer_aux_loss.py` | Extends `TRLPPOTrainer` with auxiliary loss support (reconstruction, latent alignment). |
| `PolicyAndValueWrapper` | `trl/trainer/ppo_trainer.py` | Wraps policy + value model into a single `nn.Module` for DDP-safe forward passes. |
| `ManagerEnvWrapper` | `envs/wrapper/manager_env_wrapper.py` | Bridges IsaacLab `ManagerBasedRLEnv` with the training loop. Handles obs flattening, action transforms, replay. |
| `CheckpointEvaluator` | `eval_exp.py` | Monitors experiment directory, evaluates new checkpoints, logs to W&B. |