README.md

---
language: en
tags: [robotics, 6-axis-arm, visual-policy, pytorch, imitation-learning, action-chunking]
license: apache-2.0
---

# 6Net 2.0 — 6-Axis Visual Robot Policy (~228M)

Custom transformer policy for visual 6-DoF robot arm control, v2.
Optimised for **LDX-218**, **LD-1501MG**, and **LFD-01M** hardware.

| Component | Detail | Params |
|---|---|---|
| Visual Encoder | ResNet-50 (ImageNet V2, shared) | ~23.5M |
| Visual Projection (overhead) | Linear(2048→1024) | ~2.1M |
| Visual Projection (wrist) | Linear(2048→1024) | ~2.1M |
| State Encoder | MLP(6→256→1024) | ~0.3M |
| Transformer | 16L · d=1024 · 16h · ffn=4096 | ~201.6M |
| Action Head | MLP(1024→512→K×6) | ~0.5M |
| **Total** | | **~228M** |

**Dataset:** `lerobot/pusht_image` · **Steps:** 910 · **Eff. batch:** 32  
**Hardware profile:** `LDX-218`

## Key improvements over v1
- **2× parameter count** via ResNet-50 backbone + wider/deeper transformer
- **Dual-camera** overhead + wrist tokens
- **Action chunking** (K=10): predicts 10 future steps; returns step 0 at inference
- **Hardware profiles**: joint limits, max velocity, and gravity-comp for LDX-218 / LD-1501MG / LFD-01M
- **Streaming fallback**: tries streaming download before falling back to synthetic data

## Inference

```python
import torch
from train_6net_v2 import SixNetV2, Config, HARDWARE_PROFILES
import torchvision.transforms as T
from PIL import Image

cfg   = Config(hardware="LDX-218")
model = SixNetV2(cfg)
ckpt  = torch.load("6net_v2_final.pt", map_location="cpu")
model.load_state_dict(ckpt["model_state"])
model.eval()

tf   = T.Compose([T.Resize((224,224)), T.ToTensor(),
                  T.Normalize([.485,.456,.406],[.229,.224,.225])])
img  = tf(Image.open("overhead.jpg")).unsqueeze(0)
wrist = tf(Image.open("wrist.jpg")).unsqueeze(0)
jts  = torch.zeros(1, 6)   # current joint angles (rad)

action = model.predict(img, jts, wrist=wrist, hw=HARDWARE_PROFILES["LDX-218"])
# → tensor of shape (1, 6), clamped to LDX-218 joint limits
```