# --- Add this TrainConfig entry to openpi/src/openpi/training/config.py ---
# Place it inside the _CONFIGS list (anywhere; convention is near the
# pi05_real_pk_* family).
#
# Required imports (already present in config.py):
#   - TrainConfig, DataConfig, AssetsConfig, RealRobotDataConfig
#   - pi0_config, weight_loaders, _optimizer

TrainConfig(
    name="pi05_real_pk_mixed_new_v3",
    model=pi0_config.Pi0Config(pi05=True, action_dim=32, action_horizon=10),
    data=RealRobotDataConfig(
        # Eval/finetune-data repo_id (any LeRobot repo with the same schema works for
        # inference; the model is data-agnostic at inference time once norm_stats are loaded).
        repo_id="real_pass_knife_mixed_v2_train",
        # CRITICAL: asset_id MUST be set on AssetsConfig, not on base_config.
        # base_config.asset_id is silently overwritten by self.assets.asset_id during
        # create_base_config() — see openpi-asset-id-override skill / the line
        # `base_config = dataclasses.replace(base_config, asset_id=self.assets.asset_id)`
        # in DataConfigFactory.create_base_config().
        assets=AssetsConfig(asset_id="real_pass_knife_new_mixed_train"),
        base_config=DataConfig(prompt_from_task=True, action_sequence_keys=("action",)),
    ),
    batch_size=128,
    fsdp_devices=8,
    lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500, peak_lr=5e-5, decay_steps=8000, decay_lr=5e-6,
    ),
    optimizer=_optimizer.AdamW(clip_gradient_norm=1.0),
    weight_loader=weight_loaders.CheckpointWeightLoader(
        "gs://openpi-assets/checkpoints/pi05_base/params",
    ),
    num_train_steps=8000, save_interval=1000, keep_period=1000,
    val_repo_id="real_pass_knife_mixed_v2_eval",
    val_interval=500,
    save_best_val=True,
),