Questions while reproducing the results

[Garsonjw]

Thank you very much for your great work. I encountered several questions while reproducing the results and would really appreciate your clarification:

1. Regarding reusing normalization, I would like to confirm whether this means copying the norm_stats.json from the corresponding pre-training model into
   openpi-comet/outputs/assets/train/pi05_b1k-turning_on_radio_lr2.5e-6_step20k_sft,
   and using it directly during fine-tuning.

2. When training a single task (turning_on_radio) with pi05-b1kpt12-cs32, I observed that the training loss initially increases and then gradually decreases.
   Is this behavior expected or normal in this setting?

3. I noticed that enabling
   env_wrapper._target_=behavior.learning.wrappers.RGBWrapper
   does not seem to significantly affect the resolution. From the saved videos, I cannot observe an obvious difference.
   Could you clarify what changes this wrapper is expected to introduce, and whether the effect should be visible in the rendered videos?

Thank you very much for your time and help.

[Garsonjw]

Additionally, I have another question.

I’m fine-tuning pi05-b1kpt12-cs32 on COMET-1.5k, using only the turning_on_radio task, and I noticed that the loss curve looks quite unusual during training. I’m not sure whether I might have missed something in the configuration or setup, so I’d like to ask for some advice.

Below is my exact training configuration, kept unchanged:

TrainConfig(
    name="pi05_b1k-turning_on_radio_lr2.5e-6_step20k_sft",
    exp_name="openpi",
    project_name="B1K",
    model=pi0_config.Pi0Config(pi05=True, action_horizon=32),
    data=LeRobotB1KDataConfig(
        repo_id="behavior-1k/2025-challenge-demos",
        base_config=DataConfig(
            prompt_from_task=True,
            behavior_dataset_root="/mnt/projects/datasets/comet-1.5k", 
            tasks=["turning_on_radio"],
            fine_grained_level=0,  # 0, 1, 2
        ),
    ),
    weight_loader=weight_loaders.CheckpointWeightLoader("/mnt/project/openpai_comet_model/pi05-b1kpt12-cs32/params"),
    num_train_steps=20_000,
    lr_schedule=_optimizer.CosineDecaySchedule(
        peak_lr=2.5e-6,
        decay_steps=20_000,
    ),
    freeze_filter=pi0_config.Pi0Config(pi05=True, action_horizon=32).get_freeze_filter(),
    ema_decay=None,
    checkpoint_base_dir=".",
    num_workers=8,
    batch_size=8 * 32,
)

Below are partial training logs, also kept exactly the same:

13:16:29.676 [I] Loaded norm stats from /root/projects/openpi-comet/outputs/assets/train/pi05_b1k-turning_on_radio_lr2.[107/3410] 
0k_sft/behavior-1k/2025-challenge-demos (3836088:config.py:272)                                                                  
13:16:29.687 [I] PyTorch version 2.7.1 available.                                                 (3836088:config.py:54)         
13:16:29.688 [I] Polars version 1.30.0 available.                                                 (3836088:config.py:66)         
13:16:29.689 [I] JAX version 0.5.3 available.                                                     (3836088:config.py:125)        
13:16:30.775 [I] Loaded metadata for 120 episodes.                                                (3836088:dataset.py:644)       
13:16:30.782 [I] Total episodes: 120                                                              (3836088:dataset.py:205)       
Resolving data files: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████
████████████████████████████████| 120/120 [00:00<00:00, 319160.74it/s]                                                           
prepare task with fine_grained_level 0 for /mnt/projects/datasets/comet-1.5k                            
13:17:59.108 [I] Initialized data loader:                                                                                        
[0].images['base_0_rgb']: (256, 224, 224, 3)@float32                                                                             
[0].images['left_wrist_0_rgb']: (256, 224, 224, 3)@float32                                                                       
[0].images['right_wrist_0_rgb']: (256, 224, 224, 3)@float32                                                                      
[0].image_masks['base_0_rgb']: (256,)@bool                                                                                       
[0].image_masks['left_wrist_0_rgb']: (256,)@bool                                                                                 
[0].image_masks['right_wrist_0_rgb']: (256,)@bool                                                                                
[0].state: (256, 32)@float32                                                                                                     
[0].tokenized_prompt: (256, 200)@int32                          
[0].tokenized_prompt_mask: (256, 200)@bool                      
[1]: (256, 32, 32)@float32 (3836088:train.py:253)                                                                                
13:17:59.599 [I] Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOpti
ons(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x7f2819
b2b6d0> (3836088:base_pytree_checkpoint_handler.py:334)         
13:18:00.030 [I] Restoring checkpoint from /mnt/project/openpai_comet_model/pi05-b1kpt12-cs32/params. (3836088:c
heckpointer.py:298)
13:18:13.491 [I] [process=0] /jax/checkpoint/read/bytes_per_sec: 950.3 MiB/s (total bytes: 12.5 GiB) (time elapsed: 13 seconds) (
per-host) (3836088:base_pytree_checkpoint_handler.py:114)
13:18:13.492 [I] Finished restoring checkpoint in 13.46 seconds from /mnt/project/openpai_comet_model/pi05-b1kpt
12-cs32/params. (3836088:checkpointer.py:309)
2026-01-08 13:18:13.865483: W external/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc:3021] Can't reduce memory use 
below 32.09GiB (34461143882 bytes) by rematerialization; only reduced to 37.48GiB (40241207884 bytes), down from 37.48GiB (402412
07884 bytes) originally         
13:18:25.302 [I] Initialized train state:

From my understanding, the setup itself looks reasonable, but the loss curve during training does not behave as expected. I’m not sure whether this is normal for fine-tuning on a small, single-task subset like this, or if there is something subtle I might be missing.

I’d really appreciate any insights or suggestions.
Thank you very much for your help!