From ade8ec8afbe718fe12af7003f1d236f5c0edd399 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Mon, 17 Mar 2025 15:28:47 -0700 Subject: [PATCH 01/29] add parquet export example --- examples/05_convert_to_parquet.py | 74 +++++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 examples/05_convert_to_parquet.py diff --git a/examples/05_convert_to_parquet.py b/examples/05_convert_to_parquet.py new file mode 100644 index 0000000..b5a7455 --- /dev/null +++ b/examples/05_convert_to_parquet.py @@ -0,0 +1,74 @@ +import argparse +import pandas +from reader import Reader +import numpy as np +import tqdm +import PIL.Image +import io + +parser = argparse.ArgumentParser() +parser.add_argument("recording_path") +parser.add_argument("output_path") +args = parser.parse_args() + +reader = Reader(recording_path=args.recording_path) + +index = 0 + + +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 + + +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)) + + 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(args.output_path, engine="pyarrow") \ No newline at end of file From c421cd851e4bdf86b3a76217af61ebb6d458438a Mon Sep 17 00:00:00 2001 From: John Welsh Date: Tue, 18 Mar 2025 12:33:29 -0700 Subject: [PATCH 02/29] ensure replay uses originally recorded velocities --- scripts/replay.py | 2 +- scripts/replay_directory.py | 2 +- scripts/replay_implementation.py | 9 ++++++--- 3 files changed, 8 insertions(+), 5 deletions(-) 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..d3f8e11 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) diff --git a/scripts/replay_implementation.py b/scripts/replay_implementation.py index c549156..1508245 100644 --- a/scripts/replay_implementation.py +++ b/scripts/replay_implementation.py @@ -48,7 +48,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) @@ -101,9 +101,9 @@ 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 +117,9 @@ scenario.update_state() state_dict = scenario.state_dict_common() + + state_dict.update(state_dict_original) # overwrite with original state, to ensure physics based values are correct + state_rgb = scenario.state_dict_rgb() state_segmentation = scenario.state_dict_segmentation() state_depth = scenario.state_dict_depth() From ea9549286a94e928652eaab74cdc97fa5dd927a9 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Tue, 18 Mar 2025 12:33:42 -0700 Subject: [PATCH 03/29] add world linear/angular velocity to robot state --- exts/omni.ext.mobility_gen/omni/ext/mobility_gen/robots.py | 4 ++++ 1 file changed, 4 insertions(+) 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..fd822c9 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 @@ -187,6 +189,8 @@ def update_state(self): 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): From 6fbed37f0ff42dfeebcf4ca2cb68f003b8452365 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Tue, 18 Mar 2025 12:34:08 -0700 Subject: [PATCH 04/29] add parquet export example --- examples/05_convert_to_parquet.py | 1 + examples/06_load_parquet.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+) create mode 100644 examples/06_load_parquet.py diff --git a/examples/05_convert_to_parquet.py b/examples/05_convert_to_parquet.py index b5a7455..4a49bcf 100644 --- a/examples/05_convert_to_parquet.py +++ b/examples/05_convert_to_parquet.py @@ -46,6 +46,7 @@ def numpy_array_to_jpg_columns(key: str, value: np.ndarray): 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)) + # TODO: handle normals for k, v in state_common.items(): if isinstance(v, np.ndarray): diff --git a/examples/06_load_parquet.py b/examples/06_load_parquet.py new file mode 100644 index 0000000..ab54e30 --- /dev/null +++ b/examples/06_load_parquet.py @@ -0,0 +1,18 @@ +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", columns=["robot.linear_velocity"]) + + +vel = np.stack(data['robot.linear_velocity'].to_numpy()) + + +plt.plot(vel[:, 0], 'r-') +plt.plot(vel[:, 1], 'r-') +plt.show() From 5b4ecf861698a0bb6d214dd1791fdabaa0b700e8 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Tue, 18 Mar 2025 12:45:47 -0700 Subject: [PATCH 05/29] load parquet info --- examples/06_load_parquet.py | 3 ++- examples/PARQUET_DATA_FORMAT.md | 28 ++++++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) create mode 100644 examples/PARQUET_DATA_FORMAT.md diff --git a/examples/06_load_parquet.py b/examples/06_load_parquet.py index ab54e30..ac96b96 100644 --- a/examples/06_load_parquet.py +++ b/examples/06_load_parquet.py @@ -7,9 +7,10 @@ parser.add_argument("parquet_path") args = parser.parse_args() -data = pandas.read_parquet(args.parquet_path, engine="pyarrow", columns=["robot.linear_velocity"]) +data = pandas.read_parquet(args.parquet_path, engine="pyarrow") +print(data.columns) vel = np.stack(data['robot.linear_velocity'].to_numpy()) 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 From b50f84a895ab7634d7853f6fa8f49300043fd18e Mon Sep 17 00:00:00 2001 From: John Welsh Date: Tue, 25 Mar 2025 14:11:13 -0700 Subject: [PATCH 06/29] fix segmentation overwrite issue --- scripts/replay_implementation.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/scripts/replay_implementation.py b/scripts/replay_implementation.py index 1508245..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 @@ -99,6 +100,8 @@ 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) @@ -118,7 +121,10 @@ state_dict = scenario.state_dict_common() - state_dict.update(state_dict_original) # overwrite with original state, to ensure physics based values are correct + 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() @@ -131,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 From 2966cd40ac1a824f396531a314f8ee64bf18552c Mon Sep 17 00:00:00 2001 From: John Welsh Date: Mon, 31 Mar 2025 11:35:57 -0700 Subject: [PATCH 07/29] update changelog and readme --- CHANGELOG.md | 3 +++ README.md | 2 ++ 2 files changed, 5 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a597fb1..25f726f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,9 @@ # main +- 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..f29d2d6 100644 --- a/README.md +++ b/README.md @@ -399,6 +399,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 From 58423988e37742ccff20a36c8200e6ed84948f5f Mon Sep 17 00:00:00 2001 From: John Welsh Date: Mon, 31 Mar 2025 13:29:36 -0700 Subject: [PATCH 08/29] add occupancy map to cfg to support override --- .../omni/ext/mobility_gen/build.py | 42 ++++++++++++++++--- .../omni/ext/mobility_gen/config.py | 17 +++++++- .../mobility_gen/utils/occupancy_map_utils.py | 42 ++++++++++++------- 3 files changed, 78 insertions(+), 23 deletions(-) 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..9e65a01 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 @@ -22,8 +22,8 @@ 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.config import Config, OccupancyMapConfig +from omni.ext.mobility_gen.utils.occupancy_map_utils import occupancy_map_generate_from_prim_async, compute_occupancy_bounds_from_prim_path from omni.ext.mobility_gen.utils.global_utils import new_stage, new_world, set_viewport_camera from omni.ext.mobility_gen.scenarios import Scenario, SCENARIOS from omni.ext.mobility_gen.robots import ROBOTS @@ -60,12 +60,42 @@ async def build_scenario_from_config(config: Config): add_reference_to_stage(config.scene_usd,"/World/scene") objects.GroundPlane("/World/ground_plane", visible=False) robot = robot_type.build("/World/robot") + + + if config.occupancy_map_config is None: + config.occupancy_map_config = OccupancyMapConfig() + + if config.occupancy_map_config.z_min is None: + config.occupancy_map_config.z_min = robot.occupancy_map_z_min + + if config.occupancy_map_config.z_max is None: + config.occupancy_map_config.z_max = robot.occupancy_map_z_max + + if config.occupancy_map_config.cell_size is None: + config.occupancy_map_config.cell_size = robot.occupancy_map_cell_size + + if config.occupancy_map_config.origin is None \ + or config.occupancy_map_config.lower_bound is None \ + or config.occupancy_map_config.upper_bound is None: + + origin, lower_bound, upper_bound = compute_occupancy_bounds_from_prim_path( + prim_path="/World/scene", + z_min=config.occupancy_map_config.z_min, + z_max=config.occupancy_map_config.z_max, + cell_size=config.occupancy_map_config.cell_size + ) + config.occupancy_map_config.origin = origin + config.occupancy_map_config.lower_bound = lower_bound + config.occupancy_map_config.upper_bound = upper_bound + + 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 + origin=config.occupancy_map_config.origin, + lower_bound=config.occupancy_map_config.lower_bound, + upper_bound=config.occupancy_map_config.upper_bound, + cell_size=config.occupancy_map_config.cell_size ) + chase_camera_path = robot.build_chase_camera() set_viewport_camera(chase_camera_path) scenario = scenario_type.from_robot_occupancy_map(robot, occupancy_map) 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..8afb805 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,17 +1,30 @@ import json -from typing import Literal +from typing import Literal, Optional, Tuple from dataclasses import dataclass, asdict +@dataclass +class OccupancyMapConfig: + z_min: Optional[float] = None + z_max: Optional[float] = None + cell_size: Optional[float] = None + origin: Optional[Tuple[float, float, float]] = None + lower_bound: Optional[Tuple[float, float, float]] = None + upper_bound: Optional[Tuple[float, float, float]] = None + + @dataclass class Config: scenario_type: str robot_type: str scene_usd: str + occupancy_map_config: Optional[OccupancyMapConfig] = None def to_json(self): return json.dumps(asdict(self), indent=2) @staticmethod def from_json(data: str): - return Config(**json.loads(data)) \ No newline at end of file + data = json.loads(data) + data['occupancy_map_config'] = OccupancyMapConfig(**data['occupancy_map_config']) + return Config(**data) \ No newline at end of file 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, From 7b200a9d60e9f11be2ac3421d87f25d08f049ce2 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Mon, 31 Mar 2025 15:04:58 -0700 Subject: [PATCH 09/29] add manual occ bounds to UI --- .../omni/ext/mobility_gen/build.py | 25 ++----- .../omni/ext/mobility_gen/config.py | 3 +- .../omni/ext/mobility_gen/extension.py | 68 +++++++++++++++++-- 3 files changed, 70 insertions(+), 26 deletions(-) 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 9e65a01..981b587 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 @@ -61,27 +61,12 @@ async def build_scenario_from_config(config: Config): objects.GroundPlane("/World/ground_plane", visible=False) robot = robot_type.build("/World/robot") - - if config.occupancy_map_config is None: - config.occupancy_map_config = OccupancyMapConfig() - - if config.occupancy_map_config.z_min is None: - config.occupancy_map_config.z_min = robot.occupancy_map_z_min - - if config.occupancy_map_config.z_max is None: - config.occupancy_map_config.z_max = robot.occupancy_map_z_max - - if config.occupancy_map_config.cell_size is None: - config.occupancy_map_config.cell_size = robot.occupancy_map_cell_size - - if config.occupancy_map_config.origin is None \ - or config.occupancy_map_config.lower_bound is None \ - or config.occupancy_map_config.upper_bound is None: + if config.occupancy_map_config.prim_path is not None: origin, lower_bound, upper_bound = compute_occupancy_bounds_from_prim_path( - prim_path="/World/scene", - z_min=config.occupancy_map_config.z_min, - z_max=config.occupancy_map_config.z_max, + prim_path=config.occupancy_map_config.prim_path, + z_min=config.occupancy_map_config.lower_bound[2], + z_max=config.occupancy_map_config.upper_bound[2], cell_size=config.occupancy_map_config.cell_size ) config.occupancy_map_config.origin = origin @@ -99,4 +84,4 @@ async def build_scenario_from_config(config: Config): 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 + return scenario, config \ 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 8afb805..712b738 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 @@ -5,8 +5,7 @@ @dataclass class OccupancyMapConfig: - z_min: Optional[float] = None - z_max: Optional[float] = None + prim_path: Optional[str] = None cell_size: Optional[float] = None origin: Optional[Tuple[float, float, float]] = None lower_bound: Optional[Tuple[float, float, float]] = None 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..6df0da1 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,13 +24,23 @@ import omni.ext import omni.ui as ui +from isaacsim.gui.components.ui_utils import ( + btn_builder, + cb_builder, + color_picker_builder, + dropdown_builder, + float_builder, + multi_btn_builder, + xyz_builder, +) + 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.robots import ROBOTS -from omni.ext.mobility_gen.config import Config +from omni.ext.mobility_gen.config import Config, OccupancyMapConfig from omni.ext.mobility_gen.build import build_scenario_from_config @@ -88,7 +98,17 @@ def on_startup(self, ext_id): with ui.HStack(): ui.Label("Robot Type") self.robot_combo_box = ui.ComboBox(0, *ROBOTS.names()) - + + with ui.VStack(): + self._omap_override = cb_builder( + "Manual occupancy bounds", + tooltip="If true, use manually specified occupancy map bounds", + default_val=False + ) + self._omap_origin = xyz_builder(label="Origin") + self._omap_lower_bound = xyz_builder(label="Lower Bound") + self._omap_upper_bound = xyz_builder(label="Upper Bound") + ui.Button("Build", clicked_fn=self.build_scenario) with ui.VStack(): @@ -198,6 +218,43 @@ 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 get_omap_config(self, robot_type): + z_min = robot_type.occupancy_map_z_min + z_max = robot_type.occupancy_map_z_max + cell_size = robot_type.occupancy_map_cell_size + omap_config = OccupancyMapConfig( + origin=( + self._omap_origin[0].get_value_as_float(), + self._omap_origin[1].get_value_as_float(), + 0. + ), + lower_bound=( + self._omap_lower_bound[0].get_value_as_float(), + self._omap_lower_bound[1].get_value_as_float(), + z_min + ), + upper_bound=( + self._omap_upper_bound[0].get_value_as_float(), + self._omap_upper_bound[1].get_value_as_float(), + z_max + ), + cell_size=cell_size + ) + + if not self._omap_override.get_value_as_bool(): + omap_config.prim_path = "/World/scene" + + return omap_config + + def update_ui_from_config(self): + omap_cfg: OccupancyMapConfig = self.config.occupancy_map_config + for i in range(3): + self._omap_origin[i].set_value(omap_cfg.origin[i]) + for i in range(3): + self._omap_lower_bound[i].set_value(omap_cfg.lower_bound[i]) + for i in range(3): + self._omap_upper_bound[i].set_value(omap_cfg.upper_bound[i]) + def build_scenario(self): async def _build_scenario_async(): @@ -206,9 +263,12 @@ async def _build_scenario_async(): self.clear_scenario() config = self.create_config() + robot_type = ROBOTS.get(config.robot_type) + config.occupancy_map_config = self.get_omap_config(robot_type) + + self.scenario, self.config = await build_scenario_from_config(config) - self.config = config - self.scenario = await build_scenario_from_config(config) + self.update_ui_from_config() self.draw_occ_map() From 35c8251640c4386be7507105a736f409ad088ba0 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Mon, 31 Mar 2025 16:15:21 -0700 Subject: [PATCH 10/29] move robot after occ map build --- exts/omni.ext.mobility_gen/omni/ext/mobility_gen/build.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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 981b587..3f41f87 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 @@ -58,8 +58,7 @@ async def build_scenario_from_config(config: Config): 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") + # open_stage(config.scene_usd) if config.occupancy_map_config.prim_path is not None: @@ -81,6 +80,9 @@ async def build_scenario_from_config(config: Config): cell_size=config.occupancy_map_config.cell_size ) + objects.GroundPlane("/World/ground_plane", visible=False) + robot = robot_type.build("/World/robot") + chase_camera_path = robot.build_chase_camera() set_viewport_camera(chase_camera_path) scenario = scenario_type.from_robot_occupancy_map(robot, occupancy_map) From ff837330b4477ced409c716855fdb16c974c0361 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Wed, 2 Apr 2025 14:11:18 -0700 Subject: [PATCH 11/29] added custom method for getting world pose, with major speedup --- .../omni/ext/mobility_gen/sensors.py | 8 +-- .../omni/ext/mobility_gen/utils/prim_utils.py | 53 +++++++++++++++++++ 2 files changed, 57 insertions(+), 4 deletions(-) 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/prim_utils.py b/exts/omni.ext.mobility_gen/omni/ext/mobility_gen/utils/prim_utils.py index 29cc7c6..7899752 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[Gf.Vec3d, Gf.Rotation, Gf.Vec3d]: + """ + 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[Gf.Vec3d, Gf.Rotation, Gf.Vec3d]: + """ + 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) From baa0d1ba7d5f0c16597e5001bb9b17a4aedbca24 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Wed, 2 Apr 2025 14:13:54 -0700 Subject: [PATCH 12/29] make robot use get_world_pose() method --- exts/omni.ext.mobility_gen/omni/ext/mobility_gen/robots.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 fd822c9..8f249b7 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 @@ -184,7 +184,7 @@ 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()) From c09cdae55ad6e6ecc10307abd4778749ab0f4590 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Wed, 2 Apr 2025 14:17:36 -0700 Subject: [PATCH 13/29] fix type annotations --- .../omni/ext/mobility_gen/utils/prim_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 7899752..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 @@ -282,7 +282,7 @@ def _rotation_to_np_quat(r: Gf.Rotation): return np.array([real, imag[0], imag[1], imag[2]]) -def prim_get_local_transform(prim: Usd.Prim) -> Tuple[Gf.Vec3d, Gf.Rotation, Gf.Vec3d]: +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 @@ -303,7 +303,7 @@ def prim_get_local_transform(prim: Usd.Prim) -> Tuple[Gf.Vec3d, Gf.Rotation, Gf. return _translation_to_np(translation), _rotation_to_np_quat(rotation) -def prim_get_world_transform(prim: Usd.Prim) -> Tuple[Gf.Vec3d, Gf.Rotation, Gf.Vec3d]: +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 From 86941827d526f14331eb975514ee4bd417b3a6c1 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Wed, 2 Apr 2025 15:01:36 -0700 Subject: [PATCH 14/29] add compression for path to reduce number of points for speed up --- .../omni/ext/mobility_gen/scenarios.py | 3 ++- path_planner/compress.py | 16 ++++++++++++++++ .../mobility_gen_path_planner/__init__.py | 16 +++++++++++++++- 3 files changed, 33 insertions(+), 2 deletions(-) create mode 100644 path_planner/compress.py 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..96fb52f 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 @@ -17,7 +17,7 @@ import numpy as np 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 @@ -193,6 +193,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) 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 From 3ffd751f76261b71f651a46399f1560c43973cf3 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Wed, 2 Apr 2025 15:01:56 -0700 Subject: [PATCH 15/29] use get_world_pose in get/set pose 2d methods to speed up --- exts/omni.ext.mobility_gen/omni/ext/mobility_gen/robots.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 8f249b7..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 @@ -207,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 @@ -217,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], From cdfeb65799b69bf0e6e1e4a733cecc15f589774f Mon Sep 17 00:00:00 2001 From: John Welsh Date: Wed, 2 Apr 2025 15:15:55 -0700 Subject: [PATCH 16/29] reset on goal reach for path planning --- exts/omni.ext.mobility_gen/omni/ext/mobility_gen/scenarios.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 96fb52f..5003e7d 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 @@ -229,7 +229,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) From 16b76b7f592c511379e9b5c035ddcc9c3684a708 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Thu, 3 Apr 2025 11:33:08 -0700 Subject: [PATCH 17/29] add venv to gitignore --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 From 4318aee30c5920f92c88348e7941c8ef91c26e2b Mon Sep 17 00:00:00 2001 From: John Welsh Date: Thu, 3 Apr 2025 15:16:56 -0700 Subject: [PATCH 18/29] modify build to open stage, add more fficient replay directory. TODO: fix comments --- .../omni/ext/mobility_gen/build.py | 5 +- scripts/replay_directory.py | 51 ++---- scripts/replay_directory_implementation.py | 166 ++++++++++++++++++ 3 files changed, 185 insertions(+), 37 deletions(-) create mode 100644 scripts/replay_directory_implementation.py 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 3f41f87..9be928d 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 @@ -54,10 +54,11 @@ def load_scenario(path: str) -> 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() + # new_stage() + # add_reference_to_stage(config.scene_usd,"/World/scene") + open_stage(config.scene_usd) world = new_world(physics_dt=robot_type.physics_dt) await world.initialize_simulation_context_async() - add_reference_to_stage(config.scene_usd,"/World/scene") # open_stage(config.scene_usd) if config.occupancy_map_config.prim_path is not None: diff --git a/scripts/replay_directory.py b/scripts/replay_directory.py index d3f8e11..7076fab 100644 --- a/scripts/replay_directory.py +++ b/scripts/replay_directory.py @@ -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..f6214a9 --- /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, "*")) + + count = 0 + for recording_path in recording_paths: + 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 {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 From f9a9acc9e651900da699238494aa2f97f6f0c37f Mon Sep 17 00:00:00 2001 From: John Welsh Date: Thu, 3 Apr 2025 15:21:26 -0700 Subject: [PATCH 19/29] fix double use of count variable --- scripts/replay_directory_implementation.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/scripts/replay_directory_implementation.py b/scripts/replay_directory_implementation.py index f6214a9..85df56a 100644 --- a/scripts/replay_directory_implementation.py +++ b/scripts/replay_directory_implementation.py @@ -70,9 +70,9 @@ recording_paths = glob.glob(os.path.join(args.input, "*")) - count = 0 + recording_count = 0 for recording_path in recording_paths: - count += 1 + recording_count += 1 name = os.path.basename(recording_path) output_path = os.path.join(args.output, name) @@ -113,7 +113,7 @@ writer.copy_init(recording_path) - print(f"============== Replaying {count} / {len(recording_paths)}==============") + 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}") From 95ba0e5590581713fe3c7b8550c1d7a5f37df7d3 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Fri, 4 Apr 2025 12:13:27 -0700 Subject: [PATCH 20/29] make parquet conversion work with directory --- examples/05_convert_to_parquet.py | 110 ++++++++++++++++++++---------- 1 file changed, 73 insertions(+), 37 deletions(-) diff --git a/examples/05_convert_to_parquet.py b/examples/05_convert_to_parquet.py index 4a49bcf..161c3a4 100644 --- a/examples/05_convert_to_parquet.py +++ b/examples/05_convert_to_parquet.py @@ -6,14 +6,10 @@ import PIL.Image import io -parser = argparse.ArgumentParser() -parser.add_argument("recording_path") -parser.add_argument("output_path") -args = parser.parse_args() - -reader = Reader(recording_path=args.recording_path) - -index = 0 +import os +import subprocess +import glob +import argparse def numpy_array_to_flattened_columns(key: str, value: np.ndarray): @@ -36,40 +32,80 @@ def numpy_array_to_jpg_columns(key: str, value: np.ndarray): return columns -output: pandas.DataFrame = None +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))): + for index in tqdm.tqdm(range(len(reader))): - data_dict = {} + 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)) - # 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) + # 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) + # 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()) + + # use first frame to initialize + if output is None: + output = pandas.DataFrame(columns=data_dict.keys()) - output.loc[index] = data_dict + output.loc[index] = data_dict -output.to_parquet(args.output_path, engine="pyarrow") \ No newline at end of file + output.to_parquet(output_path, engine="pyarrow") \ No newline at end of file From 3c0071d5badab1b1e7fcd8843ce034404b619644 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Tue, 8 Apr 2025 11:45:03 -0700 Subject: [PATCH 21/29] use manual occupancy map path instead --- .../omni/ext/mobility_gen/extension.py | 122 ++++++++---------- 1 file changed, 57 insertions(+), 65 deletions(-) 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 6df0da1..b7ea2e0 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,6 +24,7 @@ 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, @@ -33,15 +34,18 @@ 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, OccupancyMapConfig 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: @@ -87,10 +91,15 @@ 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("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("Scenario Type") self.scenario_combo_box = ui.ComboBox(0, *SCENARIOS.names()) @@ -99,17 +108,8 @@ def on_startup(self, ext_id): ui.Label("Robot Type") self.robot_combo_box = ui.ComboBox(0, *ROBOTS.names()) - with ui.VStack(): - self._omap_override = cb_builder( - "Manual occupancy bounds", - tooltip="If true, use manually specified occupancy map bounds", - default_val=False - ) - self._omap_origin = xyz_builder(label="Origin") - self._omap_lower_bound = xyz_builder(label="Lower Bound") - self._omap_upper_bound = xyz_builder(label="Upper Bound") - 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("") @@ -135,6 +135,7 @@ def build_occ_map_frame(self): def draw_occ_map(self): if self.scenario is not None: image = self.scenario.occupancy_map.ros_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() @@ -160,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() @@ -218,74 +220,64 @@ 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 get_omap_config(self, robot_type): - z_min = robot_type.occupancy_map_z_min - z_max = robot_type.occupancy_map_z_max - cell_size = robot_type.occupancy_map_cell_size - omap_config = OccupancyMapConfig( - origin=( - self._omap_origin[0].get_value_as_float(), - self._omap_origin[1].get_value_as_float(), - 0. - ), - lower_bound=( - self._omap_lower_bound[0].get_value_as_float(), - self._omap_lower_bound[1].get_value_as_float(), - z_min - ), - upper_bound=( - self._omap_upper_bound[0].get_value_as_float(), - self._omap_upper_bound[1].get_value_as_float(), - z_max - ), - cell_size=cell_size - ) - - if not self._omap_override.get_value_as_bool(): - omap_config.prim_path = "/World/scene" - - return omap_config - def update_ui_from_config(self): - omap_cfg: OccupancyMapConfig = self.config.occupancy_map_config - for i in range(3): - self._omap_origin[i].set_value(omap_cfg.origin[i]) - for i in range(3): - self._omap_lower_bound[i].set_value(omap_cfg.lower_bound[i]) - for i in range(3): - self._omap_upper_bound[i].set_value(omap_cfg.upper_bound[i]) - def build_scenario(self): async def _build_scenario_async(): self.clear_recording() self.clear_scenario() + self.disable_recording() - config = self.create_config() - robot_type = ROBOTS.get(config.robot_type) - config.occupancy_map_config = self.get_omap_config(robot_type) + # 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 - self.scenario, self.config = await build_scenario_from_config(config) + 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 + ) - self.update_ui_from_config() + occupancy_map = OccupancyMap.from_ros_yaml(self.omap_field_string_model.as_string) - self.draw_occ_map() + # Open stage + open_stage(scene_usd_str) + self.cached_stage_path = os.path.join(tempfile.mkdtemp(), "stage.usd") + + # Add ground plane + objects.GroundPlane("/World/ground_plane", visible=False) - world = get_world() - await world.reset_async() + # Save the stage with ground plane + save_stage(self.cached_stage_path) - self.scenario.reset() + # Initialize world + world = new_world(physics_dt=robot_type.physics_dt) + await world.initialize_simulation_context_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) - world.add_physics_callback("scenario_physics", self.on_physics) + # Draw the occupancy map + self.draw_occ_map() - # cache stage - self.cached_stage_path = os.path.join(tempfile.mkdtemp(), "stage.usd") - save_stage(self.cached_stage_path) + # Run the scenario + await world.reset_async() + self.scenario.reset() + world.add_physics_callback("scenario_physics", self.on_physics) - if self.recording_enabled: - self.start_new_recording() - # self.scenario.save(path) asyncio.ensure_future(_build_scenario_async()) \ No newline at end of file From 500d0a8cea4ee7003e9ea98b80ffe26a99ec243b Mon Sep 17 00:00:00 2001 From: John Welsh Date: Tue, 8 Apr 2025 12:09:30 -0700 Subject: [PATCH 22/29] do not save ground plane in scenario --- .../omni/ext/mobility_gen/build.py | 55 ++++--------------- .../omni/ext/mobility_gen/config.py | 17 ++---- .../omni/ext/mobility_gen/extension.py | 13 ++--- 3 files changed, 21 insertions(+), 64 deletions(-) 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 9be928d..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, OccupancyMapConfig -from omni.ext.mobility_gen.utils.occupancy_map_utils import occupancy_map_generate_from_prim_async, compute_occupancy_bounds_from_prim_path -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,41 +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() - # add_reference_to_stage(config.scene_usd,"/World/scene") - open_stage(config.scene_usd) - world = new_world(physics_dt=robot_type.physics_dt) - await world.initialize_simulation_context_async() - # open_stage(config.scene_usd) - - if config.occupancy_map_config.prim_path is not None: - - origin, lower_bound, upper_bound = compute_occupancy_bounds_from_prim_path( - prim_path=config.occupancy_map_config.prim_path, - z_min=config.occupancy_map_config.lower_bound[2], - z_max=config.occupancy_map_config.upper_bound[2], - cell_size=config.occupancy_map_config.cell_size - ) - config.occupancy_map_config.origin = origin - config.occupancy_map_config.lower_bound = lower_bound - config.occupancy_map_config.upper_bound = upper_bound - - - occupancy_map = await occupancy_map_generate_from_prim_async( - origin=config.occupancy_map_config.origin, - lower_bound=config.occupancy_map_config.lower_bound, - upper_bound=config.occupancy_map_config.upper_bound, - cell_size=config.occupancy_map_config.cell_size - ) - - objects.GroundPlane("/World/ground_plane", visible=False) - robot = robot_type.build("/World/robot") - - 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, config \ 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 712b738..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 @@ -3,21 +3,11 @@ from dataclasses import dataclass, asdict -@dataclass -class OccupancyMapConfig: - prim_path: Optional[str] = None - cell_size: Optional[float] = None - origin: Optional[Tuple[float, float, float]] = None - lower_bound: Optional[Tuple[float, float, float]] = None - upper_bound: Optional[Tuple[float, float, float]] = None - - @dataclass class Config: scenario_type: str robot_type: str scene_usd: str - occupancy_map_config: Optional[OccupancyMapConfig] = None def to_json(self): return json.dumps(asdict(self), indent=2) @@ -25,5 +15,8 @@ def to_json(self): @staticmethod def from_json(data: str): data = json.loads(data) - data['occupancy_map_config'] = OccupancyMapConfig(**data['occupancy_map_config']) - return Config(**data) \ No newline at end of file + 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 b7ea2e0..a17488c 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 @@ -43,8 +43,7 @@ from omni.ext.mobility_gen.scenarios import SCENARIOS, Scenario 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, OccupancyMapConfig -from omni.ext.mobility_gen.build import build_scenario_from_config +from omni.ext.mobility_gen.config import Config from omni.ext.mobility_gen.occupancy_map import OccupancyMap @@ -246,19 +245,17 @@ async def _build_scenario_async(): occupancy_map = OccupancyMap.from_ros_yaml(self.omap_field_string_model.as_string) - # Open stage + # Open stage and save local copy open_stage(scene_usd_str) self.cached_stage_path = os.path.join(tempfile.mkdtemp(), "stage.usd") - - # Add ground plane - objects.GroundPlane("/World/ground_plane", visible=False) - - # Save the stage with ground plane save_stage(self.cached_stage_path) # Initialize world world = new_world(physics_dt=robot_type.physics_dt) await world.initialize_simulation_context_async() + + # Add ground plane + objects.GroundPlane("/World/ground_plane", visible=False) # Add robot robot = robot_type.build("/World/robot") From febf0deb6a7e7cd9dbbfece59bfa23cb668a1b26 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Tue, 8 Apr 2025 12:11:25 -0700 Subject: [PATCH 23/29] add update state to reset --- exts/omni.ext.mobility_gen/omni/ext/mobility_gen/scenarios.py | 4 ++++ 1 file changed, 4 insertions(+) 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 5003e7d..c2b2268 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 @@ -67,6 +67,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 +106,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 +141,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): @@ -205,6 +208,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): From 65f5b4f9367856eb09a2569dd6abdde2885a17bd Mon Sep 17 00:00:00 2001 From: John Welsh Date: Tue, 8 Apr 2025 12:37:32 -0700 Subject: [PATCH 24/29] move scenario after robot --- .../omni/ext/mobility_gen/extension.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) 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 a17488c..d010b7c 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 @@ -99,13 +99,14 @@ def on_startup(self, ext_id): 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()) - with ui.HStack(): - ui.Label("Robot Type") - self.robot_combo_box = ui.ComboBox(0, *ROBOTS.names()) ui.Button("Build", clicked_fn=self.build_scenario) # ui.Button("Build", clicked_fn=self.build_scenario) From 1a670a5aff6120ee018e2ad4378e72f67227593f Mon Sep 17 00:00:00 2001 From: John Welsh Date: Tue, 8 Apr 2025 13:39:24 -0700 Subject: [PATCH 25/29] add docs for manual occ map gen --- README.md | 82 ++++++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 78 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index f29d2d6..8925115 100644 --- a/README.md +++ b/README.md @@ -163,20 +163,94 @@ 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. +1. Select ``File`` -> ``Open`` -1. Under Scene USD URL / Path copy and paste the following +2. Enter the following URL under file (change this to use your environment) ``` http://omniverse-content-production.s3-us-west-2.amazonaws.com/Assets/Isaac/4.2/Isaac/Environments/Simple_Warehouse/warehouse_multiple_shelves.usd ``` -2. Under the ``Scenario`` dropdown select ``KeyboardTeleoperationScenario`` to start +### Step 3 - Create an occupancy map + +To use MobilityGen, we first 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 + +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`` + +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. Save the image to ``~/MobilityGenData/maps/warehouse_multiple_shelves/map.png`` + + > Note: The image file name must match that specified in the YAML file saved above. + +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 + +Perform the following steps in the ``MobilityGen`` extension window. + +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 folder we created in [Step 3](#step-3---create-an-occupancy-map) + + ``` + ~/MobilityGenData/maps/warehouse_multiple_shelves/ + ``` 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. From 35e1d77acfbdc3268c7c2be123de6445c54929cd Mon Sep 17 00:00:00 2001 From: John Welsh Date: Tue, 8 Apr 2025 14:06:34 -0700 Subject: [PATCH 26/29] add occupancy map generation instructions --- README.md | 32 ++++++++++++++----- .../omni/ext/mobility_gen/extension.py | 7 ++-- 2 files changed, 28 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 8925115..32833be 100644 --- a/README.md +++ b/README.md @@ -165,22 +165,31 @@ Below details a typical workflow for collecting data with MobilityGen. ### Step 2 - Load a stage +To get started, we'll open an example warehouse stage. + 1. Select ``File`` -> ``Open`` -2. Enter the following URL under file (change this to use your environment) +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`` + + > 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 -To use MobilityGen, we first need to build an occupancy map for the environment. +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`` @@ -213,6 +222,9 @@ To do this, we'll use the Occupancy Map tool provided with Isaac Sim 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`` @@ -221,9 +233,11 @@ To do this, we'll use the Occupancy Map tool provided with Isaac Sim 11. Back in the ``Visualization`` window click ``Save Image`` -12. Save the image to ``~/MobilityGenData/maps/warehouse_multiple_shelves/map.png`` +12. In the tree explorer open the folder ``~/MobilityGenData/maps/warehouse_multiple_shelves`` + +12. Under file name enter ``map.png`` - > Note: The image file name must match that specified in the YAML file saved above. +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. @@ -233,7 +247,9 @@ That's it! You should now have a folder ``~/MobilityGenData/maps/warehouse_mult ### Step 4 - Build a scenario -Perform the following steps in the ``MobilityGen`` extension window. +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) @@ -241,10 +257,10 @@ Perform the following steps in the ``MobilityGen`` extension window. 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 folder we created in [Step 3](#step-3---create-an-occupancy-map) +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/ + ~/MobilityGenData/maps/warehouse_multiple_shelves/map.yaml ``` 3. Under the ``Robot`` dropdown select ``H1Robot`` @@ -257,7 +273,7 @@ After a few seconds, you should see the scene and occupancy map appear. ### Step 3 - 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 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 d010b7c..b1fe9cb 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 @@ -102,7 +102,7 @@ def on_startup(self, ext_id): 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()) @@ -243,8 +243,9 @@ async def _build_scenario_async(): robot_type=robot_type_str, scene_usd=scene_usd_str ) - - occupancy_map = OccupancyMap.from_ros_yaml(self.omap_field_string_model.as_string) + + occupancy_map = OccupancyMap.from_ros_yaml( + os.path.expanduser(self.omap_field_string_model.as_string)) # Open stage and save local copy open_stage(scene_usd_str) From 40a39f8af6522976c1948766cefcef4d9656e1f4 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Tue, 8 Apr 2025 14:31:15 -0700 Subject: [PATCH 27/29] add visualization --- .../omni/ext/mobility_gen/extension.py | 9 +++++---- .../omni/ext/mobility_gen/scenarios.py | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+), 4 deletions(-) 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 b1fe9cb..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 @@ -132,9 +132,9 @@ 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]) @@ -202,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): @@ -270,13 +271,13 @@ async def _build_scenario_async(): self.scenario = scenario_type.from_robot_occupancy_map(robot, occupancy_map) # Draw the occupancy map - self.draw_occ_map() + self.draw_visualization_image() # Run the scenario await world.reset_async() - self.scenario.reset() world.add_physics_callback("scenario_physics", self.on_physics) + self.reset() asyncio.ensure_future(_build_scenario_async()) \ No newline at end of file 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 c2b2268..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,6 +15,8 @@ import numpy as np +import PIL.Image +from PIL import Image, ImageDraw from typing import Tuple from mobility_gen_path_planner import generate_paths, compress_path @@ -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]() @@ -253,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 From 1e57a1a7a14e5b1945cb589fbf34b65f5eee9026 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Thu, 10 Apr 2025 12:28:49 -0700 Subject: [PATCH 28/29] fix steps in readme --- README.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 32833be..1ea7ed7 100644 --- a/README.md +++ b/README.md @@ -271,12 +271,12 @@ Perform the following steps in the ``MobilityGen`` extension window to build a n 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 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 @@ -287,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. @@ -299,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. @@ -324,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. From 9c013591b0bd2801d7ff6a6b87e816d328b36f25 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Thu, 10 Apr 2025 12:31:59 -0700 Subject: [PATCH 29/29] update changelog --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 25f726f..b929b91 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ # 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