| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
|
|
| import numpy as np |
| from rlbench.action_modes.action_mode import BimanualMoveArmThenGripper |
| from rlbench.action_modes.arm_action_modes import BimanualEndEffectorPoseViaPlanning |
| from rlbench.action_modes.gripper_action_modes import BimanualDiscrete |
| from rlbench.environment import Environment |
|
|
| from sim_rlbench.camera_spec import default_three_camera_spec |
| from sim_rlbench.obs_config import build_obs_config |
| from sim_rlbench.task_resolver import resolve_task_class |
|
|
|
|
| def _camera_shape(obs: object, camera_name: str): |
| perception_data = getattr(obs, "perception_data", None) |
| if isinstance(perception_data, dict): |
| value = perception_data.get(f"{camera_name}_rgb") |
| if value is not None: |
| return list(np.asarray(value).shape) |
| return None |
|
|
|
|
| def _noop_bimanual_action(obs: object) -> np.ndarray: |
| right_obs = getattr(obs, "right", None) |
| left_obs = getattr(obs, "left", None) |
| right = np.concatenate( |
| [ |
| np.asarray(right_obs.gripper_pose, dtype=np.float32), |
| np.array([float(right_obs.gripper_open), 1.0], dtype=np.float32), |
| ], |
| axis=0, |
| ) |
| left = np.concatenate( |
| [ |
| np.asarray(left_obs.gripper_pose, dtype=np.float32), |
| np.array([float(left_obs.gripper_open), 1.0], dtype=np.float32), |
| ], |
| axis=0, |
| ) |
| return np.concatenate([right, left], axis=0) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--task", default="bimanual_push_box") |
| parser.add_argument("--resolution", type=int, default=224) |
| parser.add_argument("--display", default=None) |
| parser.add_argument("--headless", action="store_true", default=True) |
| parser.add_argument("--visible", action="store_true") |
| args = parser.parse_args() |
|
|
| headless = args.headless and not args.visible |
| spec = default_three_camera_spec(args.resolution) |
| task_cls = resolve_task_class(args.task) |
| obs_config = build_obs_config(list(spec.upstream_cameras), args.resolution) |
| action_mode = BimanualMoveArmThenGripper( |
| BimanualEndEffectorPoseViaPlanning(absolute_mode=True, frame="world", collision_checking=False), |
| BimanualDiscrete(), |
| ) |
| env = Environment( |
| action_mode=action_mode, |
| obs_config=obs_config, |
| headless=headless, |
| robot_setup="dual_panda", |
| ) |
|
|
| try: |
| env.launch() |
| task = env.get_task(task_cls) |
| descriptions, obs = task.reset() |
| noop_action = _noop_bimanual_action(obs) |
| next_obs, reward, done = task.step(noop_action) |
| payload = { |
| "display": args.display or os.environ.get("DISPLAY"), |
| "headless": headless, |
| "task": task_cls.__name__, |
| "description": descriptions[0] if descriptions else "", |
| "front_rgb_shape": _camera_shape(obs, "front"), |
| "wrist_left_rgb_shape": _camera_shape(obs, "wrist_left"), |
| "wrist_right_rgb_shape": _camera_shape(obs, "wrist_right"), |
| "right_pose_shape": list(obs.right.gripper_pose.shape), |
| "left_pose_shape": list(obs.left.gripper_pose.shape), |
| "stepped_mode": "bimanual_noop", |
| "action_finite": bool(np.isfinite(noop_action).all()), |
| "action_dim": int(noop_action.shape[0]), |
| "reward": float(reward), |
| "done": bool(done), |
| "front_rgb_shape_after_step": _camera_shape(next_obs, "front"), |
| } |
| print(json.dumps(payload, indent=2)) |
| finally: |
| env.shutdown() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
