diff --git a/.github/workflows/build-push-containers.yml b/.github/workflows/build-push-containers.yml index aa27c728b..87ee68a8b 100644 --- a/.github/workflows/build-push-containers.yml +++ b/.github/workflows/build-push-containers.yml @@ -7,11 +7,15 @@ on: branches: - 'main' - 'og-develop' + - 'action-primitives' jobs: docker: runs-on: [self-hosted, linux, gpu] steps: + - + name: Checkout + uses: actions/checkout@v3 - name: Set up Docker Buildx uses: docker/setup-buildx-action@v2 @@ -51,17 +55,24 @@ jobs: - name: Build and push dev image uses: docker/build-push-action@v4 + id: build-dev with: + context: . push: true tags: ${{ steps.meta-dev.outputs.tags }} labels: ${{ steps.meta-dev.outputs.labels }} file: docker/dev.Dockerfile cache-from: type=gha cache-to: type=gha,mode=max + + - name: Update prod image Dockerfile with dev image tag + run: | + sed -i "s/omnigibson-dev:latest/omnigibson-dev@${{ steps.build-dev.outputs.digest }}/g" docker/prod.Dockerfile && cat docker/prod.Dockerfile - name: Build and push prod image uses: docker/build-push-action@v4 with: + context: . push: true tags: ${{ steps.meta-prod.outputs.tags }} labels: ${{ steps.meta-prod.outputs.labels }} diff --git a/docker/dev.Dockerfile b/docker/dev.Dockerfile index 113889b7d..8f897cdfb 100644 --- a/docker/dev.Dockerfile +++ b/docker/dev.Dockerfile @@ -1,4 +1,4 @@ -FROM nvcr.io/nvidia/isaac-sim:2022.2.0 +FROM nvcr.io/nvidia/isaac-sim:2022.2.1 # Set up all the prerequisites. RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \ diff --git a/docker/prod.Dockerfile b/docker/prod.Dockerfile index a0eaeef21..9feb45495 100644 --- a/docker/prod.Dockerfile +++ b/docker/prod.Dockerfile @@ -1,4 +1,4 @@ -FROM stanfordvl/omnigibson-dev +FROM stanfordvl/omnigibson-dev:latest # Copy over omnigibson source ADD . /omnigibson-src diff --git a/docker/run_docker.sh b/docker/run_docker.sh index 146ae218c..78302514b 100755 --- a/docker/run_docker.sh +++ b/docker/run_docker.sh @@ -117,9 +117,6 @@ docker run \ -e DISPLAY=${DOCKER_DISPLAY} \ -e OMNIGIBSON_HEADLESS=${OMNIGIBSON_HEADLESS} \ -v $DATA_PATH/datasets:/data \ - -v ${ICD_PATH}:/etc/vulkan/icd.d/nvidia_icd.json \ - -v ${LAYERS_PATH}:/etc/vulkan/implicit_layer.d/nvidia_layers.json \ - -v ${EGL_VENDOR_PATH}:/usr/share/glvnd/egl_vendor.d/10_nvidia.json \ -v $DATA_PATH/isaac-sim/cache/kit:/isaac-sim/kit/cache/Kit:rw \ -v $DATA_PATH/isaac-sim/cache/ov:/root/.cache/ov:rw \ -v $DATA_PATH/isaac-sim/cache/pip:/root/.cache/pip:rw \ diff --git a/omnigibson/action_primitives/starter_semantic_action_primitives.py b/omnigibson/action_primitives/starter_semantic_action_primitives.py index cb726c44d..9dd48ed3b 100644 --- a/omnigibson/action_primitives/starter_semantic_action_primitives.py +++ b/omnigibson/action_primitives/starter_semantic_action_primitives.py @@ -991,6 +991,9 @@ def _move_hand_direct_joint(self, joint_pos, control_idx, stop_on_contact=False, action = self._empty_action() controller_name = "arm_{}".format(self.arm) + + action[self.robot.controller_action_idx[controller_name]] = joint_pos + prev_eef_pos = [0.0, 0.0, 0.0] for _ in range(max_steps_for_hand_move): current_joint_pos = self.robot.get_joint_positions()[control_idx] @@ -999,12 +1002,17 @@ def _move_hand_direct_joint(self, joint_pos, control_idx, stop_on_contact=False, return if stop_on_contact and detect_robot_collision_in_sim(self.robot, ignore_obj_in_hand=False): return + if max(np.abs(np.array(self.robot.get_eef_position(self.arm) - prev_eef_pos))) < 0.0001: + raise ActionPrimitiveError( + ActionPrimitiveError.Reason.EXECUTION_ERROR, + f"Hand is stuck" + ) if use_delta: action[self.robot.controller_action_idx[controller_name]] = self._compute_delta_command(controller_name, diff_joint_pos) - else: - action[self.robot.controller_action_idx[controller_name]] = joint_pos action = self._overwrite_head_action(action, self._tracking_object) if self._tracking_object is not None else action + + prev_eef_pos = self.robot.get_eef_position(self.arm) yield self.with_context(action) if not ignore_failure: @@ -1595,13 +1603,13 @@ def _sample_pose_near_object(self, obj, pose_on_obj=None, **kwargs): - 3-array: (x,y,z) Position in the world frame - 4-array: (x,y,z,w) Quaternion orientation in the world frame """ - if pose_on_obj is None: - pos_on_obj = self._sample_position_on_aabb_face(obj) - pose_on_obj = np.array([pos_on_obj, [0, 0, 0, 1]]) - with UndoableContext(self.robot, self.robot_copy, "simplified") as context: obj_rooms = obj.in_rooms if obj.in_rooms else [self.scene._seg_map.get_room_instance_by_point(pose_on_obj[0][:2])] for _ in range(MAX_ATTEMPTS_FOR_SAMPLING_POSE_NEAR_OBJECT): + if pose_on_obj is None: + pos_on_obj = self._sample_position_on_aabb_face(obj) + pose_on_obj = [pos_on_obj, np.array([0, 0, 0, 1])] + distance = np.random.uniform(0.0, 1.0) yaw = np.random.uniform(-np.pi, np.pi) avg_arm_workspace_range = np.mean(self.robot.arm_workspace_range[self.arm]) @@ -1619,7 +1627,8 @@ def _sample_pose_near_object(self, obj, pose_on_obj=None, **kwargs): return pose_2d raise ActionPrimitiveError( - ActionPrimitiveError.Reason.SAMPLING_ERROR, "Could not find valid position near object." + ActionPrimitiveError.Reason.SAMPLING_ERROR, "Could not find valid position near object.", + {"target object": obj.name, "target pos": obj.get_position(), "pose on target": pose_on_obj} ) @staticmethod diff --git a/omnigibson/action_primitives/symbolic_semantic_action_primitives.py b/omnigibson/action_primitives/symbolic_semantic_action_primitives.py index a073b7f64..c68458d98 100644 --- a/omnigibson/action_primitives/symbolic_semantic_action_primitives.py +++ b/omnigibson/action_primitives/symbolic_semantic_action_primitives.py @@ -16,6 +16,7 @@ import gym import numpy as np +from omnigibson.transition_rules import REGISTERED_RULES, TransitionRuleAPI from scipy.spatial.transform import Rotation, Slerp from omnigibson.utils.constants import JointType from pxr import PhysxSchema @@ -270,6 +271,9 @@ def _grasp(self, obj: DatasetObject): {"target object": obj.name, "object currently in hand": obj_in_hand.name}, ) + # Get close + yield from self._navigate_if_needed(obj) + # Perform forced assisted grasp obj.set_position(self.robot.get_eef_position(self.arm)) self.robot._ag_data[self.arm] = (obj, obj.root_link) @@ -613,12 +617,18 @@ def _cut(self, obj): "The target object is not sliceable or diceable.", {"target object": obj.name} ) + + # Get close + yield from self._navigate_if_needed(obj) - # TODO: Trigger the slicing rule manually. - - - pass + # TODO: Do some more validation + added_obj_attrs = [] + removed_objs = [] + output = REGISTERED_RULES["SlicingRule"].transition({"sliceable": [obj]}) + added_obj_attrs += output.add + removed_objs += output.remove + TransitionRuleAPI.execute_transition(added_obj_attrs=added_obj_attrs, removed_objs=removed_objs) def _place_near_heating_element(self, heat_source_obj): obj_in_hand = self._get_obj_in_hand() @@ -772,14 +782,14 @@ def _sample_pose_near_object(self, obj, pose_on_obj=None, **kwargs): - 3-array: (x,y,z) Position in the world frame - 4-array: (x,y,z,w) Quaternion orientation in the world frame """ - if pose_on_obj is None: - pos_on_obj = self._sample_position_on_aabb_face(obj) - pose_on_obj = np.array([pos_on_obj, [0, 0, 0, 1]]) - with UndoableContext(self.robot, self.robot_copy, "simplified") as context: obj_rooms = obj.in_rooms if obj.in_rooms else [self.scene._seg_map.get_room_instance_by_point(pose_on_obj[0][:2])] for _ in range(MAX_ATTEMPTS_FOR_SAMPLING_POSE_NEAR_OBJECT): - distance = np.random.uniform(0.0, 1.0) + if pose_on_obj is None: + pos_on_obj = self._sample_position_on_aabb_face(obj) + pose_on_obj = [pos_on_obj, np.array([0, 0, 0, 1])] + + distance = np.random.uniform(0.0, 3.0) yaw = np.random.uniform(-np.pi, np.pi) pose_2d = np.array( [pose_on_obj[0][0] + distance * np.cos(yaw), pose_on_obj[0][1] + distance * np.sin(yaw), yaw + np.pi] @@ -796,7 +806,8 @@ def _sample_pose_near_object(self, obj, pose_on_obj=None, **kwargs): return pose_2d raise ActionPrimitiveError( - ActionPrimitiveError.Reason.SAMPLING_ERROR, "Could not find valid position near object." + ActionPrimitiveError.Reason.SAMPLING_ERROR, "Could not find valid position near object.", + {"target object": obj.name, "target pos": obj.get_position(), "pose on target": pose_on_obj} ) @staticmethod diff --git a/omnigibson/configs/fetch_behavior.yaml b/omnigibson/configs/fetch_behavior.yaml index 84da35e39..a2e57cad0 100644 --- a/omnigibson/configs/fetch_behavior.yaml +++ b/omnigibson/configs/fetch_behavior.yaml @@ -33,6 +33,15 @@ robots: default_trunk_offset: 0.365 default_arm_pose: diagonal30 reset_joint_pos: tuck + sensor_config: + VisionSensor: + sensor_kwargs: + image_height: 128 + image_width: 128 + ScanSensor: + sensor_kwargs: + min_range: 0.05 + max_range: 10.0 controller_config: base: name: DifferentialDriveController diff --git a/omnigibson/configs/homeboy.yaml b/omnigibson/configs/homeboy.yaml index 171888c2b..c28beb53e 100644 --- a/omnigibson/configs/homeboy.yaml +++ b/omnigibson/configs/homeboy.yaml @@ -32,6 +32,7 @@ robots: action_normalize: false action_type: continuous grasping_mode: sticky + disable_grasp_handling: True rigid_trunk: false default_trunk_offset: 0.365 default_arm_pose: diagonal30 diff --git a/omnigibson/configs/robots/fetch.yaml b/omnigibson/configs/robots/fetch.yaml index 42bdeb8ac..9d8054e17 100644 --- a/omnigibson/configs/robots/fetch.yaml +++ b/omnigibson/configs/robots/fetch.yaml @@ -19,6 +19,15 @@ robot: rigid_trunk: false default_trunk_offset: 0.365 default_arm_pose: vertical + sensor_config: + VisionSensor: + sensor_kwargs: + image_height: 128 + image_width: 128 + ScanSensor: + sensor_kwargs: + min_range: 0.05 + max_range: 10.0 controller_config: base: name: DifferentialDriveController diff --git a/omnigibson/configs/robots/freight.yaml b/omnigibson/configs/robots/freight.yaml index e3de96e4f..7431b464d 100644 --- a/omnigibson/configs/robots/freight.yaml +++ b/omnigibson/configs/robots/freight.yaml @@ -10,6 +10,15 @@ robot: base_name: null scale: 1.0 self_collision: false + sensor_config: + VisionSensor: + sensor_kwargs: + image_height: 128 + image_width: 128 + ScanSensor: + sensor_kwargs: + min_range: 0.05 + max_range: 10.0 controller_config: base: name: DifferentialDriveController \ No newline at end of file diff --git a/omnigibson/configs/robots/husky.yaml b/omnigibson/configs/robots/husky.yaml index 9f63cd58b..d35680b5b 100644 --- a/omnigibson/configs/robots/husky.yaml +++ b/omnigibson/configs/robots/husky.yaml @@ -10,6 +10,15 @@ robot: base_name: null scale: 1.0 self_collision: false + sensor_config: + VisionSensor: + sensor_kwargs: + image_height: 128 + image_width: 128 + ScanSensor: + sensor_kwargs: + min_range: 0.05 + max_range: 10.0 controller_config: base: name: JointController \ No newline at end of file diff --git a/omnigibson/configs/robots/locobot.yaml b/omnigibson/configs/robots/locobot.yaml index f5565fb75..c36e144f3 100644 --- a/omnigibson/configs/robots/locobot.yaml +++ b/omnigibson/configs/robots/locobot.yaml @@ -10,6 +10,15 @@ robot: base_name: null scale: 1.0 self_collision: false + sensor_config: + VisionSensor: + sensor_kwargs: + image_height: 128 + image_width: 128 + ScanSensor: + sensor_kwargs: + min_range: 0.05 + max_range: 10.0 controller_config: base: name: DifferentialDriveController \ No newline at end of file diff --git a/omnigibson/configs/robots/turtlebot.yaml b/omnigibson/configs/robots/turtlebot.yaml index 28fdc8613..0431a430f 100644 --- a/omnigibson/configs/robots/turtlebot.yaml +++ b/omnigibson/configs/robots/turtlebot.yaml @@ -10,6 +10,15 @@ robot: base_name: null scale: 1.0 self_collision: false + sensor_config: + VisionSensor: + sensor_kwargs: + image_height: 128 + image_width: 128 + ScanSensor: + sensor_kwargs: + min_range: 0.05 + max_range: 10.0 controller_config: base: name: DifferentialDriveController \ No newline at end of file diff --git a/omnigibson/configs/sensors/scan.yaml b/omnigibson/configs/sensors/scan.yaml new file mode 100644 index 000000000..318b761de --- /dev/null +++ b/omnigibson/configs/sensors/scan.yaml @@ -0,0 +1,30 @@ +# Example ScanSensor sensor config +# See omnigibson/sensors/__init__/create_sensor and omnigibson/sensors/scan_sensor for docstring of arguments +# Arguments below are the arguments that should be specified by external user (other kwargs +# used in constructor are generated automatically at runtime) +robot: + sensor_config: + ScanSensor: + modalities: [scan, occupancy_grid] # if specified, this will override the values in robots_config["obs_modalities"] + enabled: true + noise_type: null + noise_kwargs: null + sensor_kwargs: + + # Basic LIDAR kwargs + min_range: 0.05 + max_range: 10.0 + horizontal_fov: 360.0 + vertical_fov: 1.0 + yaw_offset: 0.0 + horizontal_resolution: 1.0 + vertical_resolution: 1.0 + rotation_rate: 0.0 + draw_points: false + draw_lines: false + + # Occupancy Grid kwargs + occupancy_grid_resolution: 128 + occupancy_grid_range: 5.0 + occupancy_grid_inner_radius: 0.5 + occupancy_grid_local_link: null \ No newline at end of file diff --git a/omnigibson/configs/sensors/vision.yaml b/omnigibson/configs/sensors/vision.yaml new file mode 100644 index 000000000..44fdfc6c6 --- /dev/null +++ b/omnigibson/configs/sensors/vision.yaml @@ -0,0 +1,14 @@ +# Example VisionSensor sensor config +# See omnigibson/sensors/__init__/create_sensor and omnigibson/sensors/vision_sensor for docstring of arguments +# Arguments below are the arguments that should be specified by external user (other kwargs +# used in constructor are generated automatically at runtime) +robot: + sensor_config: + VisionSensor: + modalities: [rgb, depth] # if specified, this will override the values in robots_config["obs_modalities"] + enabled: true + noise_type: null + noise_kwargs: null + sensor_kwargs: + image_height: 128 + image_width: 128 \ No newline at end of file diff --git a/omnigibson/configs/turtlebot_nav.yaml b/omnigibson/configs/turtlebot_nav.yaml index 864af5d12..36abc695f 100644 --- a/omnigibson/configs/turtlebot_nav.yaml +++ b/omnigibson/configs/turtlebot_nav.yaml @@ -29,6 +29,15 @@ robots: self_collision: false action_normalize: true action_type: continuous + sensor_config: + VisionSensor: + sensor_kwargs: + image_height: 128 + image_width: 128 + ScanSensor: + sensor_kwargs: + min_range: 0.05 + max_range: 10.0 controller_config: base: name: DifferentialDriveController diff --git a/omnigibson/maps/traversable_map.py b/omnigibson/maps/traversable_map.py index 8d1e7a6e4..83f3fb299 100644 --- a/omnigibson/maps/traversable_map.py +++ b/omnigibson/maps/traversable_map.py @@ -155,8 +155,11 @@ def build_trav_graph(map_size, maps_path, floor, trav_map): # only take the largest connected component largest_cc = max(nx.connected_components(g), key=len) g = g.subgraph(largest_cc).copy() - with open(graph_file, "wb") as pfile: - pickle.dump(g, pfile) + try: + with open(graph_file, "wb") as pfile: + pickle.dump(g, pfile) + except: + log.warning("Cannot cache traversable graph to disk possibly because dataset is read-only. Will have to recompute each time.") floor_graph.append(g) diff --git a/omnigibson/reward_functions/grasp_reward.py b/omnigibson/reward_functions/grasp_reward.py new file mode 100644 index 000000000..0a3387779 --- /dev/null +++ b/omnigibson/reward_functions/grasp_reward.py @@ -0,0 +1,72 @@ +import math +from omnigibson.reward_functions.reward_function_base import BaseRewardFunction +from omnigibson.utils.motion_planning_utils import detect_robot_collision_in_sim +import omnigibson.utils.transform_utils as T + + +class GraspReward(BaseRewardFunction): + """ + Grasp reward + """ + + def __init__(self, obj_name, dist_coeff, grasp_reward): + # Store internal vars + self.prev_grasping = False + self.obj_name = obj_name + self.obj = None + self.dist_coeff = dist_coeff + self.grasp_reward = grasp_reward + + # Run super + super().__init__() + + def _step(self, task, env, action): + self.obj = env.scene.object_registry("name", self.obj_name) if self.obj is None else self.obj + + robot = env.robots[0] + obj_in_hand = robot._ag_obj_in_hand[robot.default_arm] + current_grasping = True if obj_in_hand == self.obj else False + + # Reward varying based on combination of whether the robot was previously grasping the desired and object + # and is currently grasping the desired object + + # not grasping -> not grasping = Distance between eef and obj reward + # not grasping -> grasping = Minimizing MOI + grasp reward + # grapsing -> not grapsing = Distance between eef and obj reward + # grasping -> grasping = Minimizing MOI + grasp reward + + reward = None + + if not self.prev_grasping and not current_grasping: + eef_pos = robot.get_eef_position(robot.default_arm) + obj_center = self.obj.aabb_center + dist = T.l2_distance(eef_pos, obj_center) + reward = math.exp(-dist) * self.dist_coeff + + elif not self.prev_grasping and current_grasping: + robot_center = robot.aabb_center + obj_center = self.obj.aabb_center + dist = T.l2_distance(robot_center, obj_center) + dist_reward = math.exp(-dist) * self.dist_coeff + reward = dist_reward + self.grasp_reward + + elif self.prev_grasping and not current_grasping: + eef_pos = robot.get_eef_position(robot.default_arm) + obj_center = self.obj.aabb_center + dist = T.l2_distance(eef_pos, obj_center) + reward = math.exp(-dist) * self.dist_coeff + + elif self.prev_grasping and current_grasping: + robot_center = robot.aabb_center + obj_center = self.obj.aabb_center + dist = T.l2_distance(robot_center, obj_center) + dist_reward = math.exp(-dist) * self.dist_coeff + reward = dist_reward + self.grasp_reward + + # Overwrite reward if robot is in collision + # The one step where it grasps the object, it is in collision and this triggers + if detect_robot_collision_in_sim(robot, ignore_obj_in_hand=True): + reward = -10 + + self.prev_grasping = current_grasping + return reward, {} diff --git a/omnigibson/robots/fetch.py b/omnigibson/robots/fetch.py index cef1b1656..048bc8c20 100644 --- a/omnigibson/robots/fetch.py +++ b/omnigibson/robots/fetch.py @@ -56,6 +56,7 @@ def __init__( # Unique to BaseRobot obs_modalities="all", proprio_obs="default", + sensor_config=None, # Unique to ManipulationRobot grasping_mode="physical", @@ -99,10 +100,14 @@ def __init__( obs_modalities (str or list of str): Observation modalities to use for this robot. Default is "all", which corresponds to all modalities being used. Otherwise, valid options should be part of omnigibson.sensors.ALL_SENSOR_MODALITIES. + Note: If @sensor_config explicitly specifies `modalities` for a given sensor class, it will + override any values specified from @obs_modalities! proprio_obs (str or list of str): proprioception observation key(s) to use for generating proprioceptive observations. If str, should be exactly "default" -- this results in the default proprioception observations being used, as defined by self.default_proprio_obs. See self._get_proprioception_dict for valid key choices + sensor_config (None or dict): nested dictionary mapping sensor class name(s) to specific sensor + configurations for this object. This will override any default values specified by this class. grasping_mode (str): One of {"physical", "assisted", "sticky"}. If "physical", no assistive grasping will be applied (relies on contact friction + finger force). If "assisted", will magnetize any object touching and within the gripper's fingers. @@ -149,6 +154,7 @@ def __init__( reset_joint_pos=reset_joint_pos, obs_modalities=obs_modalities, proprio_obs=proprio_obs, + sensor_config=sensor_config, grasping_mode=grasping_mode, **kwargs, ) diff --git a/omnigibson/robots/manipulation_robot.py b/omnigibson/robots/manipulation_robot.py index af55114fe..0d6ff3ae3 100644 --- a/omnigibson/robots/manipulation_robot.py +++ b/omnigibson/robots/manipulation_robot.py @@ -109,10 +109,11 @@ def __init__( # Unique to BaseRobot obs_modalities="all", proprio_obs="default", + sensor_config=None, # Unique to ManipulationRobot grasping_mode="physical", - disable_grasp_handling=False, + disable_grasp_handling=False, # TODO: revert to False. This is for debugging purposes **kwargs, ): @@ -149,10 +150,14 @@ def __init__( obs_modalities (str or list of str): Observation modalities to use for this robot. Default is "all", which corresponds to all modalities being used. Otherwise, valid options should be part of omnigibson.sensors.ALL_SENSOR_MODALITIES. + Note: If @sensor_config explicitly specifies `modalities` for a given sensor class, it will + override any values specified from @obs_modalities! proprio_obs (str or list of str): proprioception observation key(s) to use for generating proprioceptive observations. If str, should be exactly "default" -- this results in the default proprioception observations being used, as defined by self.default_proprio_obs. See self._get_proprioception_dict for valid key choices + sensor_config (None or dict): nested dictionary mapping sensor class name(s) to specific sensor + configurations for this object. This will override any default values specified by this class. grasping_mode (str): One of {"physical", "assisted", "sticky"}. If "physical", no assistive grasping will be applied (relies on contact friction + finger force). If "assisted", will magnetize any object touching and within the gripper's fingers. @@ -200,6 +205,7 @@ def __init__( reset_joint_pos=reset_joint_pos, obs_modalities=obs_modalities, proprio_obs=proprio_obs, + sensor_config=sensor_config, **kwargs, ) diff --git a/omnigibson/robots/robot_base.py b/omnigibson/robots/robot_base.py index 8dc491e1a..9175ad335 100644 --- a/omnigibson/robots/robot_base.py +++ b/omnigibson/robots/robot_base.py @@ -7,7 +7,7 @@ from omnigibson.objects.usd_object import USDObject from omnigibson.objects.controllable_object import ControllableObject from omnigibson.utils.gym_utils import GymObservable -from omnigibson.utils.python_utils import classproperty +from omnigibson.utils.python_utils import classproperty, merge_nested_dicts from omnigibson.utils.vision_utils import segmentation_to_rgb from omnigibson.utils.constants import PrimType from pxr import PhysxSchema @@ -59,6 +59,7 @@ def __init__( # Unique to this class obs_modalities="all", proprio_obs="default", + sensor_config=None, **kwargs, ): @@ -95,10 +96,14 @@ def __init__( obs_modalities (str or list of str): Observation modalities to use for this robot. Default is "all", which corresponds to all modalities being used. Otherwise, valid options should be part of omnigibson.sensors.ALL_SENSOR_MODALITIES. + Note: If @sensor_config explicitly specifies `modalities` for a given sensor class, it will + override any values specified from @obs_modalities! proprio_obs (str or list of str): proprioception observation key(s) to use for generating proprioceptive observations. If str, should be exactly "default" -- this results in the default proprioception observations being used, as defined by self.default_proprio_obs. See self._get_proprioception_dict for valid key choices + sensor_config (None or dict): nested dictionary mapping sensor class name(s) to specific sensor + configurations for this object. This will override any default values specified by this class. kwargs (dict): Additional keyword arguments that are used for other super() calls from subclasses, allowing for flexible compositions of various object subclasses (e.g.: Robot is USDObject + ControllableObject). """ @@ -106,6 +111,7 @@ def __init__( self._obs_modalities = obs_modalities if obs_modalities == "all" else \ {obs_modalities} if isinstance(obs_modalities, str) else set(obs_modalities) # this will get updated later when we fill in our sensors self._proprio_obs = self.default_proprio_obs if proprio_obs == "default" else list(proprio_obs) + self._sensor_config = sensor_config # Process abilities robot_abilities = {"robot": {}} @@ -148,6 +154,32 @@ def _post_load(self): # Run super post load first super()._post_load() + # Load the sensors + self._load_sensors() + + def _initialize(self): + # Run super first + super()._initialize() + + # Initialize all sensors + for sensor in self._sensors.values(): + sensor.initialize() + + # Load the observation space for this robot + self.load_observation_space() + + # Validate this robot configuration + self._validate_configuration() + + def _load_sensors(self): + """ + Loads sensor(s) to retrieve observations from this object. + Stores created sensors as dictionary mapping sensor names to specific sensor + instances used by this object. + """ + # Populate sensor config + self._sensor_config = self._generate_sensor_config(custom_config=self._sensor_config) + # Search for any sensors this robot might have attached to any of its links self._sensors = dict() obs_modalities = set() @@ -158,15 +190,17 @@ def _post_load(self): if prim_type in SENSOR_PRIMS_TO_SENSOR_CLS: # Infer what obs modalities to use for this sensor sensor_cls = SENSOR_PRIMS_TO_SENSOR_CLS[prim_type] - modalities = sensor_cls.all_modalities if self._obs_modalities == "all" else \ - sensor_cls.all_modalities.intersection(self._obs_modalities) - obs_modalities = obs_modalities.union(modalities) + sensor_kwargs = self._sensor_config[sensor_cls.__name__] + if "modalities" not in sensor_kwargs: + sensor_kwargs["modalities"] = sensor_cls.all_modalities if self._obs_modalities == "all" else \ + sensor_cls.all_modalities.intersection(self._obs_modalities) + obs_modalities = obs_modalities.union(sensor_kwargs["modalities"]) # Create the sensor and store it internally sensor = create_sensor( sensor_type=prim_type, prim_path=str(prim.GetPrimPath()), name=f"{self.name}:{link_name}_{prim_type}_sensor", - modalities=modalities, + **sensor_kwargs, ) self._sensors[sensor.name] = sensor @@ -177,19 +211,28 @@ def _post_load(self): # Update our overall obs modalities self._obs_modalities = obs_modalities - def _initialize(self): - # Run super first - super()._initialize() + def _generate_sensor_config(self, custom_config=None): + """ + Generates a fully-populated sensor config, overriding any default values with the corresponding values + specified in @custom_config - # Initialize all sensors - for sensor in self._sensors.values(): - sensor.initialize() + Args: + custom_config (None or Dict[str, ...]): nested dictionary mapping sensor class name(s) to specific custom + sensor configurations for this object. This will override any default values specified by this class - # Load the observation space for this robot - self.load_observation_space() + Returns: + dict: Fully-populated nested dictionary mapping sensor class name(s) to specific sensor configurations + for this object + """ + sensor_config = {} if custom_config is None else deepcopy(custom_config) - # Validate this robot configuration - self._validate_configuration() + # Merge the sensor dictionaries + sensor_config = merge_nested_dicts( + base_dict=self._default_sensor_config, + extra_dict=sensor_config, + ) + + return sensor_config def _validate_configuration(self): """ @@ -393,6 +436,70 @@ def proprioception_dim(self): """ return len(self.get_proprioception()) + @property + def _default_sensor_config(self): + """ + Returns: + dict: default nested dictionary mapping sensor class name(s) to specific sensor + configurations for this object. See kwargs from omnigibson/sensors/__init__/create_sensor for more + details + + Expected structure is as follows: + SensorClassName1: + modalities: ... + enabled: ... + noise_type: ... + noise_kwargs: + ... + sensor_kwargs: + ... + SensorClassName2: + modalities: ... + enabled: ... + noise_type: ... + noise_kwargs: + ... + sensor_kwargs: + ... + ... + """ + return { + "VisionSensor": { + "enabled": True, + "noise_type": None, + "noise_kwargs": None, + "sensor_kwargs": { + "image_height": 128, + "image_width": 128, + }, + }, + "ScanSensor": { + "enabled": True, + "noise_type": None, + "noise_kwargs": None, + "sensor_kwargs": { + + # Basic LIDAR kwargs + "min_range": 0.05, + "max_range": 10.0, + "horizontal_fov": 360.0, + "vertical_fov": 1.0, + "yaw_offset": 0.0, + "horizontal_resolution": 1.0, + "vertical_resolution": 1.0, + "rotation_rate": 0.0, + "draw_points": False, + "draw_lines": False, + + # Occupancy Grid kwargs + "occupancy_grid_resolution": 128, + "occupancy_grid_range": 5.0, + "occupancy_grid_inner_radius": 0.5, + "occupancy_grid_local_link": None, + }, + }, + } + @property def default_proprio_obs(self): """ diff --git a/omnigibson/robots/tiago.py b/omnigibson/robots/tiago.py index 2cc6a50e7..7874f1840 100644 --- a/omnigibson/robots/tiago.py +++ b/omnigibson/robots/tiago.py @@ -71,6 +71,7 @@ def __init__( # Unique to BaseRobot obs_modalities="all", proprio_obs="default", + sensor_config=None, # Unique to ManipulationRobot grasping_mode="physical", @@ -116,10 +117,14 @@ def __init__( obs_modalities (str or list of str): Observation modalities to use for this robot. Default is "all", which corresponds to all modalities being used. Otherwise, valid options should be part of omnigibson.sensors.ALL_SENSOR_MODALITIES. + Note: If @sensor_config explicitly specifies `modalities` for a given sensor class, it will + override any values specified from @obs_modalities! proprio_obs (str or list of str): proprioception observation key(s) to use for generating proprioceptive observations. If str, should be exactly "default" -- this results in the default proprioception observations being used, as defined by self.default_proprio_obs. See self._get_proprioception_dict for valid key choices + sensor_config (None or dict): nested dictionary mapping sensor class name(s) to specific sensor + configurations for this object. This will override any default values specified by this class. grasping_mode (str): One of {"physical", "assisted", "sticky"}. If "physical", no assistive grasping will be applied (relies on contact friction + finger force). If "assisted", will magnetize any object touching and within the gripper's fingers. @@ -169,6 +174,7 @@ def __init__( reset_joint_pos=reset_joint_pos, obs_modalities=obs_modalities, proprio_obs=proprio_obs, + sensor_config=sensor_config, grasping_mode=grasping_mode, **kwargs, ) diff --git a/omnigibson/scene_graphs/graph_builder.py b/omnigibson/scene_graphs/graph_builder.py index 1a2f202bc..15255d679 100644 --- a/omnigibson/scene_graphs/graph_builder.py +++ b/omnigibson/scene_graphs/graph_builder.py @@ -164,7 +164,7 @@ def step(self, scene): # Get the bounding box. if hasattr(obj, "get_base_aligned_bbox"): - bbox_center, bbox_orn, bbox_extent, _ = obj.get_base_aligned_bbox(visual=True) + bbox_center, bbox_orn, bbox_extent, _ = obj.get_base_aligned_bbox(visual=True, fallback_to_aabb=True) bbox_pose = T.pose2mat((bbox_center, bbox_orn)) else: bbox_pose, bbox_extent = get_robot_bbox(robot) diff --git a/omnigibson/sensors/__init__.py b/omnigibson/sensors/__init__.py index 11ba84d2c..215aa0924 100644 --- a/omnigibson/sensors/__init__.py +++ b/omnigibson/sensors/__init__.py @@ -19,6 +19,7 @@ def create_sensor( prim_path, name, modalities="all", + enabled=True, sensor_kwargs=None, noise_type=None, noise_kwargs=None @@ -35,6 +36,7 @@ def create_sensor( name (str): Name for the sensor. Names need to be unique per scene. modalities (str or list of str): Modality(s) supported by this sensor. Valid options are part of sensor.all_modalities. Default is "all", which corresponds to all modalities being used + enabled (bool): Whether this sensor should be enabled or not sensor_kwargs (dict): Any keyword kwargs to pass to the constructor noise_type (str): Type of sensor to create. Should be one of REGISTERED_SENSOR_NOISES (i.e.: the string name of the desired class to create) @@ -64,6 +66,13 @@ def create_sensor( # Create the sensor sensor_kwargs = dict() if sensor_kwargs is None else sensor_kwargs - sensor = sensor_cls(prim_path=prim_path, name=name, modalities=modalities, noise=noise, **sensor_kwargs) + sensor = sensor_cls( + prim_path=prim_path, + name=name, + modalities=modalities, + enabled=enabled, + noise=noise, + **sensor_kwargs, + ) return sensor diff --git a/omnigibson/tasks/__init__.py b/omnigibson/tasks/__init__.py index 2a6bfaa8f..550a873f0 100644 --- a/omnigibson/tasks/__init__.py +++ b/omnigibson/tasks/__init__.py @@ -3,3 +3,4 @@ from omnigibson.tasks.point_navigation_task import PointNavigationTask from omnigibson.tasks.point_reaching_task import PointReachingTask from omnigibson.tasks.behavior_task import BehaviorTask +from omnigibson.tasks.grasp_task import GraspTask diff --git a/omnigibson/tasks/grasp_task.py b/omnigibson/tasks/grasp_task.py new file mode 100644 index 000000000..31a79205a --- /dev/null +++ b/omnigibson/tasks/grasp_task.py @@ -0,0 +1,140 @@ +import random +import numpy as np +import omnigibson as og +from omnigibson.action_primitives.starter_semantic_action_primitives import UndoableContext +from omnigibson.reward_functions.grasp_reward import GraspReward + +from omnigibson.tasks.task_base import BaseTask +from omnigibson.scenes.scene_base import Scene +from omnigibson.termination_conditions.grasp_goal import GraspGoal +from omnigibson.termination_conditions.timeout import Timeout +from omnigibson.utils.grasping_planning_utils import get_grasp_poses_for_object_sticky +from omnigibson.utils.motion_planning_utils import set_arm_and_detect_collision +from omnigibson.utils.python_utils import classproperty +from omnigibson.utils.sim_utils import land_object + +DIST_COEFF = 0.1 +GRASP_REWARD = 1.0 +MAX_JOINT_RANDOMIZATION_ATTEMPTS = 50 + +class GraspTask(BaseTask): + """ + Grasp task + """ + + def __init__( + self, + obj_name, + termination_config=None, + reward_config=None, + ): + self.obj_name = obj_name + super().__init__(termination_config=termination_config, reward_config=reward_config) + + def _load(self, env): + pass + + def _create_termination_conditions(self): + terminations = dict() + terminations["graspgoal"] = GraspGoal( + self.obj_name + ) + terminations["timeout"] = Timeout(max_steps=self._termination_config["max_steps"]) + + return terminations + + def _create_reward_functions(self): + rewards = dict() + rewards["grasp"] = GraspReward( + self.obj_name, + dist_coeff=self._reward_config["r_dist_coeff"], + grasp_reward=self._reward_config["r_grasp"] + ) + return rewards + + def _reset_agent(self, env): + # from IPython import embed; embed() + if hasattr(env, '_primitive_controller'): + robot = env.robots[0] + # Randomize the robots joint positions + joint_control_idx = np.concatenate([robot.trunk_control_idx, robot.arm_control_idx[robot.default_arm]]) + joint_combined_idx = np.concatenate([robot.trunk_control_idx, robot.arm_control_idx["combined"]]) + initial_joint_pos = np.array(robot.get_joint_positions()[joint_combined_idx]) + control_idx_in_joint_pos = np.where(np.in1d(joint_combined_idx, joint_control_idx))[0] + + with UndoableContext(env._primitive_controller.robot, env._primitive_controller.robot_copy, "original") as context: + for _ in range(MAX_JOINT_RANDOMIZATION_ATTEMPTS): + joint_pos, joint_control_idx = self._get_random_joint_position(robot) + initial_joint_pos[control_idx_in_joint_pos] = joint_pos + if not set_arm_and_detect_collision(context, initial_joint_pos): + robot.set_joint_positions(joint_pos, joint_control_idx) + og.sim.step() + break + + # Randomize the robot's 2d pose + obj = env.scene.object_registry("name", self.obj_name) + grasp_poses = get_grasp_poses_for_object_sticky(obj) + grasp_pose, _ = random.choice(grasp_poses) + sampled_pose_2d = env._primitive_controller._sample_pose_near_object(obj, pose_on_obj=grasp_pose) + robot_pose = env._primitive_controller._get_robot_pose_from_2d_pose(sampled_pose_2d) + robot.set_position_orientation(*robot_pose) + + # Overwrite reset by only removeing reset scene + def reset(self, env): + """ + Resets this task in the environment + + Args: + env (Environment): environment instance to reset + """ + # Reset the scene, agent, and variables + # self._reset_scene(env) + self._reset_agent(env) + self._reset_variables(env) + + # Also reset all termination conditions and reward functions + for termination_condition in self._termination_conditions.values(): + termination_condition.reset(self, env) + for reward_function in self._reward_functions.values(): + reward_function.reset(self, env) + + # Fill in low dim obs dim so we can use this to create the observation space later + self._low_dim_obs_dim = len(self.get_obs(env=env, flatten_low_dim=True)["low_dim"]) + + + def _get_random_joint_position(self, robot): + joint_positions = [] + joint_control_idx = np.concatenate([robot.trunk_control_idx, robot.arm_control_idx[robot.default_arm]]) + joints = np.array([joint for joint in robot.joints.values()]) + arm_joints = joints[joint_control_idx] + for i, joint in enumerate(arm_joints): + val = random.uniform(joint.lower_limit, joint.upper_limit) + joint_positions.append(val) + return joint_positions, joint_control_idx + + def _get_obs(self, env): + # No task-specific obs of any kind + return dict(), dict() + + def _load_non_low_dim_observation_space(self): + # No non-low dim observations so we return an empty dict + return dict() + + @classproperty + def valid_scene_types(cls): + # Any scene works + return {Scene} + + @classproperty + def default_termination_config(cls): + return { + "max_steps": 100000 + } + + @classproperty + def default_reward_config(cls): + return { + "r_dist_coeff": DIST_COEFF, + "r_grasp": GRASP_REWARD, + } + diff --git a/omnigibson/termination_conditions/grasp_goal.py b/omnigibson/termination_conditions/grasp_goal.py new file mode 100644 index 000000000..e5e8b0dd5 --- /dev/null +++ b/omnigibson/termination_conditions/grasp_goal.py @@ -0,0 +1,21 @@ +from omnigibson.termination_conditions.termination_condition_base import SuccessCondition +import omnigibson.utils.transform_utils as T + + +class GraspGoal(SuccessCondition): + """ + GraspGoal (success condition) + """ + + def __init__(self, obj_name): + self.obj_name = obj_name + self.obj = None + + # Run super init + super().__init__() + + def _step(self, task, env, action): + self.obj = env.scene.object_registry("name", self.obj_name) if self.obj is None else self.obj + robot = env.robots[0] + obj_in_hand = robot._ag_obj_in_hand[robot.default_arm] + return obj_in_hand == self.obj diff --git a/omnigibson/transition_rules.py b/omnigibson/transition_rules.py index 4fbd1d9c1..652b3e6b7 100644 --- a/omnigibson/transition_rules.py +++ b/omnigibson/transition_rules.py @@ -153,6 +153,10 @@ def step(cls): added_obj_attrs += output.add removed_objs += output.remove + cls.execute_transition(added_obj_attrs=added_obj_attrs, removed_objs=removed_objs) + + @classmethod + def execute_transition(cls, added_obj_attrs, removed_objs): # Process all transition results if len(removed_objs) > 0: disclaimer( @@ -1266,10 +1270,13 @@ def _execute_recipe(cls, container, recipe, in_volume): system.remove_all_group_particles(group=group_name) # Remove either all objects or only the recipe-relevant objects inside the container - objs_to_remove.extend(np.concatenate([ - cls._OBJECTS[np.where(in_volume[cls._CATEGORY_IDXS[obj_category]])[0]] - for obj_category in recipe["input_objects"].keys() - ]) if cls.ignore_nonrecipe_objects else cls._OBJECTS[np.where(in_volume)[0]]) + object_mask = in_volume.copy() + if cls.ignore_nonrecipe_objects: + object_category_mask = np.zeros_like(object_mask, dtype=bool) + for obj_category in recipe["input_objects"].keys(): + object_category_mask[cls._CATEGORY_IDXS[obj_category]] = True + object_mask &= object_category_mask + objs_to_remove.extend(cls._OBJECTS[object_mask]) volume += sum(obj.volume for obj in objs_to_remove) # Define callback for spawning new objects inside container diff --git a/omnigibson/utils/motion_planning_utils.py b/omnigibson/utils/motion_planning_utils.py index fe6849baf..423fe2f54 100644 --- a/omnigibson/utils/motion_planning_utils.py +++ b/omnigibson/utils/motion_planning_utils.py @@ -246,8 +246,12 @@ def state_valid_fn(q): end_conf[i] = joint.upper_limit if end_conf[i] < joint.lower_limit: end_conf[i] = joint.lower_limit - bounds.setLow(i, float(joint.lower_limit)) - bounds.setHigh(i, float(joint.upper_limit)) + if joint.upper_limit - joint.lower_limit > 2 * np.pi: + bounds.setLow(i, 0.0) + bounds.setHigh(i, 2 * np.pi) + else: + bounds.setLow(i, float(joint.lower_limit)) + bounds.setHigh(i, float(joint.upper_limit)) space.setBounds(bounds) # create a simple setup object diff --git a/tests/grasp_RL_script.py b/tests/grasp_RL_script.py new file mode 100644 index 000000000..577d163fc --- /dev/null +++ b/tests/grasp_RL_script.py @@ -0,0 +1,226 @@ +from datetime import datetime +import math +import uuid +import numpy as np +import matplotlib.pyplot as plt +import omnigibson as og +from omnigibson.macros import gm +from omnigibson.action_primitives.starter_semantic_action_primitives import StarterSemanticActionPrimitives, StarterSemanticActionPrimitiveSet +from omnigibson.sensors.scan_sensor import ScanSensor +from omnigibson.sensors.vision_sensor import VisionSensor +import omnigibson.utils.transform_utils as T +from omnigibson.objects.dataset_object import DatasetObject +import h5py + + +def set_start_pose(robot): + reset_pose_tiago = np.array([ + -1.78029833e-04, 3.20231302e-05, -1.85759447e-07, -1.16488536e-07, + 4.55182843e-08, 2.36128806e-04, 1.50000000e-01, 9.40000000e-01, + -1.10000000e+00, 0.00000000e+00, -0.90000000e+00, 1.47000000e+00, + 0.00000000e+00, 2.10000000e+00, 2.71000000e+00, 1.50000000e+00, + 1.71000000e+00, 1.30000000e+00, -1.57000000e+00, -1.40000000e+00, + 1.39000000e+00, 0.00000000e+00, 0.00000000e+00, 4.50000000e-02, + 4.50000000e-02, 4.50000000e-02, 4.50000000e-02, + ]) + robot.set_joint_positions(reset_pose_tiago) + og.sim.step() + +def step_sim(time): + for _ in range(int(time*100)): + og.sim.step() + +def execute_controller(ctrl_gen, env): + for action in ctrl_gen: + env.step(action) + +def reset_env(env, initial_poses): + objs = ["cologne", "coffee_table_fqluyq_0", "robot0"] + for o in objs: + env.scene.object_registry("name", o).set_position_orientation(*initial_poses[o]) + env.reset() + og.sim.step() + +DIST_COEFF = 0.1 +GRASP_REWARD = 0.3 +RL_ITERATIONS = 2 +h5py.get_config().track_order = True + +cfg = { + "scene": { + "type": "InteractiveTraversableScene", + "scene_model": "Rs_int", + "load_object_categories": ["floors", "coffee_table"], + }, + "robots": [ + { + "type": "Tiago", + "obs_modalities": ["rgb", "depth"], + "scale": 1.0, + "self_collisions": True, + "action_normalize": False, + "action_type": "continuous", + "grasping_mode": "sticky", + "rigid_trunk": False, + "default_arm_pose": "diagonal30", + "default_trunk_offset": 0.365, + "sensor_config": { + "VisionSensor": { + "modalities": ["rgb", "depth"], + "sensor_kwargs": { + "image_width": 224, + "image_height": 224 + } + } + }, + "controller_config": { + "base": { + "name": "JointController", + "motor_type": "velocity" + }, + "arm_left": { + "name": "JointController", + "motor_type": "position", + "command_input_limits": None, + "command_output_limits": None, + "use_delta_commands": False + }, + "arm_right": { + "name": "JointController", + "motor_type": "position", + "command_input_limits": None, + "command_output_limits": None, + "use_delta_commands": False + }, + "gripper_left": { + "name": "JointController", + "motor_type": "position", + "command_input_limits": [-1, 1], + "command_output_limits": None, + "use_delta_commands": True, + "use_single_command": True + }, + "gripper_right": { + "name": "JointController", + "motor_type": "position", + "command_input_limits": [-1, 1], + "command_output_limits": None, + "use_delta_commands": True, + "use_single_command": True + }, + "camera": { + "name": "JointController", + "motor_type": "velocity", + "use_delta_commands": False + } + } + } + ], + "task": { + "type": "GraspTask", + "obj_name": "cologne", + "termination_config": { + "max_steps": 100000, + }, + "reward_config": { + "r_dist_coeff": DIST_COEFF, + "r_grasp": GRASP_REWARD + } + }, + "objects": [ + { + "type": "DatasetObject", + "name": "cologne", + "category": "bottle_of_cologne", + "model": "lyipur", + "position": [-0.3, -0.8, 0.5], + }, + ] +} + +class Recorder(): + def __init__(self): + self.filename = "./RL_data.h5" + self.state_keys = ["robot0:eyes_Camera_sensor_rgb", "robot0:eyes_Camera_sensor_depth"] + self.reset() + + def add(self, state, action, reward): + for k in self.state_keys: + self.states[k].append(state['robot0'][k]) + self.actions.append(action) + self.rewards.append(reward) + self.ids.append(self.episode_id) + + def reset(self): + self.states = {} + for k in self.state_keys: + self.states[k] = [] + self.actions = [] + self.rewards = [] + self.ids = [] + self.episode_id = str(uuid.uuid4()) + + def _add_to_dataset(self, group, name, data): + if name in group: + dset_len = len(group[name]) + dset = group[name] + dset.resize(len(dset) + len(data), 0) + dset[dset_len:] = data + else: + if isinstance(data[0], np.ndarray): + group.create_dataset(name, data=data, maxshape=(None, *data[0].shape)) + else: + group.create_dataset(name, data=data, maxshape=(None,)) + + def save(self, group_name): + h5file = h5py.File(self.filename, 'a') + group = h5file[group_name] if group_name in h5file else h5file.create_group(group_name) + for k in self.state_keys: + self._add_to_dataset(group, k[k.find(":") + 1:], self.states[k]) + self._add_to_dataset(group, "actions", self.actions) + self._add_to_dataset(group, "rewards", self.rewards) + self._add_to_dataset(group, "ids", self.ids) + h5file.close() + +# Create the environment +env = og.Environment(configs=cfg, action_timestep=1 / 60., physics_timestep=1 / 60.) +scene = env.scene +robot = env.robots[0] +og.sim.step() + +# Place object in scene +# obj = DatasetObject( +# name="cologne", +# category="bottle_of_cologne", +# model="lyipur" +# ) +# og.sim.import_object(obj) +# obj.set_position([-0.3, -0.8, 0.5]) +# set_start_pose(robot) +# og.sim.step() +# env.scene.update_initial_state() + +controller = StarterSemanticActionPrimitives(None, scene, robot) +env._primitive_controller = controller +initial_poses = {} +for o in env.scene.objects: + initial_poses[o.name] = o.get_position_orientation() +obj = env.scene.object_registry("name", "cologne") +recorder = Recorder() +group_name = datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + +for i in range(RL_ITERATIONS): + try: + reset_env(env, initial_poses) + for action in controller.apply_ref(StarterSemanticActionPrimitiveSet.GRASP, obj): + state, reward, done, info = env.step(action) + recorder.add(state, action, reward) + if done: + for action in controller._execute_release(): + state, reward, done, info = env.step(action) + recorder.add(state, action, reward) + break + except: + print("Error in iteration: ", i) + recorder.save(group_name) + recorder.reset() \ No newline at end of file diff --git a/tests/test_reward_function.py b/tests/test_reward_function.py new file mode 100644 index 000000000..6efb94720 --- /dev/null +++ b/tests/test_reward_function.py @@ -0,0 +1,192 @@ +import math +import numpy as np +import matplotlib.pyplot as plt +import omnigibson as og +from omnigibson.macros import gm +from omnigibson.action_primitives.starter_semantic_action_primitives import StarterSemanticActionPrimitives, StarterSemanticActionPrimitiveSet +import omnigibson.utils.transform_utils as T +from omnigibson.objects.dataset_object import DatasetObject + + +def set_start_pose(robot): + reset_pose_tiago = np.array([ + -1.78029833e-04, 3.20231302e-05, -1.85759447e-07, -1.16488536e-07, + 4.55182843e-08, 2.36128806e-04, 1.50000000e-01, 9.40000000e-01, + -1.10000000e+00, 0.00000000e+00, -0.90000000e+00, 1.47000000e+00, + 0.00000000e+00, 2.10000000e+00, 2.71000000e+00, 1.50000000e+00, + 1.71000000e+00, 1.30000000e+00, -1.57000000e+00, -1.40000000e+00, + 1.39000000e+00, 0.00000000e+00, 0.00000000e+00, 4.50000000e-02, + 4.50000000e-02, 4.50000000e-02, 4.50000000e-02, + ]) + robot.set_joint_positions(reset_pose_tiago) + og.sim.step() + +def step_sim(time): + for _ in range(int(time*100)): + og.sim.step() + +def execute_controller(ctrl_gen, env): + for action in ctrl_gen: + env.step(action) + +def test_grasp_reward(): + DIST_COEFF = 0.1 + GRASP_REWARD = 0.3 + + cfg = { + "scene": { + "type": "InteractiveTraversableScene", + "scene_model": "Rs_int", + "load_object_categories": ["floors", "coffee_table"], + }, + "robots": [ + { + "type": "Tiago", + "obs_modalities": ["scan", "rgb", "depth"], + "scale": 1.0, + "self_collisions": True, + "action_normalize": False, + "action_type": "continuous", + "grasping_mode": "sticky", + "rigid_trunk": False, + "default_arm_pose": "diagonal30", + "default_trunk_offset": 0.365, + "controller_config": { + "base": { + "name": "JointController", + "motor_type": "velocity" + }, + "arm_left": { + "name": "JointController", + "motor_type": "position", + "command_input_limits": None, + "command_output_limits": None, + "use_delta_commands": False + }, + "arm_right": { + "name": "JointController", + "motor_type": "position", + "command_input_limits": None, + "command_output_limits": None, + "use_delta_commands": False + }, + "gripper_left": { + "name": "JointController", + "motor_type": "position", + "command_input_limits": [-1, 1], + "command_output_limits": None, + "use_delta_commands": True, + "use_single_command": True + }, + "gripper_right": { + "name": "JointController", + "motor_type": "position", + "command_input_limits": [-1, 1], + "command_output_limits": None, + "use_delta_commands": True, + "use_single_command": True + }, + "camera": { + "name": "JointController", + "motor_type": "velocity", + "use_delta_commands": False + } + } + } + ], + "task": { + "type": "GraspTask", + "obj_name": "cologne", + "termination_config": { + "max_steps": 100000, + }, + "reward_config": { + "r_dist_coeff": DIST_COEFF, + "r_grasp": GRASP_REWARD + } + } + } + + # Make sure sim is stopped + og.sim.stop() + + # Make sure GPU dynamics are enabled (GPU dynamics needed for cloth) and no flatcache + gm.ENABLE_OBJECT_STATES = True + gm.USE_GPU_DYNAMICS = False + gm.ENABLE_FLATCACHE = False + + # Create the environment + env = og.Environment(configs=cfg, action_timestep=1 / 60., physics_timestep=1 / 60.) + scene = env.scene + robot = env.robots[0] + env.reset() + + controller = StarterSemanticActionPrimitives(None, scene, robot) + + # Place object in scene + obj = DatasetObject( + name="cologne", + category="bottle_of_cologne", + model="lyipur" + ) + og.sim.import_object(obj) + obj.set_position([-0.3, -0.8, 0.5]) + + # Set robot position so it can grasp without movement + set_start_pose(robot) + pose = controller._get_robot_pose_from_2d_pose([-0.433881, -0.210183, -2.96118]) + robot.set_position_orientation(*pose) + og.sim.step() + + ctrl_gen = controller.apply_ref(StarterSemanticActionPrimitiveSet.GRASP, obj) + + + rewards = [0] + total_rewards = [0] + + # Check reward going from not grasping to not grasping + _, reward, _, _ = env.step(next(ctrl_gen)) + rewards.append(reward) + total_rewards.append(total_rewards[-1] + reward) + eef_pos = robot.get_eef_position(robot.default_arm) + expected_reward = math.exp(-T.l2_distance(eef_pos, obj.aabb_center)) * DIST_COEFF + assert math.isclose(reward, expected_reward, abs_tol=0.01) + + for action in ctrl_gen: + prev_obj_in_hand = robot._ag_obj_in_hand[robot.default_arm] + _, reward, _, _ = env.step(action) + rewards.append(reward) + total_rewards.append(total_rewards[-1] + reward) + obj_in_hand = robot._ag_obj_in_hand[robot.default_arm] + if prev_obj_in_hand is None and obj_in_hand is not None: + # Check reward going from not grapsing to after grasping + expected_reward = math.exp(-T.l2_distance(robot.aabb_center, obj.aabb_center)) * DIST_COEFF + GRASP_REWARD + assert math.isclose(reward, expected_reward, abs_tol=0.01) + elif prev_obj_in_hand is not None and obj_in_hand is not None: + # Check reward going from grasping to grasping + expected_reward = math.exp(-T.l2_distance(robot.aabb_center, obj.aabb_center)) * DIST_COEFF + GRASP_REWARD + # assert math.isclose(reward, expected_reward, abs_tol=0.01) + # break + + ctrl_gen = controller.apply_ref(StarterSemanticActionPrimitiveSet.RELEASE) + + for action in ctrl_gen: + prev_obj_in_hand = robot._ag_obj_in_hand[robot.default_arm] + _, reward, _, _ = env.step(action) + rewards.append(reward) + total_rewards.append(total_rewards[-1] + reward) + obj_in_hand = robot._ag_obj_in_hand[robot.default_arm] + if prev_obj_in_hand is not None and obj_in_hand is None: + # Check reward going from grapsing to not grasping + eef_pos = robot.get_eef_position(robot.default_arm) + expected_reward = math.exp(-T.l2_distance(eef_pos, obj.aabb_center)) * DIST_COEFF + assert math.isclose(reward, expected_reward, abs_tol=0.01) + + plt.figure(1) + plt.subplot(211) + plt.plot(rewards) + plt.subplot(212) + plt.plot(total_rewards) + plt.show() + +test_grasp_reward() \ No newline at end of file diff --git a/tests/test_symbolic_primitives.py b/tests/test_symbolic_primitives.py index 97c28a142..c8f45f8a7 100644 --- a/tests/test_symbolic_primitives.py +++ b/tests/test_symbolic_primitives.py @@ -3,9 +3,13 @@ import yaml import omnigibson as og +from omnigibson.macros import gm from omnigibson import object_states from omnigibson.action_primitives.symbolic_semantic_action_primitives import SymbolicSemanticActionPrimitiveSet, SymbolicSemanticActionPrimitives +gm.USE_GPU_DYNAMICS = True +gm.USE_FLATCACHE = True + def start_env(): config = { "env": { @@ -34,6 +38,7 @@ def start_env(): "action_normalize": False, "action_type": "continuous", "grasping_mode": "sticky", + "disable_grasp_handling": True, "rigid_trunk": False, "default_trunk_offset": 0.365, "default_arm_pose": "diagonal30", @@ -77,18 +82,26 @@ def start_env(): }, { "type": "DatasetObject", - "name": "steak", - "category": "steak", - "model": "ppykkp", - "position": [5.31, 10.28, 1.], + "name": "knife", + "category": "carving_knife", + "model": "awvoox", + "position": [5.31, 10.75, 1.2], }, { "type": "DatasetObject", - "name": "sponge", - "category": "sponge", - "model": "qewotb", + "name": "apple", + "category": "apple", + "model": "agveuv", "position": [4.75, 10.75, 1.], + "bounding_box": [0.098, 0.098, 0.115] }, + # { + # "type": "DatasetObject", + # "name": "sponge", + # "category": "sponge", + # "model": "qewotb", + # "position": [4.75, 10.75, 1.], + # }, ] } @@ -96,10 +109,17 @@ def start_env(): return env -@pytest.fixture -def env(): +@pytest.fixture(scope="module") +def shared_env(): + """Load the environment just once using module scope.""" return start_env() +@pytest.fixture(scope="function") +def env(shared_env): + """Reset the environment before each test function.""" + og.sim.scene.reset() + return shared_env + @pytest.fixture def prim_gen(env): scene = env.scene @@ -120,31 +140,36 @@ def sink(env): @pytest.fixture def pan(env): - return next(iter(env.scene.object_registry("category", "pan"))) + return next(iter(env.scene.object_registry("category", "frying_pan"))) @pytest.fixture def steak(env): - return next(iter(env.scene.object_registry("category", "steak"))) + # TODO: Why apple? Use steak. + return next(iter(env.scene.object_registry("category", "apple"))) @pytest.fixture def sponge(env): return next(iter(env.scene.object_registry("category", "sponge"))) +@pytest.fixture +def knife(env): + return next(iter(env.scene.object_registry("category", "carving_knife"))) + # def test_navigate(): # pass -# def test_open(env, prim_gen, fridge): -# assert not fridge.states[object_states.Open].get_value() -# for action in prim_gen.apply_ref(SymbolicSemanticActionPrimitiveSet.OPEN, fridge): -# env.step(action) -# assert fridge.states[object_states.Open].get_value() +def test_open(env, prim_gen, fridge): + assert not fridge.states[object_states.Open].get_value() + for action in prim_gen.apply_ref(SymbolicSemanticActionPrimitiveSet.OPEN, fridge): + env.step(action) + assert fridge.states[object_states.Open].get_value() -# def test_close(env, prim_gen, fridge): -# fridge.states[object_states.Open].set_value(True) -# assert fridge.states[object_states.Open].get_value() -# for action in prim_gen.apply_ref(SymbolicSemanticActionPrimitiveSet.CLOSE, fridge): -# env.step(action) -# assert not fridge.states[object_states.Open].get_value() +def test_close(env, prim_gen, fridge): + fridge.states[object_states.Open].set_value(True) + assert fridge.states[object_states.Open].get_value() + for action in prim_gen.apply_ref(SymbolicSemanticActionPrimitiveSet.CLOSE, fridge): + env.step(action) + assert not fridge.states[object_states.Open].get_value() def test_place_inside(env, prim_gen, steak, fridge): assert not steak.states[object_states.Inside].get_value(fridge) @@ -157,16 +182,20 @@ def test_place_inside(env, prim_gen, steak, fridge): env.step(action) assert steak.states[object_states.Inside].get_value(fridge) -# def test_place_ontop(env, prim_gen, steak, pan): -# assert not steak.states[object_states.OnTop].get_value(pan) -# for action in prim_gen.apply_ref(SymbolicSemanticActionPrimitiveSet.GRASP, steak): -# env.step(action) -# for action in prim_gen.apply_ref(SymbolicSemanticActionPrimitiveSet.PLACE_ON_TOP, pan): -# env.step(action) -# assert steak.states[object_states.OnTop].get_value(pan) +def test_place_ontop(env, prim_gen, steak, pan): + assert not steak.states[object_states.OnTop].get_value(pan) + for action in prim_gen.apply_ref(SymbolicSemanticActionPrimitiveSet.GRASP, steak): + env.step(action) + for action in prim_gen.apply_ref(SymbolicSemanticActionPrimitiveSet.PLACE_ON_TOP, pan): + env.step(action) + assert steak.states[object_states.OnTop].get_value(pan) + +def test_toggle_on(env, prim_gen, stove): + assert not stove.states[object_states.ToggledOn].get_value() + for action in prim_gen.apply_ref(SymbolicSemanticActionPrimitiveSet.TOGGLE_ON, stove): + env.step(action) + assert stove.states[object_states.ToggledOn].get_value() -# def test_toggle_on(): -# pass # def test_soak_under(): # pass @@ -177,8 +206,20 @@ def test_place_inside(env, prim_gen, steak, fridge): # def test_wipe(): # pass -# def test_cut(): -# pass +@pytest.mark.skip(reason="A bug with object addition causes the robot to crash after slicing.") +def test_cut(env, prim_gen, steak, knife, countertop): + # assert not steak.states[object_states.Cut].get_value(knife) + print("Grasping knife") + for action in prim_gen.apply_ref(SymbolicSemanticActionPrimitiveSet.GRASP, knife): + env.step(action) + for _ in range(60): og.sim.step() + print("Cutting steak") + for action in prim_gen.apply_ref(SymbolicSemanticActionPrimitiveSet.CUT, steak): + env.step(action) + for _ in range(60): og.sim.step() + print("Putting knife back on countertop") + for action in prim_gen.apply_ref(SymbolicSemanticActionPrimitiveSet.PLACE_ON_TOP, countertop): + env.step(action) # def test_place_near_heating_element(): # pass @@ -191,15 +232,19 @@ def main(): scene = env.scene robot = env.robots[0] prim_gen = SymbolicSemanticActionPrimitives(None, scene, robot) - steak = next(iter(env.scene.object_registry("category", "steak"))) - fridge = next(iter(env.scene.object_registry("category", "fridge"))) + steak = next(iter(env.scene.object_registry("category", "apple"))) + knife = next(iter(env.scene.object_registry("category", "carving_knife"))) + countertop = next(iter(env.scene.object_registry("category", "countertop"))) + + print("Will start in 3 seconds") + for _ in range(120): og.sim.step() try: - test_place_inside(env, prim_gen, steak, fridge) + test_cut(env, prim_gen, steak, knife, countertop) except: - pass + raise while True: - og.sim.render() + og.sim.step() if __name__ == "__main__": main() \ No newline at end of file