File size: 2,786 Bytes
e6d30a8 f7872bf e6d30a8 58570fc e6d30a8 58570fc e6d30a8 f7872bf e6d30a8 |
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: apache-2.0
datasets:
- lerobot/pusht_keypoints
pipeline_tag: robotics
---
# Model Card for Diffusion Policy / PushT (keypoints)
Diffusion Policy (as per [Diffusion Policy: Visuomotor Policy
Learning via Action Diffusion](https://arxiv.org/abs/2303.04137)) trained for the `PushT` environment from [gym-pusht](https://github.com/huggingface/gym-pusht) with keypoint-only observations.
Note: The original work trains keypoints-only with conditioning via inpainting. Here, we encode the observation along with the agent position and use the encoding as global conditioning for the denoising U-Net.
## How to Get Started with the Model
Use `python lerobot/scripts/eval.py -p lerobot/diffusion_pusht` to evaluate for 50 episodes with the outputs sent to `outputs/eval`.
For further information, please see the [LeRobot library](https://github.com/huggingface/lerobot) (particularly the [evaluation script](https://github.com/huggingface/lerobot/blob/main/lerobot/scripts/eval.py)).
## Training Details
Trained with [LeRobot@cc2f6e7](https://github.com/huggingface/lerobot/tree/cc2f6e74047bd65db0f9705fa602636b625bc28c).
The model was trained using [LeRobot's training script]([TODO link with commit hash one the PR is merged](https://github.com/huggingface/lerobot/blob/cc2f6e74047bd65db0f9705fa602636b625bc28c/lerobot/scripts/train.py)) and with the [pusht_keypoints](https://huggingface.co/datasets/lerobot/pusht_keypoints/tree/v1.5) dataset, using this command:
```bash
python lerobot/scripts/train.py \
hydra.job.name=diffusion_pusht_keypoints \
hydra.run.dir=outputs/train/2024-07-03/13-52-44_diffusion_pusht_keypoints \
env=pusht_keypoints \
policy=diffusion_pusht_keypoints \
training.save_checkpoint=true \
training.offline_steps=200000 \
training.save_freq=20000 \
training.eval_freq=10000 \
training.log_freq=50 \
training.num_workers=4 \
eval.n_episodes=50 \
eval.batch_size=50 \
wandb.enable=true \
wandb.disable_artifact=true \
device=cuda \
use_amp=true
```
The training curves may be found at https://wandb.ai/alexander-soare/lerobot/runs/5z9d8q9q/overview.
This took about 5 hours to train on an Nvida RTX H100.
## Evaluation
The model was evaluated on the `PushT` environment from [gym-pusht](https://github.com/huggingface/gym-pusht). There are two evaluation metrics on a per-episode basis:
- Maximum overlap with target (seen as `eval/avg_max_reward` in the charts above). This ranges in [0, 1].
- Success: whether or not the maximum overlap is at least 95%.
Here are the metrics for 500 episodes worth of evaluation.
Metric|Average over 500 episodes
-|-
Average max. overlap ratio | 0.97
Success rate (%) | 71.0
The results of each of the individual rollouts may be found in [eval_info.json](eval_info.json). |