# Action Heads

Action heads are the trainable decoders that convert visual-language features into robot joint commands. They're the only part of Neon that trains from scratch — everything else is frozen pre-trained models.

!!! info "Parameter Golf v2"
    Since v2, action heads use techniques from the [Parameter Golf competition](../concepts/parameter-golf.md): ReLU², RMSNorm, learnable residual scales, U-Net skip connections, and soft-capping.

---

## Design Choice: Separate Heads

Neon uses **separate MLP decoders for each joint group** instead of a single decoder:

```mermaid
graph TD
    F["Fused Features (2048-d)"]
    F --> ARM["Arms Head<br/>ActionChunkingHead<br/>→ (batch, 16, 14)"]
    F --> LOCO["Locomotion Head<br/>ActionChunkingHead<br/>→ (batch, 16, 3)"]
    F --> TH["Torso+Head<br/>MLPDecoder<br/>→ (batch, 3)"]
    F --> LEG["Legs Head<br/>MLPDecoder<br/>→ (batch, 12)"]

    ARM --> CAT["Concatenate"]
    LOCO --> CAT
    TH --> CAT
    LEG --> CAT
    CAT --> OUT["(batch, 16, 17+)"]

    style F fill:#333,color:#fff
    style ARM fill:#1b5e20,color:#fff
    style LOCO fill:#e65100,color:#fff
```

**Why separate heads?**

1. **Different dynamics** — Arms need fine motor control. Locomotion needs smooth velocity.
2. **Progressive training** — Start with arms only, add locomotion later.
3. **Debugging** — Monitor per-group loss independently.

---

## MLPDecoder (v2)

The enhanced MLP decoder with Parameter Golf techniques:

```mermaid
graph TD
    IN["Input Features"] --> PROJ["Linear(in, hidden)"]
    PROJ --> N1["RMSNorm"]
    N1 --> A1["ReLU²"]
    A1 --> |"save as skip"| D1["Dropout"]
    D1 --> N2["RMSNorm"]
    N2 --> A2["ReLU²"]
    A2 --> D2["Dropout"]
    D2 --> |"concat skip"| CAT["[h ∥ skip]"]
    CAT --> L3["Linear(2×hidden, hidden)"]
    L3 --> N3["RMSNorm"]
    N3 --> A3["ReLU²"]
    A3 --> RES["h + α·residual"]
    RES --> OUT["Linear(hidden, action_dim)"]
    OUT --> SC["soft_cap(x, 1.0)"]
    SC --> ACTIONS["Actions ∈ (-1, 1)"]

    style A1 fill:#e65100,color:#fff
    style A2 fill:#e65100,color:#fff
    style A3 fill:#e65100,color:#fff
    style SC fill:#1b5e20,color:#fff
```

```python
decoder = MLPDecoder(
    input_dim=2048,
    output_dim=14,              # Arm joints
    hidden_dim=512,
    num_layers=3,
    dropout=0.1,
    use_relu_squared=True,      # ReLU² activation
    use_rmsnorm=True,           # RMSNorm before each layer
    use_residual_scale=True,    # Learned α for residuals
    use_skip_connections=True,  # U-Net skip layer 0 → last
    soft_cap_value=1.0,         # Soft-cap output to (-1, 1)
)
```

### ActionChunkingHead (v2)

Predicts multiple future timesteps using learnable step embeddings. Also uses RMSNorm, ReLU², and soft-capping:

```python
class ActionChunkingHead(nn.Module):
    # Input:  (batch, hidden_size)
    # Output: (batch, num_steps, action_dim)
    #
    # Each step gets its own embedding:
    #   step_embed[0]  = "what to do now"
    #   step_embed[7]  = "what to do in 140ms"
    #   step_embed[15] = "what to do in 300ms"
    #
    # v2: RMSNorm + ReLU² + soft-cap throughout
```

---

## Configuration

```python
from neon.model.action_heads import ActionHeadConfig

config = ActionHeadConfig(
    hidden_size=2048,            # Must match backbone output
    action_dim=17,               # Auto-set from action space
    num_action_steps=16,         # Chunk size
    mlp_hidden=512,              # MLP width
    num_layers=3,                # MLP depth
    dropout=0.1,
    use_separate_heads=True,     # Per-group decoders
    use_action_chunking=True,    # Temporal prediction
    # Parameter Golf v2
    use_relu_squared=True,
    use_rmsnorm=True,
    use_residual_scale=True,
    use_skip_connections=True,
    soft_cap_value=1.0,
)
```

### Parameter Counts

For `arms_only` mode (17 DoF, separate chunking heads):

| Component | Parameters |
|---|---|
| Arms ActionChunkingHead | ~1.5M |
| Locomotion ActionChunkingHead | ~400K |
| Fusion layer | ~4.2M |
| Proprioception encoder | ~100K |
| RMSNorm weights | ~3K |
| Residual scales | ~10 |
| **Total trainable** | **~6M** |

This is tiny compared to the frozen 3B+ backbone.

---

## Loss Function

MSE loss with optional masking:

```python
def compute_loss(self, predicted, target, mask=None):
    loss = F.mse_loss(predicted, target, reduction="none")
    if mask is not None:
        loss = loss * mask.unsqueeze(-1)
    return loss.mean()
```

!!! tip "Per-group loss weighting"
    Arms might need 2× more loss weight than locomotion if manipulation accuracy matters more. Configurable in training.

---

## Backward Compatibility

Disable all v2 features to get the original behavior:

```python
config = ActionHeadConfig(
    use_relu_squared=False,      # → GELU
    use_rmsnorm=False,           # → Identity
    use_residual_scale=False,    # → No learned scales
    use_skip_connections=False,  # → No U-Net skip
    soft_cap_value=0.0,          # → Hard Tanh
)
```
