Add custom TWIN conversion + configs; document experimental status

openpi/README.md

@@ -321,3 +321,15 @@ We will collect common issues and their solutions here. If you encounter an issu
 | Import errors when running examples       | Make sure you've installed all dependencies with `uv sync`. Some examples may have additional requirements listed in their READMEs.                    |
 | Action dimensions mismatch                | Verify your data processing transforms match the expected input/output dimensions of your robot. Check the action space definitions in your policy classes.                                  |
 | Diverging training loss                            | Check the `q01`, `q99`, and `std` values in `norm_stats.json` for your dataset. Certain dimensions that are rarely used can end up with very small `q01`, `q99`, or `std` values, leading to huge states and actions after normalization. You can manually adjust the norm stats as a workaround. |
+
+## Multiarm/TWIN Custom Additions (Experimental)
+
+These multiarm + TWIN additions were custom-integrated for research experiments and are **not** an official upstream openpi release.
+
+- Includes a custom TWIN->LeRobot conversion script at `openpi/scripts/convert_twin_squashfs_to_lerobot.py`.
+- Includes custom training config entries in `openpi/src/openpi/training/config.py` (e.g. `pi05_twin_bimanual_parallel_finetune`).
+- Intended for rapid experimentation; edge cases may still exist.
+- Behavior is not guaranteed to be flawless across all datasets/environments without further validation.
+
+Updated: 2026-03-05
+

openpi/scripts/convert_twin_squashfs_to_lerobot.py

@@ -0,0 +1,320 @@
+#!/usr/bin/env python3
+"""
+Convert TWIN bimanual RLBench squashfs files to a LeRobot dataset.
+
+This script intentionally supports bounded conversion via --max-episodes/--max-frames
+for local stress testing before running full conversion on larger machines.
+"""
+
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import logging
+import os
+import pickle
+from pathlib import Path
+import re
+import shutil
+import subprocess
+import tempfile
+from typing import Any
+
+from lerobot.common.datasets.lerobot_dataset import HF_LEROBOT_HOME
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+import numpy as np
+from PIL import Image
+
+
+@dataclasses.dataclass(frozen=True)
+class ConverterStats:
+    episodes_seen: int = 0
+    episodes_written: int = 0
+    frames_written: int = 0
+    frames_skipped: int = 0
+
+
+class _FallbackUnpickler(pickle.Unpickler):
+    """Loads RLBench pickles even when rlbench is not installed."""
+
+    _cache: dict[tuple[str, str], type] = {}
+
+    def find_class(self, module: str, name: str) -> Any:
+        try:
+            return super().find_class(module, name)
+        except Exception:
+            key = (module, name)
+            if key not in self._cache:
+                cls = type(name, (), {})
+                cls.__module__ = module
+                self._cache[key] = cls
+            return self._cache[key]
+
+
+def _load_pickle(path: Path) -> Any:
+    with path.open("rb") as f:
+        return _FallbackUnpickler(f).load()
+
+
+def _run_unsquashfs(squashfs_path: Path, dest_dir: Path, patterns: list[str]) -> None:
+    cmd = [
+        "unsquashfs",
+        "-f",
+        "-d",
+        str(dest_dir),
+        str(squashfs_path),
+        *patterns,
+    ]
+    logging.info("Running: %s", " ".join(cmd))
+    subprocess.run(cmd, check=True)
+
+
+def _episode_sort_key(ep_name: str) -> int:
+    m = re.match(r"episode(\d+)$", ep_name)
+    if not m:
+        return 10**9
+    return int(m.group(1))
+
+
+def _collect_episode_dirs(extract_root: Path) -> list[Path]:
+    base = extract_root / "all_variations" / "episodes"
+    if not base.exists():
+        return []
+    episodes = [p for p in base.iterdir() if p.is_dir() and p.name.startswith("episode")]
+    return sorted(episodes, key=lambda p: _episode_sort_key(p.name))
+
+
+def _image_array(path: Path) -> np.ndarray:
+    return np.asarray(Image.open(path).convert("RGB"))
+
+
+def _to_state(obs: Any) -> np.ndarray:
+    left = obs.left
+    right = obs.right
+    return np.concatenate(
+        [
+            np.asarray(left.joint_positions, dtype=np.float32),
+            np.asarray([left.gripper_open], dtype=np.float32),
+            np.asarray(right.joint_positions, dtype=np.float32),
+            np.asarray([right.gripper_open], dtype=np.float32),
+        ],
+        dtype=np.float32,
+    )
+
+
+def _to_action(obs: Any) -> np.ndarray:
+    left = obs.left
+    right = obs.right
+    return np.concatenate(
+        [
+            np.asarray(left.joint_velocities, dtype=np.float32),
+            np.asarray([left.gripper_open], dtype=np.float32),
+            np.asarray(right.joint_velocities, dtype=np.float32),
+            np.asarray([right.gripper_open], dtype=np.float32),
+        ],
+        dtype=np.float32,
+    )
+
+
+def _select_prompt(variation_descriptions: Any) -> str:
+    if isinstance(variation_descriptions, (list, tuple)) and variation_descriptions:
+        for item in variation_descriptions:
+            if isinstance(item, str) and item.strip():
+                return item.strip()
+    return "perform the task"
+
+
+def convert(
+    squashfs_path: Path,
+    repo_id: str,
+    *,
+    cameras: list[str],
+    max_episodes: int | None,
+    max_frames: int | None,
+    fps: int,
+    push_to_hub: bool,
+    private: bool,
+    cleanup_output: bool,
+) -> ConverterStats:
+    output_path = HF_LEROBOT_HOME / repo_id
+    if cleanup_output and output_path.exists():
+        shutil.rmtree(output_path)
+
+    with tempfile.TemporaryDirectory(prefix="twin_unsquash_") as tmp:
+        extract_root = Path(tmp)
+
+        # 1) Extract episode metadata only.
+        _run_unsquashfs(
+            squashfs_path,
+            extract_root,
+            [
+                "all_variations/episodes/episode*/low_dim_obs.pkl",
+                "all_variations/episodes/episode*/variation_descriptions.pkl",
+                "all_variations/episodes/episode*/variation_number.pkl",
+            ],
+        )
+
+        episode_dirs = _collect_episode_dirs(extract_root)
+        if max_episodes is not None:
+            episode_dirs = episode_dirs[:max_episodes]
+        if not episode_dirs:
+            raise RuntimeError("No episodes found after metadata extraction.")
+
+        # 2) Extract RGB images only for selected episodes/cameras.
+        image_patterns: list[str] = []
+        for ep in episode_dirs:
+            for camera in cameras:
+                image_patterns.append(f"all_variations/episodes/{ep.name}/{camera}_rgb/*.png")
+        _run_unsquashfs(squashfs_path, extract_root, image_patterns)
+
+        # Determine image shapes from first valid frame.
+        first_frame = None
+        for ep in episode_dirs:
+            for camera in cameras:
+                candidate = ep / f"{camera}_rgb" / "rgb_0000.png"
+                if candidate.exists():
+                    first_frame = candidate
+                    break
+            if first_frame is not None:
+                break
+        if first_frame is None:
+            raise RuntimeError("No RGB frames found in extracted episodes.")
+        h, w, c = _image_array(first_frame).shape
+        if c != 3:
+            raise RuntimeError(f"Expected RGB images with 3 channels, got shape {(h, w, c)}")
+
+        features: dict[str, dict[str, Any]] = {
+            f"{camera}_image": {
+                "dtype": "image",
+                "shape": (h, w, 3),
+                "names": ["height", "width", "channel"],
+            }
+            for camera in cameras
+        }
+        features["state"] = {
+            "dtype": "float32",
+            "shape": (16,),
+            "names": [f"state_{i}" for i in range(16)],
+        }
+        features["action"] = {
+            "dtype": "float32",
+            "shape": (16,),
+            "names": [f"action_{i}" for i in range(16)],
+        }
+
+        dataset = LeRobotDataset.create(
+            repo_id=repo_id,
+            robot_type="rlbench_bimanual",
+            fps=fps,
+            features=features,
+            image_writer_threads=min(8, os.cpu_count() or 4),
+            image_writer_processes=1,
+        )
+
+        stats = ConverterStats()
+        for ep_dir in episode_dirs:
+            stats = dataclasses.replace(stats, episodes_seen=stats.episodes_seen + 1)
+            low_dim_obs = _load_pickle(ep_dir / "low_dim_obs.pkl")
+            variation_descriptions = _load_pickle(ep_dir / "variation_descriptions.pkl")
+            prompt = _select_prompt(variation_descriptions)
+
+            observations = getattr(low_dim_obs, "_observations", None)
+            if observations is None:
+                logging.warning("Skipping %s: missing _observations.", ep_dir.name)
+                continue
+
+            frame_limit = len(observations)
+            if max_frames is not None:
+                frame_limit = min(frame_limit, max_frames)
+
+            written_in_episode = 0
+            for i in range(frame_limit):
+                frame_data: dict[str, Any] = {}
+                missing = False
+                for camera in cameras:
+                    frame_path = ep_dir / f"{camera}_rgb" / f"rgb_{i:04d}.png"
+                    if not frame_path.exists():
+                        missing = True
+                        break
+                    frame_data[f"{camera}_image"] = _image_array(frame_path)
+                if missing:
+                    stats = dataclasses.replace(stats, frames_skipped=stats.frames_skipped + 1)
+                    continue
+
+                obs = observations[i]
+                frame_data["state"] = _to_state(obs)
+                frame_data["action"] = _to_action(obs)
+                frame_data["task"] = prompt
+
+                dataset.add_frame(frame_data)
+                stats = dataclasses.replace(stats, frames_written=stats.frames_written + 1)
+                written_in_episode += 1
+
+            if written_in_episode > 0:
+                dataset.save_episode()
+                stats = dataclasses.replace(stats, episodes_written=stats.episodes_written + 1)
+
+        if push_to_hub:
+            dataset.push_to_hub(
+                private=private,
+                push_videos=True,
+                tags=["twin", "bimanual", "rlbench", "lerobot", "openpi"],
+                license="apache-2.0",
+            )
+
+        return stats
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--squashfs-path", type=Path, required=True, help="Path to one TWIN squashfs file.")
+    parser.add_argument(
+        "--repo-id",
+        required=True,
+        help="LeRobot repo id (e.g. your_hf_username/twin_bimanual_dual_push_train).",
+    )
+    parser.add_argument(
+        "--cameras",
+        default="front,wrist_left,wrist_right",
+        help="Comma-separated camera names, without '_rgb' suffix.",
+    )
+    parser.add_argument("--max-episodes", type=int, default=None, help="Limit number of episodes to convert.")
+    parser.add_argument("--max-frames", type=int, default=None, help="Limit number of frames per episode.")
+    parser.add_argument("--fps", type=int, default=10)
+    parser.add_argument("--push-to-hub", action="store_true")
+    parser.add_argument("--private", action="store_true")
+    parser.add_argument("--no-cleanup-output", action="store_true")
+    parser.add_argument("--verbose", action="store_true")
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = _parse_args()
+    logging.basicConfig(level=logging.INFO if args.verbose else logging.WARNING, format="%(levelname)s: %(message)s")
+    cameras = [c.strip() for c in args.cameras.split(",") if c.strip()]
+    if not cameras:
+        raise ValueError("At least one camera must be specified in --cameras.")
+
+    stats = convert(
+        squashfs_path=args.squashfs_path,
+        repo_id=args.repo_id,
+        cameras=cameras,
+        max_episodes=args.max_episodes,
+        max_frames=args.max_frames,
+        fps=args.fps,
+        push_to_hub=args.push_to_hub,
+        private=args.private,
+        cleanup_output=not args.no_cleanup_output,
+    )
+    print(
+        "Conversion complete:",
+        f"episodes_seen={stats.episodes_seen}",
+        f"episodes_written={stats.episodes_written}",
+        f"frames_written={stats.frames_written}",
+        f"frames_skipped={stats.frames_skipped}",
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

openpi/src/openpi/training/config.py

@@ -462,6 +462,47 @@ class LeRobotDROIDDataConfig(DataConfigFactory):
         )
 
 
+@dataclasses.dataclass(frozen=True)
+class LeRobotTWINBimanualDataConfig(DataConfigFactory):
+    """
+    Data config for TWIN bimanual datasets converted to LeRobot format via
+    scripts/convert_twin_squashfs_to_lerobot.py.
+    """
+
+    default_prompt: str | None = None
+
+    @override
+    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
+        repack_transform = _transforms.Group(
+            inputs=[
+                _transforms.RepackTransform(
+                    {
+                        "images": {
+                            "cam_high": "front_image",
+                            "cam_left_wrist": "wrist_left_image",
+                            "cam_right_wrist": "wrist_right_image",
+                        },
+                        "state": "state",
+                        "actions": "action",
+                        "prompt": "task",
+                    }
+                )
+            ]
+        )
+        data_transforms = _transforms.Group(
+            inputs=[aloha_policy.AlohaInputs(adapt_to_pi=False)],
+            outputs=[],
+        )
+        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,
+            action_sequence_keys=("action",),
+        )
+
+
 @dataclasses.dataclass(frozen=True)
 class TrainConfig:
     # Name of the config. Must be unique. Will be used to reference this config.
@@ -938,6 +979,39 @@ _CONFIGS = [
         num_train_steps=20_000,
         batch_size=16,
     ),
+    TrainConfig(
+        # Baseline pi05 fine-tuning on TWIN bimanual LeRobot data (single action head).
+        name="pi05_twin_bimanual_finetune",
+        model=pi0_config.Pi0Config(
+            pi05=True,
+            action_dim=32,  # Keep pi05 pretraining action dimensionality.
+            action_horizon=16,
+        ),
+        data=LeRobotTWINBimanualDataConfig(
+            repo_id="your_hf_username/twin_bimanual_lerobot_train",
+            base_config=DataConfig(prompt_from_task=False),
+        ),
+        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
+        num_train_steps=20_000,
+        batch_size=16,
+    ),
+    TrainConfig(
+        # Parallel per-arm action-head pi05 fine-tuning on TWIN bimanual LeRobot data.
+        name="pi05_twin_bimanual_parallel_finetune",
+        model=pi0_config.Pi0Config(
+            pi05=True,
+            action_dim=32,
+            action_horizon=16,
+            arm_action_dims=(16, 16),
+        ),
+        data=LeRobotTWINBimanualDataConfig(
+            repo_id="your_hf_username/twin_bimanual_lerobot_train",
+            base_config=DataConfig(prompt_from_task=False),
+        ),
+        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
+        num_train_steps=20_000,
+        batch_size=16,
+    ),
     #
     # ALOHA Sim configs. This config is used to demonstrate how to train on a simple simulated environment.
     #
@@ -1010,6 +1084,33 @@ _CONFIGS = [
         wandb_enabled=False,
         pytorch_training_precision="float32",
     ),
+    TrainConfig(
+        # Local smoke-test for converted TWIN LeRobot data.
+        name="debug_pi05_twin_bimanual_parallel_local_smoke",
+        model=pi0_config.Pi0Config(
+            pi05=True,
+            paligemma_variant="dummy",
+            action_expert_variant="dummy",
+            action_dim=32,
+            action_horizon=8,
+            max_token_len=64,
+            arm_action_dims=(16, 16),
+        ),
+        data=LeRobotTWINBimanualDataConfig(
+            # This repo id is produced by scripts/convert_twin_squashfs_to_lerobot.py in local smoke mode.
+            repo_id="local/twin_bimanual_dual_push_smoke",
+            base_config=DataConfig(prompt_from_task=False),
+        ),
+        batch_size=1,
+        num_workers=0,
+        num_train_steps=2,
+        log_interval=1,
+        save_interval=1,
+        overwrite=True,
+        exp_name="debug_pi05_twin_bimanual_parallel_local_smoke",
+        wandb_enabled=False,
+        pytorch_training_precision="float32",
+    ),
     # RoboArena & PolaRiS configs.
     *roboarena_config.get_roboarena_configs(),
     *polaris_config.get_polaris_configs(),