[Group3] Submission

.gitignore

@@ -30,3 +30,5 @@ dist/
 .vscode
 
 imgui.ini
+
+output

scripts/group3/README.md

@@ -0,0 +1,34 @@
+# Group3 SimplerEnv Evaluation
+
+
+## Setup
+```
+git clone https://github.com/airoa-org/SimplerEnv.git
+cd SimplerEnv
+git checkout benchmark-v2-g3-submission
+git submodule update --init --recursive
+
+# Create a conda environment
+export REPO_ROOT="$(pwd -P)"
+conda env create -f scripts/group3/environment.yaml
+conda activate simpler-benchmark-v2-g3-submission
+
+# Downlaod the group3 simpler_env model from wasabi
+aws s3 cp s3://airoa-fm-development-competition/group3/submitted_202509291552_simpler ./g3_simpler_model/ --recursive --endpoint-url=https://s3.ap-northeast-1.wasabisys.com
+```
+
+
+## Evaluation
+
+**Google Robot**
+```
+conda activate simpler-benchmark-v2-g3-submission
+python scripts/group3/evaluate_fractal.py --ckpt-path ./g3_simpler_model
+```
+
+
+**WidowX**
+```
+conda activate simpler-benchmark-v2-g3-submission
+python scripts/group3/evaluate_bridge.py --ckpt-path ./g3_simpler_model
+```

scripts/group3/environment.yaml

@@ -0,0 +1,30 @@
+name: simpler-benchmark-v2-g3-submission
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.11
+  # system/conda packages
+  - ffmpeg
+  - libvulkan-loader
+  - libvulkan-headers
+  - libiconv=1.17
+  - libgcc-ng
+  - libstdcxx-ng
+  - awscli
+  - pip
+
+  - pip:
+      - numpy==1.25.2
+
+      - -e file://${REPO_ROOT}/ManiSkill2_real2sim
+      - -e file://${REPO_ROOT}
+
+      - torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124
+      - torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu124
+      - torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu124
+
+      - git+ssh://git@github.com/huggingface/lerobot.git@67196c9d5344cd932612cef79229f9d04134c91e#egg=lerobot[pi0]
+
+      - pytest
+      - statsmodels==0.14.5

scripts/group3/evaluate_bridge.py

@@ -0,0 +1,58 @@
+import argparse
+import time
+
+import numpy as np
+
+from scripts.group3.g3_configuration_pi0 import G3PI0Config
+from scripts.group3.g3_pi0_or_fast import G3LerobotPiFastInference
+from simpler_env.evaluation.bridge_tasks import (
+    widowx_task1_pick_object,
+    widowx_task2_stack_cube,
+    widowx_task3_put_object_on_top,
+    widowx_task4_put_object_in_basket
+)
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Run Comprehensive ManiSkill2 Evaluation")
+    parser.add_argument("--ckpt-path", type=str, required=True, help="Path to the checkpoint to evaluate.")
+    parser.add_argument("--control-freq", type=int, default=5, help="Set control frequency (default->5)")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    N_ACTION_STEPS = 2
+    ACTION_ENSEMBLE_TEMP = 0.6
+    ACTION_ENSEMBLE = True
+    STICKY_ACTION = False
+
+    args = parse_args()
+    ckpt_path = args.ckpt_path
+
+    policy = G3LerobotPiFastInference(
+        saved_model_path=ckpt_path,
+        policy_setup="widowx_bridge",
+        action_scale=1.0,
+        action_ensemble_temp=ACTION_ENSEMBLE_TEMP,
+        action_ensemble=ACTION_ENSEMBLE,
+        sticky_action=STICKY_ACTION,
+        n_action_steps=N_ACTION_STEPS,
+    )
+
+    print("Policy initialized. Starting evaluation...")
+
+    tasks = [
+        widowx_task1_pick_object, 
+        widowx_task2_stack_cube, 
+        widowx_task3_put_object_on_top, 
+        widowx_task4_put_object_in_basket
+    ]
+
+    final_scores = []
+    for task in tasks:
+        cur_scores = task(
+            env_policy=policy, ckpt_path=args.ckpt_path, control_freq=args.control_freq
+        )
+        final_scores += cur_scores
+
+    print("\nEvaluation finished.")
+    print(f"Final calculated scores: {final_scores}")

scripts/group3/evaluate_fractal.py

@@ -0,0 +1,39 @@
+import argparse
+import datetime
+
+from scripts.group3.g3_configuration_pi0 import G3PI0Config
+from scripts.group3.g3_pi0_or_fast import G3LerobotPiFastInference
+from simpler_env.evaluation.fractal_tasks import run_comprehensive_evaluation
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Run Comprehensive ManiSkill2 Evaluation")
+    parser.add_argument("--ckpt-path", type=str, required=True, help="Path to the checkpoint to evaluate.")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    N_ACTION_STEPS = 4
+    ACTION_ENSEMBLE_TEMP = 0.8
+    ACTION_ENSEMBLE = False
+    STICKY_ACTION = False
+
+    args = parse_args()
+    ckpt_path = args.ckpt_path
+
+    policy = G3LerobotPiFastInference(
+        saved_model_path=ckpt_path,
+        policy_setup="google_robot",
+        action_scale=1.0,
+        action_ensemble_temp=ACTION_ENSEMBLE_TEMP,
+        action_ensemble=ACTION_ENSEMBLE,
+        sticky_action=STICKY_ACTION,
+        n_action_steps=N_ACTION_STEPS,
+    )
+
+    print("Policy initialized. Starting evaluation...")
+
+    final_scores = run_comprehensive_evaluation(env_policy=policy, ckpt_path=args.ckpt_path)
+
+    print("\nEvaluation finished.")
+    print(f"Final calculated scores: {final_scores}")

scripts/group3/g3_configuration_pi0.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass, field
+
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.policies.pi0.configuration_pi0 import PI0Config
+
+
+@PreTrainedConfig.register_subclass("g3pi0")
+@dataclass
+class G3PI0Config(PI0Config):
+    max_ft_dim : str = ""
+    train_ft_proj : bool = True
+    action_key : str = ""
+    ft_key : str = ""
+    encoder_type : str = "seq_cnn"
+    finetune: bool = False
+    finetune_model: str = ""
+    multi_embodiment: bool = False

scripts/group3/g3_pi0_or_fast.py

@@ -0,0 +1,184 @@
+from typing import List, Optional
+import os
+from collections import deque
+
+import torch
+import numpy as np
+from PIL import Image
+from transforms3d.euler import euler2axangle
+
+from simpler_env.policies.lerobotpi.geometry import mat2euler, quat2mat
+from simpler_env.policies.lerobotpi.pi0_or_fast import LerobotPiFastInference, auto_model_fn
+from simpler_env.utils.action.action_ensemble import ActionEnsembler
+
+
+class G3LerobotPiFastInference(LerobotPiFastInference):
+    def __init__(
+        self,
+        saved_model_path: str = "pretrained/pi0",
+        unnorm_key: Optional[str] = None,
+        policy_setup: str = "widowx_bridge",
+        exec_horizon: int = 4,
+        image_size: list[int] = [224, 224],
+        action_scale: float = 1.0,
+        action_ensemble: bool = True,
+        action_ensemble_temp: float = -0.8,
+        sticky_action: bool = True,
+        n_action_steps: int = 4,
+    ) -> None:
+        gpu_idx = os.environ.get("GPU_IDX", 0)
+        self.device = f"cuda:{gpu_idx}"
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+        self.default_rot = np.array(
+            [[0, 0, 1.0], [0, 1.0, 0], [-1.0, 0, 0]]
+        )  # https://github.com/rail-berkeley/bridge_data_robot/blob/b841131ecd512bafb303075bd8f8b677e0bf9f1f/widowx_envs/widowx_controller/src/widowx_controller/widowx_controller.py#L203
+        if policy_setup == "widowx_bridge":
+            unnorm_key = "bridge_orig/1.0.0" if unnorm_key is None else unnorm_key
+            self.sticky_gripper_num_repeat = 1
+            # EE pose in Bridge data was relative to a top-down pose, instead of robot base
+        elif policy_setup == "google_robot":
+            unnorm_key = "fractal20220817_data/0.1.0" if unnorm_key is None else unnorm_key
+            self.sticky_gripper_num_repeat = 10
+        else:
+            raise NotImplementedError(
+                f"Policy setup {policy_setup} not supported for octo models. The other datasets can be found in the huggingface config.json file."
+            )
+        self.sticky_action = sticky_action
+        self.policy_setup = policy_setup
+        self.unnorm_key = unnorm_key
+
+        print(f"*** policy_setup: {policy_setup}, unnorm_key: {unnorm_key} ***")
+
+        # TODO: add pi0 loading ...
+        PI0Policy = auto_model_fn(saved_model_path)
+        self.vla = PI0Policy.from_pretrained(saved_model_path)
+        self.vla.model.paligemma_with_expert.paligemma.language_model = self.vla.model.paligemma_with_expert.paligemma.language_model.model
+        self.vla.model.paligemma_with_expert.gemma_expert.model = self.vla.model.paligemma_with_expert.gemma_expert.model.base_model
+        self.vla.config.n_action_steps = n_action_steps
+        self.vla.to(self.device)
+        self.vla.reset()
+
+        self.image_size = image_size
+        self.action_scale = action_scale
+        self.obs_horizon = 1
+        self.obs_interval = 1
+        self.pred_action_horizon = self.vla.config.n_action_steps
+        self.image_history = deque(maxlen=self.obs_horizon)
+        self.exec_horizon = exec_horizon
+
+        self.sticky_action_is_on = False
+        self.gripper_action_repeat = 0
+        self.sticky_gripper_action = 0.0
+        self.previous_gripper_action = None
+
+        self.action_ensemble = action_ensemble
+        self.action_ensemble_temp = action_ensemble_temp
+
+        if self.action_ensemble:
+            self.action_ensembler = ActionEnsembler(self.pred_action_horizon, self.action_ensemble_temp)
+        else:
+            self.action_ensembler = None
+
+        self.task = None
+        self.task_description = None
+
+
+    def step(self, image: np.ndarray, task_description: Optional[str] = None, *args, **kwargs) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
+        """
+        Input:
+            image: np.ndarray of shape (H, W, 3), uint8
+            task_description: Optional[str], task description; if different from previous task description, policy state is reset
+        Output:
+            raw_action: dict; raw policy action output
+            action: dict; processed action to be sent to the maniskill2 environment, with the following keys:
+                - 'world_vector': np.ndarray of shape (3,), xyz translation of robot end-effector
+                - 'rot_axangle': np.ndarray of shape (3,), axis-angle representation of end-effector rotation
+                - 'gripper': np.ndarray of shape (1,), gripper action
+                - 'terminate_episode': np.ndarray of shape (1,), 1 if episode should be terminated, 0 otherwise
+        """
+        if task_description is not None:
+            if task_description != self.task_description:
+                self.reset(task_description)
+
+        assert image.dtype == np.uint8
+        image = self._resize_image(image)
+        self._add_image_to_history(image)
+        images: List[Image.Image] = self._obtain_image_history()
+
+        eef_pos = kwargs.get("eef_pos", None)
+
+        state = self.preprocess_widowx_proprio(eef_pos)
+        observation = {
+            "observation.state": torch.from_numpy(state).unsqueeze(0).to(self.device).float(),
+            "observation.images.image_0": torch.from_numpy(images[0] / 255).permute(2, 0, 1).unsqueeze(0).to(self.device).float(),
+            "observation.images.image_1": torch.from_numpy(images[0] / 255).permute(2, 0, 1).unsqueeze(0).to(self.device).float(),
+            "observation.images.image_2": torch.from_numpy(images[0] / 255).permute(2, 0, 1).unsqueeze(0).to(self.device).float(),
+            "observation.images.image_3": torch.from_numpy(images[0] / 255).permute(2, 0, 1).unsqueeze(0).to(self.device).float(),
+            "task": [task_description],
+        }
+
+        actions = self.vla.select_action(observation)[0].cpu().numpy()
+
+        if self.action_ensemble:
+            action_chunk = [actions]
+            for _ in range(self.vla.config.n_action_steps-1):
+                actions = self.vla.select_action(observation)[0].cpu().numpy()
+                action_chunk.append(actions)
+            action_chunk = np.stack(action_chunk, axis=0)
+            actions = self.action_ensembler.ensemble_action(action_chunk)[None][0]
+
+        raw_action = {
+            "world_vector": np.array(actions[:3]),
+            "rotation_delta": np.array(actions[3:6]),
+            "open_gripper": np.array(actions[6:7]),  # range [0, 1]; 1 = open; 0 = close
+        }
+
+        # process raw_action to obtain the action to be sent to the maniskill2 environment
+        action = {}
+        action["world_vector"] = raw_action["world_vector"] * self.action_scale
+        action_rotation_delta = np.asarray(raw_action["rotation_delta"], dtype=np.float64)
+        roll, pitch, yaw = action_rotation_delta
+        action_rotation_ax, action_rotation_angle = euler2axangle(roll, pitch, yaw)
+        action_rotation_axangle = action_rotation_ax * action_rotation_angle
+        action["rot_axangle"] = action_rotation_axangle * self.action_scale
+
+        if self.policy_setup == "google_robot":
+            if self.sticky_action:
+                action["gripper"] = 0
+                current_gripper_action = raw_action["open_gripper"]
+                if self.previous_gripper_action is None:
+                    relative_gripper_action = np.array([0])
+                    self.previous_gripper_action = current_gripper_action
+                else:
+                    relative_gripper_action = self.previous_gripper_action - current_gripper_action
+
+                # fix a bug in the SIMPLER code here
+                # self.previous_gripper_action = current_gripper_action
+
+                if np.abs(relative_gripper_action) > 0.5 and (not self.sticky_action_is_on):
+                    self.sticky_action_is_on = True
+                    self.sticky_gripper_action = relative_gripper_action
+                    self.previous_gripper_action = current_gripper_action
+
+                if self.sticky_action_is_on:
+                    self.gripper_action_repeat += 1
+                    relative_gripper_action = self.sticky_gripper_action
+
+                if self.gripper_action_repeat == self.sticky_gripper_num_repeat:
+                    self.sticky_action_is_on = False
+                    self.gripper_action_repeat = 0
+                    self.sticky_gripper_action = 0.0
+
+                action["gripper"] = relative_gripper_action
+            else:
+                current_gripper_action = raw_action["open_gripper"]
+                current_gripper_action = (current_gripper_action * 2) - 1
+                current_gripper_action = - current_gripper_action
+                action["gripper"] = current_gripper_action
+
+        elif self.policy_setup == "widowx_bridge":
+            action["gripper"] = 2.0 * (raw_action["open_gripper"] > 0.5) - 1.0
+
+        action["terminate_episode"] = np.array([0.0])
+        return raw_action, action