# Data Soup

Multi-source data mixing for cross-embodiment training. Neon supports 7 data source types with automatic action mapping to the G1 action space.

---

## The Idea

Following [XVLA/Agibot-World](https://huggingface.co/lerobot/xvla-agibot-world), Neon trains on a **data soup** — a weighted mixture of diverse datasets, each contributing different skills:

```mermaid
graph LR
    subgraph "Source Types"
        LR["🤖 lerobot"]
        AG["🦾 agibot"]
        DG["🌌 dreamgen"]
        COS["🪐 cosmos_dreamgen"]
        OXE["🌍 oxe"]
        VC["🗣️ voice_commands"]
        S4D["📸 stereo4d"]
    end

    subgraph "Pipeline"
        MAP["ActionMapper<br/>Cross-embodiment"]
        REL["Relative Actions<br/>Cosmos EE deltas"]
        NORM["Normalize to G1"]
        MIX["Weighted Mixing"]
    end

    LR --> MAP
    AG --> MAP
    DG --> MAP
    OXE --> MAP
    COS --> REL
    MAP --> NORM
    REL --> NORM
    VC --> MIX
    S4D --> MIX
    NORM --> MIX

    MIX --> DS["NeonEpisode[]"]

    style MIX fill:#e65100,color:#fff
```

---

## Seven Source Types

| Type | Description | Action Source | Multi-Camera |
|---|---|---|---|
| `lerobot` | HuggingFace LeRobot format (Bridge, DROID) | Joint positions | ❌ |
| `agibot` | Agibot-World bimanual data | Joint positions | ✅ |
| `dreamgen` | GR00T-Dreams synthetic videos | IDM extraction | ❌ |
| `cosmos_dreamgen` | Cosmos-Predict2.5 style data | **Relative EE actions** | ✅ |
| `oxe` | Open X-Embodiment (multi-robot) | Joint positions | ❌ |
| `voice_commands` | Audio + text instructions | Language only | ❌ |
| `stereo4d` | Stereo video pairs | Visual only | ✅ (L+R) |

---

## Cross-Embodiment Mapping

The `ActionMapper` converts actions from any source embodiment to the G1 action space:

```mermaid
graph TD
    subgraph "Source Embodiments"
        FR["Franka (7 DoF)"]
        SO["SO-100 (7 DoF)"]
        AB["Agibot (14 DoF)"]
        GR["GR-1 (14 DoF)"]
        G1S["G1 (29 DoF)"]
    end

    MAP["ActionMapper<br/>Index remapping + zero-pad"]

    FR --> MAP
    SO --> MAP
    AB --> MAP
    GR --> MAP
    G1S --> MAP

    MAP --> G1["G1 Action Space<br/>(29 DoF normalized)"]

    style MAP fill:#e65100,color:#fff
    style G1 fill:#1b5e20,color:#fff
```

Example: Franka 7-DoF maps to G1 right arm indices `[7,8,9,10,11,12,13]`, with all other joints zero-padded.

---

## Cosmos DreamGen Source

The newest source type uses Cosmos-Predict2.5 relative actions for better cross-embodiment transfer:

```python
DataSourceConfig(
    name="cosmos-bridge",
    type="cosmos_dreamgen",
    path="/data/bridge",           # Local: videos/ + annotations/*.json
    # OR path="nvidia/GR1-100",    # HuggingFace dataset
    weight=2.0,
    use_relative_actions=True,
    action_scaler=20.0,            # Cosmos normalization
    camera_ids=["head", "hand_0", "hand_1"],
)
```

See [Relative Actions](relative-actions.md) for the math.

---

## Configuration

```python
from neon.data.data_soup import DataSoupConfig, DataSourceConfig

config = DataSoupConfig(
    sources=[
        DataSourceConfig(
            name="bridge",
            type="lerobot",
            path="lerobot/bridge_v2",
            weight=1.0,
        ),
        DataSourceConfig(
            name="agibot",
            type="agibot",
            path="lerobot/xvla-agibot-world",
            weight=2.0,          # 2× more agibot than bridge
            max_episodes=10000,
        ),
        DataSourceConfig(
            name="cosmos-synth",
            type="cosmos_dreamgen",
            path="nvidia/GR1-100",
            weight=1.5,
            use_relative_actions=True,
            action_scaler=20.0,
        ),
        DataSourceConfig(
            name="stereo-kitchen",
            type="stereo4d",
            path="cagataydev/strands-kitchen-stereo4d",
            weight=0.5,
            use_stereo_pair=True,
        ),
        DataSourceConfig(
            name="voice-cmds",
            type="voice_commands",
            path="cagataydev/vlm-voice-commands",
            weight=0.3,
        ),
    ],
    chunk_size=16,
    fps=15,
    shuffle=True,
)
```

---

## NeonEpisode

All source types normalize to a common `NeonEpisode` structure:

```python
@dataclass
class NeonEpisode:
    images: List[np.ndarray]          # (T, H, W, 3) RGB frames
    actions: np.ndarray               # (T, action_dim) normalized [-1, 1]
    language: str                     # Natural language instruction
    proprioception: Optional[np.ndarray]  # (T, joint_dim)
    metadata: Dict[str, Any]          # Source info, flags
```

### Episode Flags

- **`language_only=True`** — Voice command episodes (no actions, language conditioning only)
- **`visual_only=True`** — Stereo4D / DreamGen (video pretraining, no action labels)
