# API Reference

Complete reference for all public classes and functions.

---

## Model

### `NeonVLA`

::: neon.model.neon_vla.NeonVLA

The complete Vision-Language-Action model.

```python
from neon.model.neon_vla import NeonVLA, NeonConfig

model = NeonVLA(config)
model.load_backbone()
```

**Constructor:** `NeonVLA(config: NeonConfig)`

**Key Methods:**

| Method | Args | Returns | Description |
|--------|------|---------|-------------|
| `load_backbone()` | — | — | Download and load the video backbone |
| `forward(images, video_frames, text, proprioception, audio)` | See below | `(batch, steps, action_dim)` | Full forward pass |
| `predict(image, video_frames, instruction, proprioception, audio, speak)` | See below | `NeonOutput` | Inference with post-processing |
| `compute_loss(images, video_frames, text, proprioception, target_actions, action_mask, audio)` | See below | `Dict[str, Tensor]` | Training loss |
| `save_pretrained(path)` | `str` | — | Save config + trainable weights |
| `from_pretrained(path, load_backbone)` | `str`, `bool` | `NeonVLA` | Load saved model |

### `NeonConfig`

```python
@dataclass
class NeonConfig:
    backbone: BackboneConfig          # Video backbone configuration
    control_mode: str = "arms_only"   # "arms_only" | "upper_body" | "whole_body"
    include_locomotion: bool = True   # Add 3 DoF velocity output
    num_action_steps: int = 16        # Action chunk size
    mlp_hidden: int = 512             # Action head MLP width
    action_head_layers: int = 3       # Action head MLP depth
    action_head_dropout: float = 0.1
    use_separate_heads: bool = True   # Per-group action decoders
    use_action_chunking: bool = True  # Temporal action prediction
    use_proprioception: bool = True   # Include joint state encoder
    proprioception_hidden: int = 128
    learning_rate: float = 2e-4
    weight_decay: float = 0.01
    warmup_ratio: float = 0.1
    audio: Optional[Dict] = None      # AudioConfig as dict, or None
```

### `NeonOutput`

```python
@dataclass
class NeonOutput:
    actions: np.ndarray           # (steps, action_dim) normalized [-1, 1]
    raw_actions: np.ndarray       # (steps, action_dim) in physical units
    upper_body: np.ndarray        # (steps, 14) arm joint positions
    locomotion: np.ndarray        # (steps, 3) vx, vy, omega
    gripper: np.ndarray           # (steps, 2) gripper positions
    speech_path: Optional[str]    # Path to generated speech audio
```

---

## Video Backbone

### `VideoBackbone`

```python
from neon.model.video_backbone import VideoBackbone, BackboneConfig

backbone = VideoBackbone(BackboneConfig(model_id="Qwen/Qwen2.5-VL-3B-Instruct"))
backbone.load()
result = backbone.encode(images=[img], text="Pick up the cup")
```

**Key Methods:**

| Method | Returns | Description |
|--------|---------|-------------|
| `load()` | — | Load model weights from HuggingFace |
| `encode(images, video_frames, text)` | `Dict[hidden_states, pooled]` | Encode inputs to features |
| `forward(images, video_frames, text)` | `Tensor` | Returns pooled features |

**Properties:** `hidden_size: int` — Output dimension (auto-detected from model).

---

## Action Heads

### `G1ActionHead`

Orchestrates per-group action decoders.

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

head = G1ActionHead(config, action_space)
actions = head(features)  # (batch, 16, 17)
loss = head.compute_loss(predicted, target, mask)
```

### `ActionChunkingHead`

Predicts temporal action sequences with step embeddings.

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

head = ActionChunkingHead(input_dim=2048, action_dim=14, num_steps=16)
actions = head(features)  # (batch, 16, 14)
```

### `MLPDecoder`

Simple MLP with Tanh output.

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

decoder = MLPDecoder(input_dim=2048, output_dim=3, hidden_dim=256)
actions = decoder(features)  # (batch, 3)
```

---

## Audio

### `AudioEncoder`

```python
from neon.model.audio import AudioEncoder, AudioConfig

encoder = AudioEncoder(AudioConfig(encoder_type="whisper"), backbone_hidden_size=2048)
encoder.load()
features = encoder.encode(audio_tensor)  # (batch, 2048)
```

### `SpeechResponseHead`

```python
from neon.model.audio import SpeechResponseHead

head = SpeechResponseHead(hidden_size=2048)
output = head(features)
# output["response_type"]: (batch, 5) — narrate/confirm/warn/ask/silent
# output["token_logits"]: (batch, vocab_size)
```

### `PersonaPlexSpeaker`

```python
from neon.model.audio import PersonaPlexSpeaker, AudioConfig

speaker = PersonaPlexSpeaker(AudioConfig(personaplex_voice="NATM1"))
path = speaker.speak("Reaching for the cup")  # → "/tmp/neon_speech_xxx.wav"
```

---

## Action Space

### `G1ActionSpace`

```python
from neon.data.action_space import G1ActionSpace, ControlMode

space = G1ActionSpace(mode=ControlMode.ARMS_ONLY, include_locomotion=True)
```

| Property | Type | Description |
|----------|------|-------------|
| `action_dim` | `int` | Total output dimensions |
| `num_joints` | `int` | Number of controlled joints |
| `active_joints` | `List[JointInfo]` | Joints in current mode |
| `joint_names` | `List[str]` | Active joint names |

| Method | Args | Returns | Description |
|--------|------|---------|-------------|
| `normalize(actions)` | `ndarray` | `ndarray` | Raw → [-1, 1] |
| `denormalize(actions)` | `ndarray` | `ndarray` | [-1, 1] → raw |
| `split_action(action)` | `ndarray` | `Dict[str, ndarray]` | Split by joint group |
| `default_position()` | — | `ndarray` | Standing pose |
| `get_group_indices(group)` | `str` | `List[int]` | Indices for a group |

---

## Data

### `DataSoupDataset`

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

dataset = DataSoupDataset(config, action_space)
sample = dataset[0]
# sample["image"], sample["target_actions"], sample["language"], ...
```

### `ActionMapper`

```python
from neon.data.data_soup import ActionMapper

mapper = ActionMapper(target_space=action_space)
g1_actions = mapper.map_actions(franka_actions, "franka")
```

---

## Inference

### `NeonInferenceServer`

```python
from neon.inference.server import NeonInferenceServer

server = NeonInferenceServer("cagataydev/neon-g1-v1")
output = server.predict(image=frame, instruction="Pick up the cup")
server.reset()  # Clear smoothing state
```

### `G1Controller`

```python
from neon.inference.g1_controller import G1Controller, G1Config

controller = G1Controller(G1Config(robot_ip="192.168.123.10"))
controller.connect()
controller.run_control_loop(model, instruction="Pick up the cup", max_steps=200)
controller.disconnect()
```

---

## Training

### `NeonTrainer`

```python
from neon.training.train import NeonTrainer
from neon.training.config import TrainConfig

trainer = NeonTrainer(TrainConfig())
stats = trainer.train()
```

### Preset Configs

```python
from neon.training.config import default_arms_only_config, default_wholebody_config

config = default_arms_only_config()    # 3B backbone, arms_only, 3 epochs
config = default_wholebody_config()    # 7B backbone, whole_body, 5 epochs
```
