From ade8ec8afbe718fe12af7003f1d236f5c0edd399 Mon Sep 17 00:00:00 2001 From: John Welsh Date: Mon, 17 Mar 2025 15:28:47 -0700 Subject: [PATCH 1/7] 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 2/7] 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 3/7] 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 4/7] 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 5/7] 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 6/7] 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 7/7] 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