|
|
--- |
|
|
license: mit |
|
|
tags: |
|
|
- robotics |
|
|
- act |
|
|
- vla |
|
|
- gearbox-assembly |
|
|
- imitation-learning |
|
|
metrics: |
|
|
- l1_loss |
|
|
- kl_divergence |
|
|
--- |
|
|
|
|
|
# ACT Policy for Gearbox Assembly (Filtered Demos) |
|
|
|
|
|
This model is an **Action Chunking with Transformers (ACT)** policy trained to perform gearbox assembly tasks. It was trained using behavior cloning on a combination of `rocochallenge2025` and additional collected datasets. |
|
|
|
|
|
## Model Details |
|
|
|
|
|
- **Model Type**: ACT (Action Chunking with Transformers) |
|
|
- **Policy Class**: ACT |
|
|
- **Backbone**: ResNet-18 |
|
|
- **Training Dataset**: Integrated dataset (`rocochallenge2025` + `temp_new_dataset`) containing **241 episodes**. |
|
|
- **Episode Length**: Fixed to **12,600 steps** (padded/truncated) to handle variable length recordings. |
|
|
|
|
|
## Training Configuration |
|
|
- **Task Name**: `sim_gearbox_assembly_demos_filtered` |
|
|
- **Batch Size**: 32 |
|
|
- **Chunk Size (Action Horizon)**: 100 |
|
|
- **KL Weight**: 10 |
|
|
- **Hidden Dimension**: 512 |
|
|
- **Feedforward Dimension**: 3200 |
|
|
- **Learning Rate**: 1e-5 |
|
|
- **Num Epochs**: ~9500 (Early stopped/Interrupted) |
|
|
- **Seed**: 0 |
|
|
|
|
|
## Inputs and Outputs |
|
|
- **Observations**: |
|
|
- `head_rgb` (240x320) |
|
|
- `left_hand_rgb` (240x320) |
|
|
- `right_hand_rgb` (240x320) |
|
|
- `qpos` (14-dim joint positions) |
|
|
- **Actions**: |
|
|
- 14-dim combined action vector (7-dim left arm + 7-dim right arm) |
|
|
|
|
|
## Usages |
|
|
|
|
|
This model can be loaded using the `ACTPolicy` class. Ensure `dataset_stats.pkl` is loaded to normalize/unnormalize observations and actions correctly. |
|
|
|
|
|
```python |
|
|
from policy import ACTPolicy |
|
|
import pickle |
|
|
|
|
|
# Load stats |
|
|
with open('dataset_stats.pkl', 'rb') as f: |
|
|
stats = pickle.load(f) |
|
|
|
|
|
# Load policy |
|
|
policy = ACTPolicy(config) |
|
|
policy.load_state_dict(torch.load('policy_best.ckpt')) |
|
|
policy.cuda() |
|
|
policy.eval() |
|
|
``` |
|
|
|
