Add UR5 fine-tuning example#935
Open
enjoysport2022 wants to merge 1 commit into
Open
Conversation
- Add UR5Inputs and UR5Outputs transform classes in ur5_policy.py, mapping UR5 observations (6 joints + 1 gripper, base and wrist cameras) to the format expected by pi0/pi0-FAST/pi0.5 models - Add LeRobotUR5DataConfig and three training configs in config.py: pi0_ur5 (full fine-tune), pi0_ur5_low_mem_finetune (LoRA), pi0_fast_ur5 (pi0-FAST fine-tune) - Add convert_ur5_data_to_lerobot.py for converting raw UR5 recordings to LeRobot dataset format - Add unit tests covering state concatenation, image format conversion (float32 CHW -> uint8 HWC), image masks, and action slicing - Rewrite examples/ur5/README.md with step-by-step usage instructions Also fixes two issues in the original README code outline: - create_base_config() was missing the required model_config argument - UR5Inputs constructor referenced a non-existent action_dim field
enjoysport2022
requested review from
Michael-Equi,
jimmyt857 and
kvablack
as code owners
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR implements the UR5 fine-tuning example that was previously outlined in
examples/ur5/README.mdas pseudocode.Changes
UR5InputsandUR5Outputstransform classes insrc/openpi/policies/ur5_policy.py, mapping UR5 observations (6 joints + 1 gripper, base and wrist cameras) to the format expected by pi0 / pi0-FAST / pi0.5LeRobotUR5DataConfigand three ready-to-use training configs insrc/openpi/training/config.py:pi0_ur5(full fine-tune),pi0_ur5_low_mem_finetune(LoRA),pi0_fast_ur5examples/ur5/convert_ur5_data_to_lerobot.pyfor converting raw UR5 recordings to LeRobot dataset formatexamples/ur5/README.mdwith step-by-step usage instructionsBug fixes in original README outline
create_base_config()was missing the requiredmodel_configargumentUR5Inputsconstructor referenced a non-existentaction_dimfield (padding is handled byPadStatesAndActionsinModelTransformFactory)