Merge pull request #24 from suraj-nair-tri/robomimic-eval

Robomimic Policy Eval Support

droid/camera_utils/camera_readers/zed_camera.py

@@ -131,6 +131,7 @@ def _configure_camera(self, init_params):
         self._current_params = init_params
         sl_params = sl.InitParameters(**init_params)
         sl_params.set_from_serial_number(int(self.serial_number))
+        sl_params.camera_image_flip = sl.FLIP_MODE.OFF
         status = self._cam.open(sl_params)
         if status != sl.ERROR_CODE.SUCCESS:
             raise RuntimeError("Camera Failed To Open")

droid/data_processing/timestep_processing.py

@@ -13,6 +13,7 @@ def __init__(
         self,
         ignore_action=False,
         action_space="cartesian_velocity",
+        gripper_action_space=None,
         robot_state_keys=["cartesian_position", "gripper_position", "joint_positions", "joint_velocities"],
         camera_extrinsics=["hand_camera", "varied_camera", "fixed_camera"],
         state_dtype=np.float32,
@@ -22,7 +23,7 @@ def __init__(
         assert action_space in ["cartesian_position", "joint_position", "cartesian_velocity", "joint_velocity"]
 
         self.action_space = action_space
-        self.gripper_key = "gripper_velocity" if "velocity" in action_space else "gripper_position"
+        self.gripper_key = "gripper_velocity" if "velocity" in gripper_action_space else "gripper_position"
         self.ignore_action = ignore_action
 
         self.robot_state_keys = robot_state_keys
@@ -68,6 +69,27 @@ def forward(self, timestep):
         if len(extrinsics_state):
             extrinsics_state = np.concatenate(extrinsics_state)
 
+        ### Get Intrinsics ###
+        cam_intrinsics_obs = timestep["observation"]["camera_intrinsics"]
+        sorted_calibrated_ids = sorted(calibration_dict.keys())
+        intrinsics_dict = defaultdict(list)
+
+        for serial_number in sorted_camera_ids:
+            cam_type = camera_type_dict[serial_number]
+            if cam_type not in self.camera_extrinsics:
+                continue
+
+            full_cam_ids = sorted(cam_intrinsics_obs.keys())
+            for full_cam_id in full_cam_ids:
+                if serial_number in full_cam_id:
+                    intr = cam_intrinsics_obs[full_cam_id]
+                    intrinsics_dict[cam_type].append(intr)
+
+        sorted_intrinsics_keys = sorted(intrinsics_dict.keys())
+        intrinsics_state = list([np.array(intrinsics_dict[cam_type]).flatten() for cam_type in sorted_intrinsics_keys])
+        if len(intrinsics_state):
+            intrinsics_state = np.concatenate(intrinsics_state)
+
         ### Get High Dimensional State Info ###
         high_dim_state_dict = defaultdict(lambda: defaultdict(list))
 
@@ -84,7 +106,7 @@ def forward(self, timestep):
                         high_dim_state_dict[obs_type][cam_type].append(data)
 
         ### Finish Observation Portion ###
-        low_level_state = np.concatenate([robot_state, extrinsics_state], dtype=self.state_dtype)
+        low_level_state = np.concatenate([robot_state, extrinsics_state, intrinsics_state], dtype=self.state_dtype)
         processed_timestep = {"observation": {"state": low_level_state, "camera": high_dim_state_dict}}
         self.image_transformer.forward(processed_timestep)
 
@@ -95,4 +117,8 @@ def forward(self, timestep):
             action = np.concatenate([arm_action, [gripper_action]], dtype=self.action_dtype)
             processed_timestep["action"] = action
 
+        # return raw information + meta data
+        processed_timestep["extrinsics_dict"] = extrinsics_dict
+        processed_timestep["intrinsics_dict"] = intrinsics_dict
+
         return processed_timestep

droid/evaluation/eval_launcher_robomimic.py

@@ -0,0 +1,259 @@
+import json
+import os
+import numpy as np
+import torch
+from collections import OrderedDict
+from copy import deepcopy
+
+from droid.controllers.oculus_controller import VRPolicy
+from droid.evaluation.policy_wrapper import PolicyWrapperRobomimic
+from droid.robot_env import RobotEnv
+from droid.user_interface.data_collector import DataCollecter
+from droid.user_interface.gui import RobotGUI
+
+import robomimic.utils.file_utils as FileUtils
+import robomimic.utils.torch_utils as TorchUtils
+import robomimic.utils.tensor_utils as TensorUtils
+
+import cv2
+
+def eval_launcher(variant, run_id, exp_id):
+    # Get Directory #
+    dir_path = os.path.dirname(os.path.realpath(__file__))
+
+    # Prepare Log Directory #
+    variant["exp_name"] = os.path.join(variant["exp_name"], "run{0}/id{1}/".format(run_id, exp_id))
+    log_dir = os.path.join(dir_path, "../../evaluation_logs", variant["exp_name"])
+
+    # Set Random Seeds #
+    torch.manual_seed(variant["seed"])
+    np.random.seed(variant["seed"])
+
+    # Set Compute Mode #
+    use_gpu = variant.get("use_gpu", False)
+    torch.device("cuda:0" if use_gpu else "cpu")
+
+    ckpt_path = variant["ckpt_path"]
+
+    device = TorchUtils.get_torch_device(try_to_use_cuda=True)
+    ckpt_dict = FileUtils.maybe_dict_from_checkpoint(ckpt_path=ckpt_path)
+    config = json.loads(ckpt_dict["config"])
+
+    ### infer image size ###
+    for obs_key in ckpt_dict["shape_metadata"]["all_shapes"].keys():
+        if 'camera/image' in obs_key:
+            imsize = max(ckpt_dict["shape_metadata"]["all_shapes"][obs_key])
+            break
+
+    ckpt_dict["config"] = json.dumps(config)
+    policy, _ = FileUtils.policy_from_checkpoint(ckpt_dict=ckpt_dict, device=device, verbose=True)
+    policy.goal_mode = config["train"]["goal_mode"]
+    policy.eval_mode = True
+
+    # determine the action space (relative or absolute)
+    action_keys = config["train"]["action_keys"]
+    if "action/rel_pos" in action_keys:
+        action_space = "cartesian_velocity"
+        for k in action_keys:
+            assert not k.startswith("action/abs_")
+    elif "action/abs_pos" in action_keys:
+        action_space = "cartesian_position"
+        for k in action_keys:
+            assert not k.startswith("action/rel_")
+    else:
+        raise ValueError
+
+    # determine the action space for the gripper
+    if "action/gripper_velocity" in action_keys:
+        gripper_action_space = "velocity"
+    elif "action/gripper_position" in action_keys:
+        gripper_action_space = "position"
+    else:
+        raise ValueError
+
+    # determine the action space (relative or absolute)
+    action_keys = config["train"]["action_keys"]
+    if "action/rel_pos" in action_keys:
+        action_space = "cartesian_velocity"
+        for k in action_keys:
+            assert not k.startswith("action/abs_")
+    elif "action/abs_pos" in action_keys:
+        action_space = "cartesian_position"
+        for k in action_keys:
+            assert not k.startswith("action/rel_")
+    else:
+        raise ValueError
+
+    # determine the action space for the gripper
+    if "action/gripper_velocity" in action_keys:
+        gripper_action_space = "velocity"
+    elif "action/gripper_position" in action_keys:
+        gripper_action_space = "position"
+    else:
+        raise ValueError
+
+    # Prepare Policy Wrapper #
+    data_processing_kwargs = dict(
+        timestep_filtering_kwargs=dict(
+            action_space=action_space,
+            gripper_action_space=gripper_action_space,
+            robot_state_keys=["cartesian_position", "gripper_position", "joint_positions"],
+            # camera_extrinsics=[],
+        ),
+        image_transform_kwargs=dict(
+            remove_alpha=True,
+            bgr_to_rgb=True,
+            to_tensor=True,
+            augment=False,
+        ),
+    )
+    timestep_filtering_kwargs = data_processing_kwargs.get("timestep_filtering_kwargs", {})
+    image_transform_kwargs = data_processing_kwargs.get("image_transform_kwargs", {})
+
+    policy_data_processing_kwargs = {}
+    policy_timestep_filtering_kwargs = policy_data_processing_kwargs.get("timestep_filtering_kwargs", {})
+    policy_image_transform_kwargs = policy_data_processing_kwargs.get("image_transform_kwargs", {})
+
+    policy_timestep_filtering_kwargs.update(timestep_filtering_kwargs)
+    policy_image_transform_kwargs.update(image_transform_kwargs)
+
+    fs = config["train"]["frame_stack"]
+
+    wrapped_policy = PolicyWrapperRobomimic(
+        policy=policy,
+        timestep_filtering_kwargs=policy_timestep_filtering_kwargs,
+        image_transform_kwargs=policy_image_transform_kwargs,
+        frame_stack=fs,
+        eval_mode=True,
+    )
+
+    camera_kwargs = dict(
+        hand_camera=dict(image=True, concatenate_images=False, resolution=(imsize, imsize), resize_func="cv2"),
+        varied_camera=dict(image=True, concatenate_images=False, resolution=(imsize, imsize), resize_func="cv2"),
+    )
+
+    policy_camera_kwargs = {}
+    policy_camera_kwargs.update(camera_kwargs)
+
+    env = RobotEnv(
+        action_space=policy_timestep_filtering_kwargs["action_space"],
+        gripper_action_space=policy_timestep_filtering_kwargs["gripper_action_space"],
+        camera_kwargs=policy_camera_kwargs
+    )
+    controller = VRPolicy()
+
+    # Launch GUI #
+    data_collector = DataCollecter(
+        env=env,
+        controller=controller,
+        policy=wrapped_policy,
+        save_traj_dir=log_dir,
+        save_data=variant.get("save_data", True),
+    )
+    RobotGUI(robot=data_collector)
+
+
+def get_goal_im(variant, run_id, exp_id):
+    # Get Directory #
+    dir_path = os.path.dirname(os.path.realpath(__file__))
+
+    # Prepare Log Directory #
+    variant["exp_name"] = os.path.join(variant["exp_name"], "run{0}/id{1}/".format(run_id, exp_id))
+    log_dir = os.path.join(dir_path, "../../evaluation_logs", variant["exp_name"])
+
+    # Set Random Seeds #
+    torch.manual_seed(variant["seed"])
+    np.random.seed(variant["seed"])
+
+    # Set Compute Mode #
+    use_gpu = variant.get("use_gpu", False)
+    torch.device("cuda:0" if use_gpu else "cpu")
+
+    ckpt_path = variant["ckpt_path"]
+
+    device = TorchUtils.get_torch_device(try_to_use_cuda=True)
+    ckpt_dict = FileUtils.maybe_dict_from_checkpoint(ckpt_path=ckpt_path)
+    config = json.loads(ckpt_dict["config"])
+
+    ### infer image size ###
+    imsize = max(ckpt_dict["shape_metadata"]["all_shapes"]["camera/image/hand_camera_left_image"])
+
+    ckpt_dict["config"] = json.dumps(config)
+    policy, _ = FileUtils.policy_from_checkpoint(ckpt_dict=ckpt_dict, device=device, verbose=True)
+
+    # determine the action space (relative or absolute)
+    action_keys = config["train"]["action_keys"]
+    if "action/rel_pos" in action_keys:
+        action_space = "cartesian_velocity"
+        for k in action_keys:
+            assert not k.startswith("action/abs_")
+    elif "action/abs_pos" in action_keys:
+        action_space = "cartesian_position"
+        for k in action_keys:
+            assert not k.startswith("action/rel_")
+    else:
+        raise ValueError
+
+    # determine the action space for the gripper
+    if "action/gripper_velocity" in action_keys:
+        gripper_action_space = "velocity"
+    elif "action/gripper_position" in action_keys:
+        gripper_action_space = "position"
+    else:
+        raise ValueError
+
+    # Prepare Policy Wrapper #
+    data_processing_kwargs = dict(
+        timestep_filtering_kwargs=dict(
+            action_space=action_space,
+            gripper_action_space=gripper_action_space,
+            robot_state_keys=["cartesian_position", "gripper_position", "joint_positions"],
+            # camera_extrinsics=[],
+        ),
+        image_transform_kwargs=dict(
+            remove_alpha=True,
+            bgr_to_rgb=True,
+            to_tensor=True,
+            augment=False,
+        ),
+    )
+    timestep_filtering_kwargs = data_processing_kwargs.get("timestep_filtering_kwargs", {})
+    image_transform_kwargs = data_processing_kwargs.get("image_transform_kwargs", {})
+
+    policy_data_processing_kwargs = {}
+    policy_timestep_filtering_kwargs = policy_data_processing_kwargs.get("timestep_filtering_kwargs", {})
+    policy_image_transform_kwargs = policy_data_processing_kwargs.get("image_transform_kwargs", {})
+
+    policy_timestep_filtering_kwargs.update(timestep_filtering_kwargs)
+    policy_image_transform_kwargs.update(image_transform_kwargs)
+
+    wrapped_policy = PolicyWrapperRobomimic(
+        policy=policy,
+        timestep_filtering_kwargs=policy_timestep_filtering_kwargs,
+        image_transform_kwargs=policy_image_transform_kwargs,
+        frame_stack=config["train"]["frame_stack"],
+        eval_mode=True,
+    )
+
+    camera_kwargs = dict(
+        hand_camera=dict(image=True, concatenate_images=False, resolution=(imsize, imsize), resize_func="cv2"),
+        varied_camera=dict(image=True, concatenate_images=False, resolution=(imsize, imsize), resize_func="cv2"),
+    )
+
+    policy_camera_kwargs = {}
+    policy_camera_kwargs.update(camera_kwargs)
+
+    env = RobotEnv(
+        action_space=policy_timestep_filtering_kwargs["action_space"],
+        gripper_action_space=policy_timestep_filtering_kwargs["gripper_action_space"],
+        camera_kwargs=policy_camera_kwargs,
+        do_reset=False
+    )
+
+    ims = env.read_cameras()[0]["image"]
+    if not os.path.exists('eval_params'):
+        os.makedirs('eval_params')
+    for k in ims.keys():
+        image = ims[k]
+        cv2.imwrite(f'eval_params/{k}.png', image[:, :, :3])
+    return ims