diff --git a/bd_spot_wrapper/spot_wrapper/spot.py b/bd_spot_wrapper/spot_wrapper/spot.py index d8a5f4c3..f551235b 100644 --- a/bd_spot_wrapper/spot_wrapper/spot.py +++ b/bd_spot_wrapper/spot_wrapper/spot.py @@ -277,6 +277,19 @@ def rotate_gripper_with_delta(self, wrist_yaw=0.0, wrist_roll=0.0): self.set_arm_joint_positions(new_arm_joint_states) time.sleep(0.5) + def get_ee_rotation_in_body_frame_quat(self): + """ + return ee rotation in quaternion + """ + vision_T_hand = get_a_tform_b( + self.robot_state_client.get_robot_state().kinematic_state.transforms_snapshot, + "body", + "hand", + ) + quat = vision_T_hand.rotation + quat = quaternion.quaternion(quat.w, quat.x, quat.y, quat.z) + return quat + def move_gripper_to_point( self, point, rotation, seconds_to_goal=3.0, timeout_sec=10 ): diff --git a/spot_rl_experiments/configs/config.yaml b/spot_rl_experiments/configs/config.yaml index 8cf5ee2c..ac9bac73 100644 --- a/spot_rl_experiments/configs/config.yaml +++ b/spot_rl_experiments/configs/config.yaml @@ -13,6 +13,9 @@ WEIGHTS_TORCHSCRIPT: # Static place PLACE: "weights/torchscript/place_10deg_32_seed300_1649709235_ckpt.75_combined_net.torchscript" + # Open close drawer + OPEN_CLOSE_DRAWER: "weights/open_close_drawer/od14lr_30_ckpt.144.combined_net.torchscript" + WEIGHTS: # Regular Nav NAV: "weights/CUTOUT_WT_True_SD_200_ckpt.99.pvp.pth" @@ -25,6 +28,9 @@ WEIGHTS: # Static place PLACE: "weights/final_paper/place_10deg_32_seed300_1649709235_ckpt.75.pth" + + # Open close drawer + OPEN_CLOSE_DRAWER: "weights/open_close_drawer/od14lr_30_ckpt.144.combined_net.torchscript" # ASC (weight of ASC will always remain in pytorch) MIXER: "weights/final_paper/final_moe_rnn_60_1.0_SD_100_1652120928_ckpt.16_copy.pth" @@ -70,6 +76,10 @@ SUCC_XY_DIST: 0.1 SUCC_Z_DIST: 0.20 PLACE_ACTION_SPACE_LENGTH: 4 +# Open Close Drawer env +OPEN_CLOSE_DRAWER_ACTION_SPACE_LENGTH: 8 +OPEM_CLOSE_DRAWER_DISTANCE_BETWEEN_EE_HANDLE: 0.3 + # Base action params MAX_LIN_DIST: 0.25 # meters MAX_ANG_DIST: 15.0 # degrees @@ -77,12 +87,16 @@ MAX_ANG_DIST: 15.0 # degrees # Arm action params MAX_JOINT_MOVEMENT: 0.0872665 # Gaze arm speed (5 deg) MAX_JOINT_MOVEMENT_2: 0.174533 # Place arm speed (10 deg) +MAX_JOINT_MOVEMENT_OPEN_CLOSE_DRAWER: 0.06 INITIAL_ARM_JOINT_ANGLES: [0, -170, 120, 0, 75, 0] GAZE_ARM_JOINT_ANGLES: [0, -160, 100, 0, 90, 0] PLACE_ARM_JOINT_ANGLES: [0, -170, 120, 0, 75, 0] ARM_LOWER_LIMITS: [-45, -180, 0, 0, -90, 0] ARM_UPPER_LIMITS: [45, -45, 180, 0, 90, 0] +ARM_LOWER_LIMITS_OPEN_CLOSE_DRAWER: [-90, -180, 0, 0, -90, 0] +ARM_UPPER_LIMITS_OPEN_CLOSE_DRAWER: [90, 0, 180, 0, 90, 0] JOINT_BLACKLIST: [3, 5] # joints we can't control "arm0.el0", "arm0.wr1" +JOINT_BLACKLIST_OPEN_CLOSE_DRAWER: [3] ACTUALLY_MOVE_ARM: True GRASP_EVERY_STEP: False TERMINATE_ON_GRASP: False diff --git a/spot_rl_experiments/experiments/skill_test/test_open_close_drawer.py b/spot_rl_experiments/experiments/skill_test/test_open_close_drawer.py new file mode 100644 index 00000000..d21c1610 --- /dev/null +++ b/spot_rl_experiments/experiments/skill_test/test_open_close_drawer.py @@ -0,0 +1,26 @@ +# Copyright (c) Meta Platforms, Inc. and its affiliates. +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + + +import numpy as np +from spot_rl.envs.skill_manager import SpotSkillManager + +if __name__ == "__main__": + from perception_and_utils.utils.generic_utils import map_user_input_to_boolean + + # Init the skill + spotskillmanager = SpotSkillManager() + + # Using while loop + contnue = True + while contnue: + spotskillmanager.openclosedrawer(open_mode=True) + close_drawer = map_user_input_to_boolean( + "Do you want to close the drawer ? Y/N " + ) + if close_drawer: + spotskillmanager.openclosedrawer(open_mode=False) + contnue = map_user_input_to_boolean( + "Do you want to open the drawer again ? Y/N " + ) diff --git a/spot_rl_experiments/spot_rl/envs/base_env.py b/spot_rl_experiments/spot_rl/envs/base_env.py index 7a33f2ab..097a6b3e 100644 --- a/spot_rl_experiments/spot_rl/envs/base_env.py +++ b/spot_rl_experiments/spot_rl/envs/base_env.py @@ -386,7 +386,7 @@ def step( # noqa ) elif arm_action is not None: self.spot.set_arm_joint_positions( - positions=arm_action, travel_time=1 / self.ctrl_hz * 0.9 + positions=arm_action, travel_time=1 / self.ctrl_hz * 1.5 ) if self.prev_base_moved and base_action is None: @@ -499,14 +499,11 @@ def get_nav_observation(self, goal_xy, goal_heading): return observations - def get_arm_joints(self): + def get_arm_joints(self, black_list=None): # Get proprioception inputs + black_list = self.config.JOINT_BLACKLIST if black_list is None else black_list joints = np.array( - [ - j - for idx, j in enumerate(self.current_arm_pose) - if idx not in self.config.JOINT_BLACKLIST - ], + [j for idx, j in enumerate(self.current_arm_pose) if idx not in black_list], dtype=np.float32, ) diff --git a/spot_rl_experiments/spot_rl/envs/open_close_drawer_env.py b/spot_rl_experiments/spot_rl/envs/open_close_drawer_env.py new file mode 100644 index 00000000..2523979d --- /dev/null +++ b/spot_rl_experiments/spot_rl/envs/open_close_drawer_env.py @@ -0,0 +1,309 @@ +# Copyright (c) Meta Platforms, Inc. and its affiliates. +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + + +import sys +from typing import Any, Dict + +import magnum as mn +import numpy as np +import quaternion +import rospy +from spot_rl.envs.base_env import SpotBaseEnv +from spot_wrapper.spot import Spot, image_response_to_cv2, scale_depth_img + + +def get_3d_point(cam_intrinsics, pixel_uv, z): + """Helper function to compute 3D point""" + # Get camera intrinsics + fx = cam_intrinsics.focal_length.x + fy = cam_intrinsics.focal_length.y + cx = cam_intrinsics.principal_point.x + cy = cam_intrinsics.principal_point.y + + # Get 3D point + x = (pixel_uv[0] - cx) * z / fx + y = (pixel_uv[1] - cy) * z / fy + return np.array([x, y, z]) + + +class SpotOpenCloseDrawerEnv(SpotBaseEnv): + def __init__(self, config, spot: Spot): + # Select suitable keys + max_joint_movement_key = "MAX_JOINT_MOVEMENT_OPEN_CLOSE_DRAWER" + max_lin_dist_key = "MAX_LIN_DIST" + max_ang_dist_key = "MAX_ANG_DIST" + + super().__init__( + config, + spot, + stopwatch=None, + max_joint_movement_key=max_joint_movement_key, + max_lin_dist_key=max_lin_dist_key, + max_ang_dist_key=max_ang_dist_key, + ) + + self.ee_gripper_offset = mn.Vector3(config.EE_GRIPPER_OFFSET) + + # The initial joint angles is in the stow location + self.initial_arm_joint_angles = np.deg2rad([0, -180, 180, 0, 0, 0]) + + # The arm joint min max overwrite + self.arm_lower_limits = np.deg2rad(config.ARM_LOWER_LIMITS_OPEN_CLOSE_DRAWER) + self.arm_upper_limits = np.deg2rad(config.ARM_UPPER_LIMITS_OPEN_CLOSE_DRAWER) + + # Flag for done + self._success = False + + # Mode for opening or closing + self._mode = "open" + + # Distance threshold to call IK to approach the drawers + self._dis_threshold_ee_to_handle = ( + config.OPEM_CLOSE_DRAWER_DISTANCE_BETWEEN_EE_HANDLE + ) + + def reset(self, goal_dict=None, *args, **kwargs): + self.spot.open_gripper() + + # Move arm to initial configuration + cmd_id = self.spot.set_arm_joint_positions( + positions=self.initial_arm_joint_angles, travel_time=1 + ) + + # Block until arm arrives with incremental timeout for 3 attempts + timeout_sec = 1.0 + max_allowed_timeout_sec = 3.0 + status = False + while status is False and timeout_sec <= max_allowed_timeout_sec: + status = self.spot.block_until_arm_arrives(cmd_id, timeout_sec=timeout_sec) + timeout_sec += 1.0 + + # Make the arm to be in true nominal location + ee_position = self.get_gripper_position_in_base_frame_spot() + self.spot.move_gripper_to_point(ee_position, [0, 0, 0]) + + # Move arm to initial configuration again to ensure it is in the good location + cmd_id = self.spot.set_arm_joint_positions( + positions=self.initial_arm_joint_angles, travel_time=1 + ) + + # Block until arm arrives with incremental timeout for 3 attempts + timeout_sec = 1.0 + max_allowed_timeout_sec = 3.0 + status = False + while status is False and timeout_sec <= max_allowed_timeout_sec: + status = self.spot.block_until_arm_arrives(cmd_id, timeout_sec=timeout_sec) + timeout_sec += 1.0 + + self.initial_ee_orientation = self.spot.get_ee_rotation_in_body_frame_quat() + + # Update target object name as provided in config + observations = super().reset(target_obj_name="drawer handle", *args, **kwargs) + rospy.set_param("object_target", self.target_obj_name) + + # Flag for done + self._success = False + + # Get the mode: open or close drawers + self._mode = goal_dict["mode"] + + return observations + + def compute_distance_to_handle(self): + "Compute the distance in the bounding box center" + return ( + self.target_object_distance, + self.obj_center_pixel[0], + self.obj_center_pixel[1], + ) + + def bd_open_drawer_api(self): + """BD API to open the drawer in x direction""" + raise NotImplementedError + + def approach_handle_and_grasp(self, z, pixel_x, pixel_y): + """This method does IK to approach the handle and close the gripper.""" + imgs = self.spot.get_hand_image() + + # Get the camera intrinsics + cam_intrinsics = imgs[0].source.pinhole.intrinsics + + # Get the transformation + vision_T_base = self.spot.get_magnum_Matrix4_spot_a_T_b("vision", "body") + + # Get the 3D point in the hand RGB frame + point_in_hand_image_3d = get_3d_point(cam_intrinsics, (pixel_x, pixel_y), z) + + # Get the vision to hand + vision_T_hand_image: mn.Matrix4 = self.spot.get_magnum_Matrix4_spot_a_T_b( + "vision", "hand_color_image_sensor", imgs[0].shot.transforms_snapshot + ) + point_in_global_3d = vision_T_hand_image.transform_point( + mn.Vector3(*point_in_hand_image_3d) + ) + + # Get the transformation of the gripper + vision_T_hand = self.spot.get_magnum_Matrix4_spot_a_T_b("vision", "hand") + # Get the location relative to the gripper + point_in_hand_3d = vision_T_hand.inverted().transform_point(point_in_global_3d) + # Offset the x and z direction in hand frame + ee_offset_x = 0.05 + ee_offset_z = -0.05 + point_in_hand_3d[0] += ee_offset_x + point_in_hand_3d[2] += ee_offset_z + # Make it back to global frame + point_in_global_3d = vision_T_hand.transform_point(point_in_hand_3d) + + # Get the point in the base frame + point_in_base_3d = vision_T_base.inverted().transform_point(point_in_global_3d) + + # Make it to be numpy + point_in_base_3d = np.array( + [ + point_in_base_3d.x, + point_in_base_3d.y, + point_in_base_3d.z, + ] + ) + + # Get the current ee rotation in body frame + ee_rotation = self.spot.get_ee_rotation_in_body_frame_quat() + + # Move the gripper to target using current gripper pose in the body frame + # while maintaining the gripper orientation + self.spot.move_gripper_to_point( + point_in_base_3d, + [ee_rotation.w, ee_rotation.x, ee_rotation.y, ee_rotation.z], + ) + + # Close the gripper + self.spot.close_gripper() + + # Get the transformation of the gripper + vision_T_hand = self.spot.get_magnum_Matrix4_spot_a_T_b("vision", "hand") + # Get the location that we want to move to for retracting/moving forward the arm. Pull/push the drawer by 20 cm + pull_push_distance = -0.2 if self._mode == "open" else 0.25 + move_target = vision_T_hand.transform_point( + mn.Vector3([pull_push_distance, 0, 0]) + ) + # Get the move_target in base frame + move_target = vision_T_base.inverted().transform_point(move_target) + + # Retract the arm based on the current gripper location + self.spot.move_gripper_to_point( + move_target, [ee_rotation.w, ee_rotation.x, ee_rotation.y, ee_rotation.z] + ) + + # Open the gripper and retract the arm + self.spot.open_gripper() + # [0.55, 0, 0.27] is the gripper nominal location + # [0,0,0] is the roll pitch yaw + self.spot.move_gripper_to_point([0.55, 0, 0.27], [0, 0, 0]) + + # Change the flag to finish + self._success = True + + def step(self, action_dict: Dict[str, Any]): + + # Update the action_dict with place flag + action_dict["place"] = False + observations, reward, done, info = super().step( + action_dict=action_dict, + ) + + # Get bounding box + bbox = observations["handle_bbox"] + + # Compute the distance from the gripper to bounding box + # The distance is called z here + z = float("inf") + # We only compute the distance if bounding box detects something + if np.sum(bbox) > 0: + z, pixel_x, pixel_y = self.compute_distance_to_handle() + + # We close gripper here + if z != 0 and z < self._dis_threshold_ee_to_handle: + # Do IK to approach the target + self.approach_handle_and_grasp(z, pixel_x, pixel_y) + # If we can do IK, then we call it successful + done = self._success + + return observations, reward, done, info + + def angle_between_quat(self, q1, q2): + q1_inv = np.conjugate(q1) + dp = quaternion.as_float_array(q1_inv * q2) + return 2 * np.arctan2(np.linalg.norm(dp[1:]), np.abs(dp[0])) + + def angle_to_forward(self, x, y): + if np.linalg.norm(x) != 0: + x_norm = x / np.linalg.norm(x) + else: + x_norm = x + + if np.linalg.norm(y) != 0: + y_norm = y / np.linalg.norm(y) + else: + y_norm = y + + return np.arccos(np.clip(np.dot(x_norm, y_norm), -1, 1)) + + def get_angle(self, rel_pos): + """Get angle""" + forward = np.array([1.0, 0, 0]) + rel_pos = np.array(rel_pos) + forward = forward[[0, 1]] + rel_pos = rel_pos[[0, 1]] + + heading_angle = self.angle_to_forward(forward, rel_pos) + c = np.cross(forward, rel_pos) < 0 + if not c: + heading_angle = -1.0 * heading_angle + return heading_angle + + def get_cur_ee_orientation_offset(self): + # Get base to hand's transformation + ee_transform = self.spot.get_magnum_Matrix4_spot_a_T_b("vision", "hand") + # Get the base transformation + base_transform = self.spot.get_magnum_Matrix4_spot_a_T_b("vision", "body") + # Do offset: move the base center forward to be close to the gripper base + base_transform.translation = base_transform.transform_point( + mn.Vector3(0.292, 0, 0) + ) + # Get ee relative to base + ee_position = (base_transform.inverted() @ ee_transform).translation + base_T_ee_yaw = self.get_angle(ee_position) + return base_T_ee_yaw + + def get_observations(self): + # Get the depth images and handle bounding box + arm_depth, arm_depth_bbox = self.get_gripper_images() + + # Get the delta ee orientation to the initial orientation + current_ee_orientation = self.spot.get_ee_rotation_in_body_frame_quat() + delta_ee = np.array( + self.angle_between_quat( + self.initial_ee_orientation, current_ee_orientation + ), + dtype=np.float32, + ) + + # Remove the offset from the base to ee + delta_ee = np.array( + [delta_ee - abs(self.get_cur_ee_orientation_offset())], dtype=np.float32 + ) + + # Construct the observation + observations = { + "articulated_agent_arm_depth": arm_depth, + "joint": self.get_arm_joints(self.config.JOINT_BLACKLIST_OPEN_CLOSE_DRAWER), + "ee_pos": self.get_gripper_position_in_base_frame_spot(), + "handle_bbox": arm_depth_bbox, + "art_pose_delta_sensor": delta_ee, + } + return observations + + def get_success(self, observations=None): + return self._success diff --git a/spot_rl_experiments/spot_rl/envs/skill_manager.py b/spot_rl_experiments/spot_rl/envs/skill_manager.py index fc714e66..270ea37d 100644 --- a/spot_rl_experiments/spot_rl/envs/skill_manager.py +++ b/spot_rl_experiments/spot_rl/envs/skill_manager.py @@ -8,10 +8,11 @@ import numpy as np from multimethod import multimethod from perception_and_utils.utils.generic_utils import conditional_print -from spot_rl.skills.atomic_skills import Navigation, Pick, Place +from spot_rl.skills.atomic_skills import Navigation, OpenCloseDrawer, Pick, Place from spot_rl.utils.construct_configs import ( construct_config_for_gaze, construct_config_for_nav, + construct_config_for_open_close_drawer, construct_config_for_place, ) from spot_rl.utils.heuristic_nav import ( @@ -82,6 +83,7 @@ def __init__( nav_config=None, pick_config=None, place_config=None, + open_close_drawer_config=None, use_mobile_pick: bool = False, verbose: bool = True, use_policies: bool = True, @@ -101,6 +103,7 @@ def __init__( nav_config=nav_config, pick_config=pick_config, place_config=place_config, + open_close_drawer_config=open_close_drawer_config, ) # Initiate the controllers for nav, gaze, and place @@ -117,7 +120,12 @@ def __del__(self): pass def __init_spot( - self, spot: Spot = None, nav_config=None, pick_config=None, place_config=None + self, + spot: Spot = None, + nav_config=None, + pick_config=None, + place_config=None, + open_close_drawer_config=None, ): """ Initialize the Spot object, acquire lease, and construct configs @@ -148,6 +156,12 @@ def __init_spot( construct_config_for_place() if not place_config else place_config ) + self.open_close_drawer_config = ( + construct_config_for_open_close_drawer() + if not open_close_drawer_config + else open_close_drawer_config + ) + def __initiate_controllers(self, use_policies: bool = True): """ Initiate the controllers for nav, gaze, and place @@ -167,6 +181,10 @@ def __initiate_controllers(self, use_policies: bool = True): config=self.place_config, use_policies=use_policies, ) + self.open_close_drawer_controller = OpenCloseDrawer( + spot=self.spot, + config=self.open_close_drawer_config, + ) def reset(self): # Reset the policies and environments via the controllers @@ -354,6 +372,21 @@ def place( # noqa conditional_print(message=message, verbose=self.verbose) return status, message + def openclosedrawer(self, open_mode=True) -> Tuple[bool, str]: + """ + Perform open and close skills + + Returns: + bool: True if open close was successful, False otherwise + str: Message indicating the status of the open/close drawers + """ + goal_dict = { + "mode": "open" if open_mode else "close", + } # type: Dict[str, Any] + status, message = self.open_close_drawer_controller.execute(goal_dict=goal_dict) + conditional_print(message=message, verbose=self.verbose) + return status, message + def sit(self): """ Sit the robot diff --git a/spot_rl_experiments/spot_rl/real_policy.py b/spot_rl_experiments/spot_rl/real_policy.py index e8a1ad0a..40667ead 100644 --- a/spot_rl_experiments/spot_rl/real_policy.py +++ b/spot_rl_experiments/spot_rl/real_policy.py @@ -280,6 +280,50 @@ def __init__(self, checkpoint_path, device, config: CN = CN()): ) +class OpenCloseDrawerPolicy(RealPolicy): + def __init__(self, checkpoint_path, device, config: CN = CN()): + observation_space = SpaceDict( + { + "handle_bbox": spaces.Box( + low=np.finfo(np.float32).min, + high=np.finfo(np.float32).max, + shape=(240, 228, 1), + dtype=np.float32, + ), + "articulated_agent_arm_depth": spaces.Box( + low=0.0, high=1.0, shape=(240, 228, 1), dtype=np.float32 + ), + "joint": spaces.Box( + low=np.finfo(np.float32).min, + high=np.finfo(np.float32).max, + shape=(5,), + dtype=np.float32, + ), + "ee_pos": spaces.Box( + low=np.finfo(np.float32).min, + high=np.finfo(np.float32).max, + shape=(3,), + dtype=np.float32, + ), + "art_pose_delta_sensor": spaces.Box( + low=np.finfo(np.float32).min, + high=np.finfo(np.float32).max, + shape=(1,), + dtype=np.float32, + ), + "is_holding": spaces.Box( + low=0.0, high=1.0, shape=(1,), dtype=np.float32 + ), + } + ) + action_space = spaces.Box( + -1.0, 1.0, (config.get("OPEN_CLOSE_DRAWER_ACTION_SPACE_LENGTH", 8),) + ) + super().__init__( + checkpoint_path, observation_space, action_space, device, config=config + ) + + class MixerPolicy(RealPolicy): def __init__( self, diff --git a/spot_rl_experiments/spot_rl/skills/atomic_skills.py b/spot_rl_experiments/spot_rl/skills/atomic_skills.py index 8cf10ad9..49d93455 100644 --- a/spot_rl_experiments/spot_rl/skills/atomic_skills.py +++ b/spot_rl_experiments/spot_rl/skills/atomic_skills.py @@ -15,15 +15,23 @@ # Import Envs from spot_rl.envs.nav_env import SpotNavEnv +from spot_rl.envs.open_close_drawer_env import SpotOpenCloseDrawerEnv from spot_rl.envs.place_env import SpotPlaceEnv # Import policies -from spot_rl.real_policy import GazePolicy, MobileGazePolicy, NavPolicy, PlacePolicy +from spot_rl.real_policy import ( + GazePolicy, + MobileGazePolicy, + NavPolicy, + OpenCloseDrawerPolicy, + PlacePolicy, +) # Import utils and helpers from spot_rl.utils.construct_configs import ( construct_config_for_gaze, construct_config_for_nav, + construct_config_for_open_close_drawer, construct_config_for_place, ) @@ -738,3 +746,55 @@ def split_action(self, action: np.ndarray) -> Dict[str, Any]: } return action_dict + + +class OpenCloseDrawer(Skill): + """ + Open close drawer controller is used to execute open/close drawers + """ + + def __init__(self, spot, config=None) -> None: + if not config: + config = construct_config_for_open_close_drawer() + super().__init__(spot, config) + + # Setup + self.policy = OpenCloseDrawerPolicy( + self.config.WEIGHTS.OPEN_CLOSE_DRAWER, + device=self.config.DEVICE, + config=self.config, + ) + self.policy.reset() + + self.env = SpotOpenCloseDrawerEnv(self.config, spot) + + def reset_skill(self, goal_dict: Dict[str, Any]) -> Any: + """Refer to class Skill for documentation""" + # Reset the env and policy + observations = self.env.reset(goal_dict) + self.policy.reset() + + # Reset logged data at init + self.reset_logger() + + return observations + + def update_and_check_status(self, goal_dict: Dict[str, Any]) -> Tuple[bool, str]: + # Check for success and return appropriately + status = False + message = "Open/close failed to open/close the drawer" + check_open_close_success = self.env.get_success() + if check_open_close_success: + status = True + message = "Successfully opened/closed the drawer" + conditional_print(message=message, verbose=self.verbose) + return status, message + + def split_action(self, action: np.ndarray) -> Dict[str, Any]: + # Assign the action into action dict + action_dict = { + "arm_action": action[0:4], + "base_action": None, + "close_gripper": action[4], + } + return action_dict diff --git a/spot_rl_experiments/spot_rl/utils/construct_configs.py b/spot_rl_experiments/spot_rl/utils/construct_configs.py index 1512dabb..f97a038c 100644 --- a/spot_rl_experiments/spot_rl/utils/construct_configs.py +++ b/spot_rl_experiments/spot_rl/utils/construct_configs.py @@ -106,3 +106,26 @@ def construct_config_for_place(file_path=None, opts=[]): config.USE_MRCNN = False return config + + +def construct_config_for_open_close_drawer(file_path=None, opts=[]): + """ + Constructs and updates the config for open close drawer + + Args: + file_path (str): Path to the config file + opts (list): List of options to update the config + + Returns: + config (Config): Updated config object + """ + config = None + if file_path is None: + config = construct_config(opts=opts) + else: + config = construct_config(file_path=file_path, opts=opts) + + # Don't need head cameras for Gaze + config.USE_HEAD_CAMERA = False + + return config diff --git a/spot_rl_experiments/spot_rl/utils/depth_map_utils.py b/spot_rl_experiments/spot_rl/utils/depth_map_utils.py index 1e06e004..a89693f7 100644 --- a/spot_rl_experiments/spot_rl/utils/depth_map_utils.py +++ b/spot_rl_experiments/spot_rl/utils/depth_map_utils.py @@ -78,7 +78,7 @@ ) -def filter_depth(depth_img, max_depth, whiten_black=True): +def filter_depth(depth_img, max_depth, whiten_black=False): filtered_depth_img = ( fill_in_multiscale(depth_img.astype(np.float32) * (max_depth / 255.0))[0] * (255.0 / max_depth) @@ -242,7 +242,7 @@ def fill_in_multiscale( s4_blurred_depths[valid_pixels] = blurred[valid_pixels] # Calculate a top mask - top_mask = np.ones(depths_in.shape, dtype=np.bool) + top_mask = np.ones(depths_in.shape, dtype=bool) for pixel_col_idx in range(s4_blurred_depths.shape[1]): pixel_col = s4_blurred_depths[:, pixel_col_idx] top_pixel_row = np.argmax(pixel_col > 0.1) @@ -259,7 +259,7 @@ def fill_in_multiscale( # Extend highest pixel to top of image or create top mask s6_extended_depths = np.copy(s5_dilated_depths) - top_mask = np.ones(s5_dilated_depths.shape, dtype=np.bool) + top_mask = np.ones(s5_dilated_depths.shape, dtype=bool) top_row_pixels = np.argmax(s5_dilated_depths > 0.1, axis=0) top_pixel_values = s5_dilated_depths[ diff --git a/spot_rl_experiments/utils/cparamopenclosedrawer.yaml b/spot_rl_experiments/utils/cparamopenclosedrawer.yaml new file mode 100644 index 00000000..ad612cbf --- /dev/null +++ b/spot_rl_experiments/utils/cparamopenclosedrawer.yaml @@ -0,0 +1,18 @@ +# Define the model class name +MODEL_CLASS_NAME : "habitat_baselines.rl.ddppo.policy.resnet_policy.PointNavResNetPolicy" +# Define the observation dict +OBSERVATIONS_DICT: + articulated_agent_arm_depth: [[240, 228, 1], '0.0', '1.0', 'np.float32'] + joint: [[5,], 'np.finfo(np.float32).min', 'np.finfo(np.float32).max', 'np.float32'] + ee_pos: [[3,], 'np.finfo(np.float32).min', 'np.finfo(np.float32).max', 'np.float32'] + # is_holding: [[1,], '0', '1', 'np.float32'] # Some open/close drawers do not have is_holding sensor + handle_bbox: [[240, 228, 1], '0.0', '1.0', 'np.float32'] + art_pose_delta_sensor: [[1,], '0.0', '1.0', 'np.float32'] +# Define the action space output length +ACTION_SPACE_LENGTH: 8 +# The path to load and save the models +TARGET_HAB3_POLICY_PATH: "/home/jmmy/research/spot-sim2real/spot_rl_experiments/weights/open_close_drawer/od14lr_2_ckpt.129.pth" +OUTPUT_COMBINED_NET_SAVE_PATH: "/home/jmmy/research/spot-sim2real/spot_rl_experiments/weights/open_close_drawer/od14lr_2_ckpt.129.combined_net.torchscript" +# If we want to use stereo pair camera for mobile gaze +USE_STEREO_PAIR_CAMERA: False +NEW_HABITAT_LAB_POLICY_OR_OLD: 'new' diff --git a/spot_rl_experiments/utils/pytorch_to_torchscript.py b/spot_rl_experiments/utils/pytorch_to_torchscript.py index 344ac681..75139177 100644 --- a/spot_rl_experiments/utils/pytorch_to_torchscript.py +++ b/spot_rl_experiments/utils/pytorch_to_torchscript.py @@ -91,7 +91,7 @@ def forward( ) features, rnn_hidden_states = output_of_model[0], output_of_model[1] action = self.get_action(self.action_dist_ts(features), self.std) - action = action.mean + action = action.sample() return action, rnn_hidden_states @@ -288,7 +288,7 @@ def act(self, observations): self.test_recurrent_hidden_states, self.prev_actions, self.not_done_masks, - deterministic=True, + deterministic=False, ) actions = ( PolicyActionData.actions