diff --git a/.gitignore b/.gitignore index c58ab24..29b64f1 100644 --- a/.gitignore +++ b/.gitignore @@ -10,4 +10,5 @@ _*/ /.vs /app -data \ No newline at end of file +data +/.venv \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index a597fb1..b929b91 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,13 @@ # main +- Added instructions for manually building Occupancy Map with Isaac Sim Occupancy Map tool +- Modified extension to load pre-defined occupancy map rather than building on the fly + - User is now required to build the Occupancy Map before data collection. +- Added prim_get_world_transform to get world and local pose to address performance bottleneck with Isaac Sim method +- Added example for parquet conversion (to support X-Mobility training) +- Added robot linear and angular velocity to state (to support X-Mobility training) +- Fixed bug when replay rendering does not include segmentation info - Added support for surface normals image in replay rendering - Added support for instance ID segmentation rendering - Added camera world pose to state diff --git a/README.md b/README.md index c32294f..1ea7ed7 100644 --- a/README.md +++ b/README.md @@ -163,30 +163,120 @@ Below details a typical workflow for collecting data with MobilityGen. ./scripts/launch_sim.sh ``` -### Step 2 - Build a scenario +### Step 2 - Load a stage -This assumes you see the MobilityGen extension window. +To get started, we'll open an example warehouse stage. -1. Under Scene USD URL / Path copy and paste the following +1. Select ``File`` -> ``Open`` + +2. Enter the following URL under ``File name`` ``` http://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/Isaac/4.2/Isaac/Environments/Simple_Warehouse/warehouse_multiple_shelves.usd ``` +3. Click ``Open File`` -2. Under the ``Scenario`` dropdown select ``KeyboardTeleoperationScenario`` to start + > If you see a prompt ``Opening a Read Only File`` appear you can click ``Open Original File`` + +After a few seconds, you should see the stage appear. + +### Step 3 - Create an occupancy map + +Next, we need need to build an occupancy map for the environment. + +To do this, we'll use the Occupancy Map tool provided with Isaac Sim + +1. Select ``Tools`` -> ``Robotics`` -> ``Occupancy Map`` to open the Occupancy Map extension + + > You may need to also click the ``Occupancy Map`` tab in the bottom window pane to see the extension window. + +2. In the ``Occupancy Map`` window set ``Origin`` to + + - ``X``: ``2.0`` + - ``Y``: ``0.0`` + - ``Z``: ``0.0`` + +3. In the ``Occupancy Map`` window set ``Upper Bound`` to + + - ``X``: ``10.0`` + - ``Y``: ``20.0`` + - ``Z``: ``2.0`` (We'll assume the robot can move under 2 meter overpasses.) + +4. In the ``Occupancy Map`` window set ``Lower Bound`` to + + - ``X``: ``-14.0`` + - ``Y``: ``-18.0`` + - ``Z``: ``0.1`` (We'll assume we can move over 5cm bumps.) + +5. Click ``Calculate`` to generate the Occupancy Map + +7. Click ``Visualize Image`` to view the Occupancy Map + +8. In this ``Visualization`` window under ``Rotate Image`` select ``180`` + +8. In this ``Visualization`` window under ``Coordinate Type`` select ``ROS Occupancy Map Parameters File YAML`` + +10. Click ``Regenerate Image`` + +12. Copy the YAML text generated to your clipboard + +11. In a text editor of choice, create a new file named ``~/MobilityGenData/maps/warehouse_multiple_shelves/map.yaml`` + + > Note: ``~`` corresponds to your user's home directory. By default, + > we'll keep our data in ``~/MobilityGenData`` + +12. Paste the YAML text copied from the ``Visualization`` window into the created file. + +13. Edit the line ``image: warehouse_multiple_shelves.png`` to read ``image: map.png`` + +14. Save the file. + +11. Back in the ``Visualization`` window click ``Save Image`` + +12. In the tree explorer open the folder ``~/MobilityGenData/maps/warehouse_multiple_shelves`` + +12. Under file name enter ``map.png`` + +13. Click ``Save`` + +That's it! You should now have a folder ``~/MobilityGenData/maps/warehouse_multiple_shelves/`` with a file named +``map.yaml`` and ``map.png`` inside. + +> Note: For more details on generating occupancy maps, check the documents [here](https://docs.omniverse.nvidia.com/isaacsim/latest/features/ext_omni_isaac_occupancy_map.html). +> However, please note, to work with MobilityGen, you must use the rotation and coodinate type specifications detailed above. + +### Step 4 - Build a scenario + +Now that we have a ROS format Occupancy Map of our environment, we're ready to use MobilityGen! + +Perform the following steps in the ``MobilityGen`` extension window to build a new scenario. + +1. Under ``Stage`` paste the following, corresponding to the environment USD we used in [Step 3](#step-3---create-an-occupancy-map) + + ``` + http://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/Isaac/4.2/Isaac/Environments/Simple_Warehouse/warehouse_multiple_shelves.usd + ``` + +2. Under ``Occupancy Map`` enter the following corresponding to the Occupancy Map we created in [Step 3](#step-3---create-an-occupancy-map) + + ``` + ~/MobilityGenData/maps/warehouse_multiple_shelves/map.yaml + ``` 3. Under the ``Robot`` dropdown select ``H1Robot`` +2. Under the ``Scenario`` dropdown select ``KeyboardTeleoperationScenario`` to start + 4. Click ``Build`` After a few seconds, you should see the scene and occupancy map appear. -### Step 3 - Initialize / reset the scenario +### Step 5 - Initialize / reset the scenario -1. Click the ``Reset`` function to randomly initialize the scenario. Do this until the robot spawns inside the warehouse. +1. Click the ``Reset`` function to randomly initialize the scenario. Do this until the robot spawns in a desirable location. -### Step 4 - Test drive the robot +### Step 6 - Test drive the robot Before you start recording, try moving the robot around to get a feel for it @@ -197,7 +287,7 @@ To move the robot, use the following keys - ``S`` - Move Backwards - ``D`` - Turn right -### Step 5 - Start recording! +### Step 7 - Start recording! Once you're comfortable, you can record a log. @@ -209,7 +299,7 @@ Once you're comfortable, you can record a log. The data is recorded to ``~/MobilityGenData/recordings`` by default. -### Step 6 - Render data +### Step 8 - Render data If you've gotten this far, you've recorded a trajectory, but it doesn't include the rendered sensor data. @@ -234,7 +324,7 @@ Rendering the sensor data is done offline. To do this call the following That's it! Now the data with renderings should be stored in ``~/MobilityGenData/replays``. -### Step 7 - Visualize the Data +### Step 9 - Visualize the Data We provide a few examples in the [examples](./examples) folder for working with the data. @@ -399,6 +489,8 @@ The state_dict has the following schema "robot.action": np.ndarray, # [2] - Linear, angular command velocity "robot.position": np.ndarray, # [3] - XYZ "robot.orientation": np.ndarray, # [4] - Quaternion + "robot.linear_velocity": np.ndarray, # [3] - The linear velocity in world frame (As retrieved by robot.get_linear_velocity() in isaac sim) + "robot.angular_velocity": np.ndarray, # [3] - The angular velocity of the robot in the world frame. (As retrieved by robot.get_angular_velocity() in isaac sim) "robot.joint_positions": np.ndarray, # [J] - Joint positions "robot.joint_velocities": np.ndarray, # [J] - Joint velocities "robot.front_camera.left.rgb_image": np.ndarray, # [HxWx3], np.uint8 - RGB image diff --git a/examples/05_convert_to_parquet.py b/examples/05_convert_to_parquet.py new file mode 100644 index 0000000..161c3a4 --- /dev/null +++ b/examples/05_convert_to_parquet.py @@ -0,0 +1,111 @@ +import argparse +import pandas +from reader import Reader +import numpy as np +import tqdm +import PIL.Image +import io + +import os +import subprocess +import glob +import argparse + + +def numpy_array_to_flattened_columns(key: str, value: np.ndarray): + columns = { + f"{key}": value.flatten() + } + # add shape if ndim > 1 + if value.ndim > 1: + columns[f"{key}.shape"] = tuple(value.shape) + return columns + + +def numpy_array_to_jpg_columns(key: str, value: np.ndarray): + image = PIL.Image.fromarray(value) + buffer = io.BytesIO() + image.save(buffer, format="JPEG") + columns = { + key: buffer.getvalue() + } + return columns + + +if "MOBILITY_GEN_DATA" in os.environ: + DATA_DIR = os.environ['MOBILITY_GEN_DATA'] +else: + DATA_DIR = os.path.expanduser("~/MobilityGenData") + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + parser.add_argument("--input_dir", type=str, default=None) + parser.add_argument("--output_dir", type=str, default=None) + + + args = parser.parse_args() + + if args.input_dir is None: + args.input_dir = os.path.join(DATA_DIR, "replays") + + if args.output_dir is None: + args.output_dir = os.path.join(DATA_DIR, "parquet") + + if not os.path.exists(args.output_dir): + os.makedirs(args.output_dir) + + input_recordings = glob.glob(os.path.join(args.input_dir, "*")) + + processed_count = 0 + + for input_recording_path in input_recordings: + processed_count += 1 + print(f"Processing {processed_count} / {len(input_recordings)}") + + recording_name = os.path.basename(input_recording_path) + output_path = os.path.join(args.output_dir, recording_name + ".pqt") + + reader = Reader(recording_path=input_recording_path) + + index = 0 + + + output: pandas.DataFrame = None + + for index in tqdm.tqdm(range(len(reader))): + + data_dict = {} + + # Common data (saved as raw arrays) + state_common = reader.read_state_dict_common(index=index) + state_common.update(reader.read_state_dict_depth(index=index)) + state_common.update(reader.read_state_dict_segmentation(index=index)) + # state_common.update(reader.read_state_dict_depth(index=index)) + # TODO: handle normals + + for k, v in state_common.items(): + if isinstance(v, np.ndarray): + columns = numpy_array_to_flattened_columns(k, v) + else: + columns = {k: v} + data_dict.update(columns) + + # RGB data (saved as jpg) + state_rgb = reader.read_state_dict_rgb(index=index) + for k, v in state_rgb.items(): + if isinstance(v, np.ndarray): + columns = numpy_array_to_jpg_columns(k, v) + else: + columns = {k: v} + data_dict.update(columns) + + + # use first frame to initialize + if output is None: + output = pandas.DataFrame(columns=data_dict.keys()) + + output.loc[index] = data_dict + + + output.to_parquet(output_path, engine="pyarrow") \ No newline at end of file diff --git a/examples/06_load_parquet.py b/examples/06_load_parquet.py new file mode 100644 index 0000000..ac96b96 --- /dev/null +++ b/examples/06_load_parquet.py @@ -0,0 +1,19 @@ +import argparse +import pandas +import matplotlib.pyplot as plt +import numpy as np + +parser = argparse.ArgumentParser() +parser.add_argument("parquet_path") +args = parser.parse_args() + +data = pandas.read_parquet(args.parquet_path, engine="pyarrow") + + +print(data.columns) +vel = np.stack(data['robot.linear_velocity'].to_numpy()) + + +plt.plot(vel[:, 0], 'r-') +plt.plot(vel[:, 1], 'r-') +plt.show() diff --git a/examples/PARQUET_DATA_FORMAT.md b/examples/PARQUET_DATA_FORMAT.md new file mode 100644 index 0000000..7664c66 --- /dev/null +++ b/examples/PARQUET_DATA_FORMAT.md @@ -0,0 +1,28 @@ +# Parquet Data Format + +Below is a description of the fields in a typical MobilityGen recording, common to all scenarios. + +| Field | Type | Shape | Description | +|-------|------|-------|-------------| +| robot.action | array | 2 | The command linear, angular velocity in the robot frame. | +| robot.position | array | 3 | The xyz position of the robot in the world frame. | +| robot.orientation | array | 4 | The quaternion of the robot in the world frame. | +| robot.joint_positions| array | N | The joint positions of the robot. | +| robot.joint_velocities | array | N | The joint velocities of the robot. | +| robot.linear_velocity | array | 3 | The linear velocity of the robot in the world frame. (Retrieved by robot.get_linear_velocity() in isaac sim) | +| robot.angular_velocity | array | 3 | The linear velocity of the robot in the world frame. (Retrieved by robot.get_angular_velocity() in isaac sim) | +| robot.front_camera.left.segmentation_info | dict | | The segmentation info dictionary as retrieved by the Isaac Sim replicator annotator. | +| robot.front_camera.left.segmentation_image | array | (HxW) flattened | The segmentation image as retrieved by the Isaac Sim replicator annotator. Flattened | +| robot.front_camera.left.segmentation_image.shape | tuple | 2 | The segmentation image shape. | +| robot.front_camera.left.rgb_image | bytes | | The RGB camera image compressed to JPG. | +| robot.front_camera.left.depth_image | array | (HxW) | The depth image (in meters) flattened into an array. | +| robot.front_camera.left.depth_image.shape | tuple | 2 | The shape of the depth image. + +> Note, there are additional fields with similar semantics for other cameras we have excluded brevity. + +Below are fields specific to the path following scenario + +| Field | Type | Shape | Description | +|-------|------|-------|-------------| +| target_path | array | (Nx2) flattened | The target path generated by the path planner in world coordinates. This is updated whenever the path planner is called, which occurs when the robot reaches a goal (or at the beginning of a new recording) | +| target_path.shape | tuple | 2 | The shape of the target path array. | \ No newline at end of file diff --git a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/build.py b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/build.py index 51d361f..2f580d5 100644 --- a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/build.py +++ b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/build.py @@ -18,13 +18,10 @@ import isaacsim.core.utils.prims as prim_utils from isaacsim.core.utils.stage import open_stage import isaacsim.core.api.objects as objects -from isaacsim.core.utils.stage import add_reference_to_stage from omni.ext.mobility_gen.occupancy_map import OccupancyMap -from omni.ext.mobility_gen.config import Config -from omni.ext.mobility_gen.utils.occupancy_map_utils import occupancy_map_generate_from_prim_async -from omni.ext.mobility_gen.utils.global_utils import new_stage, new_world, set_viewport_camera +from omni.ext.mobility_gen.utils.global_utils import new_world, set_viewport_camera, get_stage from omni.ext.mobility_gen.scenarios import Scenario, SCENARIOS from omni.ext.mobility_gen.robots import ROBOTS from omni.ext.mobility_gen.reader import Reader @@ -37,10 +34,18 @@ def load_scenario(path: str) -> Scenario: robot_type = ROBOTS.get(config.robot_type) scenario_type = SCENARIOS.get(config.scenario_type) open_stage(os.path.join(path, "stage.usd")) - prim_utils.delete_prim("/World/robot") + + stage = get_stage() + + robot_prim_path = "/World/robot" + + if stage.GetPrimAtPath(path).IsValid(): + prim_utils.delete_prim(robot_prim_path) + + objects.GroundPlane("/World/ground_plane", visible=False) new_world(physics_dt=robot_type.physics_dt) occupancy_map = reader.read_occupancy_map() - robot = robot_type.build("/World/robot") + robot = robot_type.build(robot_prim_path) chase_camera_path = robot.build_chase_camera() set_viewport_camera(chase_camera_path) robot_type = ROBOTS.get(config.robot_type) @@ -50,23 +55,3 @@ def load_scenario(path: str) -> Scenario: scenario = scenario_type.from_robot_occupancy_map(robot, occupancy_map) return scenario - -async def build_scenario_from_config(config: Config): - robot_type = ROBOTS.get(config.robot_type) - scenario_type = SCENARIOS.get(config.scenario_type) - new_stage() - world = new_world(physics_dt=robot_type.physics_dt) - await world.initialize_simulation_context_async() - add_reference_to_stage(config.scene_usd,"/World/scene") - objects.GroundPlane("/World/ground_plane", visible=False) - robot = robot_type.build("/World/robot") - occupancy_map = await occupancy_map_generate_from_prim_async( - "/World/scene", - cell_size=robot.occupancy_map_cell_size, - z_min=robot.occupancy_map_z_min, - z_max=robot.occupancy_map_z_max - ) - chase_camera_path = robot.build_chase_camera() - set_viewport_camera(chase_camera_path) - scenario = scenario_type.from_robot_occupancy_map(robot, occupancy_map) - return scenario \ No newline at end of file diff --git a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/config.py b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/config.py index 9b87f6c..87d070c 100644 --- a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/config.py +++ b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/config.py @@ -1,5 +1,5 @@ import json -from typing import Literal +from typing import Literal, Optional, Tuple from dataclasses import dataclass, asdict @@ -14,4 +14,9 @@ def to_json(self): @staticmethod def from_json(data: str): - return Config(**json.loads(data)) \ No newline at end of file + data = json.loads(data) + return Config( + scenario_type=data['scenario_type'], + robot_type=data['robot_type'], + scene_usd=data['scene_usd'] + ) \ No newline at end of file diff --git a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/extension.py b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/extension.py index 6c48712..1ef26e1 100644 --- a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/extension.py +++ b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/extension.py @@ -24,14 +24,27 @@ import omni.ext import omni.ui as ui +from isaacsim.core.utils.stage import open_stage +from isaacsim.gui.components.ui_utils import ( + btn_builder, + cb_builder, + color_picker_builder, + dropdown_builder, + float_builder, + multi_btn_builder, + xyz_builder, +) +import isaacsim.core.api.objects as objects +import carb + from omni.ext.mobility_gen.utils.global_utils import save_stage from omni.ext.mobility_gen.writer import Writer from omni.ext.mobility_gen.inputs import GamepadDriver, KeyboardDriver from omni.ext.mobility_gen.scenarios import SCENARIOS, Scenario -from omni.ext.mobility_gen.utils.global_utils import get_world +from omni.ext.mobility_gen.utils.global_utils import get_world, new_world, set_viewport_camera from omni.ext.mobility_gen.robots import ROBOTS from omni.ext.mobility_gen.config import Config -from omni.ext.mobility_gen.build import build_scenario_from_config +from omni.ext.mobility_gen.occupancy_map import OccupancyMap if "MOBILITY_GEN_DATA" in os.environ: @@ -77,19 +90,26 @@ def on_startup(self, ext_id): with ui.VStack(): with ui.VStack(): with ui.HStack(): - ui.Label("USD Path / URL") + ui.Label("Stage") self.scene_usd_field_string_model = ui.SimpleStringModel() self.scene_usd_field = ui.StringField(model=self.scene_usd_field_string_model, height=25) with ui.HStack(): - ui.Label("Scenario Type") - self.scenario_combo_box = ui.ComboBox(0, *SCENARIOS.names()) + ui.Label("Occupancy Map") + self.omap_field_string_model = ui.SimpleStringModel() + self.omap_field = ui.StringField(model=self.omap_field_string_model, height=25) with ui.HStack(): ui.Label("Robot Type") self.robot_combo_box = ui.ComboBox(0, *ROBOTS.names()) - + + with ui.HStack(): + ui.Label("Scenario Type") + self.scenario_combo_box = ui.ComboBox(0, *SCENARIOS.names()) + + ui.Button("Build", clicked_fn=self.build_scenario) + # ui.Button("Build", clicked_fn=self.build_scenario) with ui.VStack(): self.recording_count_label = ui.Label("") @@ -112,9 +132,10 @@ def build_occ_map_frame(self): self._occupancy_map_image_provider ) - def draw_occ_map(self): + def draw_visualization_image(self): if self.scenario is not None: - image = self.scenario.occupancy_map.ros_image().copy().convert("RGBA") + image = self.scenario.get_visualization_image().copy().convert("RGBA") + carb.log_warn(image) data = list(image.tobytes()) self._occupancy_map_image_provider.set_bytes_data(data, [image.width, image.height]) self._occ_map_frame.rebuild() @@ -140,7 +161,8 @@ def on_shutdown(self): self.keyboard.disconnect() self.gamepad.disconnect() world = get_world() - world.remove_physics_callback("scenario_physics", self.on_physics) + if world is not None: + world.remove_physics_callback("scenario_physics") def start_new_recording(self): recording_name = datetime.datetime.now().isoformat() @@ -180,6 +202,7 @@ def reset(self): self.scenario.reset() if self.recording_enabled: self.start_new_recording() + self.draw_visualization_image() def on_physics(self, step_size: int): @@ -198,34 +221,63 @@ def on_physics(self, step_size: int): if self.step % 15 == 0: self.recording_step_label.text = f"Current recording duration: {self.recording_time:.2f}s" + def build_scenario(self): async def _build_scenario_async(): self.clear_recording() self.clear_scenario() + self.disable_recording() + + # Get parameters from UI + scenario_type_str = list(SCENARIOS.names())[self.scenario_combo_box.model.get_item_value_model().get_value_as_int()] + robot_type_str = list(ROBOTS.names())[self.robot_combo_box.model.get_item_value_model().get_value_as_int()] + scene_usd_str = self.scene_usd_field_string_model.as_string + + robot_type = ROBOTS.get(robot_type_str) + scenario_type = SCENARIOS.get(scenario_type_str) + + # Set config + self.config = Config( + scenario_type=scenario_type_str, + robot_type=robot_type_str, + scene_usd=scene_usd_str + ) + + occupancy_map = OccupancyMap.from_ros_yaml( + os.path.expanduser(self.omap_field_string_model.as_string)) - config = self.create_config() + # Open stage and save local copy + open_stage(scene_usd_str) + self.cached_stage_path = os.path.join(tempfile.mkdtemp(), "stage.usd") + save_stage(self.cached_stage_path) - self.config = config - self.scenario = await build_scenario_from_config(config) + # Initialize world + world = new_world(physics_dt=robot_type.physics_dt) + await world.initialize_simulation_context_async() - self.draw_occ_map() + # Add ground plane + objects.GroundPlane("/World/ground_plane", visible=False) - world = get_world() - await world.reset_async() + # Add robot + robot = robot_type.build("/World/robot") + + # Set the chase camera + chase_camera_path = robot.build_chase_camera() + set_viewport_camera(chase_camera_path) + + # Set the scenario + self.scenario = scenario_type.from_robot_occupancy_map(robot, occupancy_map) - self.scenario.reset() + # Draw the occupancy map + self.draw_visualization_image() + # Run the scenario + await world.reset_async() world.add_physics_callback("scenario_physics", self.on_physics) - # cache stage - self.cached_stage_path = os.path.join(tempfile.mkdtemp(), "stage.usd") - save_stage(self.cached_stage_path) - - if self.recording_enabled: - self.start_new_recording() + self.reset() - # self.scenario.save(path) asyncio.ensure_future(_build_scenario_async()) \ No newline at end of file diff --git a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/robots.py b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/robots.py index 2dd2355..57f02ca 100644 --- a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/robots.py +++ b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/robots.py @@ -139,6 +139,8 @@ def __init__(self, self.orientation = Buffer() self.joint_positions = Buffer() self.joint_velocities = Buffer() + self.linear_velocity = Buffer() + self.angular_velocity = Buffer() self.front_camera = front_camera @classmethod @@ -182,11 +184,13 @@ def write_action(self, step_size: float): raise NotImplementedError def update_state(self): - pos, ori = self.robot.get_local_pose() + pos, ori = self.robot.get_world_pose() self.position.set_value(pos) self.orientation.set_value(ori) self.joint_positions.set_value(self.robot.get_joint_positions()) self.joint_velocities.set_value(self.robot.get_joint_velocities()) + self.linear_velocity.set_value(self.robot.get_linear_velocity()) + self.angular_velocity.set_value(self.robot.get_angular_velocity()) super().update_state() def write_replay_data(self): @@ -203,7 +207,7 @@ def set_pose_2d(self, pose: Pose2d): self.articulation_view.initialize() self.robot.set_world_velocity(np.array([0., 0., 0., 0., 0., 0.])) self.robot.post_reset() - position, orientation = self.robot.get_local_pose() + position, orientation = self.robot.get_world_pose() position[0] = pose.x position[1] = pose.y position[2] = self.z_offset @@ -213,7 +217,7 @@ def set_pose_2d(self, pose: Pose2d): ) def get_pose_2d(self) -> Pose2d: - position, orientation = self.robot.get_local_pose() + position, orientation = self.robot.get_world_pose() theta = rot_utils.quats_to_euler_angles(orientation)[2] return Pose2d( x=position[0], diff --git a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/scenarios.py b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/scenarios.py index a4d4722..a7732ec 100644 --- a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/scenarios.py +++ b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/scenarios.py @@ -15,9 +15,11 @@ import numpy as np +import PIL.Image +from PIL import Image, ImageDraw from typing import Tuple -from mobility_gen_path_planner import generate_paths +from mobility_gen_path_planner import generate_paths, compress_path from omni.ext.mobility_gen.utils.path_utils import PathHelper, vector_angle from omni.ext.mobility_gen.utils.registry import Registry @@ -49,6 +51,9 @@ def reset(self): def step(self, step_size: float) -> bool: raise NotImplementedError + def get_visualization_image(self): + image = self.occupancy_map.ros_image() + return image SCENARIOS = Registry[Scenario]() @@ -67,6 +72,7 @@ def __init__(self, def reset(self): pose = self.pose_sampler.sample(self.buffered_occupancy_map) self.robot.set_pose_2d(pose) + self.update_state() def step(self, step_size): @@ -105,6 +111,7 @@ def __init__(self, def reset(self): pose = self.pose_sampler.sample(self.buffered_occupancy_map) self.robot.set_pose_2d(pose) + self.update_state() def step(self, step_size: float): @@ -139,6 +146,7 @@ def reset(self): pose = self.pose_sampler.sample(self.buffered_occupancy_map) self.robot.set_pose_2d(pose) self.is_alive = True + self.update_state() def step(self, step_size: float): @@ -193,6 +201,7 @@ def set_random_target_path(self): output = generate_paths(start, freespace) end = output.sample_random_end_point() path = output.unroll_path(end) + path, _ = compress_path(path) # remove redundant points path = path[:, ::-1] # y,x -> x,y coordinates path = self.occupancy_map.pixel_to_world_numpy(path) self.target_path.set_value(path) @@ -204,6 +213,7 @@ def reset(self): self.robot.set_pose_2d(pose) self.set_random_target_path() self.is_alive = True + self.update_state() def step(self, step_size: float): @@ -228,7 +238,7 @@ def step(self, step_size: float): dist_to_target = np.sqrt(np.sum((pt_robot - path_end)**2)) if dist_to_target < self.robot.path_following_stop_distance_threshold: - self.set_random_target_path() + self.is_alive = False else: vec_robot_unit = np.array([np.cos(current_pose.theta), np.sin(current_pose.theta)]) vec_target = (pt_target - pt_robot) @@ -248,3 +258,17 @@ def step(self, step_size: float): return self.is_alive + def get_visualization_image(self): + image = self.occupancy_map.ros_image().copy().convert("RGBA") + draw = ImageDraw.Draw(image) + path = self.target_path.get_value() + if path is not None: + line_coordinates = [] + path_pixels = self.occupancy_map.world_to_pixel_numpy(path) + for i in range(len(path_pixels)): + line_coordinates.append(int(path_pixels[i, 0])) + line_coordinates.append(int(path_pixels[i, 1])) + width_pixels = self.robot.occupancy_map_radius / self.occupancy_map.resolution + draw.line(line_coordinates, fill="green", width=int(width_pixels/2), joint="curve") + + return image \ No newline at end of file diff --git a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/sensors.py b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/sensors.py index 6a8a80e..dd91fb7 100644 --- a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/sensors.py +++ b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/sensors.py @@ -19,11 +19,11 @@ import omni.replicator.core as rep -from isaacsim.core.prims import SingleXFormPrim as XFormPrim from omni.ext.mobility_gen.utils.global_utils import get_stage -from omni.ext.mobility_gen.utils.stage_utils import stage_add_usd_ref +from omni.ext.mobility_gen.utils.stage_utils import stage_add_usd_ref, stage_get_prim +from omni.ext.mobility_gen.utils.prim_utils import prim_get_world_transform from omni.ext.mobility_gen.common import Module, Buffer @@ -51,7 +51,7 @@ def __init__(self, self._instance_id_segmentation_annotator = None self._normals_annotator = None self._depth_annotator = None - self._xform_prim = XFormPrim(self._prim_path) + self._prim = stage_get_prim(get_stage(), self._prim_path) self.rgb_image = Buffer(tags=["rgb"]) self.segmentation_image = Buffer(tags=["segmentation"]) @@ -166,7 +166,7 @@ def update_state(self): data = self._normals_annotator.get_data() self.normals_image.set_value(data) - position, orientation = self._xform_prim.get_world_pose() + position, orientation = prim_get_world_transform(self._prim) self.position.set_value(position) self.orientation.set_value(orientation) diff --git a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/utils/occupancy_map_utils.py b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/utils/occupancy_map_utils.py index 770cbe8..04191dc 100644 --- a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/utils/occupancy_map_utils.py +++ b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/utils/occupancy_map_utils.py @@ -64,11 +64,35 @@ def degrees(self): +def compute_occupancy_bounds_from_prim_path( + prim_path: str, + z_min: float, + z_max: float, + cell_size: float + ): + + stage = get_stage() + + prim_path = os.path.join(prim_path) + prim_path = stage.GetPrimAtPath(prim_path) + + lower_bound, upper_bound, midpoint = prim_compute_bbox(prim_path) + lower_bound = (lower_bound[0], lower_bound[1], z_min) + upper_bound = (upper_bound[0], upper_bound[1], z_max) + + width = upper_bound[0] - lower_bound[0] + height = upper_bound[1] - lower_bound[1] + origin = (lower_bound[0] - cell_size, lower_bound[1] - cell_size, 0) + lower_bound = (0, 0, z_min) + upper_bound = (width + cell_size, height + cell_size, z_max) + return origin, lower_bound, upper_bound + + async def occupancy_map_generate_from_prim_async( - prim_path: str, + origin: tp.Tuple[float, float, float], + lower_bound: tp.Tuple[float, float, float], + upper_bound: tp.Tuple[float, float, float], cell_size: float = OCCUPANCY_MAP_DEFAULT_CELL_SIZE, - z_min: float = OCCUPANCY_MAP_DEFAULT_Z_MIN, - z_max: float = OCCUPANCY_MAP_DEFAULT_Z_MAX, rotation: OccupancyMapGenerateRotation = OccupancyMapGenerateRotation.ROTATE_180, unknown_as_freespace: bool = True ) -> OccupancyMap: @@ -94,18 +118,6 @@ async def occupancy_map_generate_from_prim_async( # Compute bounds for occupancy map calculation using specified prim - prim_path = os.path.join(prim_path) - prim_path = stage.GetPrimAtPath(prim_path) - - lower_bound, upper_bound, midpoint = prim_compute_bbox(prim_path) - lower_bound = (lower_bound[0], lower_bound[1], z_min) - upper_bound = (upper_bound[0], upper_bound[1], z_max) - - width = upper_bound[0] - lower_bound[0] - height = upper_bound[1] - lower_bound[1] - origin = (lower_bound[0] - cell_size, lower_bound[1] - cell_size, 0) - lower_bound = (0, 0, z_min) - upper_bound = (width + cell_size, height + cell_size, z_max) update_location( om, diff --git a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/utils/prim_utils.py b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/utils/prim_utils.py index 29cc7c6..c9748a8 100644 --- a/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/utils/prim_utils.py +++ b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/utils/prim_utils.py @@ -270,3 +270,56 @@ def prim_rotate_z(prim: Usd.Prim, angle: float): prim_xform_op_move_end_to_front(prim) return prim + +def _translation_to_np(t: Gf.Vec3d): + return np.array(t) + + +def _rotation_to_np_quat(r: Gf.Rotation): + quat = r.GetQuaternion() + real = quat.GetReal() + imag: Gf.Vec3d = quat.GetImaginary() + return np.array([real, imag[0], imag[1], imag[2]]) + + +def prim_get_local_transform(prim: Usd.Prim) -> Tuple[np.ndarray, np.ndarray]: + """ + From: https://docs.omniverse.nvidia.com/dev-guide/latest/programmer_ref/usd/transforms/get-local-transforms.html + + Get the local transformation of a prim using Xformable. + See https://openusd.org/release/api/class_usd_geom_xformable.html + Args: + prim: The prim to calculate the local transformation. + Returns: + A tuple of: + - Translation vector. + - Rotation quaternion, i.e. 3d vector plus angle. + - Scale vector. + """ + xform = UsdGeom.Xformable(prim) + local_transformation: Gf.Matrix4d = xform.GetLocalTransformation() + translation: Gf.Vec3d = local_transformation.ExtractTranslation() + rotation: Gf.Rotation = local_transformation.ExtractRotation() + return _translation_to_np(translation), _rotation_to_np_quat(rotation) + + +def prim_get_world_transform(prim: Usd.Prim) -> Tuple[np.ndarray, np.ndarray]: + """ + From: https://docs.omniverse.nvidia.com/dev-guide/latest/programmer_ref/usd/transforms/get-world-transforms.html + + Get the local transformation of a prim using Xformable. + See https://openusd.org/release/api/class_usd_geom_xformable.html + Args: + prim: The prim to calculate the world transformation. + Returns: + A tuple of: + - Translation vector. + - Rotation quaternion, i.e. 3d vector plus angle. + - Scale vector. + """ + xform = UsdGeom.Xformable(prim) + time = Usd.TimeCode.Default() # The time at which we compute the bounding box + world_transform: Gf.Matrix4d = xform.ComputeLocalToWorldTransform(time) + translation: Gf.Vec3d = world_transform.ExtractTranslation() + rotation: Gf.Rotation = world_transform.ExtractRotation() + return _translation_to_np(translation), _rotation_to_np_quat(rotation) diff --git a/path_planner/compress.py b/path_planner/compress.py new file mode 100644 index 0000000..11bb780 --- /dev/null +++ b/path_planner/compress.py @@ -0,0 +1,16 @@ +import numpy as np + +from mobility_gen_path_planner import generate_paths, compress_path + +start = (0, 0) +freespace = np.ones((40, 40), dtype=bool) + +output = generate_paths(start, freespace) + +path = output.unroll_path(end=(20, 39)) + +cpath, ckeep = compress_path(path) + +print(path) +print(cpath) +print(ckeep) \ No newline at end of file diff --git a/path_planner/mobility_gen_path_planner/__init__.py b/path_planner/mobility_gen_path_planner/__init__.py index e0a92cf..8ee0a3a 100644 --- a/path_planner/mobility_gen_path_planner/__init__.py +++ b/path_planner/mobility_gen_path_planner/__init__.py @@ -68,4 +68,18 @@ def generate_paths(start: Tuple[int, int], freespace: np.ndarray) -> GeneratePat distance_to_start=distance_to_start, prev_i=prev_i, prev_j=prev_j - ) \ No newline at end of file + ) + +def compress_path(path: np.ndarray, eps=1e-3): + pref = path[1:-1] + pnext = path[2:] + pprev = path[:-2] + + vnext = pnext - pref + vprev = pref - pprev + + keepmask = np.ones((path.shape[0],), dtype=bool) # keep beginning / end by default + keepmask[1:-1] = np.sum((vnext - vprev)**2, axis=-1) > eps + + + return path[keepmask], keepmask \ No newline at end of file diff --git a/scripts/replay.py b/scripts/replay.py index 5c9b143..daf91e5 100644 --- a/scripts/replay.py +++ b/scripts/replay.py @@ -38,7 +38,7 @@ parser.add_argument("--rgb_enabled", type=bool, default=True) parser.add_argument("--segmentation_enabled", type=bool, default=True) parser.add_argument("--depth_enabled", type=bool, default=True) - parser.add_argument("--instance_id_segmentation_enabled", type=bool, default=True) + parser.add_argument("--instance_id_segmentation_enabled", type=bool, default=False) parser.add_argument("--normals_enabled", type=bool, default=False) parser.add_argument("--render_rt_subframes", type=int, default=1) parser.add_argument("--render_interval", type=int, default=1) diff --git a/scripts/replay_directory.py b/scripts/replay_directory.py index 8ca6956..7076fab 100644 --- a/scripts/replay_directory.py +++ b/scripts/replay_directory.py @@ -38,7 +38,7 @@ parser.add_argument("--rgb_enabled", type=bool, default=True) parser.add_argument("--segmentation_enabled", type=bool, default=True) parser.add_argument("--depth_enabled", type=bool, default=True) - parser.add_argument("--instance_id_segmentation_enabled", type=bool, default=True) + parser.add_argument("--instance_id_segmentation_enabled", type=bool, default=False) parser.add_argument("--normals_enabled", type=bool, default=False) parser.add_argument("--render_rt_subframes", type=int, default=1) parser.add_argument("--render_interval", type=int, default=1) @@ -50,38 +50,19 @@ if args.output is None: args.output = os.path.join(DATA_DIR, "replays") - args.input = os.path.expanduser(args.input) - args.output = os.path.expanduser(args.output) - - recording_paths = glob.glob(os.path.join(args.input, "*")) - - for recording_path in recording_paths: - - name = os.path.basename(recording_path) - - output_path = os.path.join(args.output, name) - - input_steps = len(glob.glob(os.path.join(recording_path, "state", "common", "*.npy"))) - output_steps = len(glob.glob(os.path.join(output_path, "state", "common", "*.npy"))) - - if input_steps == output_steps: - print(f"Skipping {name} as it is already replayed") - else: - print(f"Replaying {name}") - - subprocess.call([ - "./app/python.sh", - "scripts/replay_implementation.py", - "--ext-folder", "exts", - "--enable", "omni.ext.mobility_gen", - "--enable", "isaacsim.asset.gen.omap", - "--input_path", recording_path, - "--output_path", output_path, - "--render_interval", str(args.render_interval), - "--render_rt_subframes", str(args.render_rt_subframes), - "--rgb_enabled", str(args.rgb_enabled), - "--segmentation_enabled", str(args.segmentation_enabled), - "--instance_id_segmentation_enabled", str(args.instance_id_segmentation_enabled), - "--normals_enabled", str(args.normals_enabled), - "--depth_enabled", str(args.depth_enabled) - ]) \ No newline at end of file + subprocess.call([ + "./app/python.sh", + "scripts/replay_directory_implementation.py", + "--ext-folder", "exts", + "--enable", "omni.ext.mobility_gen", + "--enable", "isaacsim.asset.gen.omap", + "--input", args.input, + "--output", args.output, + "--render_interval", str(args.render_interval), + "--render_rt_subframes", str(args.render_rt_subframes), + "--rgb_enabled", str(args.rgb_enabled), + "--segmentation_enabled", str(args.segmentation_enabled), + "--instance_id_segmentation_enabled", str(args.instance_id_segmentation_enabled), + "--normals_enabled", str(args.normals_enabled), + "--depth_enabled", str(args.depth_enabled) + ]) \ No newline at end of file diff --git a/scripts/replay_directory_implementation.py b/scripts/replay_directory_implementation.py new file mode 100644 index 0000000..85df56a --- /dev/null +++ b/scripts/replay_directory_implementation.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" + +This script launches a simulation app for replaying and rendering +a recording. + +""" + +from isaacsim import SimulationApp + +simulation_app = SimulationApp(launch_config={"headless": True}) + +import argparse +import os +import shutil +import numpy as np +from PIL import Image +import glob +import tqdm +import time + +import omni.replicator.core as rep + +from omni.ext.mobility_gen.utils.global_utils import get_world +from omni.ext.mobility_gen.writer import Writer +from omni.ext.mobility_gen.reader import Reader +from omni.ext.mobility_gen.build import load_scenario + +if "MOBILITY_GEN_DATA" in os.environ: + DATA_DIR = os.environ['MOBILITY_GEN_DATA'] +else: + DATA_DIR = os.path.expanduser("~/MobilityGenData") + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + parser.add_argument("--input", type=str, default=None) + parser.add_argument("--output", type=str, default=None) + parser.add_argument("--rgb_enabled", type=bool, default=True) + parser.add_argument("--segmentation_enabled", type=bool, default=True) + parser.add_argument("--depth_enabled", type=bool, default=True) + parser.add_argument("--instance_id_segmentation_enabled", type=bool, default=False) + parser.add_argument("--normals_enabled", type=bool, default=False) + parser.add_argument("--render_rt_subframes", type=int, default=1) + parser.add_argument("--render_interval", type=int, default=1) + args, unknown = parser.parse_known_args() + + if args.input is None: + args.input = os.path.join(DATA_DIR, "recordings") + + if args.output is None: + args.output = os.path.join(DATA_DIR, "replays") + + args.input = os.path.expanduser(args.input) + args.output = os.path.expanduser(args.output) + + recording_paths = glob.glob(os.path.join(args.input, "*")) + + recording_count = 0 + for recording_path in recording_paths: + recording_count += 1 + name = os.path.basename(recording_path) + + output_path = os.path.join(args.output, name) + + scenario = load_scenario(recording_path) + + world = get_world() + world.reset() + + print(scenario) + + if args.rgb_enabled: + scenario.enable_rgb_rendering() + + if args.segmentation_enabled: + scenario.enable_segmentation_rendering() + + if args.depth_enabled: + scenario.enable_depth_rendering() + + # if args.instance_id_segmentation_enabled: + # scenario.enable_instance_id_segmentation_rendering() + + # if args.normals_enabled: + # scenario.enable_normals_rendering() + + simulation_app.update() + rep.orchestrator.step( + rt_subframes=args.render_rt_subframes, + delta_time=0.0, + pause_timeline=False + ) + + reader = Reader(recording_path) + num_steps = len(reader) + + writer = Writer(output_path) + writer.copy_init(recording_path) + + + print(f"============== Replaying {recording_count} / {len(recording_paths)}==============") + print(f"\tInput path: {recording_path}") + print(f"\tOutput path: {output_path}") + print(f"\tRgb enabled: {args.rgb_enabled}") + print(f"\tSegmentation enabled: {args.segmentation_enabled}") + print(f"\tRendering RT subframes: {args.render_rt_subframes}") + print(f"\tRender interval: {args.render_interval}") + + t0 = time.perf_counter() + count = 0 + for step in tqdm.tqdm(range(0, num_steps, args.render_interval)): + + state_dict_original = reader.read_state_dict(index=step) + + scenario.load_state_dict(state_dict_original) + scenario.write_replay_data() + + simulation_app.update() + + rep.orchestrator.step( + rt_subframes=args.render_rt_subframes, + delta_time=0.00, + pause_timeline=False + ) + + scenario.update_state() + + state_dict = scenario.state_dict_common() + + for k, v in state_dict_original.items(): + # overwrite with original state, to ensure physics based values are correct + if v is not None: + state_dict[k] = v # don't overwrite "None" values + + state_rgb = scenario.state_dict_rgb() + state_segmentation = scenario.state_dict_segmentation() + state_depth = scenario.state_dict_depth() + state_normals = scenario.state_dict_normals() + + writer.write_state_dict_common(state_dict, step) + writer.write_state_dict_rgb(state_rgb, step) + writer.write_state_dict_segmentation(state_segmentation, step) + writer.write_state_dict_depth(state_depth, step) + writer.write_state_dict_normals(state_normals, step) + + count += 1 + t1 = time.perf_counter() + + print(f"Process time per frame: {count / (t1 - t0)}") + + simulation_app.close() \ No newline at end of file diff --git a/scripts/replay_implementation.py b/scripts/replay_implementation.py index c549156..9d51fb2 100644 --- a/scripts/replay_implementation.py +++ b/scripts/replay_implementation.py @@ -31,6 +31,7 @@ from PIL import Image import glob import tqdm +import time import omni.replicator.core as rep @@ -48,7 +49,7 @@ parser.add_argument("--rgb_enabled", type=bool, default=True) parser.add_argument("--segmentation_enabled", type=bool, default=True) parser.add_argument("--depth_enabled", type=bool, default=True) - parser.add_argument("--instance_id_segmentation_enabled", type=bool, default=True) + parser.add_argument("--instance_id_segmentation_enabled", type=bool, default=False) parser.add_argument("--normals_enabled", type=bool, default=False) parser.add_argument("--render_rt_subframes", type=int, default=1) parser.add_argument("--render_interval", type=int, default=1) @@ -99,11 +100,13 @@ print(f"\tRendering RT subframes: {args.render_rt_subframes}") print(f"\tRender interval: {args.render_interval}") + t0 = time.perf_counter() + count = 0 for step in tqdm.tqdm(range(0, num_steps, args.render_interval)): - state_dict = reader.read_state_dict(index=step) + state_dict_original = reader.read_state_dict(index=step) - scenario.load_state_dict(state_dict) + scenario.load_state_dict(state_dict_original) scenario.write_replay_data() simulation_app.update() @@ -117,6 +120,12 @@ scenario.update_state() state_dict = scenario.state_dict_common() + + for k, v in state_dict_original.items(): + # overwrite with original state, to ensure physics based values are correct + if v is not None: + state_dict[k] = v # don't overwrite "None" values + state_rgb = scenario.state_dict_rgb() state_segmentation = scenario.state_dict_segmentation() state_depth = scenario.state_dict_depth() @@ -128,4 +137,9 @@ writer.write_state_dict_depth(state_depth, step) writer.write_state_dict_normals(state_normals, step) + count += 1 + t1 = time.perf_counter() + + print(f"Process time per frame: {count / (t1 - t0)}") + simulation_app.close() \ No newline at end of file