# Action Chunking

Instead of predicting one action at a time, Neon predicts a **chunk** of 16 future actions simultaneously. This single design decision dramatically improves control quality.

---

## The Problem with Single-Step Prediction

A naive policy outputs one action per observation:

```
Observe → Predict 1 action → Execute → Observe → Predict 1 action → ...
```

This creates several problems:

1. **Compounding errors** — Each prediction uses the previous noisy prediction as context. Small errors accumulate into large drift.
2. **Jerky motion** — Independent per-step predictions have no temporal smoothness guarantee. The robot vibrates.
3. **High latency sensitivity** — If inference takes 80ms, you're running at 12 Hz. The robot reacts slowly.

---

## How Action Chunking Works

Action chunking predicts multiple future steps at once:

```
Observe → Predict 16 actions → Execute action[0] → (next step) Execute action[1] → ...
```

After executing the full chunk (or a portion of it), re-observe and predict a new chunk:

```mermaid
sequenceDiagram
    participant Camera
    participant Model
    participant Robot

    Camera->>Model: Frame t
    Model->>Robot: Actions [t, t+1, ..., t+15]
    Robot->>Robot: Execute action[t]
    Robot->>Robot: Execute action[t+1]
    Robot->>Robot: Execute action[t+2]
    Note over Robot: Continue executing chunk...
    Camera->>Model: Frame t+K (re-observe)
    Model->>Robot: Actions [t+K, ..., t+K+15]
```

---

## Why It Works

### 1. Temporal Consistency

The model is forced to predict a coherent trajectory, not isolated snapshots. The loss function penalizes *the entire sequence*, so the model learns smooth transitions:

```python
# Loss is computed over all 16 predicted steps
loss = MSE(predicted_chunk, target_chunk)  # Shape: (batch, 16, action_dim)
```

### 2. Reduced Compounding Error

With a chunk size of 16 at 50 Hz, you only re-predict every 320ms instead of every 20ms. That's 16× fewer compounding steps per second.

### 3. Amortized Inference Cost

If your model takes 50ms to run, a single-step policy runs at 20 Hz. With chunking, you get 16 actions for that 50ms, and the robot can execute them at full 50 Hz while the *next* chunk is being computed.

---

## Neon's Implementation

The `ActionChunkingHead` uses **temporal step embeddings** to distinguish predictions at different future timesteps:

```python
class ActionChunkingHead(nn.Module):
    def __init__(self, input_dim, action_dim, num_steps=16, hidden_dim=512):
        self.step_embed = nn.Embedding(num_steps, hidden_dim)  # Learnable per-step embeddings
        self.feature_proj = nn.Linear(input_dim, hidden_dim)
        self.decoder = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),  # features + step embedding
            nn.GELU(),
            nn.Linear(hidden_dim, action_dim),
            nn.Tanh(),  # Output in [-1, 1]
        )

    def forward(self, features):
        # features: (batch, hidden_size) — pooled backbone output
        feat = self.feature_proj(features)              # (batch, hidden_dim)
        step_emb = self.step_embed(arange(16))          # (16, hidden_dim)

        # Combine features with each step embedding
        feat = feat.unsqueeze(1).expand(-1, 16, -1)     # (batch, 16, hidden_dim)
        step_emb = step_emb.unsqueeze(0).expand(B, -1, -1)  # (batch, 16, hidden_dim)

        combined = cat([feat, step_emb], dim=-1)         # (batch, 16, hidden_dim*2)
        return self.decoder(combined)                    # (batch, 16, action_dim)
```

The key insight: the **same pooled features** are decoded with **different step embeddings** for each future timestep. Step 0 learns "what to do now", step 15 learns "what to do 300ms from now".

---

## Choosing Chunk Size

| Chunk Size | Control Freq | Re-predict Interval | Trade-off |
|-----------|-------------|---------------------|-----------|
| 1 | 50 Hz | 20ms | No chunking — baseline |
| 4 | 50 Hz | 80ms | Minimal smoothing |
| **8** | **50 Hz** | **160ms** | **Good default** — Neon config uses this |
| **16** | **50 Hz** | **320ms** | **More temporal consistency** — Neon code default |
| 32 | 50 Hz | 640ms | Very smooth but less reactive |

!!! tip "Use 8 for reactive tasks, 16 for smooth manipulation"
    Pouring water needs smooth trajectories (chunk=16). Catching a thrown ball needs fast reactions (chunk=4-8).

---

## Temporal Smoothing at Inference

Neon also applies **exponential moving average (EMA)** smoothing during inference to blend successive chunks:

```python
# In NeonInferenceServer.predict():
if self._prev_actions is not None:
    output.actions = (
        alpha * output.actions +          # 70% new prediction
        (1 - alpha) * self._prev_actions  # 30% previous
    )
```

This prevents discontinuities when transitioning from one predicted chunk to the next.

---

## References

- **ACT** (Zhao et al., 2023) — *Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware*. Introduced action chunking for imitation learning.
- **Diffusion Policy** (Chi et al., 2023) — Uses diffusion to predict action sequences.
- **GR00T N1** (NVIDIA, 2025) — Uses action tokenization with VQ-BeT for temporal action prediction.
