diff --git a/.gitignore b/.gitignore index d8002eb9..99b74941 100644 --- a/.gitignore +++ b/.gitignore @@ -30,3 +30,5 @@ dist/ .vscode imgui.ini + +output diff --git a/scripts/group3/README.md b/scripts/group3/README.md new file mode 100644 index 00000000..c1b236c6 --- /dev/null +++ b/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 +``` diff --git a/scripts/group3/environment.yaml b/scripts/group3/environment.yaml new file mode 100644 index 00000000..ee79c21d --- /dev/null +++ b/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 diff --git a/scripts/group3/evaluate_bridge.py b/scripts/group3/evaluate_bridge.py new file mode 100644 index 00000000..f57e585d --- /dev/null +++ b/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}") diff --git a/scripts/group3/evaluate_fractal.py b/scripts/group3/evaluate_fractal.py new file mode 100644 index 00000000..2858a0d2 --- /dev/null +++ b/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}") diff --git a/scripts/group3/g3_configuration_pi0.py b/scripts/group3/g3_configuration_pi0.py new file mode 100644 index 00000000..e6987710 --- /dev/null +++ b/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 diff --git a/scripts/group3/g3_pi0_or_fast.py b/scripts/group3/g3_pi0_or_fast.py new file mode 100644 index 00000000..5a15c463 --- /dev/null +++ b/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 diff --git a/scripts/openpi/challenge_widowx.py b/scripts/openpi/challenge_widowx.py index 503997bc..15861e27 100644 --- a/scripts/openpi/challenge_widowx.py +++ b/scripts/openpi/challenge_widowx.py @@ -1,42 +1,299 @@ import argparse +from datetime import datetime +from types import SimpleNamespace +import json +import os +import tempfile +import time +from typing import Dict, List +import draccus +import numpy as np +import torch + +from scripts.g3_lerobotpi.g3_pi0_or_fast import G3LerobotPiFastInference +from scripts.g3_lerobotpi.submission_utils import ( + copy_run_summaries, + dump_submission_payload, + prepare_submission_dir, +) 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 + widowx_task4_put_object_in_basket, ) -from simpler_env.policies.openpi.pi0_or_fast import OpenPiFastInference +from simpler_env.evaluation.evaluate import calculate_robust_score +from simpler_env.policies.g3lerobotpi.adapter import AiroaToG3Pi0FractalBridgeAdapter +from simpler_env.policies.g3lerobotpi.policy import G3Pi0multiLerobotToAiroaPolicy +from g3_haptics.policies.g3factory import get_policy_class, get_policy_config_class 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)") + parser.add_argument( + "--policy-setup", + type=str, + default="widowx_bridge", + choices=["widowx_bridge", "google_robot"], + help="Policy setup to use when loading the G3 Pi0 checkpoint.", + ) + parser.add_argument( + "--control-freq", + type=int, + default=5, + help="Control frequency (Hz) used for evaluation (default 5)", + ) + parser.add_argument( + "--save-path-suffix", + type=str, + default="", + help="Suffix appended to the checkpoint name when saving results.", + ) + parser.add_argument( + "--multi-embodiment", + action=argparse.BooleanOptionalAction, + default=False, + help="Enable the multi-embodiment Pi0 policy pipeline.", + ) + parser.add_argument( + "--rot6d", + action=argparse.BooleanOptionalAction, + default=False, + help="Use 6D rotation state when running the multi-embodiment adapter.", + ) + parser.add_argument( + "--action-ensemble", + action=argparse.BooleanOptionalAction, + default=False, + help="Toggle action ensemble (applies to both single and multi pipelines).", + ) + parser.add_argument( + "--action-ensemble-temp", + type=float, + default=-0.8, + help="Temperature for the action ensemble module.", + ) + parser.add_argument( + "--sticky-action", + action=argparse.BooleanOptionalAction, + default=False, + help="Use sticky gripper control when policy setup is google_robot.", + ) + parser.add_argument( + "--exec-horizon", + type=int, + default=4, + help="Number of actions executed per inference step in the fast single-embodiment pipeline.", + ) + parser.add_argument( + "--add-task-prefix", + action="store_true", + help="Prepend the policy setup to the language prompt (multi-embodiment only).", + ) + parser.add_argument( + "--submission-dir", + type=str, + default="submissions", + help="Directory to write submission artifacts (score.json, summaries).", + ) + parser.add_argument( + "--submission-name", + type=str, + default=None, + help="Optional fixed name for the submission folder (defaults to timestamped checkpoint id).", + ) + parser.add_argument( + "--no-save-submission", + action="store_true", + help="Skip writing submission artifacts; only print results.", + ) return parser.parse_args() +def _build_multi_policy(args: argparse.Namespace) -> AiroaToG3Pi0FractalBridgeAdapter: + torch.backends.cudnn.benchmark = True + torch.backends.cuda.matmul.allow_tf32 = True + + max_state_dim = 10 if args.rot6d else 8 + + dataset_cfg = SimpleNamespace( + embodiments=[ + SimpleNamespace( + state_key="observation.state", + image_key=["observation.images.image"], + action_key="action", + ft_key=None, + ), + SimpleNamespace( + state_key="observation.state", + image_key=["observation.images.image"], + action_key="action", + ft_key=None, + ), + ], + max_state_dim=max_state_dim, + max_image_num=1, + max_image_shape=[3, 256, 320], + max_action_dim=7, + ) + + emb_count = 2 + dataset_stats = { + "observation.state": { + "mean": torch.zeros(emb_count, max_state_dim), + "std": torch.ones(emb_count, max_state_dim), + }, + "action": { + "mean": torch.zeros(emb_count, 7), + "std": torch.ones(emb_count, 7), + }, + } + + with open(os.path.join(args.ckpt_path, "config.json"), encoding="utf-8") as f: + raw_cfg = json.load(f) + + PolicyCfgClass = get_policy_config_class(raw_cfg.pop("type")) + with tempfile.NamedTemporaryFile("w+", encoding="utf-8") as tmp: + json.dump(raw_cfg, tmp) + tmp.flush() + with draccus.config_type("json"): + cfg = draccus.parse(PolicyCfgClass, tmp.name, args=[]) + + cfg.multi_embodiment = True + + PolicyCls = get_policy_class("g3pi0") + base_policy = PolicyCls.from_pretrained( + pretrained_name_or_path=args.ckpt_path, + config=cfg, + dataset_stats=dataset_stats, + ) + + wrapped_policy = G3Pi0multiLerobotToAiroaPolicy( + policy=base_policy, + dataset_cfg=dataset_cfg, + policy_setup=args.policy_setup, + ) + + return AiroaToG3Pi0FractalBridgeAdapter( + policy=wrapped_policy, + rot6d=args.rot6d, + policy_setup=args.policy_setup, + action_ensemble=args.action_ensemble, + action_ensemble_temp=args.action_ensemble_temp, + sticky_action=args.sticky_action, + add_taks_prefix=args.add_task_prefix, + ) + + +def _build_single_policy(args: argparse.Namespace) -> G3LerobotPiFastInference: + return G3LerobotPiFastInference( + saved_model_path=args.ckpt_path, + policy_setup=args.policy_setup, + action_scale=1.0, + action_ensemble=args.action_ensemble, + action_ensemble_temp=args.action_ensemble_temp, + sticky_action=args.sticky_action, + exec_horizon=args.exec_horizon, + ) + + +def _flatten_runs(runs: List[List[bool]]) -> List[bool]: + return [bool(item) for episode in runs for item in episode] + + if __name__ == "__main__": args = parse_args() - ckpt_path = args.ckpt_path - policy = OpenPiFastInference(saved_model_path=ckpt_path, policy_setup="widowx_bridge", action_scale=1.0) + ckpt_identifier = ( + args.ckpt_path + if not args.save_path_suffix + else args.ckpt_path + f"_{args.save_path_suffix}" + ) + should_save, submission_path = prepare_submission_dir( + args.submission_dir, + args.submission_name, + ckpt_identifier, + default_tokens=[args.policy_setup, "widowx"], + skip=args.no_save_submission, + ) + + if args.multi_embodiment: + policy = _build_multi_policy(args) + else: + policy = _build_single_policy(args) 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 + widowx_task1_pick_object, + widowx_task2_stack_cube, + widowx_task3_put_object_on_top, + widowx_task4_put_object_in_basket, ] - final_scores = [] + final_scores: List[List[bool]] = [] + task_breakdown: Dict[str, Dict[str, object]] = {} + for task in tasks: + start_time = time.time() + cur_scores = task( - env_policy=policy, ckpt_path=args.ckpt_path, control_freq=args.control_freq + env_policy=policy, + ckpt_path=ckpt_identifier, + control_freq=args.control_freq, ) final_scores += cur_scores + flat_scores = _flatten_runs(cur_scores) + success_rate = float(np.mean(flat_scores)) if flat_scores else 0.0 + num_success = int(np.sum(flat_scores)) + num_trials = int(len(flat_scores)) + + print(f"Task: {task.__name__}") + print(f"Time: {(time.time() - start_time) / 60:.2f} min") + print(f"Success Rate: {success_rate * 100:.2f}% ({num_success}/{num_trials})") + + task_breakdown[task.__name__] = { + "success_rate": success_rate, + "num_success": num_success, + "num_trials": num_trials, + "per_trial_success": [bool(x) for x in flat_scores], + } + + flattened = _flatten_runs(final_scores) + robust_score = float(calculate_robust_score(final_scores)) if final_scores else 0.0 + overall_success = float(np.mean(flattened)) if flattened else 0.0 + print("\nEvaluation finished.") - print(f"Final calculated scores: {final_scores}") + print(f"Final Success Rate: {overall_success * 100:.2f}%") + print(f"Final calculated scores: {robust_score}") + + if should_save and submission_path is not None: + payload: Dict[str, object] = { + "policy_setup": args.policy_setup, + "checkpoint": args.ckpt_path, + "timestamp": datetime.now().isoformat(), + "config": { + "control_freq": args.control_freq, + "multi_embodiment": args.multi_embodiment, + "rot6d": args.rot6d, + "action_ensemble": args.action_ensemble, + "action_ensemble_temp": args.action_ensemble_temp, + "sticky_action": args.sticky_action, + "exec_horizon": args.exec_horizon, + "add_task_prefix": args.add_task_prefix, + "save_path_suffix": args.save_path_suffix, + }, + "results": { + "robust_score": robust_score, + "overall_success_rate": overall_success, + "num_success": int(np.sum(flattened)), + "num_trials": int(len(flattened)), + "task_breakdown": task_breakdown, + }, + } + + dump_submission_payload(submission_path, payload) + copy_run_summaries(submission_path, ckpt_identifier) diff --git a/scripts/rt1/evaluate_fractal.py b/scripts/rt1/evaluate_fractal.py index ffbd9d59..396ecfc1 100644 --- a/scripts/rt1/evaluate_fractal.py +++ b/scripts/rt1/evaluate_fractal.py @@ -1,24 +1,270 @@ import argparse +from datetime import datetime +from types import SimpleNamespace +import json +import os +import tempfile +from typing import Dict -from simpler_env.evaluation.fractal_tasks import run_comprehensive_evaluation -from simpler_env.policies.rt1.rt1_model import RT1Inference +import draccus +import torch + +from scripts.g3_lerobotpi.g3_pi0_or_fast import G3LerobotPiFastInference +from scripts.g3_lerobotpi.submission_utils import ( + copy_run_summaries, + dump_submission_payload, + prepare_submission_dir, +) +from simpler_env.evaluation.evaluate import run_comprehensive_evaluation, run_partial_evaluation +from simpler_env.policies.g3lerobotpi.adapter import AiroaToG3Pi0FractalBridgeAdapter +from simpler_env.policies.g3lerobotpi.policy import G3Pi0multiLerobotToAiroaPolicy +from g3_haptics.policies.g3factory import get_policy_class, get_policy_config_class + +PARTIAL_TASKS = { + "pick_object", + "pick_object_among", + "drawer", + "move_near", + "put_in_drawer", + "calc_score", +} 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( + "--eval-task", + type=str, + default="all", + help="Evaluation task name (use 'all' for the full benchmark).", + ) + parser.add_argument( + "--policy-setup", + type=str, + default="google_robot", + choices=["google_robot", "widowx_bridge"], + help="Policy setup to use when loading the G3 Pi0 checkpoint.", + ) + parser.add_argument( + "--save-path-suffix", + type=str, + default="", + help="Suffix appended to the checkpoint name when saving results.", + ) + parser.add_argument( + "--multi-embodiment", + action="store_true", + help="Enable the multi-embodiment Pi0 policy pipeline.", + ) + parser.add_argument( + "--rot6d", + action="store_true", + help="Use 6D rotation state when running the multi-embodiment adapter.", + ) + parser.add_argument( + "--action-ensemble", + action=argparse.BooleanOptionalAction, + default=False, + help="Toggle action ensemble (applies to both single and multi pipelines).", + ) + parser.add_argument( + "--action-ensemble-temp", + type=float, + default=-0.8, + help="Temperature for the action ensemble module.", + ) + parser.add_argument( + "--sticky-action", + action=argparse.BooleanOptionalAction, + default=False, + help="Use sticky gripper control when policy setup is google_robot.", + ) + parser.add_argument( + "--exec-horizon", + type=int, + default=4, + help="Number of actions executed per inference step in the fast single-embodiment pipeline.", + ) + parser.add_argument( + "--add-task-prefix", + action="store_true", + help="Prepend the policy setup to the language prompt (multi-embodiment only).", + ) + parser.add_argument( + "--submission-dir", + type=str, + default="submissions", + help="Directory to write submission artifacts (score.json, summaries).", + ) + parser.add_argument( + "--submission-name", + type=str, + default=None, + help="Optional fixed name for the submission folder (defaults to timestamped checkpoint id).", + ) + parser.add_argument( + "--no-save-submission", + action="store_true", + help="Skip writing submission artifacts; only print results.", + ) return parser.parse_args() +def _build_multi_policy(args: argparse.Namespace) -> AiroaToG3Pi0FractalBridgeAdapter: + torch.backends.cudnn.benchmark = True + torch.backends.cuda.matmul.allow_tf32 = True + + max_state_dim = 10 if args.rot6d else 8 + + dataset_cfg = SimpleNamespace( + embodiments=[ + SimpleNamespace( + state_key="observation.state", + image_key=["observation.images.image"], + action_key="action", + ft_key=None, + ), + SimpleNamespace( + state_key="observation.state", + image_key=["observation.images.image"], + action_key="action", + ft_key=None, + ), + ], + max_state_dim=max_state_dim, + max_image_num=1, + max_image_shape=[3, 256, 320], + max_action_dim=7, + ) + + emb_count = 2 + dataset_stats = { + "observation.state": { + "mean": torch.zeros(emb_count, max_state_dim), + "std": torch.ones(emb_count, max_state_dim), + }, + "action": { + "mean": torch.zeros(emb_count, 7), + "std": torch.ones(emb_count, 7), + }, + } + + with open(os.path.join(args.ckpt_path, "config.json"), encoding="utf-8") as f: + raw_cfg = json.load(f) + + PolicyCfgClass = get_policy_config_class(raw_cfg.pop("type")) + with tempfile.NamedTemporaryFile("w+", encoding="utf-8") as tmp: + json.dump(raw_cfg, tmp) + tmp.flush() + with draccus.config_type("json"): + cfg = draccus.parse(PolicyCfgClass, tmp.name, args=[]) + + if args.multi_embodiment: + cfg.multi_embodiment = True + + PolicyCls = get_policy_class("g3pi0") + base_policy = PolicyCls.from_pretrained( + pretrained_name_or_path=args.ckpt_path, + config=cfg, + dataset_stats=dataset_stats, + ) + + wrapped_policy = G3Pi0multiLerobotToAiroaPolicy( + policy=base_policy, + dataset_cfg=dataset_cfg, + policy_setup=args.policy_setup, + ) + + return AiroaToG3Pi0FractalBridgeAdapter( + policy=wrapped_policy, + rot6d=args.rot6d, + policy_setup=args.policy_setup, + action_ensemble=args.action_ensemble, + action_ensemble_temp=args.action_ensemble_temp, + sticky_action=args.sticky_action, + add_taks_prefix=args.add_task_prefix, + ) + + +def _build_single_policy(args: argparse.Namespace) -> G3LerobotPiFastInference: + return G3LerobotPiFastInference( + saved_model_path=args.ckpt_path, + policy_setup=args.policy_setup, + action_scale=1.0, + action_ensemble=args.action_ensemble, + action_ensemble_temp=args.action_ensemble_temp, + sticky_action=args.sticky_action, + exec_horizon=args.exec_horizon, + ) + + +def _normalize_scores(scores: Dict[str, float]) -> Dict[str, float]: + return {key: float(value) for key, value in scores.items()} + + if __name__ == "__main__": args = parse_args() - ckpt_path = args.ckpt_path - policy = RT1Inference(saved_model_path=ckpt_path, policy_setup="google_robot") + ckpt_identifier = ( + args.ckpt_path + if not args.save_path_suffix + else args.ckpt_path + f"_{args.save_path_suffix}" + ) + should_save, submission_path = prepare_submission_dir( + args.submission_dir, + args.submission_name, + ckpt_identifier, + default_tokens=[args.policy_setup, args.eval_task], + skip=args.no_save_submission, + ) + + if args.multi_embodiment: + policy = _build_multi_policy(args) + else: + policy = _build_single_policy(args) print("Policy initialized. Starting evaluation...") - final_scores = run_comprehensive_evaluation(env_policy=policy, ckpt_path=args.ckpt_path) + if args.eval_task == "all": + final_scores = run_comprehensive_evaluation(env_policy=policy, ckpt_path=ckpt_identifier) + elif args.eval_task in PARTIAL_TASKS: + final_scores = run_partial_evaluation( + env_policy=policy, + ckpt_path=ckpt_identifier, + task=args.eval_task, + ) + else: + raise ValueError( + f"Unknown eval-task '{args.eval_task}'. Use 'all' or one of {sorted(PARTIAL_TASKS)}." + ) print("\nEvaluation finished.") print(f"Final calculated scores: {final_scores}") + + if should_save and submission_path is not None: + results = _normalize_scores(final_scores) if final_scores else {} + payload: Dict[str, object] = { + "policy_setup": args.policy_setup, + "eval_task": args.eval_task, + "checkpoint": args.ckpt_path, + "timestamp": datetime.now().isoformat(), + "config": { + "multi_embodiment": args.multi_embodiment, + "rot6d": args.rot6d, + "action_ensemble": args.action_ensemble, + "action_ensemble_temp": args.action_ensemble_temp, + "sticky_action": args.sticky_action, + "exec_horizon": args.exec_horizon, + "add_task_prefix": args.add_task_prefix, + "save_path_suffix": args.save_path_suffix, + }, + "results": results, + } + if not results: + payload["note"] = ( + "Partial evaluation completed. Run with --eval-task all or --eval-task calc_score to compute final benchmark metrics." + ) + + dump_submission_payload(submission_path, payload) + copy_run_summaries(submission_path, ckpt_identifier) diff --git a/simpler_env/evaluation/maniskill2_evaluator.py b/simpler_env/evaluation/maniskill2_evaluator.py index 4ece8452..d8ebcd52 100644 --- a/simpler_env/evaluation/maniskill2_evaluator.py +++ b/simpler_env/evaluation/maniskill2_evaluator.py @@ -269,6 +269,7 @@ def maniskill2_evaluator(model, args): def _run_single_evaluation(model, args, control_mode, robot_init_x, robot_init_y, robot_init_quat): + success_arr = [] kwargs = dict( model=model, task_name=args.task_name, @@ -316,7 +317,6 @@ def _run_single_evaluation(model, args, control_mode, robot_init_x, robot_init_y obj_init_y = rng.uniform(args.obj_init_y_range[0], args.obj_init_y_range[1]) success = run_maniskill2_eval_single_episode(obj_init_x=obj_init_x, obj_init_y=obj_init_y, **kwargs) success_arr.append(success) - else: raise NotImplementedError() diff --git a/simpler_env/policies/g3lerobotpi/adapter.py b/simpler_env/policies/g3lerobotpi/adapter.py new file mode 100644 index 00000000..4d7304d8 --- /dev/null +++ b/simpler_env/policies/g3lerobotpi/adapter.py @@ -0,0 +1,335 @@ +from collections import deque +import os +from typing import List, Dict, Optional, Sequence, Tuple + +import matplotlib.pyplot as plt +import numpy as np +import cv2 +import logging +import torch +from PIL import Image + +from simpler_env.utils.action.action_ensemble import ActionEnsembler + +from simpler_env.utils.geometry import euler2axangle, mat2euler, quat2mat + +from simpler_env.policies.g3lerobotpi.geometry import quat_to_rot6d + +class BaseAdapter: + def __init__(self, policy): + self.policy = policy + + def reset(self, task_description): + pass + + def preprocess(self, image: np.ndarray, eef_pos: np.ndarray, prompt: str) -> Dict: + pass + + def postprocess(self, outputs: Dict) -> Dict: + pass + + def step(self, image: np.ndarray, eef_pos: np.ndarray, prompt: Optional[str] = None) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: + inputs = self.preprocess(image, eef_pos, prompt) + outputs = self.policy.step(inputs) + state_gripper = inputs["state"][-1] + action_gripper = outputs["actions"][-1] + # print(f"state: {state_gripper} action: {action_gripper}") + final_outputs = self.postprocess(outputs) + simpler_outputs = { + "world_vector": outputs["actions"][:3], + "rot_axangle": outputs["actions"][3:6], + "gripper": outputs["actions"][6:], + "terminate_episode": outputs["terminate_episode"], + } + final_simpler_outputs = { + "world_vector": final_outputs["actions"][:3], + "rot_axangle": final_outputs["actions"][3:6], + "gripper": final_outputs["actions"][6:], + "terminate_episode": final_outputs["terminate_episode"], + } + return simpler_outputs, final_simpler_outputs + + def _resize_image(self, image: np.ndarray) -> np.ndarray: + # Resize to 256x256 using Lanczos (approx via OpenCV's INTER_LANCZOS4) + img = cv2.resize(image, (256, 256), interpolation=cv2.INTER_LANCZOS4) + # If float image likely in [0,1], scale to [0,255] + if np.issubdtype(img.dtype, np.floating): + maxv = np.nanmax(img) + if maxv <= 1.0 + 1e-6: + img = img * 255.0 + # Clip, round and cast to uint8 + img = np.clip(np.rint(img), 0, 255).astype(np.uint8) + return img + + def visualize_epoch(self, predicted_raw_actions: Sequence[np.ndarray], images: Sequence[np.ndarray], save_path: str) -> None: + images = [self._resize_image(image) for image in images] + ACTION_DIM_LABELS = ["x", "y", "z", "roll", "pitch", "yaw", "grasp"] + + img_strip = np.concatenate(np.array(images[::3]), axis=1) + + # set up plt figure + figure_layout = [["image"] * len(ACTION_DIM_LABELS), ACTION_DIM_LABELS] + plt.rcParams.update({"font.size": 12}) + fig, axs = plt.subplot_mosaic(figure_layout) + fig.set_size_inches([45, 10]) + + # plot actions + pred_actions = np.array([np.concatenate([a["world_vector"], a["rotation_delta"], a["open_gripper"]], axis=-1) for a in predicted_raw_actions]) + for action_dim, action_label in enumerate(ACTION_DIM_LABELS): + # actions have batch, horizon, dim, in this example we just take the first action for simplicity + axs[action_label].plot(pred_actions[:, action_dim], label="predicted action") + axs[action_label].set_title(action_label) + axs[action_label].set_xlabel("Time in one episode") + + axs["image"].imshow(img_strip) + axs["image"].set_xlabel("Time in one episode (subsampled)") + plt.legend() + plt.savefig(save_path) + + +class AiroaToG3Pi0FractalBridgeAdapter(BaseAdapter): + def __init__( + self, + policy, + policy_setup: str = "widowx_bridge", + rot6d: bool = False, + exec_horizon: int = 4, + action_scale: float = 1.0, + action_ensemble: bool = True, + action_ensemble_temp: float = -0.8, + sticky_action: bool = False, + add_taks_prefix: bool = False + ) -> None: + super().__init__(policy) + self.sticky_gripper_num_repeat = 10 # same to lerobotpi0 + self.policy = policy + self.policy_setup = policy_setup + self.rot6d = rot6d + self.sticky_action_is_on = False + self.gripper_action_repeat = 0 + self.sticky_gripper_action = 0.0 + self.previous_gripper_action = None + self.action_scale = action_scale + self.action_ensemble_temp = action_ensemble_temp + self.obs_horizon = 1 + self.obs_interval = 1 + self.pred_action_horizon = 5 + self.image_history = deque(maxlen=self.obs_horizon) + self.exec_horizon = exec_horizon + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.sticky_action = sticky_action + self.add_taks_prefix = add_taks_prefix + + if self.policy_setup == "widowx_bridge": + self.action_ensemble = True + self.sticky_gripper_num_repeat = 1 + # EE pose in Bridge data was relative to a top-down pose, instead of robot base + 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 + elif self.policy_setup == "google_robot": + self.action_ensemble = True + 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." + ) + + 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 reset(self, task_description: str) -> None: + self.image_history.clear() + if self.action_ensemble: + self.action_ensembler.reset() + self.task_description = task_description + self.sticky_action_is_on = False + self.gripper_action_repeat = 0 + self.sticky_gripper_action = 0.0 + self.previous_gripper_action = None + self.action_plan = deque() + + def preprocess_widowx_proprio(self, eef_pos: np.ndarray, prompt: str): + """convert ee rotation to the frame of top-down + https://github.com/allenzren/open-pi-zero/blob/c3df7fb062175c16f69d7ca4ce042958ea238fb7/src/agent/env_adapter/simpler.py#L167 + """ + # StateEncoding.POS_EULER: xyz + rpy + pad + gripper(openness) + + proprio = eef_pos.copy() + if self.rot6d: + rpy_bridge_converted = quat_to_rot6d(torch.from_numpy(eef_pos).unsqueeze(0)).squeeze(0).numpy() + gripper_openness = proprio[7] # from simpler, 0 for close, 1 for open + state = np.concatenate( + [ + rpy_bridge_converted[:9], + [gripper_openness], + ] + ) + else: + proprio = eef_pos + rm_bridge = quat2mat(proprio[3:7]) + rpy_bridge_converted = mat2euler(rm_bridge @ self.default_rot.T) + gripper_openness = proprio[7] # from simpler, 0 for close, 1 for open + state = np.concatenate( + [ + proprio[:3], + rpy_bridge_converted, + np.zeros(1), + [gripper_openness], + ] + ) + + if self.add_taks_prefix: + prompts = ["google robot: ", "widowx: "] + prompt = prompts[1] + prompt + + return state, prompt + + def preprocess_google_robot_proprio(self, eef_pos: np.ndarray, prompt: str) -> dict: + """convert wxyz quat from simpler to xyzw used in fractal""" + + if self.rot6d: + gripper_width = eef_pos[-1] # from simpler, 0 for close, 1 for open continuous + gripper_closedness = gripper_width + quat_wxyz = quat_to_rot6d(torch.from_numpy(eef_pos).unsqueeze(0)).squeeze(0).numpy() + state = np.concatenate( + ( + quat_wxyz[:9], + [gripper_closedness], + ) + ) + + else: + gripper_width = eef_pos[-1] # from simpler, 0 for close, 1 for open continuous + gripper_closedness = ( + 1 - gripper_width + ) + quat_xyzw = np.roll(eef_pos[3:7], -1) + state = np.concatenate( + ( + eef_pos[:3], + quat_xyzw, + [gripper_closedness], + ) + ) + + + if self.add_taks_prefix: + prompts = ["google robot: ", "widowx: "] + prompt = prompts[0] + prompt + + return state, prompt + def step(self, image: np.ndarray, task_description: Optional[str] = None, *args, **kwargs): + """ + 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 + self._add_image_to_history(image) + images: List[Image.Image] = self._obtain_image_history() + + eef_pos = kwargs.get("eef_pos", None) + if self.policy_setup == "google_robot": + state, task_description = self.preprocess_google_robot_proprio(eef_pos, task_description) + image_key = "observation.images.image" + elif self.policy_setup == "widowx_bridge": + state, task_description = self.preprocess_widowx_proprio(eef_pos, task_description) + image_key = "observation.images.image" + + if not self.action_plan: + observation = { + "observation.state": torch.from_numpy(state).unsqueeze(0).to(self.device).float(), + image_key: images[0], + "task": [task_description], + } + + # model output gripper action, +1 = open, 0 = close + action_chunk = self.policy.step(observation)[: self.pred_action_horizon].cpu().numpy() + self.action_plan.extend(action_chunk[: self.exec_horizon]) + + raw_actions = self.action_plan.popleft() + + raw_action = { + "world_vector": np.array(raw_actions[:3]), + "rotation_delta": np.array(raw_actions[3:6]), + "open_gripper": np.array(raw_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"].astype(np.float32) * 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 + + def _add_image_to_history(self, image: np.ndarray) -> None: + if len(self.image_history) == 0: + self.image_history.extend([image] * self.obs_horizon) + else: + self.image_history.append(image) + + def _obtain_image_history(self) -> List[np.ndarray]: + image_history = list(self.image_history) + images = image_history[:: self.obs_interval] + # images = [Image.fromarray(image).convert("RGB") for image in images] + return images diff --git a/simpler_env/policies/g3lerobotpi/geometry.py b/simpler_env/policies/g3lerobotpi/geometry.py new file mode 100644 index 00000000..1f780346 --- /dev/null +++ b/simpler_env/policies/g3lerobotpi/geometry.py @@ -0,0 +1,119 @@ +import math + +import numpy as np +import torch + +def matrix_to_6d(R: torch.Tensor) -> torch.Tensor: + """ + Extract the first two columns of a 3x3 rotation matrix => 6D representation. + R: shape [B, 3, 3] + returns: shape [B, 6] + """ + # columns 0 and 1 => shape [B, 3, 2] + col0 = R[..., :, 0] # [B, 3] + col1 = R[..., :, 1] # [B, 3] + return torch.cat([col0, col1], dim=-1) # [B, 6] + +def quaternion_to_matrix(quat: torch.Tensor) -> torch.Tensor: + """ + Convert quaternion to a rotation matrix. + Args: + quat: [B,4] tensor with order (w, x, y, z) + Returns: + [B, 3, 3] rotation matrix. + """ + w = quat[:, 0] + x = quat[:, 1] + y = quat[:, 2] + z = quat[:, 3] + B = quat.shape[0] + R = torch.zeros(B, 3, 3, device=quat.device, dtype=quat.dtype) + R[:, 0, 0] = 1 - 2*(y*y + z*z) + R[:, 0, 1] = 2*(x*y - z*w) + R[:, 0, 2] = 2*(x*z + y*w) + R[:, 1, 0] = 2*(x*y + z*w) + R[:, 1, 1] = 1 - 2*(x*x + z*z) + R[:, 1, 2] = 2*(y*z - x*w) + R[:, 2, 0] = 2*(x*z - y*w) + R[:, 2, 1] = 2*(y*z + x*w) + R[:, 2, 2] = 1 - 2*(x*x + y*y) + return R + +def quat_to_rot6d(tcp_pose: torch.Tensor) -> torch.Tensor: + """ + Convert tcp_pose from quaternion representation to 6D rotation representation. + + Args: + tcp_pose: [B, 8] tensor with format [x, y, z, w, x, y, z, gripper] + (i.e. position (3) + quaternion (4) + gripper (1)). + Returns: + [B, 10] tensor with format [x, y, z, r6d0, r6d1, r6d2, r6d3, r6d4, r6d5, gripper]. + """ + xyz = tcp_pose[:, 0:3] # [B,3] + quat = tcp_pose[:, 3:7] # [B,4] assumed order: (w, x, y, z) + gripper = tcp_pose[:, 7] # [B] + + # quaternion -> rotation matrix + R = quaternion_to_matrix(quat) # [B,3,3] + # Extract 6D representation from the first two columns + r6d = matrix_to_6d(R) # [B,6] + + # Concatenate: [x,y,z] + [r6d (6)] + [gripper] => [B,10] + state_10d = torch.cat([xyz, r6d, gripper.unsqueeze(-1)], dim=-1) + return state_10d + +def euler_to_matrix(roll: torch.Tensor, pitch: torch.Tensor, yaw: torch.Tensor) -> torch.Tensor: + """ + Convert Euler angles (roll, pitch, yaw) in ZYX convention to a 3x3 rotation matrix. + roll, pitch, yaw: shape [B] (batch) + Returns: shape [B, 3, 3] + """ + sr, cr = torch.sin(roll), torch.cos(roll) + sp, cp = torch.sin(pitch), torch.cos(pitch) + sy, cy = torch.sin(yaw), torch.cos(yaw) + + # Rz(yaw) * Ry(pitch) * Rx(roll) + R00 = cy * cp + R01 = cy * sp * sr - sy * cr + R02 = cy * sp * cr + sy * sr + + R10 = sy * cp + R11 = sy * sp * sr + cy * cr + R12 = sy * sp * cr - cy * sr + + R20 = -sp + R21 = cp * sr + R22 = cp * cr + + # Stack into [B, 3, 3] + R = torch.stack([ + torch.stack([R00, R01, R02], dim=-1), + torch.stack([R10, R11, R12], dim=-1), + torch.stack([R20, R21, R22], dim=-1), + ], dim=-2) + return R + +def rpy_to_rot6d(tcp_pose: torch.Tensor) -> torch.Tensor: + """ + tcp_pose: shape [B, 8] => [x, y, z, roll, pitch, yaw, pad, gripper] + returns: shape [B, 10] => + [x, y, z, r6d0, r6d1, r6d2, r6d3, r6d4, r6d5, gripper] + """ + x = tcp_pose[:, 0] + y = tcp_pose[:, 1] + z = tcp_pose[:, 2] + roll = tcp_pose[:, 3] + pitch = tcp_pose[:, 4] + yaw = tcp_pose[:, 5] + pad = tcp_pose[:, 6] + gr = tcp_pose[:, 7] # gripper + + # Convert RPY -> 3x3 matrix -> 6D + R = euler_to_matrix(roll, pitch, yaw) # [B, 3, 3] + r6d = matrix_to_6d(R) # [B, 6] + + # Combine into 10D => [x, y, z, r6d(6), gripper] + xyz = tcp_pose[:, 0:3] # [B, 3] + xyz_r6d = torch.cat([xyz, r6d], dim=-1) # [B, 9] + state_10d = torch.cat([xyz_r6d, gr.unsqueeze(-1)], dim=-1) # [B, 10] + return state_10d \ No newline at end of file diff --git a/simpler_env/policies/g3lerobotpi/policy.py b/simpler_env/policies/g3lerobotpi/policy.py new file mode 100644 index 00000000..7ea7da76 --- /dev/null +++ b/simpler_env/policies/g3lerobotpi/policy.py @@ -0,0 +1,79 @@ + +import numpy as np +import cv2 +import torch +from typing import List, Dict, Optional, Sequence, Tuple + +from simpler_env.policies.base import AiroaBasePolicy +from g3_haptics.utils.lerobot_dataset_utils import create_g3multi_embodiment +from g3_haptics.datasets.embodiment import EmbodimentTag + +class G3Pi0multiLerobotToAiroaPolicy(AiroaBasePolicy): + def __init__( + self, + policy, + dataset_cfg=None, + policy_setup: str = "widowx_bridge", + ): + self.policy = policy + self.policy.eval() + if policy_setup == "widowx_bridge": + h, w = 256, 256 + self.image_size = (w, h) # cv2.resize は (W, H) + elif policy_setup == "google_robot": + self.image_size = (320, 256) + + self.embodiment = ( + create_g3multi_embodiment(dataset_cfg) if dataset_cfg is not None else None + ) + + self.policy_setup = policy_setup + if self.policy_setup == "google_robot": + self.tag = 0 + elif policy_setup == "widowx_bridge": + self.tag = 1 + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + def step(self, obs: Dict) -> Dict: + image = self._resize_image(obs["observation.images.image"]) + prompt = obs["task"] + state = obs["observation.state"] + + observation = { + "observation.state": obs["observation.state"], + "observation.images.image": torch.from_numpy(image / 255.) + .permute(2, 0, 1) + .unsqueeze(0) + .to(self.device) + .float(), + "task": obs["task"], + } + + if self.embodiment is not None: + device = next( + (v.device for v in observation.values() if torch.is_tensor(v)), None + ) + tag = ( + torch.tensor(self.tag, dtype=torch.long, device=device) + if device is not None + else torch.tensor(self.tag, dtype=torch.long) + ) + observation = self.embodiment.pad_item( + observation + | { + # FIXME: We cannot know embodiment + EmbodimentTag: tag + } + ) + + with torch.inference_mode(): + actions = self.policy.select_action(observation) + + return actions + + def reset(self) -> None: + self.policy.reset() + + def _resize_image(self, image: np.ndarray) -> np.ndarray: + image = cv2.resize(image, tuple(self.image_size), interpolation=cv2.INTER_AREA) + return image \ No newline at end of file