Issues with Reproducing pi0-SF on Real Hardware

[CortexSphere]

Previously, I successfully trained and deployed pi0.5-SF on the simulation dataset LIBERO PLUS, improving the success rate from 61.1% for pi0.5 to 75.2%. Now, when attempting to deploy this algorithm on real hardware using the default simulation parameters, I encounter situations where tasks that pi0.5 can handle perfectly result in a 0% success rate for pi0.5-SF. I would like to know if you and your team have any suggestions on this, such as parameter selection, dataset scale, etc.?

[haofuly]

Hi @CortexSphere

1. Have you ever update your code to branch 9a7e590 or the newer one? There is an important bug before, which may lead to collapse.
2. Have the action_loss converged to 0.01 or smaller, together with the grad_norm maintaining smaller than 1?
3. BTW, would you like to share your pi0.5-SF config and codebase modification for LIBERO PLUS? It is quite meaningful for our community, thanks!

[CortexSphere]

1.Of course, as mentioned earlier, I have successfully trained on LIBERO PLUS and achieved a good improvement in success rate.
2.Information related to real hardware training:

3.My pi05-sf-libero_plus configuration config.py:

@dataclasses.dataclass(frozen=True)
class LeRobotLiberoPlusDataConfig(LeRobotLiberoDataConfig):@override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        # Start with the base repack from parent, but override the prompt mapping to use 'task' from LIBERO-Plus dataset
        repack_transform = _transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "observation/image": "image",
                        "observation/wrist_image": "wrist_image",
                        "observation/state": "state",
                        "actions": "actions",
                        "prompt": "task",  # Key change: LIBERO-Plus uses 'task' for language instruction
                    }
                )
            ]
        )

        # Reuse the rest from parent (data_transforms, model_transforms, etc.)
        base_config = super().create(assets_dirs, model_config)
        return dataclasses.replace(
            base_config,
            repack_transforms=repack_transform,
        )

    TrainConfig(
        name="pi05_sf_libero_plus",
        fsdp_devices=4,
        model=pi0_config.Pi0Config(pi05=True, action_horizon=10, discrete_state_input=False),
        data=LeRobotLiberoPlusDataConfig(  # Use the new Plus config class
            repo_id="/mnt/disk1/BigBaobei/mynvme1/cyc/hf_cache/lerobot/cyc/libero_plus",  # Your local converted dataset path
            base_config=DataConfig(prompt_from_task=True),
            extra_delta_transform=False,  # Actions in LIBERO-Plus are already deltas, like original
        ),
        batch_size=32,
        lr_schedule=_optimizer.CosineDecaySchedule(
            warmup_steps=10_000,
            peak_lr=5e-5,
            decay_steps=1_000_000,
            decay_lr=5e-5,
        ),
        optimizer=_optimizer.AdamW(clip_gradient_norm=1.0),
        ema_decay=0.999,
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        pytorch_weight_path="/home/robotlab/.cache/openpi/openpi-assets/checkpoints/pi05_base.pth",
        num_train_steps=30_000,
        save_interval=5_000,
        #VGGT
        vggt_weight_path='/mnt/mynvme0/cyc/hf_cache/hub/models--facebook--VGGT-1B/snapshots/860abec7937da0a4c03c41d3c269c366e82abdf9',
        vla_layers_align=12,
        vggt_layers_align=-1,
        pooling_func='bilinear',
        use_vggt_pe=True,
        use_vlm_norm=True,
        align_loss_coeff=0.5,
    ),

Regarding how to integrate LIBERO PLUS, you can refer to https://github.com/sylvestf/LIBERO-plus. I simplified the original benchmark by taking 1/10 of it.

[haofuly]

As I use agilex cobot as my real-world robot, my dataset scale for each task is 40 demonstrations, which works quite well. At step0 during training period, mostly the action_loss≈0.2 and grad_norm≈3.

Your training curve seems stable, so I am not sure where the bug lies. Could you make sure that the pi05 and pi05-sf start at and end at the similiar level action_loss and grad_norm? Besides, is the batch_size setting of pi05-sf the same as that of pi05?

Thanks for your LIBERO PLUS contribution.