Skip to content

A TrainConfig about PI0.5 when use the JointPos as actions on Lerobot dataset #933

@XLINYIN

Description

@XLINYIN

This is a trainconfig I modified myself to further fine-tune based on PI0.5_Droid_JointPos.
I used the Lerobot dataset format and followed the tutorial in example/Droid on customizing a small-scale droid dataset, simply changing velocity in the action to joint_position.
However, I found that PI0.5Droid_jointpos didn't perform well initially on the small dataset generated in the examples, with a loss of around 3. After 150 epochs of training, it decreased to 1.5. I don't know if this is due to my settings or if the model parameters themselves were not sufficiently trained. Because on the same dataset, the initial loss using joint velocity parameters is around 1.3.
Below is the configuration I wrote. The JointPosFlag is an extra parameter I added to LeRobotDROIDDataConfig to control whether to use bool_mask(7,-1).

TrainConfig(
name="pi05_droid_jointpos_Lerobot_trans",
model=pi0_config.Pi0Config(action_dim=32, action_horizon=15, pi05=True),
data=LeRobotDROIDDataConfig(
# Replace with your custom DROID LeRobot dataset repo id.
repo_id="HF_TEST_LEROBOT/my_droid_dataset_jointpos",
base_config=DataConfig(prompt_from_task=True),
assets=AssetsConfig(
# Important: reuse the original DROID norm stats during fine-tuning!
assets_dir="gs://openpi-assets/checkpoints/pi05_droid_jointpos/assets",
asset_id="droid",
),
JointPosFlag=True
),
weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_droid_jointpos/params"),
num_train_steps=20_000,
batch_size=32
)

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions