File size: 1,745 Bytes
91802ec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | ---
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()
```
|