Add ManiSkill to LeRobot dataset converter

[dancher00]

Converts .npz files to LeRobot format

Closes #545

[StoneT2000]

npz files? We don't have those kind?

thanks for the PR though! Would be great to have this integration so people can use Lerobot built in algorithms on maniskill sim and data

[dancher00]

@StoneT2000 Thanks for the feedback!

I initially converted HDF5 → NPZ for convenience during development, but I can update the converter to work directly with HDF5 files. I can update the PR accordingly though. The core conversion logic stays the same, just need to change the data loading part

[StoneT2000]

Yeah we would just need something for the default .h5 files generated by our RecordEpisode wrapper. Feel free to request a review from me once done.

[sijitu666]

@StoneT2000 Thanks for the feedback!

I initially converted HDF5 → NPZ for convenience during development, but I can update the converter to work directly with HDF5 files. I can update the PR accordingly though. The core conversion logic stays the same, just need to change the data loading part

Hi @dancher00,

Just checking in on this PR. This converter looks incredibly useful! Is the updated version that works directly with HDF5 (.h5) files available somewhere? I'm eager to try it out for my project.

Thanks for the great contribution!

[dancher00]

@StoneT2000 Apologies for the long delay! I've updated the converter to work with HDF5 files as you requested.

[dancher00]

@sijitu666 Sorry for the delay! Thanks for checking in on this - really appreciate the interest.

The HDF5 converter is ready now. I'm currently using it to train policies for the SO-101 arm and it's working well.

Check the updated code and README in this PR for usage instructions. Let me know if you have any questions!

[StoneT2000]

Thank you @dancher00

A few requests:

Can this code be moved to the same folder as our trajectory replay tool? https://github.com/haosulab/ManiSkill/tree/main/mani_skill/trajectory

And delete the tools folder you made.

Moreover, we should move the (excellent) readme to the docs folder. We should have a page on our docs discussing trajectory replay, just add a section for lerobot format exports and add your content there.

Lastly it would be great if you could share a sample training command to test the lerobot format exports.

[dancher00]

Hi @StoneT2000 !

I've added docs to ManiSkill/docs/source/user_guide/datasets/replay.md

Here is quick pipline for test

1. Download demo data

python -m mani_skill.utils.download_demo "PickCube-v1"

2. Replay trajectory with RGB observations

python -m mani_skill.trajectory.replay_trajectory \
  --traj-path ~/.maniskill/demos/PickCube-v1/motionplanning/trajectory.h5 \
  --use-first-env-state \
  -c pd_joint_pos \
  -o rgbd \
  --save-traj \
  -n 1 \
  --count 10

3. Convert to LeRobot format

python -m mani_skill.trajectory.convert_to_lerobot \
  ~/.maniskill/demos/PickCube-v1/motionplanning/trajectory.rgbd.pd_joint_pos.physx_cpu.h5 \
  /tmp/lerobot_pickcube \
  --task-name "Pick up the red cube"
  --robot-type "panda"

You will get this - ManiSkill Lerobot Dataset

4. Training Test

lerobot-train \
  --policy.path=lerobot/smolvla_base \
  --dataset.repo_id=maniskill-pickcube \
  --dataset.root=/tmp/lerobot_pickcube \
  --batch_size=4 \
  --steps=100 \
  --output_dir=/tmp/test_training \
  --policy.device=cuda \
  --wandb.enable=false

[dancher00]

@StoneT2000 Hi! Just a friendly ping on this PR when you have a moment. Thanks!

[StoneT2000]

Sorry, lot of stuff happening recently! Promise to review and test today, thanks again!