Question about fine-tuning.

[Yyb-XJTU]

Hi, thanks for the amazing work!

I followed your repository to fine-tune our pi05 model on our custom dataset.

I wanted to double-check if the training is proceeding correctly. Here is the loss curve I got:

Does this convergence trend look normal to you?

Thanks for your time!

[Sakits]

Hi @Yyb-XJTU ,

Thank you very much for your interest in our work!

From the curves you attached, the training loss trend doesn’t look quite normal. In our experiments, training for around 30k–50k steps is usually already enough on most datasets.

Here is an example:

Could you please share your full training config so we can help?

Best,
Jiaming

[Yyb-XJTU]

Hi @Sakits

Thanks for the guidance! It's good to know that 30k-50k steps should be sufficient.

To give you more context on my workflow, I am working with the official OpenPi implementation. Since OpenPi is JAX-based, I converted the relevant parts from JAX to PyTorch, modified the YAML config in your repository, and then trained directly using your training code.

I suspect the issue might stem from this conversion or the configuration adaptation.

Here is the full config I used:

# VLASH Training Configuration for PI05 with Async Inference

# Policy configuration
policy:
  type: pi05
  pretrained_path: [path]
  push_to_hub: false
  dtype: bfloat16
  device: cuda
  state_cond: true

# Dataset configuration
dataset:
  repo_id: [repo_id]
  video_backend: pyav

# Training parameters
output_dir: outputs/train/debug
job_name: debug
batch_size: 8
steps: 100000
num_workers: 8
seed: 1000

# Optimizer configuration
use_policy_training_preset: false
optimizer:
  type: adamw
  lr: 5.0e-5
  betas: [0.9, 0.95]
  weight_decay: 1.0e-10

# Learning rate scheduler
scheduler:
  type: cosine_decay_with_warmup
  num_warmup_steps: 1000
  peak_lr: 5.0e-5
  decay_lr: 2.5e-6
  num_decay_steps: 50000

# Checkpointing
save_checkpoint: true
save_freq: 10000

# Logging
log_freq: 200
wandb:
  enable: false
  project: vlash
  entity: null

# Maximum temporal delay steps
# Set to N to enable random delays from [0, N]
max_delay_steps: 8

Given this setup, do you see any obvious issues with the config, or should I double-check the data processing/conversion?

Best,
Yibo

[Sakits]

Hi @Yyb-XJTU ,

Thanks for sharing the full config and the details of your setup!

The configuration looks good to me. The only minor things I would suggest to modify are:

• Batch size is a bit small: if your GPU allows, increasing it to 16 / 32 usually gives more stable actions for the tuned model.
• video_backend: using torchcodec is significantly faster than pyav.

A few more information might help find the reason:
Is your dataset collected using the LeRobot repo?
Could you try running a training job with the default config from our repo (without the JAX→PyTorch conversion) on your dataset and see whether the loss curve looks normal?

The PI0.5 weights from LeRobot in the config is already converted from the official PI weights and cross-checked, so in principle the logic should be fully consistent with OpenPi and shouldn’t require any manual conversion.
If training with the default setup works fine, then the problem is likely in the custom conversion; if not, we can dig deeper into the data or training logs. :)

Best,
Jiaming

[SkieyFly]

Hi @Yyb-XJTU ,
And may I ask whether you are using absolute actions or delta actions in your dataset?

[Yyb-XJTU]

Hi @SkieyFly ,
I am using absolute actions in my dataset, and for context, I adapted the official OpenPi implementation from JAX to PyTorch.

[Yyb-XJTU]

Hi @Sakits ,
Here is my ckpt's config.json for reference:

{
  "action_dim": 32,
  "action_horizon": 50,
  "paligemma_variant": "gemma_2b",
  "action_expert_variant": "gemma_300m",
  "precision": "bfloat16"
}

[Yyb-XJTU]

If convenient, I’d love to add your WeChat to discuss this further. But if you prefer to keep the conversation here on GitHub/Email, that is completely fine too!

My WeChat ID: yyb1013150902