"""See _CONFIGS for the list of available configs."""

import abc
from collections.abc import Sequence
import dataclasses
import difflib
import logging
import pathlib
from typing import Any, Literal, Protocol, TypeAlias

import etils.epath as epath
import flax.nnx as nnx
from typing_extensions import override
import tyro

import openpi.models.model as _model
import openpi.models.pi0_config as pi0_config
import openpi.models.pi0_fast as pi0_fast
import openpi.models.tokenizer as _tokenizer
import openpi.policies.aloha_policy as aloha_policy
import openpi.policies.b601_policy as b601_policy
import openpi.policies.droid_policy as droid_policy
import openpi.policies.libero_policy as libero_policy
import openpi.policies.ur3_policy as ur3_policy
import openpi.policies.ur3_bimanual_policy as ur3_bimanual_policy
import openpi.policies.unitree_policy as unitree_policy
import openpi.policies.unitree_eef_policy as unitree_eef_policy
import openpi.policies.unitree_eef6d_policy as unitree_eef6d_policy
import openpi.shared.download as _download
import openpi.shared.normalize as _normalize
import openpi.training.droid_rlds_dataset as droid_rlds_dataset
import openpi.training.misc.polaris_config as polaris_config
import openpi.training.misc.roboarena_config as roboarena_config
import openpi.training.optimizer as _optimizer
import openpi.training.weight_loaders as weight_loaders
import openpi.transforms as _transforms

ModelType: TypeAlias = _model.ModelType
# Work around a tyro issue with using nnx.filterlib.Filter directly.
Filter: TypeAlias = nnx.filterlib.Filter


@dataclasses.dataclass(frozen=True)
class AssetsConfig:
    """Determines the location of assets (e.g., norm stats) that will be used to set up the data pipeline.

    These assets will be replicated inside the checkpoint under the `assets/asset_id` directory.

    This can be used to load assets from a different checkpoint (e.g., base model checkpoint) or some other
    centralized location. For example, to load the norm stats for the Trossen robot from the base model checkpoint
    during fine-tuning, use:

    ```
    AssetsConfig(
        assets_dir="gs://openpi-assets/checkpoints/pi0_base/assets",
        asset_id="trossen",
    )
    ```
    """

    # Assets directory. If not provided, the config assets_dirs will be used. This is useful to load assets from
    # a different checkpoint (e.g., base model checkpoint) or some other centralized location.
    assets_dir: str | None = None

    # Asset id. If not provided, the repo id will be used. This allows users to reference assets that describe
    # different robot platforms.
    asset_id: str | None = None


@dataclasses.dataclass(frozen=True)
class DataConfig:
    # LeRobot repo id. If None, fake data will be created.
    repo_id: str | None = None
    # Directory within the assets directory containing the data assets.
    asset_id: str | None = None
    # Contains precomputed normalization stats. If None, normalization will not be performed.
    norm_stats: dict[str, _transforms.NormStats] | None = None

    # Used to adopt the inputs from a dataset specific format to a common format
    # which is expected by the data transforms.
    repack_transforms: _transforms.Group = dataclasses.field(default_factory=_transforms.Group)
    # Data transforms, typically include robot specific transformations. Will be applied
    # before the data is normalized. See `model.Observation` and `model.Actions` to learn about the
    # normalized data.
    data_transforms: _transforms.Group = dataclasses.field(default_factory=_transforms.Group)
    # Model specific transforms. Will be applied after the data is normalized.
    model_transforms: _transforms.Group = dataclasses.field(default_factory=_transforms.Group)
    # If true, will use quantile normalization. Otherwise, normal z-score normalization will be used.
    use_quantile_norm: bool = False

    # Names of keys that will be used by the data loader to generate the action sequence. The length of the
    # sequence is defined by the `action_horizon` field in the model config. This should be adjusted if your
    # LeRobot dataset is using different keys to represent the action.
    action_sequence_keys: Sequence[str] = ("actions",)

    # If true, will use the LeRobot dataset task to define the prompt.
    prompt_from_task: bool = False

    # Local root directory for LeRobot datasets stored on disk.
    # When set, LeRobotDataset loads from `local_root / repo_id` instead of
    # downloading from HuggingFace Hub.  Useful during development when the
    # dataset has not been uploaded yet.
    local_root: pathlib.Path | None = None

    # Only used for RLDS data loader (ie currently only used for DROID).
    rlds_data_dir: str | None = None
    # Action space for DROID dataset.
    action_space: droid_rlds_dataset.DroidActionSpace | None = None
    # List of datasets to sample from: name, version, weight, and optionally filter_dict_path
    datasets: Sequence[droid_rlds_dataset.RLDSDataset] = ()


class GroupFactory(Protocol):
    def __call__(self, model_config: _model.BaseModelConfig) -> _transforms.Group:
        """Create a group."""


@dataclasses.dataclass(frozen=True)
class ModelTransformFactory(GroupFactory):
    """Creates model transforms for standard pi0 models."""

    # If provided, will determine the default prompt that be used by the model.
    default_prompt: str | None = None

    def __call__(self, model_config: _model.BaseModelConfig) -> _transforms.Group:
        match model_config.model_type:
            case _model.ModelType.PI0:
                return _transforms.Group(
                    inputs=[
                        _transforms.InjectDefaultPrompt(self.default_prompt),
                        _transforms.ResizeImages(224, 224),
                        _transforms.TokenizePrompt(
                            _tokenizer.PaligemmaTokenizer(model_config.max_token_len),
                        ),
                        _transforms.PadStatesAndActions(model_config.action_dim),
                    ],
                )
            case _model.ModelType.PI05:
                assert isinstance(model_config, pi0_config.Pi0Config)
                return _transforms.Group(
                    inputs=[
                        _transforms.InjectDefaultPrompt(self.default_prompt),
                        _transforms.ResizeImages(224, 224),
                        _transforms.TokenizePrompt(
                            _tokenizer.PaligemmaTokenizer(model_config.max_token_len),
                            discrete_state_input=model_config.discrete_state_input,
                        ),
                        _transforms.PadStatesAndActions(model_config.action_dim),
                    ],
                )
            case _model.ModelType.PI0_FAST:
                tokenizer_cls = (
                    _tokenizer.FASTTokenizer
                    if model_config.fast_model_tokenizer is None
                    else model_config.fast_model_tokenizer
                )
                tokenizer_kwargs = (
                    {} if model_config.fast_model_tokenizer_kwargs is None else model_config.fast_model_tokenizer_kwargs
                )
                return _transforms.Group(
                    inputs=[
                        _transforms.InjectDefaultPrompt(self.default_prompt),
                        _transforms.ResizeImages(224, 224),
                        _transforms.TokenizeFASTInputs(
                            tokenizer_cls(model_config.max_token_len, **tokenizer_kwargs),
                        ),
                    ],
                    outputs=[
                        _transforms.ExtractFASTActions(
                            tokenizer_cls(model_config.max_token_len, **tokenizer_kwargs),
                            action_horizon=model_config.action_horizon,
                            action_dim=model_config.action_dim,
                        )
                    ],
                )


@dataclasses.dataclass(frozen=True)
class DataConfigFactory(abc.ABC):
    # The LeRobot repo id.
    repo_id: str = tyro.MISSING
    # Determines how the assets will be loaded.
    assets: AssetsConfig = dataclasses.field(default_factory=AssetsConfig)
    # Base config that will be updated by the factory.
    base_config: tyro.conf.Suppress[DataConfig | None] = None

    @abc.abstractmethod
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        """Create a data config."""

    def create_base_config(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        repo_id = self.repo_id if self.repo_id is not tyro.MISSING else None
        asset_id = self.assets.asset_id or repo_id
        return dataclasses.replace(
            self.base_config or DataConfig(),
            repo_id=repo_id,
            asset_id=asset_id,
            norm_stats=self._load_norm_stats(epath.Path(self.assets.assets_dir or assets_dirs), asset_id),
            use_quantile_norm=model_config.model_type != ModelType.PI0,
        )

    def _load_norm_stats(self, assets_dir: epath.Path, asset_id: str | None) -> dict[str, _transforms.NormStats] | None:
        if asset_id is None:
            return None
        try:
            data_assets_dir = str(assets_dir / asset_id)
            norm_stats = _normalize.load(_download.maybe_download(data_assets_dir))
            logging.info(f"Loaded norm stats from {data_assets_dir}")
            return norm_stats
        except FileNotFoundError:
            logging.info(f"Norm stats not found in {data_assets_dir}, skipping.")
        return None


@dataclasses.dataclass(frozen=True)
class FakeDataConfig(DataConfigFactory):
    repo_id: str = "fake"

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        return DataConfig(repo_id=self.repo_id)


@dataclasses.dataclass(frozen=True)
class SimpleDataConfig(DataConfigFactory):
    # Factory for the data transforms.
    data_transforms: tyro.conf.Suppress[GroupFactory] = dataclasses.field(default_factory=GroupFactory)
    # Factory for the model transforms.
    model_transforms: tyro.conf.Suppress[GroupFactory] = dataclasses.field(default_factory=ModelTransformFactory)

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            data_transforms=self.data_transforms(model_config),
            model_transforms=self.model_transforms(model_config),
        )


@dataclasses.dataclass(frozen=True)
class LeRobotAlohaDataConfig(DataConfigFactory):
    # If true, will convert joint dimensions to deltas with respect to the current state before passing to the model.
    # Gripper dimensions will remain in absolute values.
    use_delta_joint_actions: bool = True
    # If provided, will be injected into the input data if the "prompt" key is not present.
    default_prompt: str | None = None
    # If true, this will convert the joint and gripper values from the standard Aloha space to
    # the space used by the pi internal runtime which was used to train the base model. People who
    # use standard Aloha data should set this to true.
    adapt_to_pi: bool = True

    # Repack transforms.
    repack_transforms: tyro.conf.Suppress[_transforms.Group] = dataclasses.field(
        default=_transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "images": {"cam_high": "observation.images.top"},
                        "state": "observation.state",
                        "actions": "action",
                    }
                )
            ]
        )
    )
    # Action keys that will be used to read the action sequence from the dataset.
    action_sequence_keys: Sequence[str] = ("action",)

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        data_transforms = _transforms.Group(
            inputs=[aloha_policy.AlohaInputs(adapt_to_pi=self.adapt_to_pi)],
            outputs=[aloha_policy.AlohaOutputs(adapt_to_pi=self.adapt_to_pi)],
        )
        if self.use_delta_joint_actions:
            delta_action_mask = _transforms.make_bool_mask(6, -1, 6, -1)
            data_transforms = data_transforms.push(
                inputs=[_transforms.DeltaActions(delta_action_mask)],
                outputs=[_transforms.AbsoluteActions(delta_action_mask)],
            )

        model_transforms = ModelTransformFactory(default_prompt=self.default_prompt)(model_config)

        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            repack_transforms=self.repack_transforms,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
            action_sequence_keys=self.action_sequence_keys,
        )


@dataclasses.dataclass(frozen=True)
class LeRobotLiberoDataConfig(DataConfigFactory):
    """
    This config is used to configure transforms that are applied at various parts of the data pipeline.
    For your own dataset, you can copy this class and modify the transforms to match your dataset based on the
    comments below.
    """

    extra_delta_transform: bool = False

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        # The repack transform is *only* applied to the data coming from the dataset,
        # and *not* during inference. We can use it to make inputs from the dataset look
        # as close as possible to those coming from the inference environment (e.g. match the keys).
        # Below, we match the keys in the dataset (which we defined in the data conversion script) to
        # the keys we use in our inference pipeline (defined in the inference script for libero).
        # For your own dataset, first figure out what keys your environment passes to the policy server
        # and then modify the mappings below so your dataset's keys get matched to those target keys.
        # The repack transform simply remaps key names here.
        repack_transform = _transforms.Group(
          inputs=[
            _transforms.RepackTransform(
              {
                "observation/image":       "observation.images.image",
                "observation/wrist_image": "observation.images.wrist_image",
                "observation/state":       "observation.state",
                "actions":                 "actions",
                "prompt":                  "prompt",
              }
            )
        ]
      )

        # The data transforms are applied to the data coming from the dataset *and* during inference.
        # Below, we define the transforms for data going into the model (``inputs``) and the transforms
        # for data coming out of the model (``outputs``) (the latter is only used during inference).
        # We defined these transforms in `libero_policy.py`. You can check the detailed comments there for
        # how to modify the transforms to match your dataset. Once you created your own transforms, you can
        # replace the transforms below with your own.
        data_transforms = _transforms.Group(
            inputs=[libero_policy.LiberoInputs(model_type=model_config.model_type)],
            outputs=[libero_policy.LiberoOutputs()],
        )

        # One additional data transform: pi0 models are trained on delta actions (relative to the first
        # state in each action chunk). IF your data has ``absolute`` actions (e.g. target joint angles)
        # you can uncomment the following line to convert the actions to delta actions. The only exception
        # is for the gripper actions which are always absolute.
        # In the example below, we would apply the delta conversion to the first 6 actions (joints) and
        # leave the 7th action (gripper) unchanged, i.e. absolute.
        # In Libero, the raw actions in the dataset are already delta actions, so we *do not* need to
        # apply a separate delta conversion (that's why it's commented out). Choose whether to apply this
        # transform based on whether your dataset uses ``absolute`` or ``delta`` actions out of the box.

        # LIBERO already represents actions as deltas, but we have some old Pi0 checkpoints that are trained with this
        # extra delta transform.
        if self.extra_delta_transform:
            delta_action_mask = _transforms.make_bool_mask(6, -1)
            data_transforms = data_transforms.push(
                inputs=[_transforms.DeltaActions(delta_action_mask)],
                outputs=[_transforms.AbsoluteActions(delta_action_mask)],
            )

        # Model transforms include things like tokenizing the prompt and action targets
        # You do not need to change anything here for your own dataset.
        model_transforms = ModelTransformFactory()(model_config)

        # We return all data transforms for training and inference. No need to change anything here.
        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            repack_transforms=repack_transform,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
        )

@dataclasses.dataclass(frozen=True)
class LeRobotUnitreeG1DataConfig(DataConfigFactory):
    """Unitree G1 Dex1 bimanual dataset config.
    16-DoF: L_arm(7) + R_arm(7) + L_grip(1) + R_grip(1).
    Three cameras: left_high (base), left_wrist, right_wrist.

    Set action_dim=12 for G1_23 (5-DoF arms: L_arm(5) + R_arm(5) + L_grip + R_grip)."""

    action_dim: int = unitree_policy.ACTION_DIM

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        repack_transform = _transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "observation/image":             "observation.images.cam_left_high",
                        "observation/left_wrist_image":  "observation.images.cam_left_wrist",
                        "observation/right_wrist_image": "observation.images.cam_right_wrist",
                        "observation/state":             "observation.state",
                        "actions":                       "action",   # single, matches dataset
                        "prompt":                        "prompt",
                    }
                )
            ]
        )
        data_transforms = _transforms.Group(
            inputs=[unitree_policy.UnitreeG1Inputs(model_type=model_config.model_type)],
            outputs=[unitree_policy.UnitreeG1Outputs(action_dim=self.action_dim)],
        )
        model_transforms = ModelTransformFactory()(model_config)
        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            repack_transforms=repack_transform,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
            action_sequence_keys=("action",),   # single
        )


@dataclasses.dataclass(frozen=True)
class LeRobotUnitreeG1EEFDataConfig(DataConfigFactory):
    """Unitree G1 Dex1 bimanual EEF-space dataset config.
    16-dim: L_eef xyz+quat(7) + R_eef xyz+quat(7) + L_grip(1) + R_grip(1),
    pelvis frame, quaternion order (qx, qy, qz, qw).
    Three cameras: left_high (base), left_wrist, right_wrist.

    If `delta_position_actions` is True, the two EEF POSITION triplets (dims
    0:3 and 7:10 of the 16-dim quaternion layout) are converted to deltas
    relative to the current state before the model, and added back after.
    Quaternions and grippers stay absolute (delta of a quaternion is not
    meaningful). Same spatial-generalization rationale as the 6D delta config;
    the on-robot client contract is unchanged (AbsoluteActions re-adds the
    current pose, so returned actions are still absolute quaternion poses)."""

    # Delta the L/R EEF position triplets (quaternions/grippers stay absolute).
    # Needs its own norm_stats -> use a distinct config name / assets dir.
    delta_position_actions: bool = False

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        repack_transform = _transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "observation/image":             "observation.images.cam_left_high",
                        "observation/left_wrist_image":  "observation.images.cam_left_wrist",
                        "observation/right_wrist_image": "observation.images.cam_right_wrist",
                        "observation/state":             "observation.state",
                        "actions":                       "action",   # single, matches dataset
                        "prompt":                        "prompt",
                    }
                )
            ]
        )
        input_transforms = [unitree_eef_policy.UnitreeG1EEFInputs(model_type=model_config.model_type)]
        output_transforms = [unitree_eef_policy.UnitreeG1EEFOutputs()]
        if self.delta_position_actions:
            # 16-dim quat layout: [Lpos(3) | Lquat(4) | Rpos(3) | Rquat(4) | Lgrip | Rgrip].
            # Delta only the two position triplets; quaternions and grippers absolute.
            pos_delta_mask = _transforms.make_bool_mask(3, -4, 3, -4, -2)
            input_transforms.append(_transforms.DeltaActions(pos_delta_mask))
            output_transforms.insert(0, _transforms.AbsoluteActions(pos_delta_mask))
        data_transforms = _transforms.Group(inputs=input_transforms, outputs=output_transforms)
        model_transforms = ModelTransformFactory()(model_config)
        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            repack_transforms=repack_transform,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
            action_sequence_keys=("action",),   # single
        )


@dataclasses.dataclass(frozen=True)
class LeRobotUnitreeG1EEF6DDataConfig(DataConfigFactory):
    """Unitree G1 Dex1 bimanual EEF-space dataset config, 6D rotation action space.

    On disk the dataset is the same 16-dim quaternion layout as
    LeRobotUnitreeG1EEFDataConfig (L/R xyz+quat + grippers, pelvis frame). The
    data transform converts quat->6D on the way into the model (state + actions,
    so norm_stats are computed on the 20-dim 6D representation) and 6D->quat on
    the way out, so the on-robot client contract stays 16-dim quaternion.

    If `delta_position_actions` is True, the two EEF POSITION triplets (dims
    0:3 and 9:12 of the 20-dim 6D layout) are converted to deltas relative to
    the current state before the model, and added back after. Rotations (6D)
    and grippers stay absolute. This makes the trajectory largely independent
    of the robot's absolute standing pose / object position and is the
    single most effective change for spatial generalization when collection and
    deployment are NOT at fixed positions (see UMI / action-space studies).
    The on-robot client contract is unchanged: AbsoluteActions re-adds the
    current EEF pose, so the client still receives absolute quaternion poses."""

    # Convert the L/R EEF position triplets to delta actions (rotations/grippers
    # stay absolute). Requires its own norm_stats (delta action distribution
    # differs from absolute), so use a distinct config name / assets dir.
    delta_position_actions: bool = False
    # When delta_position_actions is True, restrict the delta to the HORIZONTAL
    # (x, y) axes and keep z ABSOLUTE. Rationale: delta helps the relative
    # approach/grasp (x, y) and is robust to absolute-frame bias, but on z it
    # loses the absolute-height reference and accumulates drift — bad for
    # stacking to a fixed top-of-stack height. So delta x,y (generalize the
    # approach) + absolute z (precise placement height) is the best of both.
    delta_horizontal_only: bool = False

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        repack_transform = _transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "observation/image":             "observation.images.cam_left_high",
                        "observation/left_wrist_image":  "observation.images.cam_left_wrist",
                        "observation/right_wrist_image": "observation.images.cam_right_wrist",
                        "observation/state":             "observation.state",
                        "actions":                       "action",   # single, matches dataset
                        "prompt":                        "prompt",
                    }
                )
            ]
        )
        input_transforms = [unitree_eef6d_policy.UnitreeG1EEF6DInputs(model_type=model_config.model_type)]
        output_transforms = [unitree_eef6d_policy.UnitreeG1EEF6DOutputs()]
        if self.delta_position_actions:
            # 20-dim 6D layout: [Lx Ly Lz | Lrot6d(6) | Rx Ry Rz | Rrot6d(6) | Lgrip | Rgrip].
            if self.delta_horizontal_only:
                # delta x,y only; z, rotation-6D, grippers absolute.
                pos_delta_mask = _transforms.make_bool_mask(2, -1, -6, 2, -1, -6, -2)
            else:
                # delta the full position triplets; rotation-6D and grippers absolute.
                pos_delta_mask = _transforms.make_bool_mask(3, -6, 3, -6, -2)
            # DeltaActions runs AFTER the quat->6D input transform (needs the 20-dim
            # state+actions); AbsoluteActions runs BEFORE the 6D->quat output transform
            # (must act on the 20-dim layout the mask indexes).
            input_transforms.append(_transforms.DeltaActions(pos_delta_mask))
            output_transforms.insert(0, _transforms.AbsoluteActions(pos_delta_mask))
        data_transforms = _transforms.Group(inputs=input_transforms, outputs=output_transforms)
        model_transforms = ModelTransformFactory()(model_config)
        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            repack_transforms=repack_transform,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
            action_sequence_keys=("action",),   # single
        )


@dataclasses.dataclass(frozen=True)
class LeRobotUR3MergedDataConfig(DataConfigFactory):
    """UR3 merged dataset config. Same as LiberoDataConfig but with
    correct column names for UR3 merged dataset (action singular,
    observation.images.image / wrist_image)."""

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        repack_transform = _transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "observation/image":       "observation.images.fixed",
                        "observation/wrist_image": "observation.images.cam_wrist",
                        "observation/state":       "observation.state",
                        "actions":                 "action",     # single to match dataset
                        "prompt":                  "prompt",
                    }
                )
            ]
        )

        data_transforms = _transforms.Group(
            inputs=[libero_policy.LiberoInputs(model_type=model_config.model_type)],
            outputs=[libero_policy.LiberoOutputs()],
        )

        model_transforms = ModelTransformFactory()(model_config)

        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            repack_transforms=repack_transform,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
            action_sequence_keys=("action",),  # single
        )

        
@dataclasses.dataclass(frozen=True)
class OldLeRobotUR3MergedDataConfig(DataConfigFactory):
    """UR3 merged dataset config. Same as LiberoDataConfig but with
    correct column names for UR3 merged dataset (action singular,
    observation.images.image / wrist_image)."""

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        repack_transform = _transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "observation/image":       "observation.images.fixed",
                        "observation/wrist_image": "observation.images.cam_wrist",
                        "observation/state":       "observation.state",
                        "actions":                 "action",     # single to mathc dataset
                        "prompt":                  "prompt",
                    }
                )
            ]
        )

        data_transforms = _transforms.Group(
            inputs=[libero_policy.LiberoInputs(model_type=model_config.model_type)],
            outputs=[libero_policy.LiberoOutputs()],
        )

        model_transforms = ModelTransformFactory()(model_config)

        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            repack_transforms=repack_transform,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
            action_sequence_keys=("action",),  # ← 单数，匹配数据集
        )



@dataclasses.dataclass(frozen=True)
class LeRobotB601DataConfig(DataConfigFactory):
    """Data config for fine-tuning on a Seeed B601 (single-arm follower) LeRobot v2.1 dataset.

    The B601 dataset stores a 7-dim absolute-joint-position state/action vector (degrees) and only two
    cameras (top, wrist); b601_policy zero-fills and masks the third view. Actions are *absolute* joint
    positions, so we do NOT apply a delta transform.
    """

    # Injected as the prompt for every sample. Leave None to take the prompt from the dataset's tasks.jsonl.
    default_prompt: str | None = None

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        # Remap raw LeRobot column names to the flat keys expected by B601Inputs (and by the inference client).
        repack_transform = _transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "observation/top": "observation.images.top",
                        "observation/wrist": "observation.images.wrist",
                        "observation/state": "observation.state",
                        "actions": "action",
                        # PromptFromLeRobotTask runs before repack and repack drops unlisted keys, so the
                        # prompt has to be carried through explicitly when prompt_from_task is set.
                        "prompt": "prompt",
                    }
                )
            ]
        )

        data_transforms = _transforms.Group(
            inputs=[b601_policy.B601Inputs(model_type=model_config.model_type)],
            outputs=[b601_policy.B601Outputs()],
        )

        model_transforms = ModelTransformFactory(default_prompt=self.default_prompt)(model_config)

        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            repack_transforms=repack_transform,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
        )


@dataclasses.dataclass(frozen=True)
class LeRobotRM65DataConfig(DataConfigFactory):
    """Data config for fine-tuning on the RealMan RM65 right-arm LeRobot v2.1 dataset.

    Same layout as B601 -- 7-dim state/action (right_q1..q6 in rad, right_gripper 0=closed..100=open)
    and two cameras (top, wrist) -- so the B601 policy transforms are reused as-is. Unlike B601,
    observation.state and action share one coordinate convention, so the joint actions can be
    trained as deltas relative to the current state.
    """

    # Injected as the prompt for every sample. Leave None to take the prompt from the dataset's tasks.jsonl.
    default_prompt: str | None = None
    # Train the 6 joints as deltas relative to the current state; the gripper stays absolute.
    use_delta_joint_actions: bool = False

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        repack_transform = _transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "observation/top": "observation.images.top",
                        "observation/wrist": "observation.images.wrist",
                        "observation/state": "observation.state",
                        "actions": "action",
                        "prompt": "prompt",
                    }
                )
            ]
        )

        data_transforms = _transforms.Group(
            inputs=[b601_policy.B601Inputs(model_type=model_config.model_type)],
            outputs=[b601_policy.B601Outputs()],
        )
        if self.use_delta_joint_actions:
            delta_action_mask = _transforms.make_bool_mask(6, -1)
            data_transforms = data_transforms.push(
                inputs=[_transforms.DeltaActions(delta_action_mask)],
                outputs=[_transforms.AbsoluteActions(delta_action_mask)],
            )

        model_transforms = ModelTransformFactory(default_prompt=self.default_prompt)(model_config)

        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            repack_transforms=repack_transform,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
        )


@dataclasses.dataclass(frozen=True)
class RLDSDroidDataConfig(DataConfigFactory):
    """
    Config for training on DROID, using RLDS data format (for efficient training on larger datasets).
    """

    rlds_data_dir: str | None = None
    action_space: droid_rlds_dataset.DroidActionSpace | None = None

    # Filtering options. Can pass a path to a dictionary that maps episodes to timestep ranges
    # to tuples denoting ranges of time steps to keep (start, end). Episodes are uniquely identified with
    # f"{recording_folderpath}--{file_path}", both of which are present in the RLDS episode metadata.

    # List of datasets to sample from: name, version, weight, and optionally filter_dict_path
    datasets: Sequence[droid_rlds_dataset.RLDSDataset] = (
        droid_rlds_dataset.RLDSDataset(
            name="droid",
            version="1.0.1",
            weight=1.0,
            filter_dict_path="gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json",
        ),
    )

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        repack_transform = _transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "observation/exterior_image_1_left": "observation/image",
                        "observation/wrist_image_left": "observation/wrist_image",
                        "observation/joint_position": "observation/joint_position",
                        "observation/gripper_position": "observation/gripper_position",
                        "actions": "actions",
                        "prompt": "prompt",
                    }
                )
            ]
        )

        data_transforms = _transforms.Group(
            inputs=[droid_policy.DroidInputs(model_type=model_config.model_type)],
            outputs=[droid_policy.DroidOutputs()],
        )

        if self.action_space == droid_rlds_dataset.DroidActionSpace.JOINT_POSITION:
            # Data loader returns absolute joint position actions -- convert to delta actions for training.
            delta_action_mask = _transforms.make_bool_mask(7, -1)
            data_transforms = data_transforms.push(
                inputs=[_transforms.DeltaActions(delta_action_mask)],
                outputs=[_transforms.AbsoluteActions(delta_action_mask)],
            )

        model_transforms = ModelTransformFactory()(model_config)

        assert self.rlds_data_dir is not None, "Need to set rlds data dir for RLDS data loader."

        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            repack_transforms=repack_transform,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
            rlds_data_dir=self.rlds_data_dir,
            action_space=self.action_space,
            datasets=self.datasets,
        )


@dataclasses.dataclass(frozen=True)
class LeRobotDROIDDataConfig(DataConfigFactory):
    """
    Example data config for custom DROID dataset in LeRobot format.
    To convert your custom DROID dataset (<10s of hours) to LeRobot format, see examples/droid/convert_droid_data_to_lerobot.py
    """

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        repack_transform = _transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "observation/exterior_image_1_left": "exterior_image_1_left",
                        "observation/exterior_image_2_left": "exterior_image_2_left",
                        "observation/wrist_image_left": "wrist_image_left",
                        "observation/joint_position": "joint_position",
                        "observation/gripper_position": "gripper_position",
                        "actions": "actions",
                        "prompt": "prompt",
                    }
                )
            ]
        )
        # We assume joint *velocity* actions, so we should *not* apply an additional delta transform.
        data_transforms = _transforms.Group(
            inputs=[droid_policy.DroidInputs(model_type=model_config.model_type)],
            outputs=[droid_policy.DroidOutputs()],
        )
        model_transforms = ModelTransformFactory()(model_config)

        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            repack_transforms=repack_transform,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
        )

@dataclasses.dataclass(frozen=True)
class LeRobotUR3DataConfig(DataConfigFactory):

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        repack_transform = _transforms.Group(
            inputs=[_transforms.RepackTransform(ur3_policy.REPACK_STRUCTURE)]
        )
        data_transforms = _transforms.Group(
            inputs=[ur3_policy.UR3Inputs(model_type=model_config.model_type)],
            outputs=[ur3_policy.UR3Outputs()],
        )
        # Actions are absolute joint positions → convert to delta for training.
        # Gripper (last dim) stays absolute, hence -1 at the end.
        delta_action_mask = _transforms.make_bool_mask(6, -1)
        data_transforms = data_transforms.push(
            inputs=[_transforms.DeltaActions(delta_action_mask)],
            outputs=[_transforms.AbsoluteActions(delta_action_mask)],
        )
        model_transforms = ModelTransformFactory()(model_config)
        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            repack_transforms=repack_transform,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
        )


@dataclasses.dataclass(frozen=True)
class LeRobotUR3BimanualDataConfig(DataConfigFactory):
    """Dual-arm UR3 (14-dim: per arm 6 joints + gripper, 3 cameras), see ur3_bimanual_policy."""

    @override
    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
        repack_transform = _transforms.Group(
            inputs=[_transforms.RepackTransform(ur3_bimanual_policy.REPACK_STRUCTURE)]
        )
        data_transforms = _transforms.Group(
            inputs=[ur3_bimanual_policy.UR3BimanualInputs(model_type=model_config.model_type)],
            outputs=[ur3_bimanual_policy.UR3BimanualOutputs()],
        )
        # Actions are absolute joint positions -> delta for training; both grippers stay absolute.
        delta_action_mask = _transforms.make_bool_mask(6, -1, 6, -1)
        data_transforms = data_transforms.push(
            inputs=[_transforms.DeltaActions(delta_action_mask)],
            outputs=[_transforms.AbsoluteActions(delta_action_mask)],
        )
        model_transforms = ModelTransformFactory()(model_config)
        return dataclasses.replace(
            self.create_base_config(assets_dirs, model_config),
            repack_transforms=repack_transform,
            data_transforms=data_transforms,
            model_transforms=model_transforms,
            action_sequence_keys=("action",),
        )



@dataclasses.dataclass(frozen=True)
class TrainConfig:
    # Name of the config. Must be unique. Will be used to reference this config.
    name: tyro.conf.Suppress[str]
    # Project name.
    project_name: str = "openpi"
    # Experiment name. Will be used to name the metadata and checkpoint directories.
    exp_name: str = tyro.MISSING

    # Defines the model config. Some attributes (action_dim, action_horizon, and max_token_len) are shared by all models
    # -- see BaseModelConfig. Specific model implementations (e.g., Pi0Config) inherit from BaseModelConfig and may
    # define additional attributes.
    model: _model.BaseModelConfig = dataclasses.field(default_factory=pi0_config.Pi0Config)

    # A weight loader can optionally load (possibly partial) weights from disk after the model is initialized.
    weight_loader: weight_loaders.WeightLoader = dataclasses.field(default_factory=weight_loaders.NoOpWeightLoader)

    # Optional path to a PyTorch checkpoint to load weights from.
    pytorch_weight_path: str | None = None

    # Precision for PyTorch training.
    pytorch_training_precision: Literal["bfloat16", "float32"] = "bfloat16"

    lr_schedule: _optimizer.LRScheduleConfig = dataclasses.field(default_factory=_optimizer.CosineDecaySchedule)
    optimizer: _optimizer.OptimizerConfig = dataclasses.field(default_factory=_optimizer.AdamW)
    ema_decay: float | None = 0.99

    # Specifies which weights should be frozen.
    freeze_filter: tyro.conf.Suppress[Filter] = dataclasses.field(default_factory=nnx.Nothing)

    # Determines the data to be trained on.
    data: DataConfigFactory = dataclasses.field(default_factory=FakeDataConfig)

    # Base directory for config assets (e.g., norm stats).
    assets_base_dir: str = "./assets"
    # Base directory for checkpoints.
    checkpoint_base_dir: str = "./checkpoints"

    # Random seed that will be used by random generators during training.
    seed: int = 42
    # Global batch size.
    batch_size: int = 32
    # Number of workers to use for the data loader. Increasing this number will speed up data loading but
    # will increase memory and CPU usage.
    num_workers: int = 2
    # Number of train steps (batches) to run.
    num_train_steps: int = 30_000

    # How often (in steps) to log training metrics.
    log_interval: int = 100
    # How often (in steps) to save checkpoints.
    save_interval: int = 1000
    # If set, any existing checkpoints matching step % keep_period == 0 will not be deleted.
    keep_period: int | None = 5000

    # If true, will overwrite the checkpoint directory if it already exists.
    overwrite: bool = False
    # If true, will resume training from the last checkpoint.
    resume: bool = False

    # If true, will enable wandb logging.
    wandb_enabled: bool = True

    # Used to pass metadata to the policy server.
    policy_metadata: dict[str, Any] | None = None

    # If the value is greater than 1, FSDP will be enabled and shard across number of specified devices; overall
    # device memory will be reduced but training could potentially be slower.
    # eg. if total device is 4 and fsdp devices is 2; then the model will shard to 2 devices and run
    # data parallel between 2 groups of devices.
    fsdp_devices: int = 1

    # --- Periodic open-loop eval (used by scripts/train_eval.py) ---
    # How often (in steps) to run eval. 0 disables eval.
    eval_interval: int = 0
    # Number of episodes (from the start of the dataset) to evaluate.
    eval_episodes: int = 1
    # Evaluate every Nth frame within an episode.
    eval_stride: int = 1
    # Optional cap on the number of frames evaluated per episode.
    eval_max_frames: int | None = None
    # Dataset directory for eval. If None, uses the training dataset.
    eval_data_path: str | None = None

    @property
    def assets_dirs(self) -> pathlib.Path:
        """Get the assets directory for this config."""
        return (pathlib.Path(self.assets_base_dir) / self.name).resolve()

    @property
    def checkpoint_dir(self) -> pathlib.Path:
        """Get the checkpoint directory for this config."""
        if not self.exp_name:
            raise ValueError("--exp_name must be set")
        return (pathlib.Path(self.checkpoint_base_dir) / self.name / self.exp_name).resolve()

    @property
    def trainable_filter(self) -> nnx.filterlib.Filter:
        """Get the filter for the trainable parameters."""
        return nnx.All(nnx.Param, nnx.Not(self.freeze_filter))

    def __post_init__(self) -> None:
        if self.resume and self.overwrite:
            raise ValueError("Cannot resume and overwrite at the same time.")

#modify here lyh
# Use `get_config` if you need to get a config by name in your code.
_CONFIGS = [
    #
    # Inference Aloha configs.
    #
    TrainConfig(
        name="pi05_libero_ur3",
        #model=pi0_config.Pi0Config(pi05=True, action_horizon=10, discrete_state_input=False, paligemma_variant="gemma_2b_lora", action_expert_variant="gemma_300m_lora"),
        model=pi0_config.Pi0Config(pi05=True, action_horizon=10, discrete_state_input=False, paligemma_variant="gemma_2b_lora", action_expert_variant="gemma_300m_lora"),
        data=LeRobotLiberoDataConfig(
            repo_id="/home/ur3-exp/pi/outcomes/tmp/ur3-pink-cylinder_6D_final/",
            base_config=DataConfig(prompt_from_task=True),
            extra_delta_transform=False,
#            assets=AssetsConfig(
#              assets_dir="/home/ur3-exp/pi/openpi/assets",   # 本地目录
#              asset_id="ur3_pick_place",                      # 自己起个名
#    ),
        ),
        batch_size=64,
        freeze_filter=pi0_config.Pi0Config(
            paligemma_variant="gemma_2b_lora", action_expert_variant="gemma_300m_lora"
        ).get_freeze_filter(),
        ema_decay=None,
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        num_train_steps=1500,
        keep_period=250,
        save_interval=250,
    ),
    
    TrainConfig(
    # Unitree G1 Dex1 bimanual pick-and-place, 16-DoF.
    # T2 strategy: PaliGemma LoRA + Action Expert full fine-tuning.
    # Small dataset (~20k frames, 50 eps) → smaller bs, fewer steps than UR3 5task.
      name="pi05_unitree_g1",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=32,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1DataConfig(
        repo_id="stack-cube-new",
        base_config=DataConfig(
            local_root=pathlib.Path("/home/ur3-exp/unitree/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),

      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
    #   keep_last=3,
      keep_period=2000,
    #   fsdp_devices=4,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
          # How often (in steps) to run eval. 0 disables eval.
    eval_interval = 500,
    # Number of episodes (from the start of the dataset) to evaluate.
    eval_episodes = 5,
    # Evaluate every Nth frame within an episode.
    eval_stride= 20,
    # Optional cap on the number of frames evaluated per episode.
    eval_max_frames = None,
    # Dataset directory for eval. If None, uses the training dataset.
    eval_data_path = "/home/ur3-exp/unitree/data/stack-cube-new-eval",

    ),

    TrainConfig(
    # Serving alias for checkpoints/pi05_g1_eef/v1 (hang-tool, trained 2026-07-28,
    # wandb k6xv7l34). That run used pi05_g1_eef with repo_id="hang-tool-eef";
    # pi05_g1_eef has since been repointed to stack-cube-eef, so serving v1 under
    # that name looks up the wrong norm stats. Identical recipe otherwise.
    # Serve with --policy.config=pi05_g1_hang_tool --policy.dir=checkpoints/pi05_g1_eef/v1/19999
      name="pi05_g1_hang_tool",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=32,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEFDataConfig(
        repo_id="hang-tool-eef",
        base_config=DataConfig(
            local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
    ),

    TrainConfig(
    # Unitree G1 Dex1 bimanual pick-and-place, EEF action space.
    # Same recipe as pi05_unitree_g1 but state/action are EEF poses
    # (L/R xyz+quat + grippers, 16-dim) converted by ~/unitree/data/joint_to_eef.py.
    # T2 strategy: PaliGemma LoRA + Action Expert full fine-tuning.
      name="pi05_g1_eef",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=32,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEFDataConfig(
        repo_id="stack-cube-eef",
        base_config=DataConfig(
            local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
    #   weight_loader=weight_loaders.CheckpointWeightLoader("/home/ur3-exp/pi/openpi/checkpoints/pi05_g1_eef/eef_v1/19999/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
    # item-classification has no separate eval split; None -> eval on the training dataset.
      eval_data_path = None,
    ),
    TrainConfig(
    # Unitree G1 Dex1 bimanual pick-and-place, EEF action space.
    # Same recipe as pi05_unitree_g1 but state/action are EEF poses
    # (L/R xyz+quat + grippers, 16-dim) converted by ~/unitree/data/joint_to_eef.py.
    # T2 strategy: PaliGemma LoRA + Action Expert full fine-tuning.
      name="pi05_g1_eef2",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=32,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEFDataConfig(
        repo_id="stack-v3-eef",
        base_config=DataConfig(
            local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
    #   weight_loader=weight_loaders.CheckpointWeightLoader("/home/ur3-exp/pi/openpi/checkpoints/pi05_g1_eef/eef_v1/19999/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
      eval_data_path = "/shared/user64/workspace/yuhao/pi/data/sort-tools-mered-eef-eval",
    ),
    TrainConfig(
    # Unitree G1 Dex1 bimanual block-stacking, EEF pose action space with a 6D
    # rotation representation (see unitree_eef6d_policy.py). Same recipe as
    # pi05_g1_eef; the ONLY difference is quat->6D on model in/out, which
    # removes the quaternion double-cover ambiguity that made the left-arm
    # orientation the worst-tracked channel. On-disk data stays 16-dim quat.
    # Trains on stack-merged-eef (stack-cube-eef + stack-the-cube-v2-eef merged,
    # both relabeled to the color-stacking prompt).
      name="pi05_g1_eef6d",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=48,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEF6DDataConfig(
        repo_id="stack-merged-eef",
        base_config=DataConfig(
            local_root=pathlib.Path("/home/ur3-exp/unitree/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
      eval_data_path = "/home/ur3-exp/unitree/data/stack-cube-new-eval-eef",
    ),

    TrainConfig(
    # Same as pi05_g1_eef6d, but the two EEF position triplets are DELTA actions
    # (relative to the current EEF pose); rotations (6D) and grippers stay
    # absolute. Highest-leverage change for spatial generalization when object /
    # robot positions vary between collection and deployment (delta trajectories
    # are largely independent of absolute start pose). Needs its OWN norm_stats
    # (delta action distribution differs) -> distinct config name / assets dir.
    # Client contract is unchanged: AbsoluteActions re-adds the current pose, so
    # the returned actions are still absolute quaternion poses.
      name="pi05_g1_eef6d_delta",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=48,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEF6DDataConfig(
        repo_id="stack-merged-eef",
        delta_position_actions=True,
        base_config=DataConfig(
            local_root=pathlib.Path("/home/ur3-exp/unitree/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
      eval_data_path = "/home/ur3-exp/unitree/data/stack-cube-eval-eef",
    ),

    TrainConfig(
    # VARIABLE-ISOLATION configs: identical recipe to pi05_g1_eef6d(_delta) but
    # trained on stack-cube-eef ONLY (95 eps, correct stacking prompt), NOT the
    # merged set. Use these to disentangle whether the "grasp-then-oscillate"
    # regression is caused by (a) the added stack-the-cube-v2 data (conflicting
    # placement height) or (b) the delta representation:
    #   pi05_g1_eef6d_scube        : 6D absolute, stack-cube-eef only
    #   pi05_g1_eef6d_delta_scube  : 6D + position-delta, stack-cube-eef only
    # If _scube places cleanly and the merged model does not -> v2 is the cause.
    # If _delta_scube oscillates but _scube does not -> delta feedback is the cause.
      name="pi05_g1_eef6d_scube",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=48,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEF6DDataConfig(
        repo_id="stack-cube-eef",
        base_config=DataConfig(
            local_root=pathlib.Path("/home/ur3-exp/unitree/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
      eval_data_path = "/home/ur3-exp/unitree/data/stack-cube-eval-eef",
    ),

    TrainConfig(
    # 6D + position-delta, stack-cube-eef ONLY. See pi05_g1_eef6d_scube header.
      name="pi05_g1_eef6d_delta_scube",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=48,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEF6DDataConfig(
        repo_id="stack-cube-eef",
        delta_position_actions=True,
        base_config=DataConfig(
            local_root=pathlib.Path("/home/ur3-exp/unitree/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
      eval_data_path = "/home/ur3-exp/unitree/data/stack-cube-eval-eef",
    ),

    TrainConfig(
    # 7D (quaternion) EEF representation + position-delta, stack-cube-eef ONLY.
    # Same as pi05_g1_eef6d_delta_scube but keeps the ORIGINAL 16-dim quaternion
    # action space (L/R xyz+quat + grippers) instead of 6D rotation. Position
    # triplets are delta; quaternions and grippers stay absolute. Use it to
    # isolate the effect of the delta representation on your known-good 7D setup,
    # and to compare 7D-delta vs 6D-delta head to head.
      name="pi05_g1_eef_delta_scube",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=48,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEFDataConfig(
        repo_id="stack-cube-eef",
        delta_position_actions=True,
        base_config=DataConfig(
            local_root=pathlib.Path("/home/ur3-exp/unitree/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
      eval_data_path = "/home/ur3-exp/unitree/data/stack-cube-new-eval-eef",
    ),

    TrainConfig(
    # MIXED action space: 6D, position-delta on x,y only, z/rotation/grippers
    # ABSOLUTE. Motivated by the empirical split observed on the robot:
    #   - full-delta improved GRASP precision (relative approach on x,y, robust
    #     to absolute-frame bias, closed-loop convergence) but
    #   - full-delta learned HEIGHT poorly (delta loses the absolute-z reference
    #     and accumulates drift -> bad for stacking to a fixed top-of-stack z).
    # Delta x,y keeps the grasp benefit; absolute z restores placement-height
    # precision. Rotation (6D) and grippers stay absolute. stack-cube-eef only.
      name="pi05_g1_eef6d_deltaxy_scube",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=48,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEF6DDataConfig(
        repo_id="stack-cube-eef",
        delta_position_actions=True,
        delta_horizontal_only=True,
        base_config=DataConfig(
            local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
      eval_data_path = "/shared/user64/workspace/yuhao/pi/data/stack-cube-eef-eval",
    ),
    TrainConfig(
    # Unitree G1 Dex1 bimanual color-matching, EEF pose action space with 6D
    # rotation (same recipe as pi05_g1_eef6d). Trains on biomanual-eef
    # (LGG100/biomanual, 50 eps / 19951 frames, converted by
    # scripts/joint_to_eef.py). Task: put each cup on the coaster of the same
    # color, blue first and then green onto the coaster the blue cup vacated; the
    # two arms act in sequence rather than together (measured gripper overlap is
    # 0.000 across all 50 episodes), and the layout is randomized per episode.
    #
    # Steps: at batch_size 32 this set is 623 step/epoch, so 20000 steps is ~32
    # epochs. pi05_g1_tennis and pi05_g1_simple_pick use the same 20000 on sets of
    # nearly identical size (19008 / 18463 frames), which keeps the three directly
    # comparable; pi05_g1_drawer keeps the same step budget but on 2.4x the data,
    # so it sees only ~13 epochs -- this family is step-matched, not epoch-matched.
      name="pi05_g1_biomanual",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=48,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEF6DDataConfig(
        repo_id="biomanual-eef",
        base_config=DataConfig(
            local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
    # No separate eval split; None -> open-loop eval on the training dataset.
      eval_data_path = None,
    ),
    TrainConfig(
    # Unitree G1 Dex1 bimanual drawer task, EEF pose action space with 6D
    # rotation (same recipe as pi05_g1_eef6d). Trains on drawer-eef
    # (LGG100/drawer, 100 eps / 47741 frames, converted by
    # scripts/joint_to_eef.py). Task: the arm nearer to the drawer pulls it open,
    # the other puts the purple cube inside, then the drawer is closed -- a long,
    # sequential two-arm episode (~477 frames avg), hence the 48-step chunk.
    #
    # Steps: at batch_size 32 this set is 1492 step/epoch, so 20000 steps is only
    # ~13 epochs -- the same step budget as the other three G1 EEF tasks but 2.4x
    # the data, so proportionally far less exposure per episode. Raise to ~45000
    # if you want epoch parity with them; left at 20000 so the four runs cost the
    # same and stay comparable step-for-step.
      name="pi05_g1_drawer",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=48,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEF6DDataConfig(
        repo_id="drawer-eef",
        base_config=DataConfig(
            local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
    # No separate eval split; None -> open-loop eval on the training dataset.
      eval_data_path = None,
    ),
    TrainConfig(
    # Unitree G1 Dex1 tennis-ball-into-tube, EEF pose action space with 6D rotation
    # (same recipe as pi05_g1_biomanual / pi05_g1_drawer). Trains on tennis-eef
    # (51 eps / 19008 frames, converted by scripts/joint_to_eef.py). Task: one arm
    # holds the ball tube while the other drops a tennis ball in, then the tube is
    # stood upright on the empty coaster -- asymmetric bimanual, one arm keeps a
    # static grasp for most of the episode.
    #
    # Steps: at batch_size 32 this set is 594 step/epoch, so 20000 steps is ~34
    # epochs -- the same budget pi05_g1_biomanual gets on a set of nearly identical
    # size (19951 frames, 623 step/epoch, ~32 epochs), which keeps the two directly
    # comparable. save_interval 4000 + keep_period 2000 leave intermediate
    # checkpoints to fall back on if it overfits.
      name="pi05_g1_tennis",
      model=pi0_config.Pi0Config(
        pi05=True,
        # 48 covers ~13% of a 373-frame episode here, matching biomanual/drawer.
        action_horizon=48,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEF6DDataConfig(
        repo_id="tennis-eef",
        base_config=DataConfig(
            local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
    # No separate eval split; None -> open-loop eval on the training dataset.
      eval_data_path = None,
    ),
    TrainConfig(
    # Unitree G1 Dex1 single-arm pick-and-drop, EEF pose action space with 6D
    # rotation (same recipe as pi05_g1_biomanual / pi05_g1_drawer). Trains on
    # simple-pick-eef (91 eps / 18463 frames, converted by scripts/joint_to_eef.py).
    # Task: pick the pink cube up with the nearer arm and drop it into the green
    # bin -- the simplest set of this family, one reach-grasp-drop and no second
    # phase.
    #
    # action_horizon is 32 here rather than the 48 the other three G1 EEF tasks use:
    # episodes run 144-286 frames (mean 202), about half of tennis-eef's and a
    # third of drawer-eef's, so a 48-step chunk would span a quarter of a whole
    # episode and pad heavily near the end. 32 matches pi05_g1_eef.
    #
    # Steps: at batch_size 32 this set is 577 step/epoch, so 20000 steps is ~35
    # epochs -- same budget as the sibling tasks. If it overfits (the likeliest of
    # the four, being the simplest), compare the 12000/14000/16000 checkpoints
    # before shortening the schedule.
      name="pi05_g1_simple_pick",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=32,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEF6DDataConfig(
        repo_id="simple-pick-eef",
        base_config=DataConfig(
            local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
    # No separate eval split; None -> open-loop eval on the training dataset.
      eval_data_path = None,
    ),
    TrainConfig(
    # Unitree G1_23 Dex1 (real robot, 5-DoF arms) "pick up the leftmost X bottle",
    # JOINT action space, 12-dim: L_arm(5) + R_arm(5) + L_grip + R_grip. Trains on
    # pick-leftmost-g1-23 (181 eps / 50699 frames): the four v3.0 sets from
    # pick_leftmost_g1_23_dex1_lerobot_only_20260922 (cola, water, two hand-soap
    # sessions -> 3 prompts) merged to v2.1 by scripts/convert_g1_23_v30_to_v21.py,
    # with handsoap_094922Z ep28 dropped (right arm sagged, state diverged from
    # action). Left-arm-only task; the right arm and gripper barely move.
    #
    # Not interchangeable with the G1_29 configs above: different arm DoF and a
    # real-robot camera setup, so initialise from pi05_base, not from those runs.
    #
    # Episodes run 145-586 frames (mean 280), so action_horizon 32 as in
    # pi05_g1_simple_pick. At batch_size 32 this is ~1584 step/epoch, so 20000
    # steps is ~12.6 epochs -- the pi05_g1_drawer budget on a set of similar size.
      name="pi05_g1_23_pick_leftmost",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=32,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1DataConfig(
        repo_id="pick-leftmost-g1-23",
        action_dim=12,
        base_config=DataConfig(
            local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
    # No separate eval split; None -> open-loop eval on the training dataset.
      eval_data_path = None,
    ),
    TrainConfig(
    # Unitree G1_23 Dex1 "pick up the leftmost X", v2 collection: same robot, action
    # space and recipe as pi05_g1_23_pick_leftmost, but a new scene (shelf instead of
    # table) and a different initial arm pose, so it trains on v2 alone from pi05_base.
    # pick-leftmost-v2-g1-23 (251 eps / 96200 frames): the raw sessions in
    # pickup_leftmost_v2_g1_23_20261001 (chips can 131, cola 120 -> 2 prompts)
    # converted by scripts/convert_g1_23_raw_to_v21.py, no episodes dropped.
    # Left-arm-only task; the right arm is held still. 42 episodes (chips can
    # ep72-112, cola ep111) keep the right gripper half closed throughout -- kept.
    #
    # Episodes run 205-576 frames (mean 383). At batch_size 32 this is ~3006
    # step/epoch, so 20000 steps is ~6.7 epochs.
      name="pi05_g1_23_pick_leftmost_v2",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=32,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1DataConfig(
        repo_id="pick-leftmost-v2-g1-23",
        action_dim=12,
        base_config=DataConfig(
            local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
    # No separate eval split; None -> open-loop eval on the training dataset.
      eval_data_path = None,
    ),
    TrainConfig(
    # Unitree G1_23 Dex1 "place the X bottle at the red X", same robot, action
    # space and recipe as pi05_g1_23_pick_leftmost. Trains on place-at-red-x-g1-23
    # (99 eps / 22091 frames): the four v3.0 sets from
    # place_at_red_x_g1_23_dex1_lerobot_only_20260923 (two cola sessions, hand
    # soap, water -> 3 prompts) merged by scripts/convert_g1_23_v30_to_v21.py, no
    # episodes dropped. Every episode starts with the bottle already in the closed
    # left gripper and opens it once at ~60% -- the place half of pick-and-place.
    #
    # Episodes run 128-389 frames (mean 223), so action_horizon 32. At batch_size
    # 32 this is ~690 step/epoch, so 20000 steps is ~29 epochs, close to
    # pi05_g1_simple_pick's ~35 on a set of similar size; if it overfits, compare
    # the 12000/16000 checkpoints before shortening the schedule.
      name="pi05_g1_23_place_red_x",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=32,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1DataConfig(
        repo_id="place-at-red-x-g1-23",
        action_dim=12,
        base_config=DataConfig(
            local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
    # No separate eval split; None -> open-loop eval on the training dataset.
      eval_data_path = None,
    ),
    TrainConfig(
    # Unitree G1 Dex1 three-colour sorting, EEF pose action space with 6D rotation
    # (same recipe as pi05_g1_biomanual / pi05_g1_drawer). Trains on
    # item-classification-new-eef (100 eps / 66933 frames). Task: red object into
    # the red box, then blue into blue, then yellow into yellow -- three sequential
    # pick-and-places, the longest episodes of this family (486-881 frames, mean
    # 669, vs 477 for drawer).
    #
    # The prompt is byte-identical to the first task of merge2 / item-sort-merged,
    # so those datasets' text assets apply here too; this is a re-collection of the
    # same task, not a new instruction.
    #
    # Steps: at batch_size 32 this set is 2092 step/epoch, so 20000 steps is ~9.6
    # epochs -- the lowest of the family, since it holds the most frames. Kept at
    # 20000 so all six G1 EEF runs cost the same and stay comparable step-for-step
    # (see pi05_g1_biomanual for the step-matched-not-epoch-matched rationale);
    # raise to ~65000 if you want epoch parity with the ~19k-frame tasks.
      name="pi05_g1_item_classification_new",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=48,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEF6DDataConfig(
        repo_id="item-classification-new-eef",
        base_config=DataConfig(
            local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
    # No separate eval split; None -> open-loop eval on the training dataset.
      eval_data_path = None,
    ),
    TrainConfig(
    # Unitree G1 Dex1 screwdriver-into-toolbox, EEF pose action space with 6D
    # rotation (same recipe as pi05_g1_biomanual / pi05_g1_drawer). Trains on
    # sort-tools-new-eef (100 eps / 32746 frames, episodes 241-435, mean 327).
    # Task: place the yellow screwdriver into the empty slot of the toolbox with
    # the left hand -- a fixed-arm single-phase task, unlike the "nearer arm"
    # phrasing the other G1 sets use.
    #
    # The prompt is byte-identical to the second task of hang-sort-merged /
    # item-sort-merged, so this is a re-collection of that task rather than a new
    # instruction.
    #
    # Steps: at batch_size 32 this set is 1023 step/epoch, so 20000 steps is ~19.5
    # epochs -- between drawer (~13) and the ~19k-frame tasks (~32-35). Kept at
    # 20000 for step parity with the rest of the family.
      name="pi05_g1_sort_tools_new",
      model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=48,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ),
      data=LeRobotUnitreeG1EEF6DDataConfig(
        repo_id="sort-tools-new-eef",
        base_config=DataConfig(
            local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            prompt_from_task=True,
        ),
      ),
      weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
      batch_size=32,
      num_workers=8,
      num_train_steps=20000,
      save_interval=4000,
      log_interval=100,
      keep_period=2000,
      wandb_enabled=True,
      lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=500,
        peak_lr=2.5e-5,
        decay_steps=15_000,
        decay_lr=2.5e-6,
      ),
      optimizer=_optimizer.AdamW(
        b1=0.9,
        b2=0.95,
        eps=1e-8,
        weight_decay=1e-4,
        clip_gradient_norm=1.0,
      ),
      freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m",
      ).get_freeze_filter(),
      ema_decay=0.99,
      eval_interval = 500,
      eval_episodes = 5,
      eval_stride= 20,
      eval_max_frames = None,
    # No separate eval split; None -> open-loop eval on the training dataset.
      eval_data_path = None,
    ),
    #
    # Fine-tuning Seeed B601 configs.
    #
    # Single-arm 7-dim pick-and-place, delivered as LeRobot v2.1 by the collection party. Only two
    # cameras, so b601_policy zero-fills and masks the right wrist view. The action expert is adapted
    # with LoRA while PaliGemma (vision + LLM) is fully fine-tuned; note that get_freeze_filter only
    # matches ".*llm.*", so the SigLIP tower stays trainable.
    #
    # Actions are absolute joint positions in degrees, and observation.state does NOT share the
    # action's coordinate convention: shoulder_pan, shoulder_lift, wrist_roll and gripper are
    # sign-flipped, and the gripper uses a different unit (state spans [-270, 0], action [0, 57]).
    # The state (follower) -> action (leader) conversion is NOT implemented here; a real-robot
    # client has to do it. State normalization is minmax, so a state outside the training range
    # saturates pi0.5's 256-bin state discretization.
    TrainConfig(
        name="pi05_b601_pick_red_cube",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=30,  # 30 fps -> ~1s action chunk
            action_expert_variant="gemma_300m_lora",
        ),
        data=LeRobotB601DataConfig(
            # Keep the slash: asset_id defaults to repo_id, and the delivered checkpoints bundle their
            # norm stats under assets/b601/pick_red_cube_v1/. Renaming this breaks checkpoint loading.
            repo_id="b601/pick_red_cube_v1",
            base_config=DataConfig(
                local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
                prompt_from_task=True,
                action_sequence_keys=("action",),
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        freeze_filter=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=30,
            action_expert_variant="gemma_300m_lora",
        ).get_freeze_filter(),
        num_train_steps=30_000,
        batch_size=32,
        ema_decay=None,  # EMA is turned off for LoRA finetuning.
    ),
    # Same robot and data pipeline as above, but with the capacity split the other way round: LoRA on
    # the Gemma-2B LLM only, so its 2.5B of pretrained language weights stay frozen, while the action
    # expert -- the one part that has to learn B601's novel action space -- trains in full. The SigLIP
    # tower also trains in full; openpi has no mechanism to freeze or LoRA it (get_freeze_filter only
    # matches ".*llm.*"). This is the branch get_freeze_filter documents as "If only freeze gemma
    # params, exclude action expert params".
    #
    # Trainable 872.8M of 3.381B (25.8%): SigLIP 414.8M + action expert 427.9M + LLM LoRA 27.9M +
    # flow heads 2.2M. Frozen: the 2508.5M Gemma-2B base.
    #
    # Registered as a separate config rather than editing the one above because the two have different
    # parameter trees -- reusing the name would make the other's checkpoints unloadable. Norm stats are
    # shared via AssetsConfig so both configs read one file.
    TrainConfig(
        name="pi05_b601_pick_red_cube_llm_lora",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=30,  # 30 fps -> ~1s action chunk
            paligemma_variant="gemma_2b_lora",
            # action_expert_variant stays at the gemma_300m default: full fine-tuning.
        ),
        data=LeRobotB601DataConfig(
            repo_id="b601/pick_red_cube_v1",
            assets=AssetsConfig(
                assets_dir="./assets/pi05_b601_pick_red_cube",
                asset_id="b601/pick_red_cube_v1",
            ),
            base_config=DataConfig(
                local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
                prompt_from_task=True,
                action_sequence_keys=("action",),
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        freeze_filter=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=30,
            paligemma_variant="gemma_2b_lora",
        ).get_freeze_filter(),
        num_train_steps=30_000,
        batch_size=32,
        # The default of 2 starves the GPUs: each sample decodes two AV1 streams, and splitting a
        # batch of 32 over four devices makes each step fast enough that the loader becomes the
        # bottleneck (observed ~30% idle time on all four GPUs).
        num_workers=8,
        # Kept at None so this run differs from pi05_b601_pick_red_cube in the capacity split alone.
        ema_decay=None,
    ),
    #
    # Fine-tuning RealMan RM65 configs.
    #
    # Right-arm 7-dim pick-and-hold over 4 objects (yellow Lay's, Nongfu Spring, hand soap, Coca-Cola),
    # 187 episodes / 66,121 frames, merged from the 4 per-task deliveries with scripts/merge_datasets.py.
    # The task string in tasks.jsonl names the object, so prompt_from_task selects the task.
    # Capacity split follows pi05_b601_pick_red_cube_llm_lora: SigLIP full + LLM LoRA + action expert full.
    #
    # Two configs differing only in the action representation; each keeps its own norm stats under
    # ./assets/<config name>/rm65/4tasks_v1/, because the delta variant's action stats are in delta space.
    TrainConfig(
        name="pi05_rm65_4tasks",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=30,  # 30 fps -> ~1s action chunk
            paligemma_variant="gemma_2b_lora",
        ),
        data=LeRobotRM65DataConfig(
            repo_id="rm65/4tasks_v1",
            base_config=DataConfig(
                local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
                prompt_from_task=True,
                action_sequence_keys=("action",),
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        freeze_filter=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=30,
            paligemma_variant="gemma_2b_lora",
        ).get_freeze_filter(),
        num_train_steps=30_000,
        batch_size=32,
        num_workers=8,
        ema_decay=None,
    ),
    # Same as above, but the 6 joints are trained as deltas relative to the current state (gripper absolute).
    TrainConfig(
        name="pi05_rm65_4tasks_delta",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=30,
            paligemma_variant="gemma_2b_lora",
        ),
        data=LeRobotRM65DataConfig(
            repo_id="rm65/4tasks_v1",
            use_delta_joint_actions=True,
            base_config=DataConfig(
                local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
                prompt_from_task=True,
                action_sequence_keys=("action",),
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        freeze_filter=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=30,
            paligemma_variant="gemma_2b_lora",
        ).get_freeze_filter(),
        num_train_steps=30_000,
        batch_size=32,
        num_workers=8,
        ema_decay=None,
    ),
    # Same as pi05_rm65_4tasks, but on rm65/4tasks_v1_trim (scripts/trim_static_frames.py): the static head
    # of every episode is dropped and the static tail is cut to 10 frames, 66,121 -> 58,399 frames.
    # Short mid-episode pauses are kept.
    TrainConfig(
        name="pi05_rm65_4tasks_trim",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=30,
            paligemma_variant="gemma_2b_lora",
        ),
        data=LeRobotRM65DataConfig(
            repo_id="rm65/4tasks_v1_trim",
            base_config=DataConfig(
                local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
                prompt_from_task=True,
                action_sequence_keys=("action",),
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        freeze_filter=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=30,
            paligemma_variant="gemma_2b_lora",
        ).get_freeze_filter(),
        num_train_steps=30_000,
        batch_size=32,
        num_workers=8,
        ema_decay=None,
    ),
    # 2026-10-05 delivery: rightmost-object pick-and-hold over 4 new objects (yellow original / cucumber /
    # purple tomato chips, black Coca-Cola), 350 episodes. The 4 per-task lingbot_v033 folders were merged
    # with scripts/merge_datasets.py into rm65/4tasks_350_v2 and trimmed with scripts/trim_static_frames.py
    # (130,578 -> 123,698 frames). Training recipe copied from pi05_g1_eef: LLM LoRA + full action expert,
    # horizon 32, cosine LR 2.5e-5, EMA 0.99, 20k steps. Joint-space actions (absolute), as in pi05_rm65_4tasks.
    TrainConfig(
        name="pi05_rm65_4tasks_350",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=32,
            discrete_state_input=False,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m",
        ),
        data=LeRobotRM65DataConfig(
            repo_id="rm65/4tasks_350_v2_trim",
            base_config=DataConfig(
                local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
                prompt_from_task=True,
                action_sequence_keys=("action",),
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        batch_size=32,
        num_workers=8,
        num_train_steps=20000,
        save_interval=4000,
        log_interval=100,
        keep_period=2000,
        wandb_enabled=True,
        lr_schedule=_optimizer.CosineDecaySchedule(
            warmup_steps=500,
            peak_lr=2.5e-5,
            decay_steps=15_000,
            decay_lr=2.5e-6,
        ),
        optimizer=_optimizer.AdamW(
            b1=0.9,
            b2=0.95,
            eps=1e-8,
            weight_decay=1e-4,
            clip_gradient_norm=1.0,
        ),
        freeze_filter=pi0_config.Pi0Config(
            pi05=True,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m",
        ).get_freeze_filter(),
        ema_decay=0.99,
        eval_interval=500,
        eval_episodes=5,
        eval_stride=20,
        eval_max_frames=None,
        # No separate eval split; None -> eval on the training dataset.
        eval_data_path=None,
    ),


    TrainConfig(
     # This config is for fine-tuning pi05-base-mutitask on a custom multitask dataset.
     # Here, we use LeRobot data format (like for all other fine-tuning examples)
    name="pi05_ur3_multitask",
    model=pi0_config.Pi0Config(pi05=True, action_horizon=10, discrete_state_input=False, paligemma_variant="gemma_2b_lora", 
                                action_expert_variant="gemma_300m_lora",),
    data=LeRobotUR3MergedDataConfig(
        repo_id="zestcode5/ur3-multiple-task",
        base_config=DataConfig(
            prompt_from_task=True,
            # read from hub
            ),
        ),
#    weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
     weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),

    batch_size=96,
    num_train_steps=8000,
    save_interval=1000,
    log_interval=50,
    keep_period=2000,
    wandb_enabled=True,
    freeze_filter=pi0_config.Pi0Config(pi05=True,paligemma_variant="gemma_2b_lora",action_expert_variant="gemma_300m_lora",
    ).get_freeze_filter(),
    ema_decay=None,
    ),
    
    TrainConfig(
    #  for test not train
    # cp  ./assets/pi05_ur3_multitask → ./assets/pi0_5_multitask_old
    name="pi0_5_multitask_old",
    model=pi0_config.Pi0Config(
        pi05=True,
        action_horizon=10,
        discrete_state_input=False,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m_lora",
    ),
    data=OldLeRobotUR3MergedDataConfig(
        repo_id="zestcode5/ur3-multiple-task-trimmed-idle",
        base_config=DataConfig(
            prompt_from_task=True,
        ),
    ),
    weight_loader=weight_loaders.CheckpointWeightLoader(
        "gs://openpi-assets/checkpoints/pi05_base/params"
    ),
    # log from wandb
    batch_size=96,
    num_workers=2,
    num_train_steps=8000,
    save_interval=1000,
    log_interval=50,
    keep_period=2000,
    fsdp_devices=1,
    wandb_enabled=True,              
    lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=1_000,
        peak_lr=2.5e-5,
        decay_steps=30_000,
        decay_lr=2.5e-6,
    ),
    optimizer=_optimizer.AdamW(
        b1=0.9, b2=0.95, eps=1e-8,
        weight_decay=1e-10,
        clip_gradient_norm=1.0,
    ),
    freeze_filter=pi0_config.Pi0Config(
        pi05=True,
        paligemma_variant="gemma_2b_lora",
        action_expert_variant="gemma_300m_lora",
    ).get_freeze_filter(),
    ema_decay=None,
    ),
    
    
    TrainConfig(
        # T2 strategy: PaliGemma LoRA + Action Expert full fine-tuning.
        # Addresses the dual-LoRA capacity bottleneck (loss plateau at ~0.011
        # observed in pi05_ur3_multitask).
        # Dataset: UR3 multitask, ~612k frames. Hardware: 4x A100 80G with FSDP.
        name="pi05_ur3_multitask_full_action",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=10,
            discrete_state_input=False,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m",
        ),
        data=LeRobotUR3MergedDataConfig(
            repo_id="zestcode5/ur3-merged-3tasks-v1",
            base_config=DataConfig(
                prompt_from_task=True,
                # read from hub
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        
        batch_size=128,
        num_workers=16,
        num_train_steps=15000,
        save_interval=250,
        log_interval=10,
        keep_period=1000,
        fsdp_devices=4,
        wandb_enabled=True,
        lr_schedule=_optimizer.CosineDecaySchedule(
            warmup_steps=1_000,
            peak_lr=2.5e-5,
            decay_steps=30_000,
            decay_lr=2.5e-6,
        ),
        optimizer=_optimizer.AdamW(
            b1=0.9,
            b2=0.95,
            eps=1e-8,
            weight_decay=1e-4,
            clip_gradient_norm=1.0,
        ),
        freeze_filter=pi0_config.Pi0Config(
            pi05=True,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m",
        ).get_freeze_filter(),
        ema_decay=0.99,
    ),
    
    TrainConfig(
        # T2 strategy: PaliGemma LoRA + Action Expert full fine-tuning.
        # Addresses the dual-LoRA capacity bottleneck (loss plateau at ~0.011
        # observed in pi05_ur3_multitask).
        # Dataset: UR3 multitask, ~612k frames. Hardware: 4x A100 80G with FSDP.
        name="pi05_ur3_multitask_test_full_action",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=10,
            discrete_state_input=False,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m",
        ),
        data=LeRobotUR3MergedDataConfig(
            repo_id="zestcode5/ur3-multiple-task-trimmed-idle",
            base_config=DataConfig(
                prompt_from_task=True,
                # read from hub
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        
        batch_size=128,
        num_workers=16,
        num_train_steps=15000,
        save_interval=250,
        log_interval=10,
        keep_period=1000,
        fsdp_devices=4,
        wandb_enabled=True,
        lr_schedule=_optimizer.CosineDecaySchedule(
            warmup_steps=1_000,
            peak_lr=2.5e-5,
            decay_steps=30_000,
            decay_lr=2.5e-6,
        ),
        optimizer=_optimizer.AdamW(
            b1=0.9,
            b2=0.95,
            eps=1e-8,
            weight_decay=1e-4,
            clip_gradient_norm=1.0,
        ),
        freeze_filter=pi0_config.Pi0Config(
            pi05=True,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m",
        ).get_freeze_filter(),
        ema_decay=0.99,
    ),
    TrainConfig(
        # T2 strategy: PaliGemma LoRA + Action Expert full fine-tuning.
        # Addresses the dual-LoRA capacity bottleneck (loss plateau at ~0.011
        # observed in pi05_ur3_multitask).
        # Dataset: UR3 multitask, ~612k frames. Hardware: 4x A100 80G with FSDP.
        name="pi05_ur3_5task",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=10,
            discrete_state_input=False,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m",
        ),
        data=LeRobotUR3MergedDataConfig(
            repo_id="zestcode5/ur3-merged-5tasks-v1",
            base_config=DataConfig(
                prompt_from_task=True,
                # read from hub
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        
        batch_size=128,
        num_workers=16,
        num_train_steps=15000,
        save_interval=250,
        log_interval=10,
        keep_period=1000,
        fsdp_devices=4,
        wandb_enabled=True,
        lr_schedule=_optimizer.CosineDecaySchedule(
            warmup_steps=1_000,
            peak_lr=2.5e-5,
            decay_steps=30_000,
            decay_lr=2.5e-6,
        ),
        optimizer=_optimizer.AdamW(
            b1=0.9,
            b2=0.95,
            eps=1e-8,
            weight_decay=1e-4,
            clip_gradient_norm=1.0,
        ),
        freeze_filter=pi0_config.Pi0Config(
            pi05=True,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m",
        ).get_freeze_filter(),
        ema_decay=0.99,
    ),
    
    
        TrainConfig(
        # T2 strategy: PaliGemma LoRA + Action Expert full fine-tuning.
        # Addresses the dual-LoRA capacity bottleneck (loss plateau at ~0.011
        # observed in pi05_ur3_multitask).
        # Dataset: UR3 multitask, ~612k frames. Hardware: 4x A100 80G with FSDP.
        name="pi05_ur3_4task",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=10,
            discrete_state_input=False,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m",
        ),
        data=LeRobotUR3MergedDataConfig(
            repo_id="zestcode5/ur3-merged-4tasks-v1",
            base_config=DataConfig(
                prompt_from_task=True,
                # read from hub
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        
        batch_size=128,
        num_workers=16,
        num_train_steps=15000,
        save_interval=250,
        log_interval=10,
        keep_period=1000,
        fsdp_devices=4,
        wandb_enabled=True,
        lr_schedule=_optimizer.CosineDecaySchedule(
            warmup_steps=1_000,
            peak_lr=2.5e-5,
            decay_steps=30_000,
            decay_lr=2.5e-6,
        ),
        optimizer=_optimizer.AdamW(
            b1=0.9,
            b2=0.95,
            eps=1e-8,
            weight_decay=1e-4,
            clip_gradient_norm=1.0,
        ),
        freeze_filter=pi0_config.Pi0Config(
            pi05=True,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m",
        ).get_freeze_filter(),
        ema_decay=0.99,
    ),
    
    TrainConfig(
     # This config is for fine-tuning pi05-base-mutitask on a custom multitask dataset.
     # Here, we use LeRobot data format (like for all other fine-tuning examples)
    name="pi05_ur3_push_glass",
    model=pi0_config.Pi0Config(pi05=True, action_horizon=10, discrete_state_input=False, paligemma_variant="gemma_2b_lora", 
                                action_expert_variant="gemma_300m_lora",),
    data=LeRobotUR3MergedDataConfig(
        repo_id="poweredshine/ur3-single-finger-pushing-glass-to-marker",
        base_config=DataConfig(
            prompt_from_task=True,
            # read from hub
            ),
        ),
#    weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
     weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),

    batch_size=96,
    num_train_steps=8000,
    save_interval=1000,
    log_interval=50,
    keep_period=2000,
    wandb_enabled=True,
    freeze_filter=pi0_config.Pi0Config(pi05=True,paligemma_variant="gemma_2b_lora",action_expert_variant="gemma_300m_lora",
    ).get_freeze_filter(),
    ema_decay=None,
    ),
    
    
    TrainConfig(
        # Dual-arm UR3 (LGG100/ur3-test cleaned: 22 eps / 3620 frames @ 15 Hz, 14-dim joints, 3 cameras),
        # see scripts/clean_ur3_bimanual_lerobot.py.
        # PaliGemma LoRA + Action Expert LoRA, following pi05_ur3_multitask / pi05_ur3_push_glass.
        name="pi05_ur3_bimanual_lora",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=10,
            discrete_state_input=False,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m_lora",
        ),
        data=LeRobotUR3BimanualDataConfig(
            repo_id="ur3-test-clean",
            base_config=DataConfig(
                local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
                prompt_from_task=True,
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        batch_size=32,
        num_workers=8,
        num_train_steps=10000,
        save_interval=1000,
        log_interval=50,
        keep_period=2000,
        wandb_enabled=True,
        lr_schedule=_optimizer.CosineDecaySchedule(
            warmup_steps=500,
            peak_lr=2.5e-5,
            decay_steps=10_000,
            decay_lr=2.5e-6,
        ),
        optimizer=_optimizer.AdamW(
            b1=0.9,
            b2=0.95,
            eps=1e-8,
            weight_decay=1e-4,
            clip_gradient_norm=1.0,
        ),
        freeze_filter=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=10,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m_lora",
        ).get_freeze_filter(),
        ema_decay=None,
    ),

    TrainConfig(
        # Bimanual UR3 (ur3-test-clean), config/policy from https://huggingface.co/LGG100/ur3-test-lora
        # (pi05_ur3_test_lora: aloha_policy with adapt_to_pi=False), action_horizon 10 -> 32.
        # LoRA on both PaliGemma and the action expert, single GPU.
        name="pi05_ur3_test_lora_h32",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=32,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m_lora",
        ),
        data=LeRobotAlohaDataConfig(
            repo_id="ur3-test-clean",
            base_config=DataConfig(
                prompt_from_task=True,
                local_root=pathlib.Path("/shared/user64/workspace/yuhao/pi/data"),
            ),
            adapt_to_pi=False,
            repack_transforms=_transforms.Group(
                inputs=[
                    _transforms.RepackTransform(
                        {
                            "images": {
                                "cam_high": "observation.images.cam_high",
                                "cam_left_wrist": "observation.images.cam_left_wrist",
                                "cam_right_wrist": "observation.images.cam_right_wrist",
                            },
                            "state": "observation.state",
                            "actions": "action",
                            "prompt": "prompt",
                        }
                    )
                ]
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        batch_size=32,
        num_workers=8,
        num_train_steps=4000,
        save_interval=1000,
        keep_period=1000,
        log_interval=50,
        lr_schedule=_optimizer.CosineDecaySchedule(
            warmup_steps=200,
            peak_lr=2.5e-5,
            decay_steps=4000,
            decay_lr=2.5e-6,
        ),
        freeze_filter=pi0_config.Pi0Config(
            pi05=True,
            action_horizon=32,
            paligemma_variant="gemma_2b_lora",
            action_expert_variant="gemma_300m_lora",
        ).get_freeze_filter(),
        ema_decay=None,
    ),

    TrainConfig(
        name="pi0_aloha",
        model=pi0_config.Pi0Config(),
        data=LeRobotAlohaDataConfig(
            assets=AssetsConfig(asset_id="trossen"),
        ),
        policy_metadata={"reset_pose": [0, -1.5, 1.5, 0, 0, 0]},
    ),
    TrainConfig(
        name="pi05_aloha",
        model=pi0_config.Pi0Config(pi05=True),
        data=LeRobotAlohaDataConfig(
            assets=AssetsConfig(asset_id="trossen"),
        ),
        policy_metadata={"reset_pose": [0, -1.5, 1.5, 0, 0, 0]},
    ),
    TrainConfig(
        name="pi0_aloha_towel",
        model=pi0_config.Pi0Config(),
        data=LeRobotAlohaDataConfig(
            assets=AssetsConfig(asset_id="trossen"),
            default_prompt="fold the towel",
        ),
        policy_metadata={"reset_pose": [0, -1.5, 1.5, 0, 0, 0]},
    ),
    TrainConfig(
        name="pi0_aloha_tupperware",
        model=pi0_config.Pi0Config(),
        data=LeRobotAlohaDataConfig(
            assets=AssetsConfig(asset_id="trossen"),
            default_prompt="open the tupperware and put the food on the plate",
        ),
        policy_metadata={"reset_pose": [0, -1.5, 1.5, 0, 0, 0]},
    ),
    #
    # Inference DROID configs.
    #
    TrainConfig(
        name="pi0_droid",
        model=pi0_config.Pi0Config(action_horizon=10),
        data=SimpleDataConfig(
            assets=AssetsConfig(asset_id="droid"),
            data_transforms=lambda model: _transforms.Group(
                inputs=[droid_policy.DroidInputs(model_type=ModelType.PI0)],
                outputs=[droid_policy.DroidOutputs()],
            ),
            base_config=DataConfig(
                prompt_from_task=True,
            ),
        ),
    ),
    TrainConfig(
        name="pi0_fast_droid",
        model=pi0_fast.Pi0FASTConfig(action_dim=8, action_horizon=10),
        data=SimpleDataConfig(
            assets=AssetsConfig(asset_id="droid"),
            data_transforms=lambda model: _transforms.Group(
                inputs=[droid_policy.DroidInputs(model_type=ModelType.PI0_FAST)],
                outputs=[droid_policy.DroidOutputs()],
            ),
            base_config=DataConfig(
                prompt_from_task=True,
            ),
        ),
    ),
    TrainConfig(
        name="pi05_droid",
        model=pi0_config.Pi0Config(action_horizon=15, pi05=True),
        data=SimpleDataConfig(
            assets=AssetsConfig(asset_id="droid"),
            data_transforms=lambda model: _transforms.Group(
                inputs=[droid_policy.DroidInputs(model_type=ModelType.PI05)],
                outputs=[droid_policy.DroidOutputs()],
            ),
            base_config=DataConfig(
                prompt_from_task=True,
            ),
        ),
    ),
    #
    # Fine-tuning Libero configs.
    #
    # These train configs define the hyperparameters for fine-tuning the base model on your own dataset.
    # They are used to define key elements like the dataset you are training on, the base checkpoint you
    # are using, and other hyperparameters like how many training steps to run or what learning rate to use.
    # For your own dataset, you can copy this class and modify the dataset name, and data transforms based on
    # the comments below.
    TrainConfig(
        # Change the name to reflect your model and dataset.
        name="pi0_libero",
        # Here you define the model config -- In this example we use pi0 as the model
        # architecture and perform *full* finetuning. in the examples below we show how to modify
        # this to perform *low-memory* (LORA) finetuning and use pi0-FAST as an alternative architecture.
        model=pi0_config.Pi0Config(),
        # Here you define the dataset you are training on. In this example we use the Libero
        # dataset. For your own dataset, you can change the repo_id to point to your dataset.
        # Also modify the DataConfig to use the new config you made for your dataset above.
        data=LeRobotLiberoDataConfig(
            repo_id="physical-intelligence/libero",
            base_config=DataConfig(
                # This flag determines whether we load the prompt (i.e. the task instruction) from the
                # ``task`` field in the LeRobot dataset. If set to True, the prompt will show up in
                # a field called ``prompt`` in the input dict. The recommended setting is True.
                prompt_from_task=True,
            ),
            extra_delta_transform=True,
        ),
        # Here you define which pre-trained checkpoint you want to load to initialize the model.
        # This should match the model config you chose above -- i.e. in this case we use the pi0 base model.
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_base/params"),
        # Below you can define other hyperparameters like the learning rate, number of training steps, etc.
        # Check the base TrainConfig class for a full list of available hyperparameters.
        num_train_steps=30_000,
    ),
    TrainConfig(
        name="pi0_libero_low_mem_finetune",
        # Here is an example of loading a pi0 model for LoRA fine-tuning.
        model=pi0_config.Pi0Config(paligemma_variant="gemma_2b_lora", action_expert_variant="gemma_300m_lora"),
        data=LeRobotLiberoDataConfig(
            repo_id="physical-intelligence/libero",
            base_config=DataConfig(prompt_from_task=True),
            extra_delta_transform=True,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_base/params"),
        num_train_steps=30_000,
        # The freeze filter defines which parameters should be frozen during training.
        # We have a convenience function in the model config that returns the default freeze filter
        # for the given model config for LoRA finetuning. Just make sure it matches the model config
        # you chose above.
        freeze_filter=pi0_config.Pi0Config(
            paligemma_variant="gemma_2b_lora", action_expert_variant="gemma_300m_lora"
        ).get_freeze_filter(),
        # Turn off EMA for LoRA finetuning.
        ema_decay=None,
    ),
    TrainConfig(
        name="pi0_fast_libero",
        # Here is an example of loading a pi0-FAST model for full finetuning.
        # Modify action_dim and action_horizon to match your dataset (action horizon is equal to
        # the desired action chunk length).
        # The max_token_len is the maximum number of (non-image) tokens the model can handle.
        # This includes the tokenized prompt, proprioceptive state, and (FAST-tokenized) action tokens.
        # Choosing this value too small may chop off tokens at the end of your sequence (the code will throw
        # a warning), while choosing it too large will waste memory (since we pad each batch element to the
        # max_token_len). A good rule of thumb is to use approx 180 for single-arm robots, and approx 250 for
        # two-arm robots. Generally, err on the lower side here first, and potentially increase the value if
        # you see many warnings being thrown during training.
        model=pi0_fast.Pi0FASTConfig(action_dim=7, action_horizon=10, max_token_len=180),
        data=LeRobotLiberoDataConfig(
            repo_id="physical-intelligence/libero",
            base_config=DataConfig(prompt_from_task=True),
            extra_delta_transform=True,
        ),
        # Note that we load the pi0-FAST base model checkpoint here.
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_fast_base/params"),
        num_train_steps=30_000,
    ),
    TrainConfig(
        name="pi0_fast_libero_low_mem_finetune",
        # Here is an example of loading a pi0-FAST model for LoRA finetuning.
        # For setting action_dim, action_horizon, and max_token_len, see the comments above.
        model=pi0_fast.Pi0FASTConfig(
            action_dim=7, action_horizon=10, max_token_len=180, paligemma_variant="gemma_2b_lora"
        ),
        data=LeRobotLiberoDataConfig(
            repo_id="physical-intelligence/libero",
            base_config=DataConfig(prompt_from_task=True),
            extra_delta_transform=True,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_fast_base/params"),
        num_train_steps=30_000,
        # Again, make sure to match the model config above when extracting the freeze filter
        # that specifies which parameters should be frozen during LoRA finetuning.
        freeze_filter=pi0_fast.Pi0FASTConfig(
            action_dim=7, action_horizon=10, max_token_len=180, paligemma_variant="gemma_2b_lora"
        ).get_freeze_filter(),
        # Turn off EMA for LoRA finetuning.
        ema_decay=None,
    ),
    TrainConfig(
        name="pi05_libero",
        model=pi0_config.Pi0Config(pi05=True, action_horizon=10, discrete_state_input=False),
        data=LeRobotLiberoDataConfig(
            repo_id="physical-intelligence/libero",
            base_config=DataConfig(prompt_from_task=True),
            extra_delta_transform=False,
        ),
        batch_size=256,
        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="/path/to/your/pytorch_weight_path",
        num_train_steps=30_000,
    ),
    #
    # Fine-tuning Aloha configs.
    #
    # This is a test config that is used to illustate how train on a custom LeRobot dataset.
    # For instructions on how to convert and train on your own Aloha dataset see examples/aloha_real/README.md
    TrainConfig(
        name="pi0_aloha_pen_uncap",
        model=pi0_config.Pi0Config(),
        data=LeRobotAlohaDataConfig(
            repo_id="physical-intelligence/aloha_pen_uncap_diverse",
            assets=AssetsConfig(
                assets_dir="gs://openpi-assets/checkpoints/pi0_base/assets",
                asset_id="trossen",
            ),
            default_prompt="uncap the pen",
            repack_transforms=_transforms.Group(
                inputs=[
                    _transforms.RepackTransform(
                        {
                            "images": {
                                "cam_high": "observation.images.cam_high",
                                "cam_left_wrist": "observation.images.cam_left_wrist",
                                "cam_right_wrist": "observation.images.cam_right_wrist",
                            },
                            "state": "observation.state",
                            "actions": "action",
                        }
                    )
                ]
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_base/params"),
        num_train_steps=20_000,
    ),
    TrainConfig(
        name="pi05_aloha_pen_uncap",
        model=pi0_config.Pi0Config(pi05=True),
        data=LeRobotAlohaDataConfig(
            repo_id="physical-intelligence/aloha_pen_uncap_diverse",
            assets=AssetsConfig(
                assets_dir="gs://openpi-assets/checkpoints/pi05_base/assets",
                asset_id="trossen",
            ),
            default_prompt="uncap the pen",
            repack_transforms=_transforms.Group(
                inputs=[
                    _transforms.RepackTransform(
                        {
                            "images": {
                                "cam_high": "observation.images.cam_high",
                                "cam_left_wrist": "observation.images.cam_left_wrist",
                                "cam_right_wrist": "observation.images.cam_right_wrist",
                            },
                            "state": "observation.state",
                            "actions": "action",
                        }
                    )
                ]
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        num_train_steps=20_000,
        batch_size=64,
    ),
    #
    # Fine-tuning DROID configs.
    #
    TrainConfig(
        # This config is for fine-tuning pi0-FAST-base on the *full* DROID dataset.
        # We use RLDS data loading to make training on this large dataset tractable.
        # For fine-tuning on your own DROID dataset, see below.
        name="pi0_fast_full_droid_finetune",
        model=pi0_fast.Pi0FASTConfig(
            action_dim=8,
            action_horizon=16,
            max_token_len=180,
        ),
        data=RLDSDroidDataConfig(
            repo_id="droid",
            # Set this to the path to your DROID RLDS dataset (the parent directory of the `droid` directory).
            rlds_data_dir="<path_to_droid_rlds_dataset>",
            action_space=droid_rlds_dataset.DroidActionSpace.JOINT_POSITION,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_fast_base/params"),
        lr_schedule=_optimizer.CosineDecaySchedule(
            warmup_steps=1_000,
            peak_lr=5e-5,
            decay_steps=1_000_000,
            decay_lr=5e-5,
        ),
        num_train_steps=100_000,  # 100k steps should be sufficient, takes ~2 days on 8x H100s
        batch_size=256,
        log_interval=100,
        save_interval=5000,
        keep_period=20_000,
        num_workers=0,  # Important: RLDS DataLoader requires num_workers=0, handles multi-processing internally
    ),
    TrainConfig(
        # This config is for fine-tuning pi05 on the *full* DROID dataset.
        # We use RLDS data loading to make training on this large dataset tractable.
        # For fine-tuning on your own DROID dataset, see below.
        name="pi05_full_droid_finetune",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_dim=32,
            action_horizon=16,
        ),
        data=RLDSDroidDataConfig(
            repo_id="droid",
            # Set this to the path to your DROID RLDS dataset (the parent directory of the `droid` directory).
            rlds_data_dir="/mnt/pi-data/kevin",
            action_space=droid_rlds_dataset.DroidActionSpace.JOINT_POSITION,
            assets=AssetsConfig(
                assets_dir="gs://openpi-assets/checkpoints/pi05_base/assets/",
                asset_id="droid",
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
        lr_schedule=_optimizer.CosineDecaySchedule(
            warmup_steps=1_000,
            peak_lr=5e-5,
            decay_steps=1_000_000,
            decay_lr=5e-5,
        ),
        num_train_steps=100_000,
        batch_size=256,
        log_interval=100,
        save_interval=5000,
        keep_period=10_000,
        num_workers=0,  # Important: RLDS DataLoader requires num_workers=0, handles multi-processing internally
    ),
    TrainConfig(
        # This config is for fine-tuning pi05-DROID on a custom (smaller) DROID dataset.
        # Here, we use LeRobot data format (like for all other fine-tuning examples)
        # To convert your custom DROID dataset (<10s of hours) to LeRobot format, see examples/droid/convert_droid_data_to_lerobot.py
        name="pi05_droid_finetune",
        model=pi0_config.Pi0Config(
            pi05=True,
            action_dim=32,  # pi05 is trained with 32-dim actions
            action_horizon=16,
        ),
        data=LeRobotDROIDDataConfig(
            # Replace with your custom DROID LeRobot dataset repo id.
            repo_id="your_hf_username/my_droid_dataset",
            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/assets",
                asset_id="droid",
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_droid/params"),
        num_train_steps=20_000,
        batch_size=32,
    ),
    
    
    #
    # ALOHA Sim configs. This config is used to demonstrate how to train on a simple simulated environment.
    #
    TrainConfig(
        name="pi0_aloha_sim",
        model=pi0_config.Pi0Config(),
        data=LeRobotAlohaDataConfig(
            repo_id="lerobot/aloha_sim_transfer_cube_human",
            default_prompt="Transfer cube",
            use_delta_joint_actions=False,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_base/params"),
        num_train_steps=20_000,
    ),
    #
    # Debugging configs.
    #
    TrainConfig(
        name="debug",
        data=FakeDataConfig(),
        batch_size=2,
        model=pi0_config.Pi0Config(paligemma_variant="dummy", action_expert_variant="dummy"),
        save_interval=100,
        overwrite=True,
        exp_name="debug",
        num_train_steps=10,
        wandb_enabled=False,
    ),
    TrainConfig(
        name="debug_restore",
        data=FakeDataConfig(),
        batch_size=2,
        model=pi0_config.Pi0Config(paligemma_variant="dummy", action_expert_variant="dummy"),
        weight_loader=weight_loaders.CheckpointWeightLoader("./checkpoints/debug/debug/9/params"),
        overwrite=True,
        exp_name="debug",
        num_train_steps=10,
        wandb_enabled=False,
    ),
    TrainConfig(
        name="debug_pi05",
        model=pi0_config.Pi0Config(pi05=True, paligemma_variant="dummy", action_expert_variant="dummy"),
        data=FakeDataConfig(),
        batch_size=2,
        num_train_steps=10,
        overwrite=True,
        exp_name="debug_pi05",
        wandb_enabled=False,
    ),
    # RoboArena & PolaRiS configs.
    *roboarena_config.get_roboarena_configs(),
    *polaris_config.get_polaris_configs(),
]

if len({config.name for config in _CONFIGS}) != len(_CONFIGS):
    raise ValueError("Config names must be unique.")
_CONFIGS_DICT = {config.name: config for config in _CONFIGS}


def cli() -> TrainConfig:
    return tyro.extras.overridable_config_cli({k: (k, v) for k, v in _CONFIGS_DICT.items()})


def get_config(config_name: str) -> TrainConfig:
    """Get a config by name."""
    if config_name not in _CONFIGS_DICT:
        closest = difflib.get_close_matches(config_name, _CONFIGS_DICT.keys(), n=1, cutoff=0.0)
        closest_str = f" Did you mean '{closest[0]}'? " if closest else ""
        raise ValueError(f"Config '{config_name}' not found.{closest_str}")

    return _CONFIGS_DICT[config_name]