IliaLarchenko/dot_pusht_keypoints

Model Card for "Decoder Only Transformer (DOT) Policy" for PushT keypoints dataset

Read more about the model and implementation details in the DOT Policy repository.

This model is trained using the LeRobot library and achieves state-of-the-art results on behavior cloning on the PushT keypoints dataset. It achieves 88.1% success rate (and 0.969 average max reward) vs. ~78% for the previous state-of-the-art model or 69% that I managed to reproduce using VQ-BET implementation in LeRobot.

This result is achieved without the checkpoint selection. If you are interested in an even better model with a success rate of ~94% (but harder to reproduce as it requires some parameters tuning and checkpoint selection), please refer to this model

You can use this model by installing LeRobot from this branch

To train the model:

python lerobot/scripts/train.py \
    --policy.type=dot \
    --dataset.repo_id=lerobot/pusht_keypoints \
    --env.type=pusht \
    --env.task=PushT-v0 \
    --output_dir=outputs/train/pusht_keyponts \
    --batch_size=24  \
    --log_freq=1000 \
    --eval_freq=10000 \
    --save_freq=50000 \
    --offline.steps=1000000 \
    --seed=100000 \
    --wandb.enable=true \
    --num_workers=24 \
    --use_amp=true \
    --device=cuda \
    --policy.return_every_n=2

To evaluate the model:

python lerobot/scripts/eval.py \
    --policy.path=IliaLarchenko/dot_pusht_keypoints \
    --env.type=pusht \
    --env.task=PushT-v0 \
    --eval.n_episodes=1000 \
    --eval.batch_size=100 \
    --env.obs_type=environment_state_agent_pos \
    --seed=1000000

Model size:

• Total parameters: 14.1m
• Trainable parameters: 2.9m

Note: the results are even slightly better than ones reported in the repository. There was a big update in the LeRobot library, I did all the original training and evaluation using the older version of LeRobot. This model was trained using the new version of the library with the same parameters. The older version of the library that was used for the original experiments is available in the this branch