From 0d7161b0b2e79c006a830ac60cd2867504c38df4 Mon Sep 17 00:00:00 2001 From: xavierpuig user Date: Sun, 26 Jan 2025 23:33:54 +0000 Subject: [PATCH 1/7] update rearrange env --- .../articulated_agents/robots/spot_robot.py | 5 +- .../isaac_sim/isaac_mobile_manipulator.py | 73 +- .../habitat/isaac_sim/isaac_spot_robot.py | 14 +- .../tasks/rearrange/actions/actions.py | 17 + .../rearrange/articulated_agent_manager.py | 109 +- .../tasks/rearrange/isaac_rearrange_sim.py | 1231 +++++++++++++++++ test_rearrange_env.py | 156 +++ 7 files changed, 1598 insertions(+), 7 deletions(-) create mode 100644 habitat-lab/habitat/tasks/rearrange/isaac_rearrange_sim.py create mode 100644 test_rearrange_env.py diff --git a/habitat-lab/habitat/articulated_agents/robots/spot_robot.py b/habitat-lab/habitat/articulated_agents/robots/spot_robot.py index 52a8c5942d..7e85513d56 100644 --- a/habitat-lab/habitat/articulated_agents/robots/spot_robot.py +++ b/habitat-lab/habitat/articulated_agents/robots/spot_robot.py @@ -87,7 +87,10 @@ class SpotParams: class SpotRobot(MobileManipulator): - def _get_spot_params(self): + @classmethod + + + def _get_spot_params(cls): return SpotParams( arm_joints=list(range(0, 7)), gripper_joints=[7], diff --git a/habitat-lab/habitat/isaac_sim/isaac_mobile_manipulator.py b/habitat-lab/habitat/isaac_sim/isaac_mobile_manipulator.py index ed58671e02..0d598a3c23 100644 --- a/habitat-lab/habitat/isaac_sim/isaac_mobile_manipulator.py +++ b/habitat-lab/habitat/isaac_sim/isaac_mobile_manipulator.py @@ -29,12 +29,26 @@ def __init__( params: MobileManipulatorParams, agent_cfg, isaac_service, + sim=None # limit_robo_joints: bool = True, # fixed_base: bool = True, # maintain_link_order: bool = False, # base_type="mobile", ): + self.params = params + self._sim = sim self._robot_wrapper = SpotRobotWrapper(isaac_service=isaac_service, instance_id=0) + # TODO: this should move later, cameras should not be attached to agents + # @alexclegg + self._cameras = None + if hasattr(self.params, "cameras"): + from collections import defaultdict + + self._cameras = defaultdict(list) + for camera_prefix in self.params.cameras: + for sensor_name in self._sim._sensors: + if sensor_name.startswith(camera_prefix): + self._cameras[camera_prefix].append(sensor_name) def reconfigure(self) -> None: @@ -46,8 +60,63 @@ def update(self) -> None: """Updates the camera transformations and performs necessary checks on joint limits and sleep states. """ - # todo - pass + """Updates the camera transformations and performs necessary checks on + joint limits and sleep states. + """ + if self._cameras is not None: + # get the transformation + agent_node = self._sim._default_agent.scene_node + inv_T = agent_node.transformation.inverted() + # update the cameras + sim = self._sim + look_up = mn.Vector3(0,1,0) + + for cam_prefix, sensor_names in self._cameras.items(): + for sensor_name in sensor_names: + sens_obj = self._sim._sensors[sensor_name]._sensor_object + cam_info = self.params.cameras[cam_prefix] + + look_at = sim.agents_mgr._all_agent_data[0].articulated_agent.base_pos + camera_pos = look_at + mn.Vector3(-0.7, 1.5, -0.7) + + + + # if cam_info.attached_link_id == -1: + # link_trans = self.sim_obj.transformation + # else: + # link_trans = self.sim_obj.get_link_scene_node( + # cam_info.attached_link_id + # ).transformation + + # if cam_info.cam_look_at_pos == mn.Vector3(0, 0, 0): + # pos = cam_info.cam_offset_pos + # ori = cam_info.cam_orientation + # Mt = mn.Matrix4.translation(pos) + # Mz = mn.Matrix4.rotation_z(mn.Rad(ori[2])) + # My = mn.Matrix4.rotation_y(mn.Rad(ori[1])) + # Mx = mn.Matrix4.rotation_x(mn.Rad(ori[0])) + # cam_transform = Mt @ Mz @ My @ Mx + # else: + # cam_transform = mn.Matrix4.look_at( + # cam_info.cam_offset_pos, + # cam_info.cam_look_at_pos, + # mn.Vector3(0, 1, 0), + # ) + # cam_transform = ( + # link_trans + # @ cam_transform + # @ cam_info.relative_transform + # ) + # cam_transform = inv_T @ cam_transform + + # sens_obj.node.transformation = ( + # orthonormalize_rotation_shear(cam_transform) + # ) + sens_obj.node.rotation = mn.Quaternion.from_matrix( + mn.Matrix4.look_at(camera_pos, look_at, look_up).rotation() + ) + sens_obj.node.translation = camera_pos + print(sens_obj.node.translation) def reset(self) -> None: """Reset the joints on the existing robot. diff --git a/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py b/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py index 6e58a37758..b69be6c494 100644 --- a/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py +++ b/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py @@ -8,11 +8,19 @@ import numpy as np import quaternion +from habitat.articulated_agents.mobile_manipulator import ArticulatedAgentCameraParams from habitat.articulated_agents.mobile_manipulator import MobileManipulatorParams from habitat.isaac_sim.isaac_mobile_manipulator import IsaacMobileManipulator from habitat.articulated_agents.robots.spot_robot import SpotRobot + + class IsaacSpotRobot(IsaacMobileManipulator): + """Isaac-internal wrapper for a robot. + + + The goal with this wrapper is convenience but not encapsulation. See also (public) IsaacMobileManipulator, which has the goal of exposing a minimal public interface to the rest of Habitat-lab. + """ # todo: put most of this logic in IsaacMobileManipulator @property @@ -58,10 +66,12 @@ def get_ee_local_pose( def __init__( self, agent_cfg, - isaac_service + isaac_service, + sim=None ): super().__init__( - SpotRobot.get_spot_params(), + SpotRobot._get_spot_params(), agent_cfg, isaac_service, + sim=sim ) diff --git a/habitat-lab/habitat/tasks/rearrange/actions/actions.py b/habitat-lab/habitat/tasks/rearrange/actions/actions.py index 06ea299e63..a29d121711 100644 --- a/habitat-lab/habitat/tasks/rearrange/actions/actions.py +++ b/habitat-lab/habitat/tasks/rearrange/actions/actions.py @@ -428,6 +428,8 @@ def step(self, delta_pos, *args, **kwargs): self.cur_articulated_agent.arm_motor_pos = set_arm_pos + + @registry.register_task_action class BaseVelAction(ArticulatedAgentAction): """ @@ -535,6 +537,21 @@ def step(self, *args, **kwargs): self.update_base() +@registry.register_task_action +class BaseVelIsaacAction(BaseVelAction): + def step(self, *args, **kwargs): + lin_vel, ang_vel = kwargs[self._action_arg_prefix + "base_vel"] + lin_vel = np.clip(lin_vel, -1, 1) * self._lin_speed + ang_vel = np.clip(ang_vel, -1, 1) * self._ang_speed + if not self._allow_back: + lin_vel = np.maximum(lin_vel, 0) + + self.base_vel_ctrl.linear_velocity = mn.Vector3(lin_vel, 0, 0) + self.base_vel_ctrl.angular_velocity = mn.Vector3(0, ang_vel, 0) + self.cur_articulated_agent._robot_wrapper._robot.set_angular_velocity([0, 0, ang_vel]) + self.cur_articulated_agent._robot_wrapper._robot.set_linear_velocity([lin_vel, 0, 0]) + + @registry.register_task_action class BaseVelNonCylinderAction(ArticulatedAgentAction): """ diff --git a/habitat-lab/habitat/tasks/rearrange/articulated_agent_manager.py b/habitat-lab/habitat/tasks/rearrange/articulated_agent_manager.py index dbc5e3e95f..65fdb35a8b 100644 --- a/habitat-lab/habitat/tasks/rearrange/articulated_agent_manager.py +++ b/habitat-lab/habitat/tasks/rearrange/articulated_agent_manager.py @@ -3,8 +3,7 @@ # LICENSE file in the root directory of this source tree. from dataclasses import dataclass -from typing import TYPE_CHECKING, Iterator, List, Optional - +from typing import TYPE_CHECKING, Iterator, List, Optional, Any import magnum as mn import numpy as np @@ -62,6 +61,18 @@ def ik_helper(self): return self._ik_helper +@dataclass +class IsaacAgentData: + """ + Data needed to manage an agent instance. + """ + articulated_agent: Any + start_js: np.ndarray + cfg: Any + + + + class ArticulatedAgentManager: """ Handles creating, updating and managing all agent instances. @@ -236,3 +247,97 @@ def update_debug(self): for agent_data in self._all_agent_data: for grasp_mgr in agent_data.grasp_mgrs: grasp_mgr.update_debug() + + +class IsaacArticulatedAgentManager(ArticulatedAgentManager): + """ + Handles creating, updating and managing all agent instances. + """ + + def update_agents(self): + """ + Update all agent instance managers. + """ + for agent_data in self._all_agent_data: + agent_data.articulated_agent.update() + + def __init__(self, cfg, sim): + self._sim = sim + self._all_agent_data = [] + self._is_pb_installed = is_pb_installed() + self.agent_names = cfg.agents + from habitat.isaac_sim.isaac_spot_robot import IsaacSpotRobot + + for agent_name in cfg.agents_order: + + agent_cfg = cfg.agents[agent_name] + # TODO: put this later into a config + agent = IsaacSpotRobot(agent_cfg=agent_cfg, isaac_service=sim._isaac_wrapper.service, sim=sim) + + + # TODO: correct this + use_arm_init = np.array([0.0]) # (agent.params.arm_init_params) + self._all_agent_data.append( + IsaacAgentData( + articulated_agent=agent, + cfg=agent_cfg, + start_js=use_arm_init, + ) + ) + + + @property + def grasp_iter(self) -> Iterator[RearrangeGraspManager]: + return iter(()) + + def on_new_scene(self): + pass + + def pre_obj_clear(self) -> None: + + pass + + + def agent(self): + return self._all_agent_data[0].articulated_agent + + + @add_perf_timing_func() + def post_obj_load_reconfigure(self): + """ + Called at the end of the simulator reconfigure method. Used to set the starting configurations of the robots if specified in the task config. + """ + for agent_data in self._all_agent_data: + target_arm_init_params = ( + agent_data.start_js + + agent_data.cfg.joint_start_noise + * np.random.randn(len(agent_data.start_js)) + ) + + # We only randomly set the location of the particular joint if that joint can be controlled + # and given joint_that_can_control value. + if agent_data.cfg.joint_that_can_control is not None: + assert len(agent_data.start_js) == len( + agent_data.cfg.joint_that_can_control + ) + for i in range(len(agent_data.cfg.joint_that_can_control)): + # We cannot control this joint + if agent_data.cfg.joint_that_can_control[i] == 0: + # The initial parameter for this joint should be the original angle + target_arm_init_params[i] = agent_data.start_js[i] + + # TODO: reset? + # agent_data.articulated_agent.params.arm_init_params = ( + # target_arm_init_params + # ) + # agent_data.articulated_agent.reset() + + # consume a fixed position from SIMUALTOR.agent_0 if configured + if agent_data.cfg.is_set_start_state: + agent_data.articulated_agent.base_pos = mn.Vector3( + agent_data.cfg.start_position + ) + agent_rot = agent_data.cfg.start_rotation + agent_data.articulated_agent.sim_obj.rotation = mn.Quaternion( + mn.Vector3(agent_rot[:3]), agent_rot[3] + ) \ No newline at end of file diff --git a/habitat-lab/habitat/tasks/rearrange/isaac_rearrange_sim.py b/habitat-lab/habitat/tasks/rearrange/isaac_rearrange_sim.py new file mode 100644 index 0000000000..75a70ec67e --- /dev/null +++ b/habitat-lab/habitat/tasks/rearrange/isaac_rearrange_sim.py @@ -0,0 +1,1231 @@ +#!/usr/bin/env python3 + +# 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 os +import os.path as osp +import time +from collections import defaultdict +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Optional, + Tuple, + Union, + cast, +) + +import magnum as mn +import numpy as np +import numpy.typing as npt + +import habitat_sim + +# flake8: noqa +from habitat.articulated_agents.robots import FetchRobot, FetchRobotNoWheels +from habitat.config import read_write +from habitat.core.registry import registry +from habitat.core.simulator import AgentState, Observations +from habitat.datasets.rearrange.navmesh_utils import get_largest_island_index +from habitat.datasets.rearrange.rearrange_dataset import RearrangeEpisode +from habitat.datasets.rearrange.samplers.receptacle import ( + AABBReceptacle, + find_receptacles, +) +from habitat.sims.habitat_simulator.habitat_simulator import HabitatSim +from habitat.tasks.rearrange.articulated_agent_manager import ( + ArticulatedAgentData, + ArticulatedAgentManager, + IsaacArticulatedAgentManager +) +from habitat.tasks.rearrange.marker_info import MarkerInfo +from habitat.tasks.rearrange.rearrange_grasp_manager import ( + RearrangeGraspManager, +) +from habitat.tasks.rearrange.utils import ( + add_perf_timing_func, + get_rigid_aabb, + make_render_only, + rearrange_collision, + rearrange_logger, +) +from habitat_sim.logging import logger +from habitat_sim.nav import NavMeshSettings +from habitat_sim.physics import CollisionGroups, JointMotorSettings, MotionType +from habitat_sim.sim import SimulatorBackend +from habitat_sim.utils.common import quat_from_magnum + +if TYPE_CHECKING: + from omegaconf import DictConfig + +def bind_physics_material_to_hierarchy(stage, root_prim, material_name, static_friction, dynamic_friction, restitution): + + from pxr import UsdShade, UsdPhysics + from omni.isaac.core.materials.physics_material import PhysicsMaterial + + # material_path = f"/PhysicsMaterials/{material_name}" + # material_prim = stage.DefinePrim(material_path, "PhysicsMaterial") + # material = UsdPhysics.MaterialAPI(material_prim) + + # material.CreateStaticFrictionAttr().Set(static_friction) + # material.CreateDynamicFrictionAttr().Set(dynamic_friction) + # material.CreateRestitutionAttr().Set(restitution) + + physics_material = PhysicsMaterial( + prim_path=f"/PhysicsMaterials/{material_name}", + name=material_name, + static_friction=static_friction, + dynamic_friction=dynamic_friction, + restitution=restitution) + + binding_api = UsdShade.MaterialBindingAPI.Apply(root_prim) + binding_api.Bind( + physics_material.material, + bindingStrength=UsdShade.Tokens.strongerThanDescendants, + materialPurpose="physics", + ) + +@registry.register_simulator(name="IsaacRearrangeSim-v0") +class IsaacRearrangeSim(HabitatSim): + def __init__(self, config: "DictConfig"): + config.scene = "NONE" + if len(config.agents) > 1: + with read_write(config): + for agent_name, agent_cfg in config.agents.items(): + # using list to create a copy of the sim_sensors keys since we will be + # editing the sim_sensors config + sensor_keys = list(agent_cfg.sim_sensors.keys()) + for sensor_key in sensor_keys: + sensor_config = agent_cfg.sim_sensors.pop(sensor_key) + sensor_config.uuid = ( + f"{agent_name}_{sensor_config.uuid}" + ) + agent_cfg.sim_sensors[ + f"{agent_name}_{sensor_key}" + ] = sensor_config + + + super().__init__(config) + from habitat.isaac_sim.isaac_app_wrapper import IsaacAppWrapper + + self._isaac_wrapper = IsaacAppWrapper(self, headless=True) + + isaac_world = self._isaac_wrapper.service.world + self._usd_visualizer = self._isaac_wrapper.service.usd_visualizer + + self._isaac_physics_dt = 1.0 / 180 + # beware goofy behavior if physics_dt doesn't equal rendering_dt + isaac_world.set_simulation_dt(physics_dt = self._isaac_physics_dt, rendering_dt = self._isaac_physics_dt) + + # asset_path = "/home/eric/projects/habitat-lab/data/usd/scenes/102817140.usda" + asset_path = "/fsx-siro/xavierpuig/projects/habitat_isaac/habitat-lab/data/usd/scenes/102344193_with_stage.usda" + # asset_path = "/home/eric/projects/habitat-lab/data/usd/scenes/102344193_with_stage.usda" + from omni.isaac.core.utils.stage import add_reference_to_stage + add_reference_to_stage(usd_path=asset_path, prim_path="/World/test_scene") + self._usd_visualizer.on_add_reference_to_stage(usd_path=asset_path, prim_path="/World/test_scene") + + + + + + self._rigid_objects = [] + self.add_or_reset_rigid_objects() + self._pick_target_rigid_object_idx = None + + stage = self._isaac_wrapper.service.world.stage + prim = stage.GetPrimAtPath("/World") + bind_physics_material_to_hierarchy(stage=stage, root_prim=prim, material_name="my_material", static_friction=1.0, dynamic_friction=1.0, restitution=0.0) + self.agents_mgr = IsaacArticulatedAgentManager(self.habitat_config, self) + + isaac_world.reset() + self._isaac_rom.post_reset() + + for agent in self.agents_mgr.articulated_agents_iter: + agent._robot_wrapper.post_reset() + + + self.first_setup = True + self.ep_info: Optional[RearrangeEpisode] = None + self.prev_loaded_navmesh = None + self.prev_scene_id: Optional[str] = None + + # Number of physics updates per action + self.ac_freq_ratio = self.habitat_config.ac_freq_ratio + # The physics update time step. + self.ctrl_freq = self.habitat_config.ctrl_freq + # Effective control speed is (ctrl_freq/ac_freq_ratio) + + self.art_objs: List[habitat_sim.physics.ManagedArticulatedObject] = [] + self._start_art_states: Dict[ + habitat_sim.physics.ManagedArticulatedObject, + Tuple[List[float], mn.Matrix4], + ] = {} + self._prev_obj_names: Optional[List[str]] = None + self._scene_obj_ids: List[int] = [] + # The receptacle information cached between all scenes. + self._receptacles_cache: Dict[str, Dict[str, mn.Range3D]] = {} + # The per episode receptacle information. + self._receptacles: Dict[str, mn.Range3D] = {} + # Used to get data from the RL environment class to sensors. + self._goal_pos = None + self.viz_ids: Dict[Any, Any] = defaultdict(lambda: None) + self._handle_to_object_id: Dict[str, int] = {} + self._markers: Dict[str, MarkerInfo] = {} + + self._viz_templates: Dict[str, Any] = {} + self._viz_handle_to_template: Dict[str, float] = {} + self._viz_objs: Dict[str, Any] = {} + self._draw_bb_objs: List[int] = [] + + + # Setup config options. + self._debug_render_articulated_agent = ( + self.habitat_config.debug_render_articulated_agent + ) + self._debug_render_goal = self.habitat_config.debug_render_goal + self._debug_render = self.habitat_config.debug_render + self._concur_render = self.habitat_config.concur_render + self._batch_render = config.renderer.enable_batch_renderer + self._enable_gfx_replay_save = ( + self.habitat_config.habitat_sim_v0.enable_gfx_replay_save + ) + self._needs_markers = self.habitat_config.needs_markers + self._update_articulated_agent = ( + self.habitat_config.update_articulated_agent + ) + self._step_physics = self.habitat_config.step_physics + self._auto_sleep = self.habitat_config.auto_sleep + self._load_objs = self.habitat_config.load_objs + self._additional_object_paths = ( + self.habitat_config.additional_object_paths + ) + self._kinematic_mode = self.habitat_config.kinematic_mode + + self._extra_runtime_perf_stats: Dict[str, float] = defaultdict(float) + self._perf_logging_enabled = False + self.cur_runtime_perf_scope: List[str] = [] + self._should_setup_semantic_ids = ( + self.habitat_config.should_setup_semantic_ids + ) + self._isaac_wrapper.step(num_steps=1) + + + + + def enable_perf_logging(self): + """ + Will turn on the performance logging (by default this is off). + """ + self._perf_logging_enabled = True + + @property + def receptacles(self) -> Dict[str, AABBReceptacle]: + return self._receptacles + + @property + def handle_to_object_id(self) -> Dict[str, int]: + """ + Maps a handle name to the relative position of an object in `self._scene_obj_ids`. + """ + return self._handle_to_object_id + + @property + def draw_bb_objs(self) -> List[int]: + """ + Simulator object indices of objects to draw bounding boxes around if + debug render is enabled. By default, this is populated with all target + objects. + """ + return self._draw_bb_objs + + @property + def scene_obj_ids(self) -> List[int]: + """ + The simulator rigid body IDs of all objects in the scene. + """ + return self._scene_obj_ids + + @property + def articulated_agent(self): + if len(self.agents_mgr) > 1: + raise ValueError( + f"Cannot access `sim.articulated_agent` with multiple articulated agents" + ) + return self.agents_mgr[0].articulated_agent + + @property + def grasp_mgr(self): + if len(self.agents_mgr) > 1: + raise ValueError( + f"Cannot access `sim.grasp_mgr` with multiple articulated_agents" + ) + return self.agents_mgr[0].grasp_mgr + + @property + def grasp_mgrs(self): + if len(self.agents_mgr) > 1: + raise ValueError( + f"Cannot access `sim.grasp_mgr` with multiple articulated_agents" + ) + return self.agents_mgr[0].grasp_mgrs + + def _get_target_trans(self): + """ + This is how the target transforms should be accessed since + multiprocessing does not allow pickling. + """ + # Preprocess the ep_info making necessary datatype conversions. + target_trans = [] + rom = self.get_rigid_object_manager() + for target_handle, trans in self._targets.items(): + targ_idx = self._scene_obj_ids.index( + rom.get_object_by_handle(target_handle).object_id + ) + target_trans.append((targ_idx, trans)) + return target_trans + + @add_perf_timing_func() + def _try_acquire_context(self): + if self.renderer and self._concur_render: + self.renderer.acquire_gl_context() + + @add_perf_timing_func() + def _sleep_all_objects(self): + """ + De-activate (sleep) all rigid objects in the scene, assuming they are already in a dynamically stable state. + """ + rom = self.get_rigid_object_manager() + for _, ro in rom.get_objects_by_handle_substring().items(): + ro.awake = False + + aom = self.get_articulated_object_manager() + for _, ao in aom.get_objects_by_handle_substring().items(): + ao.awake = False + + def _add_markers(self, ep_info: RearrangeEpisode): + self._markers = {} + aom = self.get_articulated_object_manager() + for marker in ep_info.markers: + p = marker["params"] + ao = aom.get_object_by_handle(p["object"]) + name_to_link = {} + name_to_link_id = {} + for i in range(ao.num_links): + name = ao.get_link_name(i) + link = ao.get_link_scene_node(i) + name_to_link[name] = link + name_to_link_id[name] = i + + self._markers[marker["name"]] = MarkerInfo( + p["offset"], + name_to_link[p["link"]], + ao, + name_to_link_id[p["link"]], + ) + + def get_marker(self, name: str) -> MarkerInfo: + return self._markers[name] + + def get_all_markers(self): + return self._markers + + def _update_markers(self) -> None: + for m in self._markers.values(): + m.update() + + @add_perf_timing_func() + def reset(self): + SimulatorBackend.reset(self) + asset_path = "/fsx-siro/xavierpuig/projects/habitat_isaac/habitat-lab/data/usd/scenes/102344193_with_stage.usda" + # asset_path = "/home/eric/projects/habitat-lab/data/usd/scenes/102344193_with_stage.usda" + from omni.isaac.core.utils.stage import add_reference_to_stage + isaac_world = self._isaac_wrapper.service.world + + add_reference_to_stage(usd_path=asset_path, prim_path="/World/test_scene") + self._usd_visualizer.on_add_reference_to_stage(usd_path=asset_path, prim_path="/World/test_scene") + + + + + + self._rigid_objects = [] + self.add_or_reset_rigid_objects() + self._pick_target_rigid_object_idx = None + + stage = self._isaac_wrapper.service.world.stage + prim = stage.GetPrimAtPath("/World") + bind_physics_material_to_hierarchy(stage=stage, root_prim=prim, material_name="my_material", static_friction=1.0, dynamic_friction=1.0, restitution=0.0) + # self.agents_mgr = IsaacArticulatedAgentManager(self.habitat_config, self) + + isaac_world.reset() + self._isaac_rom.post_reset() + + for agent in self.agents_mgr.articulated_agents_iter: + agent._robot_wrapper.post_reset() + + for i in range(len(self.agents)): + self.reset_agent(i) + return None + + @add_perf_timing_func() + def reconfigure(self, config: "DictConfig", ep_info: RearrangeEpisode): + self._handle_to_goal_name = ep_info.info["object_labels"] + + self.ep_info = ep_info + + new_scene = self.prev_scene_id != ep_info.scene_id + if new_scene: + self._prev_obj_names = None + + # Only remove and re-add objects if we have a new set of objects. + ep_info.rigid_objs = sorted(ep_info.rigid_objs, key=lambda x: x[0]) + obj_names = [x[0] for x in ep_info.rigid_objs] + # Only remove and re-add objects if we have a new set of objects. + should_add_objects = self._prev_obj_names != obj_names + self._prev_obj_names = obj_names + + self.agents_mgr.pre_obj_clear() + self._clear_objects(should_add_objects, new_scene) + + is_hard_reset = new_scene or should_add_objects + return + if is_hard_reset: + with read_write(config): + config["scene"] = ep_info.scene_id + t_start = time.time() + super().reconfigure(config, should_close_on_new_scene=False) + self.add_perf_timing("super_reconfigure", t_start) + # The articulated object handles have changed. + self._start_art_states = {} + + if new_scene: + self.agents_mgr.on_new_scene() + + self.prev_scene_id = ep_info.scene_id + self._viz_templates = {} + self._viz_handle_to_template = {} + + # Set the default articulated object joint state. + for ao, (set_joint_state, set_T) in self._start_art_states.items(): + ao.clear_joint_states() + ao.joint_positions = set_joint_state + if not is_hard_reset: + # [Andrew Szot 2023-08-22]: If we don't correct for this, some + # articulated objects may "drift" over time when the scene + # reset is skipped. + ao.transformation = set_T + + # Load specified articulated object states from episode config + self._set_ao_states_from_ep(ep_info) + + self.agents_mgr.post_obj_load_reconfigure() + + # add episode clutter objects additional to base scene objects + if self._load_objs: + self._add_objs(ep_info, should_add_objects, new_scene) + self._setup_targets(ep_info) + + self._add_markers(ep_info) + + # auto-sleep rigid objects as optimization + if self._auto_sleep: + self._sleep_all_objects() + + rom = self.get_rigid_object_manager() + self._obj_orig_motion_types = { + handle: ro.motion_type + for handle, ro in rom.get_objects_by_handle_substring().items() + } + + if new_scene: + self._load_navmesh(ep_info) + + # Get the starting positions of the target objects. + scene_pos = self.get_scene_pos() + self.target_start_pos = np.array( + [ + scene_pos[ + self._scene_obj_ids.index( + rom.get_object_by_handle(t_handle).object_id + ) + ] + for t_handle, _ in self._targets.items() + ] + ) + + self._draw_bb_objs = [ + rom.get_object_by_handle(obj_handle).object_id + for obj_handle in self._targets + ] + + if self.first_setup: + self.first_setup = False + self.agents_mgr.first_setup() + # Capture the starting art states + self._start_art_states = { + ao: (ao.joint_positions, ao.transformation) + for ao in self.art_objs + } + + if self._should_setup_semantic_ids: + self._setup_semantic_ids() + + @add_perf_timing_func() + def _setup_semantic_ids(self): + # Add the rigid object id for the semantic map + rom = self.get_rigid_object_manager() + for i, handle in enumerate(rom.get_object_handles()): + obj = rom.get_object_by_handle(handle) + for node in obj.visual_scene_nodes: + node.semantic_id = ( + obj.object_id + self.habitat_config.object_ids_start + ) + + def get_agent_data(self, agent_idx: Optional[int]) -> ArticulatedAgentData: + if agent_idx is None: + return self.agents_mgr[0] + else: + return self.agents_mgr[agent_idx] + + @property + def num_articulated_agents(self): + return len(self.agents_mgr) + + def set_articulated_agent_base_to_random_point( + self, + max_attempts: int = 50, + agent_idx: Optional[int] = None, + filter_func: Optional[Callable[[np.ndarray, float], bool]] = None, + ) -> Tuple[np.ndarray, float]: + """ + :param filter_func: If specified, takes as input the agent base + position and angle and returns if the sampling point should be + allowed (true for allowed, false for rejected). + + :returns: The set base position and rotation + """ + articulated_agent = self.get_agent_data(agent_idx).articulated_agent + + for attempt_i in range(max_attempts): + # start_pos = self.pathfinder.get_random_navigable_point( + # island_index=self._largest_indoor_island_idx + # ) + start_pos = mn.Vector3([-3.39, 0.8, -4.8]) + + # start_pos = self.safe_snap_point(start_pos) + start_rot = np.random.uniform(0, 2 * np.pi) + + if filter_func is not None and not filter_func( + start_pos, start_rot + ): + continue + + articulated_agent.base_pos = start_pos + articulated_agent.base_rot = start_rot + self.perform_discrete_collision_detection() + + # Check for collisions + # did_collide, _ = rearrange_collision( + # self, True, ignore_base=False, agent_idx=agent_idx + # ) + # if not did_collide: + # break + if attempt_i == max_attempts - 1: + rearrange_logger.warning( + f"Could not find a collision free start for {self.ep_info.episode_id}" + ) + return start_pos, start_rot + + def _setup_targets(self, ep_info): + self._targets = {} + for target_handle, transform in ep_info.targets.items(): + self._targets[target_handle] = mn.Matrix4( + [[transform[j][i] for j in range(4)] for i in range(4)] + ) + + @add_perf_timing_func() + def _load_navmesh(self, ep_info): + scene_name = ep_info.scene_id.split("/")[-1].split(".")[0] + base_dir = osp.join(*ep_info.scene_id.split("/")[:2]) + + navmesh_path = osp.join(base_dir, "navmeshes", scene_name + ".navmesh") + + if osp.exists(navmesh_path): + self.pathfinder.load_nav_mesh(navmesh_path) + logger.info(f"Loaded navmesh from {navmesh_path}") + else: + logger.warning( + f"Requested navmesh to load from {navmesh_path} does not exist. Recomputing from configured values and caching." + ) + navmesh_settings = NavMeshSettings() + navmesh_settings.set_defaults() + + agent_config = None + if hasattr(self.habitat_config.agents, "agent_0"): + agent_config = self.habitat_config.agents.agent_0 + elif hasattr(self.habitat_config.agents, "main_agent"): + agent_config = self.habitat_config.agents.main_agent + else: + raise ValueError(f"Cannot find agent parameters.") + navmesh_settings.agent_radius = agent_config.radius + navmesh_settings.agent_height = agent_config.height + navmesh_settings.agent_max_climb = agent_config.max_climb + navmesh_settings.agent_max_slope = agent_config.max_slope + navmesh_settings.include_static_objects = True + self.recompute_navmesh(self.pathfinder, navmesh_settings) + os.makedirs(osp.dirname(navmesh_path), exist_ok=True) + self.pathfinder.save_nav_mesh(navmesh_path) + + # NOTE: allowing indoor islands only + self._largest_indoor_island_idx = get_largest_island_index( + self.pathfinder, self, allow_outdoor=False + ) + + @property + def largest_island_idx(self) -> int: + """ + The path finder index of the indoor island that has the largest area. + """ + return self._largest_indoor_island_idx + + @add_perf_timing_func() + def _clear_objects( + self, should_add_objects: bool, new_scene: bool + ) -> None: + rom = self.get_rigid_object_manager() + + # Clear all the rigid objects. + if should_add_objects: + for scene_obj_id in self._scene_obj_ids: + if not rom.get_library_has_id(scene_obj_id): + continue + rom.remove_object_by_id(scene_obj_id) + self._scene_obj_ids = [] + + # Reset all marker visualization points + for obj_id in self.viz_ids.values(): + if rom.get_library_has_id(obj_id): + rom.remove_object_by_id(obj_id) + self.viz_ids = defaultdict(lambda: None) + + # Remove all object mesh visualizations. + for viz_obj in self._viz_objs.values(): + if rom.get_library_has_id(viz_obj.object_id): + rom.remove_object_by_id(viz_obj.object_id) + self._viz_objs = {} + + if new_scene: + # Do not remove the articulated objects from the scene, these are + # managed by the underlying sim. + self.art_objs = [] + + @add_perf_timing_func() + def _set_ao_states_from_ep(self, ep_info: RearrangeEpisode) -> None: + """ + Sets the ArticulatedObject states for the episode which are differ from base scene state. + """ + aom = self.get_articulated_object_manager() + for aoi_handle, joint_states in ep_info.ao_states.items(): + ao = aom.get_object_by_handle(aoi_handle) + ao_pose = ao.joint_positions + for link_ix, joint_state in joint_states.items(): + joint_position_index = ao.get_link_joint_pos_offset( + int(link_ix) + ) + ao_pose[joint_position_index] = joint_state + ao.joint_positions = ao_pose + + def is_point_within_bounds(self, pos): + # NOTE: This check is loose: really we want the island bounds, not the full navmesh + lower_bound, upper_bound = self.pathfinder.get_bounds() + return all(lower_bound <= pos) and all(upper_bound >= pos) + + def safe_snap_point(self, pos: np.ndarray) -> np.ndarray: + """ + Returns the 3D coordinates corresponding to a point belonging + to the biggest navmesh island in the scenee and closest to pos. + When that point returns NaN, computes a navigable point at increasing + distances to it. + """ + new_pos = self.pathfinder.snap_point( + pos, self._largest_indoor_island_idx + ) + + max_iter = 10 + offset_distance = 1.5 + distance_per_iter = 0.5 + num_sample_points = 1000 + + regen_i = 0 + while np.isnan(new_pos[0]) and regen_i < max_iter: + # Increase the search radius + new_pos = self.pathfinder.get_random_navigable_point_near( + pos, + offset_distance + regen_i * distance_per_iter, + num_sample_points, + island_index=self._largest_indoor_island_idx, + ) + regen_i += 1 + + assert not np.isnan( + new_pos[0] + ), f"The snap position is NaN. scene_id: {self.ep_info.scene_id}, new position: {new_pos}, original position: {pos}" + + return new_pos + + @add_perf_timing_func() + def _add_objs( + self, + ep_info: RearrangeEpisode, + should_add_objects: bool, + new_scene: bool, + ) -> None: + return + # Load clutter objects: + rom = self.get_rigid_object_manager() + obj_counts: Dict[str, int] = defaultdict(int) + + self._handle_to_object_id = {} + if should_add_objects: + self._scene_obj_ids = [] + + # Get Object template manager + otm = self.get_object_template_manager() + + for i, (obj_handle, transform) in enumerate(ep_info.rigid_objs): + t_start = time.time() + if should_add_objects: + # Get object path + object_template = otm.get_templates_by_handle_substring( + obj_handle + ) + + # Exit if template is invalid + if not object_template: + raise ValueError( + f"Template not found for object with handle {obj_handle}" + ) + + # Get object path + object_path = list(object_template.keys())[0] + + # Get rigid object from the path + ro = rom.add_object_by_template_handle(object_path) + else: + ro = rom.get_object_by_id(self._scene_obj_ids[i]) + self.add_perf_timing("create_asset", t_start) + + # The saved matrices need to be flipped when reloading. + ro.transformation = mn.Matrix4( + [[transform[j][i] for j in range(4)] for i in range(4)] + ) + ro.angular_velocity = mn.Vector3.zero_init() + ro.linear_velocity = mn.Vector3.zero_init() + + other_obj_handle = ( + obj_handle.split(".")[0] + f"_:{obj_counts[obj_handle]:04d}" + ) + if self._kinematic_mode: + ro.motion_type = habitat_sim.physics.MotionType.KINEMATIC + ro.collidable = False + + if should_add_objects: + self._scene_obj_ids.append(ro.object_id) + rel_idx = self._scene_obj_ids.index(ro.object_id) + self._handle_to_object_id[other_obj_handle] = rel_idx + + if other_obj_handle in self._handle_to_goal_name: + ref_handle = self._handle_to_goal_name[other_obj_handle] + self._handle_to_object_id[ref_handle] = rel_idx + + obj_counts[obj_handle] += 1 + + if new_scene: + self._receptacles = self._create_recep_info( + ep_info.scene_id, list(self._handle_to_object_id.keys()) + ) + + ao_mgr = self.get_articulated_object_manager() + # Make all articulated objects (including the robots) kinematic + for aoi_handle in ao_mgr.get_object_handles(): + ao = ao_mgr.get_object_by_handle(aoi_handle) + if self._kinematic_mode: + ao.motion_type = habitat_sim.physics.MotionType.KINEMATIC + # remove any existing motors when converting to kinematic AO + for motor_id in ao.existing_joint_motor_ids: + ao.remove_joint_motor(motor_id) + self.art_objs.append(ao) + + def _create_recep_info( + self, scene_id: str, ignore_handles: List[str] + ) -> Dict[str, mn.Range3D]: + if scene_id not in self._receptacles_cache: + receps = {} + all_receps = find_receptacles( + self, + ignore_handles=ignore_handles, + ) + for recep in all_receps: + recep = cast(AABBReceptacle, recep) + local_bounds = recep.bounds + global_T = recep.get_global_transform(self) + # Some coordinates may be flipped by the global transformation, + # mixing the minimum and maximum bound coordinates. + bounds = np.stack( + [ + global_T.transform_point(local_bounds.min), + global_T.transform_point(local_bounds.max), + ], + axis=0, + ) + receps[recep.unique_name] = mn.Range3D( + np.min(bounds, axis=0), np.max(bounds, axis=0) + ) + self._receptacles_cache[scene_id] = receps + return self._receptacles_cache[scene_id] + + def _create_obj_viz(self): + """ + Adds a visualization of the goal for each of the target objects in the + scene. This is the same as the target object, but is a render only + object. This also places dots around the bounding box of the object to + further distinguish the goal from the target object. + """ + for marker_name, m in self._markers.items(): + m_T = m.get_current_transform() + self.viz_ids[marker_name] = self.visualize_position( + m_T.translation, self.viz_ids[marker_name] + ) + + rom = self.get_rigid_object_manager() + obj_attr_mgr = self.get_object_template_manager() + + # Enable BB render for the debug render call. + for obj_id in self._draw_bb_objs: + self.set_object_bb_draw(True, obj_id) + + if self._debug_render_goal: + for target_handle, transform in self._targets.items(): + # Visualize the goal of the object + new_target_handle = ( + target_handle.split("_:")[0] + ".object_config.json" + ) + matching_templates = ( + obj_attr_mgr.get_templates_by_handle_substring( + new_target_handle + ) + ) + ro = rom.add_object_by_template_handle( + list(matching_templates.keys())[0] + ) + self.set_object_bb_draw(True, ro.object_id) + ro.transformation = transform + make_render_only(ro, self) + bb = get_rigid_aabb(ro.object_id, self, True) + bb_viz_name1 = target_handle + "_bb1" + bb_viz_name2 = target_handle + "_bb2" + viz_r = 0.01 + self.viz_ids[bb_viz_name1] = self.visualize_position( + bb.front_bottom_right, self.viz_ids[bb_viz_name1], viz_r + ) + self.viz_ids[bb_viz_name2] = self.visualize_position( + bb.back_top_left, self.viz_ids[bb_viz_name2], viz_r + ) + + self._viz_objs[target_handle] = ro + + def capture_state(self, with_articulated_agent_js=False) -> Dict[str, Any]: + """ + Record and return a dict of state info. + + :param with_articulated_agent_js: If true, state dict includes articulated_agent joint positions in addition. + + State info dict includes: + - Robot transform + - a list of ArticulatedObject transforms + - a list of RigidObject transforms + - a list of ArticulatedObject joint states + - the object id of currently grasped object (or None) + - (optionally) the articulated_agent's joint positions + """ + # Don't need to capture any velocity information because this will + # automatically be set to 0 in `set_state`. + articulated_agent_T = [ + articulated_agent.sim_obj.transformation + for articulated_agent in self.agents_mgr.articulated_agents_iter + ] + art_T = [ao.transformation for ao in self.art_objs] + rom = self.get_rigid_object_manager() + + rigid_T, rigid_V = [], [] + for i in self._scene_obj_ids: + obj_i = rom.get_object_by_id(i) + rigid_T.append(obj_i.transformation) + rigid_V.append((obj_i.linear_velocity, obj_i.angular_velocity)) + + art_pos = [ao.joint_positions for ao in self.art_objs] + + articulated_agent_js = [ + articulated_agent.sim_obj.joint_positions + for articulated_agent in self.agents_mgr.articulated_agents_iter + ] + + ret = { + "articulated_agent_T": articulated_agent_T, + "art_T": art_T, + "rigid_T": rigid_T, + "rigid_V": rigid_V, + "art_pos": art_pos, + "obj_hold": [ + grasp_mgr.snap_idx for grasp_mgr in self.agents_mgr.grasp_iter + ], + } + if with_articulated_agent_js: + ret["articulated_agent_js"] = articulated_agent_js + return ret + + def set_state(self, state: Dict[str, Any], set_hold=False) -> None: + """ + Sets the simulation state from a cached state info dict. See capture_state(). + + :param set_hold: If true this will set the snapped object from the `state`. + + TODO: This should probably be True by default, but I am not sure the effect + it will have. + """ + rom = self.get_rigid_object_manager() + + if state["articulated_agent_T"] is not None: + for articulated_agent_T, robot in zip( + state["articulated_agent_T"], + self.agents_mgr.articulated_agents_iter, + ): + robot.sim_obj.transformation = articulated_agent_T + n_dof = len(robot.sim_obj.joint_forces) + robot.sim_obj.joint_forces = np.zeros(n_dof) + robot.sim_obj.joint_velocities = np.zeros(n_dof) + + if "articulated_agent_js" in state: + for articulated_agent_js, robot in zip( + state["articulated_agent_js"], + self.agents_mgr.articulated_agents_iter, + ): + robot.sim_obj.joint_positions = articulated_agent_js + + for T, ao in zip(state["art_T"], self.art_objs): + ao.transformation = T + + for T, V, i in zip( + state["rigid_T"], state["rigid_V"], self._scene_obj_ids + ): + # reset object transform + obj = rom.get_object_by_id(i) + obj.transformation = T + obj.linear_velocity = V[0] + obj.angular_velocity = V[1] + + for p, ao in zip(state["art_pos"], self.art_objs): + ao.joint_positions = p + + if set_hold: + if state["obj_hold"] is not None: + for obj_hold_state, grasp_mgr in zip( + state["obj_hold"], self.agents_mgr.grasp_iter + ): + self.internal_step(-1) + grasp_mgr.snap_to_obj(obj_hold_state) + else: + for grasp_mgr in self.agents_mgr.grasp_iter: + grasp_mgr.desnap(True) + + def get_agent_state(self, agent_id: int = 0) -> habitat_sim.AgentState: + articulated_agent = self.get_agent_data(agent_id).articulated_agent + rotation = mn.Quaternion.rotation( + mn.Rad(articulated_agent.base_rot) - mn.Rad(0 * np.pi / 2), + mn.Vector3(0, 1, 0), + ) + rot_offset = mn.Quaternion.rotation( + mn.Rad(-np.pi / 2), mn.Vector3(0, 1, 0) + ) + return AgentState( + articulated_agent.base_pos, + quat_from_magnum(articulated_agent.sim_obj.rotation * rot_offset), + ) + + @add_perf_timing_func() + def step(self, action: Union[str, int]) -> Observations: + + rom = self.get_rigid_object_manager() + self._isaac_wrapper.step(num_steps=1) + + if self._debug_render: + if self._debug_render_articulated_agent: + self.agents_mgr.update_debug() + rom = self.get_rigid_object_manager() + self._try_acquire_context() + + # Disable BB drawing for observation render + for obj_id in self._draw_bb_objs: + self.set_object_bb_draw(False, obj_id) + + # Remove viz objects + for obj in self._viz_objs.values(): + if obj is not None and rom.get_library_has_id(obj.object_id): + rom.remove_object_by_id(obj.object_id) + self._viz_objs = {} + + # Remove all visualized positions + add_back_viz_objs = {} + for name, viz_id in self.viz_ids.items(): + if viz_id is None: + continue + viz_obj = rom.get_object_by_id(viz_id) + before_pos = viz_obj.translation + rom.remove_object_by_id(viz_id) + r = self._viz_handle_to_template[viz_id] + add_back_viz_objs[name] = (before_pos, r) + self.viz_ids = defaultdict(lambda: None) + + self.maybe_update_articulated_agent() + + if self._batch_render: + for _ in range(self.ac_freq_ratio): + self.internal_step(-1, update_articulated_agent=False) + + obs = self.get_sensor_observations() + self.add_keyframe_to_observations(obs) + elif self._concur_render: + self.start_async_render() + + for _ in range(self.ac_freq_ratio): + self.internal_step(-1, update_articulated_agent=False) + + t_start = time.time() + obs = self._sensor_suite.get_observations( + self.get_sensor_observations_async_finish() + ) + self.add_perf_timing("get_sensor_observations", t_start) + else: + for _ in range(self.ac_freq_ratio): + self.internal_step(-1, update_articulated_agent=False) + + t_start = time.time() + obs = self._sensor_suite.get_observations( + self.get_sensor_observations() + ) + self.add_perf_timing("get_sensor_observations", t_start) + + # TODO: Support recording while batch rendering + if self._enable_gfx_replay_save and not self._batch_render: + self.gfx_replay_manager.save_keyframe() + + if self._needs_markers: + self._update_markers() + + # TODO: Make debug cameras more flexible + if "third_rgb" in obs and self._debug_render: + self._try_acquire_context() + for k, (pos, r) in add_back_viz_objs.items(): + viz_id = self.viz_ids[k] + + self.viz_ids[k] = self.visualize_position( + pos, self.viz_ids[k], r=r + ) + + # Also render debug information + self._create_obj_viz() + + debug_obs = self.get_sensor_observations() + obs["third_rgb"] = debug_obs["third_rgb"][:, :, :3] + + return obs + + def maybe_update_articulated_agent(self): + """ + Calls the update agents method on the articulated agent manager if the + `update_articulated_agent` configuration is set to True. Among other + things, this will set the articulated agent's sensors' positions to their new + positions. + """ + if self._update_articulated_agent: + self.agents_mgr.update_agents() + + def visualize_position( + self, + position: np.ndarray, + viz_id: Optional[int] = None, + r: float = 0.05, + ) -> int: + """Adds the sphere object to the specified position for visualization purpose.""" + + template_mgr = self.get_object_template_manager() + rom = self.get_rigid_object_manager() + viz_obj = None + if viz_id is None: + if r not in self._viz_templates: + template = template_mgr.get_template_by_handle( + template_mgr.get_template_handles("sphere")[0] + ) + template.scale = mn.Vector3(r, r, r) + self._viz_templates[str(r)] = template_mgr.register_template( + template, "ball_new_viz_" + str(r) + ) + viz_obj = rom.add_object_by_template_id( + self._viz_templates[str(r)] + ) + make_render_only(viz_obj, self) + self._viz_handle_to_template[viz_obj.object_id] = r + else: + viz_obj = rom.get_object_by_id(viz_id) + + viz_obj.translation = mn.Vector3(*position) + return viz_obj.object_id + + @add_perf_timing_func() + def internal_step( + self, dt: Union[int, float], update_articulated_agent: bool = True + ) -> None: + """Step the world and update the articulated_agent. + + :param dt: Timestep by which to advance the world. Multiple physics substeps can be executed within a single timestep. -1 indicates a single physics substep. + + Never call sim.step_world directly or miss updating the articulated_agent. + """ + # Optionally step physics and update the articulated_agent for benchmarking purposes + if self._step_physics: + self.step_world(dt) + + def get_targets(self) -> Tuple[np.ndarray, np.ndarray]: + """Get a mapping of object ids to goal positions for rearrange targets. + + :return: ([idx: int], [goal_pos: list]) The index of the target object + in self._scene_obj_ids and the 3D goal position, rotation is IGNORED. + Note that goal_pos is the desired position of the object, not the + starting position. + """ + target_trans = self._get_target_trans() + if len(target_trans) == 0: + return np.array([]), np.array([]) + targ_idx, targ_trans = list(zip(*self._get_target_trans())) + + a, b = np.array(targ_idx), [ + np.array(x.translation) for x in targ_trans + ] + return a, np.array(b) + + def get_n_targets(self) -> int: + """Get the number of rearrange targets.""" + return len(self.ep_info.targets) + + def get_target_objs_start(self) -> np.ndarray: + """Get the initial positions of all objects targeted for rearrangement as a numpy array.""" + return self.target_start_pos + + def get_scene_pos(self) -> np.ndarray: + """Get the positions of all clutter RigidObjects in the scene as a numpy array.""" + rom = self.get_rigid_object_manager() + return np.array( + [ + rom.get_object_by_id(idx).translation + for idx in self._scene_obj_ids + ] + ) + + def add_perf_timing(self, desc: str, t_start: float) -> None: + """ + Records a duration since `t_start` into the perf stats. Note that this + is additive, so times between successive calls accumulate, not reset. + Also note that this will only log if `self._perf_logging_enabled=True`. + """ + if not self._perf_logging_enabled: + return + + name = ".".join(self.cur_runtime_perf_scope) + if desc != "": + name += "." + desc + self._extra_runtime_perf_stats[name] += time.time() - t_start + + def get_runtime_perf_stats(self) -> Dict[str, float]: + stats_dict = {} + for name, value in self._extra_runtime_perf_stats.items(): + stats_dict[name] = value + # clear this dict so we don't accidentally collect these twice + self._extra_runtime_perf_stats = defaultdict(float) + + return stats_dict + + def add_or_reset_rigid_objects(self): + + # on dining table + drop_pos = mn.Vector3(-3.6, 0.8, -7.22) # mn.Vector3(-7.4, 0.8, -7.5) + offset_vec = mn.Vector3(1.3, 0.0, 0.0) + + # above coffee table + # drop_pos = mn.Vector3(-8.1, 0.5, -3.9) + + # middle of room + # drop_pos = mn.Vector3(-5.4, 1.2, -3.9) + + up_vec = mn.Vector3(0.0, 1.0, 0.0) + path_to_configs = "data/objects/ycb/configs" + + do_add = len(self._rigid_objects) == 0 + + # for coffee table + if False: + objects_to_add = [] + object_names = ["024_bowl", "013_apple", "011_banana", "010_potted_meat_can", "077_rubiks_cube", "036_wood_block", "004_sugar_box"] + next_obj_idx = 0 + sp = 0.25 + for cell_y in range(5): + for cell_x in range(3): + for cell_z in range(3): + offset_vec = mn.Vector3(cell_x * sp - sp, cell_y * sp, cell_z * sp - sp) + objects_to_add.append((f"{path_to_configs}/{object_names[next_obj_idx]}.object_config.json", drop_pos + offset_vec)) + next_obj_idx = (next_obj_idx + 1) % len(object_names) + + if True: + # for dining table + objects_to_add = [ + (f"{path_to_configs}/024_bowl.object_config.json", drop_pos + offset_vec * 0.0 + up_vec * 0.0), + # (f"{path_to_configs}/011_banana.object_config.json", drop_pos + offset_vec * 0.01 + up_vec * 0.05), + (f"{path_to_configs}/013_apple.object_config.json", drop_pos + offset_vec * -0.01 + up_vec * 0.05), + # (f"{path_to_configs}/011_banana.object_config.json", drop_pos + offset_vec * 0.02 + up_vec * 0.12), + (f"{path_to_configs}/013_apple.object_config.json", drop_pos + offset_vec * 0.01 + up_vec * 0.1), + + (f"{path_to_configs}/010_potted_meat_can.object_config.json", drop_pos + offset_vec * 0.3 + up_vec * 0.0), + + (f"{path_to_configs}/077_rubiks_cube.object_config.json", drop_pos + offset_vec * 0.6 + up_vec * 0.1), + (f"{path_to_configs}/036_wood_block.object_config.json", drop_pos + offset_vec * 0.6 + up_vec * 0.0), + + (f"{path_to_configs}/004_sugar_box.object_config.json", drop_pos + offset_vec * 0.9), + + (f"{path_to_configs}/004_sugar_box.object_config.json", drop_pos + offset_vec * 1.0), + (f"{path_to_configs}/004_sugar_box.object_config.json", drop_pos + offset_vec * 1.1), + (f"{path_to_configs}/004_sugar_box.object_config.json", drop_pos + offset_vec * 0.8), + + (f"{path_to_configs}/010_potted_meat_can.object_config.json", drop_pos + offset_vec * 0.22 + up_vec * 0.0), + (f"{path_to_configs}/010_potted_meat_can.object_config.json", drop_pos + offset_vec * 0.38 + up_vec * 0.0), + ] + + + from habitat.isaac_sim.isaac_rigid_object_manager import IsaacRigidObjectManager + self._isaac_rom = IsaacRigidObjectManager(self._isaac_wrapper.service) + rigid_obj_mgr = self._isaac_rom + + for i, (handle, position) in enumerate(objects_to_add): + if do_add: + ro = rigid_obj_mgr.add_object_by_template_handle(handle) + self._rigid_objects.append(ro) + else: + ro = self._rigid_objects[i] + + rotation = mn.Quaternion.rotation(-mn.Deg(90), mn.Vector3.x_axis()) + trans = mn.Matrix4.from_(rotation.to_matrix(), position) + ro.transformation = trans + # breakpoint() \ No newline at end of file diff --git a/test_rearrange_env.py b/test_rearrange_env.py new file mode 100644 index 0000000000..bccd55b5ce --- /dev/null +++ b/test_rearrange_env.py @@ -0,0 +1,156 @@ +import habitat_sim +import magnum as mn +import warnings +from habitat.tasks.rearrange.isaac_rearrange_sim import IsaacRearrangeSim +warnings.filterwarnings('ignore') +from habitat_sim.utils.settings import make_cfg +from matplotlib import pyplot as plt +from habitat_sim.utils import viz_utils as vut +from omegaconf import DictConfig +import numpy as np +from habitat.articulated_agents.robots import FetchRobot +from habitat.config.default import get_agent_config +from habitat.config.default_structured_configs import ThirdRGBSensorConfig, HeadRGBSensorConfig, HeadPanopticSensorConfig +from habitat.config.default_structured_configs import SimulatorConfig, HabitatSimV0Config, AgentConfig +from habitat.config.default import get_agent_config +import habitat +from habitat_sim.physics import JointMotorSettings, MotionType +from omegaconf import OmegaConf +import os +from habitat.isaac_sim.isaac_app_wrapper import IsaacAppWrapper +from habitat.isaac_sim import isaac_prim_utils +import random +from habitat.config.default_structured_configs import TaskConfig, EnvironmentConfig, DatasetConfig, HabitatConfig +from habitat.config.default_structured_configs import ArmActionConfig, BaseVelocityActionConfig, OracleNavActionConfig, ActionConfig +import imageio +from habitat.core.env import Env + +data_path = "/fsx-siro/xavierpuig/projects/habitat_isaac/habitat-lab/data/" + + +def make_sim_cfg(agent_dict): + # Start the scene config + sim_cfg = SimulatorConfig(type="IsaacRearrangeSim-v0") + + # This is for better graphics + sim_cfg.habitat_sim_v0.enable_hbao = True + sim_cfg.habitat_sim_v0.enable_physics = False + + + # Set up an example scene + sim_cfg.scene = "NONE" # os.path.join(data_path, "hab3_bench_assets/hab3-hssd/scenes/103997919_171031233.scene_instance.json") + # sim_cfg.scene_dataset = os.path.join(data_path, "hab3_bench_assets/hab3-hssd/hab3-hssd.scene_dataset_config.json") + # sim_cfg.additional_object_paths = [os.path.join(data_path, 'objects/ycb/configs/')] + + + cfg = OmegaConf.create(sim_cfg) + + # Set the scene agents + cfg.agents = agent_dict + cfg.agents_order = list(cfg.agents.keys()) + return cfg + + +def make_hab_cfg(agent_dict, action_dict): + sim_cfg = make_sim_cfg(agent_dict) + task_cfg = TaskConfig(type="RearrangeEmptyTask-v0") + task_cfg.actions = action_dict + env_cfg = EnvironmentConfig() + dataset_cfg = DatasetConfig(type="RearrangeDataset-v0", data_path="data/hab3_bench_assets/episode_datasets/small_large.json.gz") + + + hab_cfg = HabitatConfig() + hab_cfg.environment = env_cfg + hab_cfg.task = task_cfg + + hab_cfg.dataset = dataset_cfg + hab_cfg.simulator = sim_cfg + hab_cfg.simulator.seed = hab_cfg.seed + + return hab_cfg + + + +def init_rearrange_env(agent_dict, action_dict): + hab_cfg = make_hab_cfg(agent_dict, action_dict) + res_cfg = OmegaConf.create(hab_cfg) + return Env(res_cfg) + + + +def main(): + # Define the agent configuration + main_agent_config = AgentConfig() + + urdf_path = os.path.join(data_path, "robots/hab_spot_arm/urdf/hab_spot_arm.urdf") + main_agent_config.articulated_agent_urdf = urdf_path + main_agent_config.articulated_agent_type = "SpotRobot" + + # Define sensors that will be attached to this agent, here a third_rgb sensor and a head_rgb. + # We will later talk about why we are giving the sensors these names + main_agent_config.sim_sensors = { + "third_rgb": ThirdRGBSensorConfig(), + "head_rgb": HeadRGBSensorConfig(), + } + + # We create a dictionary with names of agents and their corresponding agent configuration + agent_dict = {"main_agent": main_agent_config} + + action_dict = { + "base_velocity_action": BaseVelocityActionConfig(type="BaseVelIsaacAction"), + } + env = init_rearrange_env(agent_dict, action_dict) + aux = env.reset() + + writer = imageio.get_writer( + "output_env.mp4", + fps=30, + ) + action = {'action': 'base_velocity_action', 'action_args': {'base_vel': np.array([ 10.0, 0], dtype=np.float32)}} + for i in range(100): + + obs = env.step(action) + im = obs["third_rgb"] + writer.append_data(im) + writer.close() + breakpoint() + # def get_pick_target_pos(): + # ro = isaac_viewer._rigid_objects[isaac_viewer._pick_target_rigid_object_idx] + # com_world = isaac_prim_utils.get_com_world(ro._rigid_prim) + # # self.draw_axis(0.05, mn.Matrix4.translation(com_world)) + # return com_world + + # isaac_viewer._spot_state_machine.set_pick_target(get_pick_target_pos) + + # # breakpoint() + # for it in range(100): + # isaac_viewer.update_isaac({}) + # look_up = mn.Vector3(0,1,0) + # isaac_viewer.update_spot_pre_step(0.01) + + # look_at = sim.agents_mgr._all_agent_data[0].articulated_agent.base_pos + # print(look_at) + # camera_pos = look_at + mn.Vector3(-0.7, 1.5, -0.7) + + # sim._sensors["third_rgb"]._sensor_object.node.rotation = mn.Quaternion.from_matrix( + # mn.Matrix4.look_at(camera_pos, look_at, look_up).rotation() + # ) + + # sim._sensors["third_rgb"]._sensor_object.node.translation = camera_pos + # # import cv2 + # sim.reset() + + # res = sim.get_sensor_observations() + + # im = res["third_rgb"][:,:,[0,1,2]] + # import cv2 + # cv2.imwrite("third2.png", im) + # breakpoint() + # writer.append_data(im) + + # writer.close() + # breakpoint() + + +if __name__ == "__main__": + main() \ No newline at end of file From 12b9f9ceb2f055ec328e00be898c3a652b4ecc80 Mon Sep 17 00:00:00 2001 From: xavierpuig user Date: Tue, 28 Jan 2025 05:35:19 +0000 Subject: [PATCH 2/7] added path planner --- .../isaac_sim/isaac_mobile_manipulator.py | 89 ++++++---- .../habitat/isaac_sim/isaac_spot_robot.py | 17 +- .../tasks/rearrange/isaac_rearrange_sim.py | 41 ++--- test_rearrange_env.py | 168 +++++++++++++----- 4 files changed, 203 insertions(+), 112 deletions(-) diff --git a/habitat-lab/habitat/isaac_sim/isaac_mobile_manipulator.py b/habitat-lab/habitat/isaac_sim/isaac_mobile_manipulator.py index 0d598a3c23..ceeb6c8e15 100644 --- a/habitat-lab/habitat/isaac_sim/isaac_mobile_manipulator.py +++ b/habitat-lab/habitat/isaac_sim/isaac_mobile_manipulator.py @@ -75,48 +75,63 @@ def update(self) -> None: for sensor_name in sensor_names: sens_obj = self._sim._sensors[sensor_name]._sensor_object cam_info = self.params.cameras[cam_prefix] - + agent = sim.agents_mgr._all_agent_data[0].articulated_agent look_at = sim.agents_mgr._all_agent_data[0].articulated_agent.base_pos - camera_pos = look_at + mn.Vector3(-0.7, 1.5, -0.7) - - - - # if cam_info.attached_link_id == -1: - # link_trans = self.sim_obj.transformation - # else: - # link_trans = self.sim_obj.get_link_scene_node( - # cam_info.attached_link_id - # ).transformation - - # if cam_info.cam_look_at_pos == mn.Vector3(0, 0, 0): - # pos = cam_info.cam_offset_pos - # ori = cam_info.cam_orientation - # Mt = mn.Matrix4.translation(pos) - # Mz = mn.Matrix4.rotation_z(mn.Rad(ori[2])) - # My = mn.Matrix4.rotation_y(mn.Rad(ori[1])) - # Mx = mn.Matrix4.rotation_x(mn.Rad(ori[0])) - # cam_transform = Mt @ Mz @ My @ Mx - # else: - # cam_transform = mn.Matrix4.look_at( - # cam_info.cam_offset_pos, - # cam_info.cam_look_at_pos, - # mn.Vector3(0, 1, 0), - # ) - # cam_transform = ( - # link_trans - # @ cam_transform - # @ cam_info.relative_transform - # ) + camera_pos = look_at + mn.Vector3(-0.5, 5.0, 0) + + if cam_info.attached_link_id == -1: + link_trans = agent.base_transformation + else: + link_trans = agent.get_link_transform(cam_info.attached_link_id) + + if cam_info.cam_look_at_pos == mn.Vector3(0, 0, 0): + pos = cam_info.cam_offset_pos + ori = cam_info.cam_orientation + Mt = mn.Matrix4.translation(pos) + Mz = mn.Matrix4.rotation_z(mn.Rad(ori[2])) + My = mn.Matrix4.rotation_y(mn.Rad(ori[1])) + Mx = mn.Matrix4.rotation_x(mn.Rad(ori[0])) + cam_transform_rel = Mt @ Mz @ My @ Mx + else: + cam_transform_rel = mn.Matrix4.look_at( + cam_info.cam_offset_pos, + cam_info.cam_look_at_pos, + mn.Vector3(0, 1, 0), + ) + gt_look_at = mn.Matrix4.look_at(camera_pos, look_at, look_up) + # print(look_at) + # print(gt_look_at) + # print("----") + # print(cam_transform) + # breakpoint() + + cam_transform = ( + link_trans + @ cam_transform_rel + @ cam_info.relative_transform + ) + + camera_pos = link_trans.transform_point(cam_info.cam_offset_pos) + look_at = link_trans.transform_point(cam_info.cam_look_at_pos) + + sens_obj.node.transformation = mn.Matrix4.look_at(camera_pos, look_at, look_up) + + # sens_obj.node.translation = camera_pos + # cam_transform = inv_T @ cam_transform # sens_obj.node.transformation = ( - # orthonormalize_rotation_shear(cam_transform) + # cam_transform # ) - sens_obj.node.rotation = mn.Quaternion.from_matrix( - mn.Matrix4.look_at(camera_pos, look_at, look_up).rotation() - ) - sens_obj.node.translation = camera_pos - print(sens_obj.node.translation) + if "third" in cam_prefix: + # breakpoint() + pass + # breakpoint() + # sens_obj.node.rotation = mn.Quaternion.from_matrix( + # mn.Matrix4.look_at(camera_pos, look_at, look_up).rotation() + # ) + # sens_obj.node.translation = camera_pos + # print(sens_obj.node.translation) def reset(self) -> None: """Reset the joints on the existing robot. diff --git a/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py b/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py index b69be6c494..0863c496aa 100644 --- a/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py +++ b/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py @@ -26,13 +26,20 @@ class IsaacSpotRobot(IsaacMobileManipulator): @property def base_transformation(self): add_rot = mn.Matrix4.rotation( - mn.Rad(-np.pi / 2), mn.Vector3(1.0, 0, 0) + mn.Rad(np.pi / 2), mn.Vector3(1.0, 0, 0) ) - base_position, base_orientation = self._robot_wrapper.robot.get_world_pose() - # todo: get Hab transform from pos and orient - assert False - return None + base_position, base_rotation = self._robot_wrapper.get_root_pose() + pose = mn.Matrix4.from_(base_rotation.to_matrix(), base_position) + return pose @ add_rot + def get_link_transform(self, link_id): + link_positions, link_rotations = self._robot_wrapper.get_link_world_poses() + position, rotation = link_positions[link_id], link_rotations[link_id] + # breakpoint() + + pose = mn.Matrix4.from_(rotation.to_matrix(), position) + return pose + def get_ee_local_pose( self, ee_index: int = 0 ) -> Tuple[np.ndarray, np.ndarray]: diff --git a/habitat-lab/habitat/tasks/rearrange/isaac_rearrange_sim.py b/habitat-lab/habitat/tasks/rearrange/isaac_rearrange_sim.py index 75a70ec67e..aa41a0ba27 100644 --- a/habitat-lab/habitat/tasks/rearrange/isaac_rearrange_sim.py +++ b/habitat-lab/habitat/tasks/rearrange/isaac_rearrange_sim.py @@ -60,13 +60,12 @@ from habitat_sim.sim import SimulatorBackend from habitat_sim.utils.common import quat_from_magnum + if TYPE_CHECKING: from omegaconf import DictConfig def bind_physics_material_to_hierarchy(stage, root_prim, material_name, static_friction, dynamic_friction, restitution): - from pxr import UsdShade, UsdPhysics - from omni.isaac.core.materials.physics_material import PhysicsMaterial # material_path = f"/PhysicsMaterials/{material_name}" # material_prim = stage.DefinePrim(material_path, "PhysicsMaterial") @@ -75,7 +74,10 @@ def bind_physics_material_to_hierarchy(stage, root_prim, material_name, static_f # material.CreateStaticFrictionAttr().Set(static_friction) # material.CreateDynamicFrictionAttr().Set(dynamic_friction) # material.CreateRestitutionAttr().Set(restitution) - + from pxr import UsdShade, UsdPhysics + from omni.isaac.core.materials.physics_material import PhysicsMaterial + + physics_material = PhysicsMaterial( prim_path=f"/PhysicsMaterials/{material_name}", name=material_name, @@ -381,6 +383,10 @@ def reconfigure(self, config: "DictConfig", ep_info: RearrangeEpisode): new_scene = self.prev_scene_id != ep_info.scene_id if new_scene: self._prev_obj_names = None + + if new_scene: + self._load_navmesh(ep_info) + # Only remove and re-add objects if we have a new set of objects. ep_info.rigid_objs = sorted(ep_info.rigid_objs, key=lambda x: x[0]) @@ -550,36 +556,13 @@ def _setup_targets(self, ep_info): @add_perf_timing_func() def _load_navmesh(self, ep_info): - scene_name = ep_info.scene_id.split("/")[-1].split(".")[0] - base_dir = osp.join(*ep_info.scene_id.split("/")[:2]) - - navmesh_path = osp.join(base_dir, "navmeshes", scene_name + ".navmesh") - + # TODO: later will work in other scenes + navmesh_path = "/home/xavierpuig/habitat_llm/habitat-llm-planner-2/habitat-llm/102344193/navmeshes/102344193.navmesh" if osp.exists(navmesh_path): self.pathfinder.load_nav_mesh(navmesh_path) logger.info(f"Loaded navmesh from {navmesh_path}") else: - logger.warning( - f"Requested navmesh to load from {navmesh_path} does not exist. Recomputing from configured values and caching." - ) - navmesh_settings = NavMeshSettings() - navmesh_settings.set_defaults() - - agent_config = None - if hasattr(self.habitat_config.agents, "agent_0"): - agent_config = self.habitat_config.agents.agent_0 - elif hasattr(self.habitat_config.agents, "main_agent"): - agent_config = self.habitat_config.agents.main_agent - else: - raise ValueError(f"Cannot find agent parameters.") - navmesh_settings.agent_radius = agent_config.radius - navmesh_settings.agent_height = agent_config.height - navmesh_settings.agent_max_climb = agent_config.max_climb - navmesh_settings.agent_max_slope = agent_config.max_slope - navmesh_settings.include_static_objects = True - self.recompute_navmesh(self.pathfinder, navmesh_settings) - os.makedirs(osp.dirname(navmesh_path), exist_ok=True) - self.pathfinder.save_nav_mesh(navmesh_path) + raise Exception # NOTE: allowing indoor islands only self._largest_indoor_island_idx = get_largest_island_index( diff --git a/test_rearrange_env.py b/test_rearrange_env.py index bccd55b5ce..c5655102bf 100644 --- a/test_rearrange_env.py +++ b/test_rearrange_env.py @@ -10,7 +10,7 @@ import numpy as np from habitat.articulated_agents.robots import FetchRobot from habitat.config.default import get_agent_config -from habitat.config.default_structured_configs import ThirdRGBSensorConfig, HeadRGBSensorConfig, HeadPanopticSensorConfig +from habitat.config.default_structured_configs import ThirdRGBSensorConfig, HeadRGBSensorConfig, ArmDepthSensorConfig, HeadPanopticSensorConfig from habitat.config.default_structured_configs import SimulatorConfig, HabitatSimV0Config, AgentConfig from habitat.config.default import get_agent_config import habitat @@ -35,6 +35,7 @@ def make_sim_cfg(agent_dict): # This is for better graphics sim_cfg.habitat_sim_v0.enable_hbao = True sim_cfg.habitat_sim_v0.enable_physics = False + sim_cfg.habitat_sim_v0.frustum_culling = False # Set up an example scene @@ -77,6 +78,103 @@ def init_rearrange_env(agent_dict, action_dict): return Env(res_cfg) +from habitat.tasks.utils import get_angle +from habitat.datasets.rearrange.navmesh_utils import compute_turn +class OracleNavSkill(): + def __init__(self, env, target_pos): + self.env = env + self.target_pos = target_pos + self.target_base_pos = target_pos + self.dist_thresh = 0.1 + self.turn_velocity = 2 + + self.forward_velocity = 10 + self.turn_thresh = 0.2 + self.articulated_agent = self.env.sim.articulated_agent + + def _path_to_point(self, point): + """ + Obtain path to reach the coordinate point. If agent_pos is not given + the path starts at the agent base pos, otherwise it starts at the agent_pos + value + :param point: Vector3 indicating the target point + """ + agent_pos = self.articulated_agent.base_pos + + path = habitat_sim.ShortestPath() + path.requested_start = agent_pos + path.requested_end = point + found_path = self.env.sim.pathfinder.find_path(path) + if not found_path: + return [agent_pos, point] + return path.points + + def get_step(self): + + obj_targ_pos = np.array(self.target_pos) + base_T = self.articulated_agent.base_transformation + + curr_path_points = self._path_to_point(self.target_base_pos) + robot_pos = np.array(self.articulated_agent.base_pos) + if len(curr_path_points) == 1: + curr_path_points += curr_path_points + + cur_nav_targ = np.array(curr_path_points[1]) + forward = np.array([1.0, 0, 0]) + robot_forward = np.array(base_T.transform_vector(forward)) + + # Compute relative target + rel_targ = cur_nav_targ - robot_pos + # Compute heading angle (2D calculation) + robot_forward = robot_forward[[0, 2]] + rel_targ = rel_targ[[0, 2]] + rel_pos = (obj_targ_pos - robot_pos)[[0, 2]] + dist_to_final_nav_targ = np.linalg.norm( + (np.array(self.target_base_pos) - robot_pos)[[0, 2]], + ) + angle_to_target = get_angle(robot_forward, rel_targ) + angle_to_obj = get_angle(robot_forward, rel_pos) + + # Compute the distance + at_goal = ( + dist_to_final_nav_targ < self.dist_thresh + and angle_to_obj < self.turn_thresh + ) + + # Planning to see if the robot needs to do back-up + + if not at_goal: + if dist_to_final_nav_targ < self.dist_thresh: + # TODO: this does not account for the sampled pose's final rotation + # Look at the object target position when getting close + vel = compute_turn( + rel_pos, + self.turn_velocity, + robot_forward, + ) + elif angle_to_target < self.turn_thresh: + # Move forward towards the target + vel = [self.forward_velocity, 0] + else: + # Look at the target waypoint + vel = compute_turn( + rel_targ, + self.turn_velocity, + robot_forward, + ) + else: + vel = [0, 0] + vel2 = compute_turn( + rel_targ, + self.turn_velocity, + robot_forward, + ) + vel2 = vel + + # print(vel, dist_to_final_nav_targ, angle_to_obj, angle_to_target) + action = {'action': 'base_velocity_action', 'action_args': {'base_vel': np.array([ vel[0], vel[1]], dtype=np.float32)}} + return action + def main(): # Define the agent configuration @@ -90,7 +188,7 @@ def main(): # We will later talk about why we are giving the sensors these names main_agent_config.sim_sensors = { "third_rgb": ThirdRGBSensorConfig(), - "head_rgb": HeadRGBSensorConfig(), + "articulated_agent_arm_depth": ArmDepthSensorConfig(), } # We create a dictionary with names of agents and their corresponding agent configuration @@ -106,50 +204,38 @@ def main(): "output_env.mp4", fps=30, ) - action = {'action': 'base_velocity_action', 'action_args': {'base_vel': np.array([ 10.0, 0], dtype=np.float32)}} - for i in range(100): + + writer2 = imageio.get_writer( + "output_env_head.mp4", + fps=30, + ) + action1 = {'action': 'base_velocity_action', 'action_args': {'base_vel': np.array([ 5.0, 0], dtype=np.float32)}} + action2 = {'action': 'base_velocity_action', 'action_args': {'base_vel': np.array([ 0, 5], dtype=np.float32)}} + action3 = {'action': 'base_velocity_action', 'action_args': {'base_vel': np.array([ 0, 0], dtype=np.float32)}} + + first_obj = env.sim._rigid_objects[0].translation + nav_point = env.sim.pathfinder.get_random_navigable_point_near(circle_center=first_obj, radius=1) + curr_pos = env.sim.articulated_agent.base_pos + dist = np.linalg.norm((np.array(curr_pos) - nav_point) * np.array([1,0,1])) + nav_planner = OracleNavSkill(env, nav_point) + i = 0 + while dist > 0.10 or i < 200: - obs = env.step(action) + i += 1 + action_planner = nav_planner.get_step() + obs = env.step(action_planner) im = obs["third_rgb"] + writer.append_data(im) + + curr_pos = env.sim.articulated_agent.base_pos + dist = np.linalg.norm((np.array(curr_pos) - nav_point) * np.array([1,0,1])) + print(dist) writer.close() breakpoint() - # def get_pick_target_pos(): - # ro = isaac_viewer._rigid_objects[isaac_viewer._pick_target_rigid_object_idx] - # com_world = isaac_prim_utils.get_com_world(ro._rigid_prim) - # # self.draw_axis(0.05, mn.Matrix4.translation(com_world)) - # return com_world - - # isaac_viewer._spot_state_machine.set_pick_target(get_pick_target_pos) - - # # breakpoint() - # for it in range(100): - # isaac_viewer.update_isaac({}) - # look_up = mn.Vector3(0,1,0) - # isaac_viewer.update_spot_pre_step(0.01) - - # look_at = sim.agents_mgr._all_agent_data[0].articulated_agent.base_pos - # print(look_at) - # camera_pos = look_at + mn.Vector3(-0.7, 1.5, -0.7) - - # sim._sensors["third_rgb"]._sensor_object.node.rotation = mn.Quaternion.from_matrix( - # mn.Matrix4.look_at(camera_pos, look_at, look_up).rotation() - # ) - - # sim._sensors["third_rgb"]._sensor_object.node.translation = camera_pos - # # import cv2 - # sim.reset() - - # res = sim.get_sensor_observations() - - # im = res["third_rgb"][:,:,[0,1,2]] - # import cv2 - # cv2.imwrite("third2.png", im) - # breakpoint() - # writer.append_data(im) - - # writer.close() - # breakpoint() + + writer2.close() + if __name__ == "__main__": From 6b911918db8b9dc7895d3f667da3e298c555b1cd Mon Sep 17 00:00:00 2001 From: xavierpuig user Date: Tue, 28 Jan 2025 16:53:46 +0000 Subject: [PATCH 3/7] update --- .../isaac_sim/isaac_mobile_manipulator.py | 39 +++++-------------- 1 file changed, 9 insertions(+), 30 deletions(-) diff --git a/habitat-lab/habitat/isaac_sim/isaac_mobile_manipulator.py b/habitat-lab/habitat/isaac_sim/isaac_mobile_manipulator.py index ceeb6c8e15..e983c3e80a 100644 --- a/habitat-lab/habitat/isaac_sim/isaac_mobile_manipulator.py +++ b/habitat-lab/habitat/isaac_sim/isaac_mobile_manipulator.py @@ -35,11 +35,15 @@ def __init__( # maintain_link_order: bool = False, # base_type="mobile", ): - self.params = params self._sim = sim self._robot_wrapper = SpotRobotWrapper(isaac_service=isaac_service, instance_id=0) + # Modify here the params: + + self.params = params + # TODO: this should move later, cameras should not be attached to agents # @alexclegg + self._cameras = None if hasattr(self.params, "cameras"): from collections import defaultdict @@ -77,13 +81,11 @@ def update(self) -> None: cam_info = self.params.cameras[cam_prefix] agent = sim.agents_mgr._all_agent_data[0].articulated_agent look_at = sim.agents_mgr._all_agent_data[0].articulated_agent.base_pos - camera_pos = look_at + mn.Vector3(-0.5, 5.0, 0) if cam_info.attached_link_id == -1: link_trans = agent.base_transformation else: - link_trans = agent.get_link_transform(cam_info.attached_link_id) - + link_trans = agent.get_link_transform(cam_info.attached_link_id+1) if cam_info.cam_look_at_pos == mn.Vector3(0, 0, 0): pos = cam_info.cam_offset_pos ori = cam_info.cam_orientation @@ -98,12 +100,6 @@ def update(self) -> None: cam_info.cam_look_at_pos, mn.Vector3(0, 1, 0), ) - gt_look_at = mn.Matrix4.look_at(camera_pos, look_at, look_up) - # print(look_at) - # print(gt_look_at) - # print("----") - # print(cam_transform) - # breakpoint() cam_transform = ( link_trans @@ -111,28 +107,11 @@ def update(self) -> None: @ cam_info.relative_transform ) - camera_pos = link_trans.transform_point(cam_info.cam_offset_pos) - look_at = link_trans.transform_point(cam_info.cam_look_at_pos) - sens_obj.node.transformation = mn.Matrix4.look_at(camera_pos, look_at, look_up) + sens_obj.node.transformation = ( + cam_transform + ) - # sens_obj.node.translation = camera_pos - - # cam_transform = inv_T @ cam_transform - - # sens_obj.node.transformation = ( - # cam_transform - # ) - if "third" in cam_prefix: - # breakpoint() - pass - # breakpoint() - # sens_obj.node.rotation = mn.Quaternion.from_matrix( - # mn.Matrix4.look_at(camera_pos, look_at, look_up).rotation() - # ) - # sens_obj.node.translation = camera_pos - # print(sens_obj.node.translation) - def reset(self) -> None: """Reset the joints on the existing robot. NOTE: only arm and gripper joint motors (not gains) are reset by default, derived class should handle any other changes. From b1420b81fd5da526b45be80d305e2f0db57af05a Mon Sep 17 00:00:00 2001 From: xavierpuig user Date: Tue, 28 Jan 2025 17:01:05 +0000 Subject: [PATCH 4/7] update --- .../habitat/isaac_sim/isaac_spot_robot.py | 14 ++++++++++++- .../tasks/rearrange/actions/actions.py | 12 ++++++++--- test_rearrange_env.py | 20 ++++++++++--------- 3 files changed, 33 insertions(+), 13 deletions(-) diff --git a/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py b/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py index 0863c496aa..3bfdc89ac8 100644 --- a/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py +++ b/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py @@ -38,6 +38,9 @@ def get_link_transform(self, link_id): # breakpoint() pose = mn.Matrix4.from_(rotation.to_matrix(), position) + add_rot = mn.Matrix4.rotation( + mn.Rad(-np.pi / 2), mn.Vector3(1.0, 0, 0) + ) return pose def get_ee_local_pose( @@ -76,8 +79,17 @@ def __init__( isaac_service, sim=None ): + # TODO: This should be obtained from _target_arm_joint_positions but it is not intialized here yet. + ee_index = 19 + arm_joints = [0, 5, 10, 15, 16, 17, 18] + leg_joints = [jid for jid in range(19) if jid not in arm_joints] + + spot_params = SpotRobot._get_spot_params() + spot_params.arm_joints = arm_joints + spot_params.gripper_joints = [ee_index] + spot_params.leg_joints = leg_joints super().__init__( - SpotRobot._get_spot_params(), + spot_params, agent_cfg, isaac_service, sim=sim diff --git a/habitat-lab/habitat/tasks/rearrange/actions/actions.py b/habitat-lab/habitat/tasks/rearrange/actions/actions.py index a29d121711..9fe1427ca9 100644 --- a/habitat-lab/habitat/tasks/rearrange/actions/actions.py +++ b/habitat-lab/habitat/tasks/rearrange/actions/actions.py @@ -548,9 +548,15 @@ def step(self, *args, **kwargs): self.base_vel_ctrl.linear_velocity = mn.Vector3(lin_vel, 0, 0) self.base_vel_ctrl.angular_velocity = mn.Vector3(0, ang_vel, 0) - self.cur_articulated_agent._robot_wrapper._robot.set_angular_velocity([0, 0, ang_vel]) - self.cur_articulated_agent._robot_wrapper._robot.set_linear_velocity([lin_vel, 0, 0]) - + if ang_vel != 0: + self.cur_articulated_agent._robot_wrapper._robot.set_angular_velocity([0, 0, ang_vel]) + self.cur_articulated_agent._robot_wrapper._robot.set_linear_velocity([lin_vel, 0, 0]) + else: + lin_vel = 0 + ang_vel = 0 + self.cur_articulated_agent._robot_wrapper._robot.set_angular_velocity([0, 0, ang_vel]) + self.cur_articulated_agent._robot_wrapper._robot.set_linear_velocity([lin_vel, 0, 0]) + self.cur_articulated_agent.base_pos = self.cur_articulated_agent.base_transformation.transform_point(mn.Vector3(0.1,0,0)) @registry.register_task_action class BaseVelNonCylinderAction(ArticulatedAgentAction): diff --git a/test_rearrange_env.py b/test_rearrange_env.py index c5655102bf..4f861cbe50 100644 --- a/test_rearrange_env.py +++ b/test_rearrange_env.py @@ -10,7 +10,7 @@ import numpy as np from habitat.articulated_agents.robots import FetchRobot from habitat.config.default import get_agent_config -from habitat.config.default_structured_configs import ThirdRGBSensorConfig, HeadRGBSensorConfig, ArmDepthSensorConfig, HeadPanopticSensorConfig +from habitat.config.default_structured_configs import ThirdRGBSensorConfig, HeadRGBSensorConfig, ArmDepthSensorConfig, ArmRGBSensorConfig, HeadPanopticSensorConfig from habitat.config.default_structured_configs import SimulatorConfig, HabitatSimV0Config, AgentConfig from habitat.config.default import get_agent_config import habitat @@ -35,13 +35,13 @@ def make_sim_cfg(agent_dict): # This is for better graphics sim_cfg.habitat_sim_v0.enable_hbao = True sim_cfg.habitat_sim_v0.enable_physics = False + + # TODO: disable this, causes performance issues sim_cfg.habitat_sim_v0.frustum_culling = False # Set up an example scene sim_cfg.scene = "NONE" # os.path.join(data_path, "hab3_bench_assets/hab3-hssd/scenes/103997919_171031233.scene_instance.json") - # sim_cfg.scene_dataset = os.path.join(data_path, "hab3_bench_assets/hab3-hssd/hab3-hssd.scene_dataset_config.json") - # sim_cfg.additional_object_paths = [os.path.join(data_path, 'objects/ycb/configs/')] cfg = OmegaConf.create(sim_cfg) @@ -188,7 +188,7 @@ def main(): # We will later talk about why we are giving the sensors these names main_agent_config.sim_sensors = { "third_rgb": ThirdRGBSensorConfig(), - "articulated_agent_arm_depth": ArmDepthSensorConfig(), + "articulated_agent_arm_rgb": ArmRGBSensorConfig(), } # We create a dictionary with names of agents and their corresponding agent configuration @@ -209,9 +209,7 @@ def main(): "output_env_head.mp4", fps=30, ) - action1 = {'action': 'base_velocity_action', 'action_args': {'base_vel': np.array([ 5.0, 0], dtype=np.float32)}} - action2 = {'action': 'base_velocity_action', 'action_args': {'base_vel': np.array([ 0, 5], dtype=np.float32)}} - action3 = {'action': 'base_velocity_action', 'action_args': {'base_vel': np.array([ 0, 0], dtype=np.float32)}} + action_example = {'action': 'base_velocity_action', 'action_args': {'base_vel': np.array([ 5.0, 0], dtype=np.float32)}} first_obj = env.sim._rigid_objects[0].translation nav_point = env.sim.pathfinder.get_random_navigable_point_near(circle_center=first_obj, radius=1) @@ -219,19 +217,23 @@ def main(): dist = np.linalg.norm((np.array(curr_pos) - nav_point) * np.array([1,0,1])) nav_planner = OracleNavSkill(env, nav_point) i = 0 - while dist > 0.10 or i < 200: + while dist > 0.10 and i < 200: i += 1 action_planner = nav_planner.get_step() obs = env.step(action_planner) im = obs["third_rgb"] + im2 = obs["articulated_agent_arm_rgb"] writer.append_data(im) + writer2.append_data(im2) curr_pos = env.sim.articulated_agent.base_pos dist = np.linalg.norm((np.array(curr_pos) - nav_point) * np.array([1,0,1])) - print(dist) + writer.close() + writer2.close() + breakpoint() writer2.close() From 9b42c243a25f1d3a3ca47b282a4f6dfdb150a1de Mon Sep 17 00:00:00 2001 From: xavierpuig user Date: Tue, 28 Jan 2025 17:01:58 +0000 Subject: [PATCH 5/7] update --- test_rearrange_env.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/test_rearrange_env.py b/test_rearrange_env.py index 4f861cbe50..3c27d6e32a 100644 --- a/test_rearrange_env.py +++ b/test_rearrange_env.py @@ -234,9 +234,6 @@ def main(): writer.close() writer2.close() - breakpoint() - - writer2.close() From d8497a3ea353e60a3704961dbe66e37f4f23f8fc Mon Sep 17 00:00:00 2001 From: xavierpuig user Date: Wed, 29 Jan 2025 06:23:05 +0000 Subject: [PATCH 6/7] add ik action --- .../isaac_sim/_internal/spot_robot_wrapper.py | 12 +++- habitat-lab/habitat/isaac_sim/actions.py | 62 +++++++++++++++++++ .../habitat/isaac_sim/isaac_spot_robot.py | 12 ++++ .../tasks/rearrange/actions/actions.py | 21 ------- test_rearrange_env.py | 43 ++++++++++--- 5 files changed, 118 insertions(+), 32 deletions(-) create mode 100644 habitat-lab/habitat/isaac_sim/actions.py diff --git a/habitat-lab/habitat/isaac_sim/_internal/spot_robot_wrapper.py b/habitat-lab/habitat/isaac_sim/_internal/spot_robot_wrapper.py index d52e9c02ee..b16672c671 100644 --- a/habitat-lab/habitat/isaac_sim/_internal/spot_robot_wrapper.py +++ b/habitat-lab/habitat/isaac_sim/_internal/spot_robot_wrapper.py @@ -100,6 +100,15 @@ def robot(self) -> Robot: def get_prim_path(self): return self._robot_prim_path + + + def set_root_pose(self, pos, rot): + + rot = [rot.scalar] + list(rot.vector) + pos_usd = isaac_prim_utils.habitat_to_usd_position(pos) + rot_usd = isaac_prim_utils.habitat_to_usd_rotation(rot) + self._robot.set_world_pose(pos_usd, rot_usd) + def get_root_pose(self): @@ -167,7 +176,8 @@ def post_reset(self): arm_joint_indices.append(dof_names.index(arm_joint_name)) self._arm_joint_indices = np.array(arm_joint_indices) - self._target_arm_joint_positions = None + # self._target_arm_joint_positions = None + self._target_arm_joint_positions = [0.0, -2.36, 0.0, 2.25, 0.0, 1.67, 0.0, -1.67] def scale_prim_mass_and_inertia(self, path, scale): diff --git a/habitat-lab/habitat/isaac_sim/actions.py b/habitat-lab/habitat/isaac_sim/actions.py new file mode 100644 index 0000000000..5b6fa704d6 --- /dev/null +++ b/habitat-lab/habitat/isaac_sim/actions.py @@ -0,0 +1,62 @@ +from habitat.core.registry import registry +from habitat.tasks.rearrange.actions.actions import BaseVelAction, ArticulatedAgentAction +import magnum as mn +import numpy as np +import habitat_sim +from examples.hitl.isaacsim_viewer.isaacsim_viewer import SpotPickHelper +@registry.register_task_action +class BaseVelIsaacAction(BaseVelAction): + def step(self, *args, **kwargs): + lin_vel, ang_vel = kwargs[self._action_arg_prefix + "base_vel"] + lin_vel = np.clip(lin_vel, -1, 1) * self._lin_speed + ang_vel = np.clip(ang_vel, -1, 1) * self._ang_speed + if not self._allow_back: + lin_vel = np.maximum(lin_vel, 0) + + self.base_vel_ctrl.linear_velocity = mn.Vector3(lin_vel, 0, 0) + self.base_vel_ctrl.angular_velocity = mn.Vector3(0, ang_vel, 0) + self.cur_articulated_agent._robot_wrapper._robot.set_angular_velocity([0, 0, ang_vel]) + self.cur_articulated_agent._robot_wrapper._robot.set_linear_velocity([lin_vel, 0, 0]) + + + +@registry.register_task_action +class ArmReachAction(ArticulatedAgentAction): + def __init__(self, *args, **kwargs): + super().__init__(self, *args, **kwargs) + + self._spot_wrapper = self.cur_articulated_agent._robot_wrapper + self._spot_pick_helper = SpotPickHelper(len(self._spot_wrapper._arm_joint_indices)) + + def step(self, *args, **kwargs): + target_pos = kwargs[self._action_arg_prefix + "target_pos"] + base_pos, base_rot = self._spot_wrapper.get_root_pose() + def inverse_transform(pos_a, rot_b, pos_b): + inv_pos = rot_b.inverted().transform_vector(pos_a - pos_b) + return inv_pos + target_rel_pos = inverse_transform(target_pos, base_rot, base_pos) + + dt = 0.5 + self._spot_wrapper._target_arm_joint_positions = self._spot_pick_helper.update(dt, target_rel_pos) + + + +@registry.register_task_action +class BaseVelKinematicIsaacAction(BaseVelAction): + + def update_base(self): + ctrl_freq = self._sim.ctrl_freq + trans = self.cur_articulated_agent.base_transformation + rigid_state = habitat_sim.RigidState( + mn.Quaternion.from_matrix(trans.rotation()), trans.translation + ) + + target_rigid_state = self.base_vel_ctrl.integrate_transform( + 1 / ctrl_freq, rigid_state + ) + target_trans = mn.Matrix4.from_( + target_rigid_state.rotation.to_matrix(), target_rigid_state.translation + ) + self.cur_articulated_agent.base_transformation = target_trans + + diff --git a/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py b/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py index 3bfdc89ac8..75822c678a 100644 --- a/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py +++ b/habitat-lab/habitat/isaac_sim/isaac_spot_robot.py @@ -32,6 +32,18 @@ def base_transformation(self): pose = mn.Matrix4.from_(base_rotation.to_matrix(), base_position) return pose @ add_rot + + @base_transformation.setter + def base_transformation(self, base_transformation): + rot = mn.Matrix4.rotation( + mn.Rad(-np.pi / 2), mn.Vector3(1.0, 0, 0) + ) + base_transformation = base_transformation @ rot + rot = mn.Quaternion.from_matrix(base_transformation.rotation()) + self._robot_wrapper.set_root_pose(base_transformation.translation, rot) + # pose = mn.Matrix4.from_(base_rotation.to_matrix(), base_position + + def get_link_transform(self, link_id): link_positions, link_rotations = self._robot_wrapper.get_link_world_poses() position, rotation = link_positions[link_id], link_rotations[link_id] diff --git a/habitat-lab/habitat/tasks/rearrange/actions/actions.py b/habitat-lab/habitat/tasks/rearrange/actions/actions.py index 9fe1427ca9..3c29923717 100644 --- a/habitat-lab/habitat/tasks/rearrange/actions/actions.py +++ b/habitat-lab/habitat/tasks/rearrange/actions/actions.py @@ -537,27 +537,6 @@ def step(self, *args, **kwargs): self.update_base() -@registry.register_task_action -class BaseVelIsaacAction(BaseVelAction): - def step(self, *args, **kwargs): - lin_vel, ang_vel = kwargs[self._action_arg_prefix + "base_vel"] - lin_vel = np.clip(lin_vel, -1, 1) * self._lin_speed - ang_vel = np.clip(ang_vel, -1, 1) * self._ang_speed - if not self._allow_back: - lin_vel = np.maximum(lin_vel, 0) - - self.base_vel_ctrl.linear_velocity = mn.Vector3(lin_vel, 0, 0) - self.base_vel_ctrl.angular_velocity = mn.Vector3(0, ang_vel, 0) - if ang_vel != 0: - self.cur_articulated_agent._robot_wrapper._robot.set_angular_velocity([0, 0, ang_vel]) - self.cur_articulated_agent._robot_wrapper._robot.set_linear_velocity([lin_vel, 0, 0]) - else: - lin_vel = 0 - ang_vel = 0 - self.cur_articulated_agent._robot_wrapper._robot.set_angular_velocity([0, 0, ang_vel]) - self.cur_articulated_agent._robot_wrapper._robot.set_linear_velocity([lin_vel, 0, 0]) - self.cur_articulated_agent.base_pos = self.cur_articulated_agent.base_transformation.transform_point(mn.Vector3(0.1,0,0)) - @registry.register_task_action class BaseVelNonCylinderAction(ArticulatedAgentAction): """ diff --git a/test_rearrange_env.py b/test_rearrange_env.py index 3c27d6e32a..1d7a18b417 100644 --- a/test_rearrange_env.py +++ b/test_rearrange_env.py @@ -21,9 +21,10 @@ from habitat.isaac_sim import isaac_prim_utils import random from habitat.config.default_structured_configs import TaskConfig, EnvironmentConfig, DatasetConfig, HabitatConfig -from habitat.config.default_structured_configs import ArmActionConfig, BaseVelocityActionConfig, OracleNavActionConfig, ActionConfig +from habitat.config.default_structured_configs import ArmActionConfig, BaseVelocityActionConfig, ActionConfig, OracleNavActionConfig, ActionConfig import imageio from habitat.core.env import Env +from habitat.isaac_sim import actions data_path = "/fsx-siro/xavierpuig/projects/habitat_isaac/habitat-lab/data/" @@ -189,26 +190,27 @@ def main(): main_agent_config.sim_sensors = { "third_rgb": ThirdRGBSensorConfig(), "articulated_agent_arm_rgb": ArmRGBSensorConfig(), + "articulated_agent_arm_depth": ArmDepthSensorConfig(), } # We create a dictionary with names of agents and their corresponding agent configuration agent_dict = {"main_agent": main_agent_config} action_dict = { - "base_velocity_action": BaseVelocityActionConfig(type="BaseVelIsaacAction"), + "base_velocity_action": BaseVelocityActionConfig(type="BaseVelKinematicIsaacAction"), + "arm_reach_action": ActionConfig(type="ArmReachAction") } env = init_rearrange_env(agent_dict, action_dict) + + + aux = env.reset() - writer = imageio.get_writer( "output_env.mp4", fps=30, ) - writer2 = imageio.get_writer( - "output_env_head.mp4", - fps=30, - ) + action_example = {'action': 'base_velocity_action', 'action_args': {'base_vel': np.array([ 5.0, 0], dtype=np.float32)}} first_obj = env.sim._rigid_objects[0].translation @@ -224,15 +226,36 @@ def main(): obs = env.step(action_planner) im = obs["third_rgb"] im2 = obs["articulated_agent_arm_rgb"] + im3 = (255 * obs["articulated_agent_arm_depth"]).astype(np.uint8) + imt = np.zeros(im.shape) + imt[:im2.shape[0], :im2.shape[1], :] = im2 + imt[im2.shape[0]:, :im2.shape[1], 0] = im3[:, :, 0] + imt[im2.shape[0]:, :im2.shape[1], 1] = im3[:, :, 0] + imt[im2.shape[0]:, :im2.shape[1], 2] = im3[:, :, 0] + im = np.concatenate([im, imt], 1) writer.append_data(im) - writer2.append_data(im2) - curr_pos = env.sim.articulated_agent.base_pos dist = np.linalg.norm((np.array(curr_pos) - nav_point) * np.array([1,0,1])) + + print(env.sim.articulated_agent._robot_wrapper._target_arm_joint_positions) + for i in range(200): + arm_reach = {'action': 'arm_reach_action', 'action_args': {'target_pos': np.array(first_obj, dtype=np.float32)}} + + obs = env.step(arm_reach) + im = obs["third_rgb"] + im2 = obs["articulated_agent_arm_rgb"] + im3 = (255 * obs["articulated_agent_arm_depth"]).astype(np.uint8) + imt = np.zeros(im.shape) + imt[:im2.shape[0], :im2.shape[1], :] = im2 + imt[im2.shape[0]:, :im2.shape[1], 0] = im3[:, :, 0] + imt[im2.shape[0]:, :im2.shape[1], 1] = im3[:, :, 0] + imt[im2.shape[0]:, :im2.shape[1], 2] = im3[:, :, 0] + im = np.concatenate([im, imt], 1) + writer.append_data(im) writer.close() - writer2.close() + From 72f2276fa8bc5a989130d6ddfd42a484fd3b219c Mon Sep 17 00:00:00 2001 From: xavierpuig user Date: Fri, 7 Feb 2025 19:45:32 +0000 Subject: [PATCH 7/7] update --- examples/hitl/isaacsim_viewer/isaacsim_viewer.py | 4 ++-- habitat-lab/habitat/isaac_sim/spot_arm_ik_helper.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/examples/hitl/isaacsim_viewer/isaacsim_viewer.py b/examples/hitl/isaacsim_viewer/isaacsim_viewer.py index ab2484654e..f3779fb31c 100644 --- a/examples/hitl/isaacsim_viewer/isaacsim_viewer.py +++ b/examples/hitl/isaacsim_viewer/isaacsim_viewer.py @@ -219,7 +219,7 @@ def update(self, dt): class SpotPickHelper: APPROACH_DIST = 0.16 - APPROACH_DURATION = 2.0 + APPROACH_DURATION = 50.0 def __init__(self, num_dof): @@ -320,7 +320,7 @@ def update(self, dt, target_rel_pos): # target_arm_joint_positions[7] = -1.0 if approach_progress > 0.0: - print(f"approach_progress: {approach_progress}") + print(f"approach_progress: {approach_progress}", target_arm_joint_positions[7]) return target_arm_joint_positions diff --git a/habitat-lab/habitat/isaac_sim/spot_arm_ik_helper.py b/habitat-lab/habitat/isaac_sim/spot_arm_ik_helper.py index 73035d2005..bd7a983d05 100644 --- a/habitat-lab/habitat/isaac_sim/spot_arm_ik_helper.py +++ b/habitat-lab/habitat/isaac_sim/spot_arm_ik_helper.py @@ -138,4 +138,5 @@ def is_valid_triangle(a, b, h): result[0] = yaw_angle_to_target is_ik_active = not is_out_of_range - return is_ik_active, result \ No newline at end of file + return is_ik_active, result +