models/DreamDojo/docs/ARCHITECTURE.md

# DreamDojo Architecture

DreamDojo is an action-conditioned video world model built on the **Diffusion Transformer (DiT)** backbone with **Rectified Flow** formulation. It generates future video frames conditioned on an input image/video, text prompt, and robot action sequences.

## Model Variants

| Variant | Blocks | Heads | Channels | Head Dim | Params |
|---------|--------|-------|----------|----------|--------|
| 2B      | 28     | 16    | 2048     | 128      | ~2B    |
| 7B      | 28     | 32    | 4096     | 128      | ~7B    |
| 14B     | 36     | 40    | 5120     | 128      | ~14B   |

Config definitions: `cosmos_predict2/_src/predict2/action/configs/action_conditioned/net.py`

## Core Architecture

### DiT Backbone: `MiniTrainDIT`

**File**: `cosmos_predict2/_src/predict2/networks/minimal_v4_dit.py`

The backbone is a standard Vision Transformer adapted for video:

```
Input (B, C_in, T, H, W)
  --> Patchify (patch_spatial=2, patch_temporal=1)
  --> Linear projection to model_channels
  --> + 3D RoPE positional embeddings
  --> N x TransformerBlock:
      |-- AdaLN (adaptive layer norm from timestep embedding)
      |-- Self-Attention (Q/K/V with RMSNorm + RoPE)
      |-- Cross-Attention (to text embeddings)
      |-- FFN (GPT-2 style: Linear -> GELU -> Linear)
  --> Unpatchify
  --> Output (B, C_out, T, H, W)
```

**Key classes:**
- `MiniTrainDIT` (line 700+): Main model with `forward()` and `forward_with_cfg()`
- `Attention` (line 389+): Multi-head attention with configurable backends
- `GPT2FeedForward` (line 238+): MLP block
- `RMSNorm` (line 220+): Root mean square normalization

### Action Conditioning

**File**: `cosmos_predict2/_src/predict2/action/networks/action_conditioned_minimal_v1_lvg_dit.py`

Actions are injected into the model via the timestep embedding pathway:

```
action: (B, T-1, action_dim=384) 
  --> ActionEmbedder (Linear projection)
  --> Add to timestep embedding (t_emb)
  --> t_emb conditions AdaLN in each transformer block
```

**Action dimension layout (384-dim):**

| Range       | Robot/Type           |
|-------------|---------------------|
| `[0, 29)`   | Fourier GR-1        |
| `[29, 58)`  | Retargeted GR-1     |
| `[58, 101)` | Unitree G1          |
| `[101, 147)`| Bimanual YAM        |
| `[147, 169)`| AgiBot              |
| `[169, 220)`| Reserved            |
| `[220, 352)`| MANO hand actions   |
| `[352, 384)`| Latent actions      |

Each action frame contains: `[delta_xyz(3), delta_rotation(3), gripper_state(1)]` for the relevant robot, with unused dimensions zeroed.

### Attention Mechanisms

**File**: `cosmos_predict2/_src/predict2/networks/minimal_v4_dit.py` (line 389+)

The `Attention` class supports multiple backends:

| Backend | Description | Default For |
|---------|-------------|-------------|
| `torch` | `F.scaled_dot_product_attention` | General |
| `torch-flex` | FlexAttention with BlockMask | Sparse patterns |
| `minimal_a2a` | Custom A2A attention | 14B model |
| `i4` | Imaginaire4 attention | Alternative |
| `transformer_engine` | NVIDIA TE attention | Training |

**Attention flow:**
```
Input x: (B, S, D)  where S = T*H*W (flattened video tokens)
  --> Q = q_proj(x): (B, S, n_heads*head_dim)
  --> K = k_proj(x): (B, S, n_heads*head_dim)  [self-attn]
       K = k_proj(context): (B, M, n_heads*head_dim)  [cross-attn]
  --> V = v_proj(...)
  --> Reshape: (B, S, n_heads, head_dim)
  --> Q, K = RMSNorm(Q), RMSNorm(K)
  --> Q, K = apply_RoPE_3D(Q, K)  [self-attn only]
  --> output = attn_op(Q, K, V)
  --> output = output_proj(output)
```

### Rectified Flow Diffusion

**File**: `cosmos_predict2/_src/predict2/models/text2world_model_rectified_flow.py`

DreamDojo uses Rectified Flow (RF), a straight-path ODE formulation:

- **Forward process**: `x_t = (1 - t) * x_0 + t * noise`, where `t in [0, 1]`
- **Training objective**: Predict velocity `v = noise - x_0`
- **Inference**: Solve ODE from `t=1` (noise) to `t=0` (clean) using Euler steps
- **Scheduler**: `FlowUniPCMultistepScheduler` (2nd-order predictor-corrector)
- **Default steps**: ~35 Euler steps with `shift=5.0`
- **CFG**: Classifier-free guidance with `guidance_scale` parameter

### VAE Tokenizer

**File**: `cosmos_predict2/_src/predict2/tokenizers/cosmos.py`

The Cosmos VAE compresses video to latent space:

| Property | Value |
|----------|-------|
| Spatial compression | 8x (via patch_spatial=2 in model) |
| Temporal compression | 4x |
| Latent channels | 16 |
| Input format | `(B, 3, T, H, W)` uint8 [0, 255] |
| Latent format | `(B, 16, T/4, H/8, W/8)` bf16 |

For 480x640 input with 13 frames: latent shape = `(1, 16, 4, 60, 80)`

### Text Encoder

Uses T5-based text encoder for computing text embeddings from prompts. The embeddings condition the model through cross-attention in each transformer block.

For distilled models, pre-computed CR1 (empty-string) embeddings can be used for efficiency.

## Config System

DreamDojo uses a layered config system:

1. **Base configs** (`cosmos_predict2/_src/predict2/action/configs/action_conditioned/config.py`): Define model, net, optimizer, scheduler defaults
2. **Network configs** (`net.py`): Register model architectures (2B, 7B, 14B)
3. **Experiment configs** (`cosmos_predict2/experiments/base/action.py`): Auto-register from YAML files
4. **YAML overrides** (`configs/*.yaml`): Per-experiment overrides (e.g., `14b_480_640_gr1.yaml`)

Config loading: `load_model_from_checkpoint()` in `cosmos_predict2/_src/predict2/utils/model_loader.py`

## Distilled Model (Student)

**File**: `cosmos_predict2/_src/predict2/interactive/networks/dit_action_causal.py`

The distilled model differs from the teacher:

| Aspect | Teacher | Student (Distilled) |
|--------|---------|---------------------|
| Denoising steps | ~35 | 4 (DMD2) |
| Attention | Bidirectional | Temporal causal |
| KV cache | No | Yes (frame-indexed) |
| torch.compile | Optional | Recommended |
| Inference mode | Batch | Streaming (per-frame) |

**Distillation pipeline** (3 stages):
1. **Teacher generation** (`launch_teacher_gen.sh`): Generate training data with teacher
2. **Warmup** (`launch_warmup.sh`): Train student to match teacher outputs
3. **Self-forcing** (`launch_self_forcing.sh`): Finetune with autoregressive self-predictions

## Key File Map

```
DreamDojo/
  cosmos_predict2/_src/predict2/
    networks/
      minimal_v4_dit.py          # Core DiT backbone + Attention class
      a2a_cp.py                  # A2A and NATTEN attention ops
    action/
      networks/                  # Action-conditioned model wrappers
      inference/
        inference.py             # Standard inference script
        inference_batch.py       # Batch inference (custom benchmark)
        inference_pipeline.py    # ActionVideo2WorldInference class
      configs/                   # Action-conditioned configs
    interactive/
      networks/dit_action_causal.py  # Distilled causal DiT
      inference/action_video2world.py # Distilled streaming inference
    models/
      text2world_model_rectified_flow.py  # Rectified flow model
    tokenizers/                  # VAE tokenizer
    schedulers/                  # ODE solvers
    utils/
      model_loader.py            # Checkpoint loading
      kv_cache.py                # KV cache utilities
  configs/                       # YAML experiment configs
  docs/                          # Documentation
  scripts/                       # Utilities (checkpoint conversion, etc.)
```