| --- |
| 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() |
| ``` |
|
|
