From e0de921e6cbe087e049de14c3f27b9c84ccf13b7 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Fri, 13 Feb 2026 16:07:25 +0000 Subject: [PATCH 1/5] wip --- embodichain/__init__.py | 2 ++ embodichain/data/dataset.py | 17 ++++++++++------ embodichain/lab/gym/envs/embodied_env.py | 26 +++++++++++++++++++----- 3 files changed, 34 insertions(+), 11 deletions(-) diff --git a/embodichain/__init__.py b/embodichain/__init__.py index 3ce011bc..d0960679 100644 --- a/embodichain/__init__.py +++ b/embodichain/__init__.py @@ -31,3 +31,5 @@ def _get_version(): __version__ = _get_version() + +import robot_challenge_tasks.data diff --git a/embodichain/data/dataset.py b/embodichain/data/dataset.py index 5b3d4a82..97be04d8 100644 --- a/embodichain/data/dataset.py +++ b/embodichain/data/dataset.py @@ -118,11 +118,18 @@ def calculate_md5(self, file_path, chunk_size=8192): return hash_md5.hexdigest() -def get_data_class(dataset_name: str): +DEFAULT_DATA_MODULES = [ + "embodichain.data", + "embodichain.data.assets", +] + + +def get_data_class(dataset_name: str, extra_modules: list[str] | None = None): """Retrieve the dataset class from the available modules. Args: dataset_name (str): The name of the dataset class. + extra_modules (list[str] | None): Optional list of additional module names to search for the dataset class. Returns: type: The dataset class. @@ -130,11 +137,9 @@ def get_data_class(dataset_name: str): Raises: AttributeError: If the dataset class is not found in any module. """ - module_names = [ - "embodichain.data", - "embodichain.data.assets", - __name__, - ] + module_names = DEFAULT_DATA_MODULES + ( + extra_modules if extra_modules is not None else [] + ) for module_name in module_names: try: diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index cd1e4df3..f858674a 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -519,7 +519,12 @@ def _setup_interactive_objects(self) -> None: self.sim.add_rigid_object_group(cfg=cfg) def preview_sensor_data( - self, name: str, data_type: str = "color", env_ids: int = 0, method: str = "plt" + self, + name: str, + data_type: str = "color", + env_ids: int = 0, + method: str = "cv2", + save: bool = False, ) -> None: """Preview the sensor data by matplotlib @@ -531,6 +536,7 @@ def preview_sensor_data( data_type (str): type of the sensor data to preview. env_ids (int): index of the arena to preview. Defaults to 0. method (str): method to preview the sensor data. Currently support "plt" and "cv2". Defaults to "plt". + save (bool): whether to save the preview image. Defaults to False. """ # TODO: this function need to be improved to support more sensor types and data types. @@ -556,15 +562,25 @@ def preview_sensor_data( if method == "cv2": import cv2 - cv2.imshow( - f"sensor_data_{data_type}", cv2.cvtColor(view, cv2.COLOR_RGB2BGR) - ) - cv2.waitKey(0) + if save: + cv2.imwrite( + f"sensor_data_{data_type}.png", + cv2.cvtColor(view, cv2.COLOR_RGB2BGR), + ) + else: + cv2.imshow( + f"sensor_data_{data_type}", cv2.cvtColor(view, cv2.COLOR_RGB2BGR) + ) + cv2.waitKey(0) elif method == "plt": from matplotlib import pyplot as plt plt.imshow(view) plt.savefig(f"sensor_data_{data_type}.png") + if not save: + plt.show() + else: + plt.close() def create_demo_action_list(self, *args, **kwargs) -> Sequence[EnvAction] | None: """Create a demonstration action list for the environment. From c600b482371d04222fd11ef95fad4b85693af205 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 14 Feb 2026 07:15:38 +0000 Subject: [PATCH 2/5] wip --- embodichain/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/embodichain/__init__.py b/embodichain/__init__.py index d0960679..3ce011bc 100644 --- a/embodichain/__init__.py +++ b/embodichain/__init__.py @@ -31,5 +31,3 @@ def _get_version(): __version__ = _get_version() - -import robot_challenge_tasks.data From 0d9b22c949f06cb9cc6b17a608afb7ab992162e2 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 14 Feb 2026 15:24:34 +0800 Subject: [PATCH 3/5] wip --- embodichain/lab/sim/sim_manager.py | 8 + embodichain/lab/sim/utility/gizmo_utils.py | 7 +- embodichain/lab/sim/utility/keyboard_utils.py | 286 +++++++++++++++++- 3 files changed, 295 insertions(+), 6 deletions(-) diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 1eb607ad..eb29fea6 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -803,6 +803,14 @@ def get_light(self, uid: str) -> Light | None: return None return self._lights[uid] + def get_light_uid_list(self) -> List[str]: + """Get current light uid list + + Returns: + List[str]: list of light uid. + """ + return list(self._lights.keys()) + def add_rigid_object( self, cfg: RigidObjectCfg, diff --git a/embodichain/lab/sim/utility/gizmo_utils.py b/embodichain/lab/sim/utility/gizmo_utils.py index b2a2fafc..5038fb9e 100644 --- a/embodichain/lab/sim/utility/gizmo_utils.py +++ b/embodichain/lab/sim/utility/gizmo_utils.py @@ -51,7 +51,7 @@ def gizmo_transform_callback(node, translation, rotation, flag): def run_gizmo_robot_control_loop( - robot: "Robot", control_part: str = "arm", end_link_name: str | None = None + robot: object | str, control_part: str = "arm", end_link_name: str | None = None ): """Run a control loop for testing gizmo controls on a robot. @@ -59,7 +59,7 @@ def run_gizmo_robot_control_loop( using gizmo controls with keyboard input for additional commands. Args: - robot (Robot): The robot to control with the gizmo. + robot (Robot | str): The robot to control with the gizmo. control_part (str, optional): The part of the robot to control. Defaults to "arm". end_link_name (str | None, optional): The name of the end link for FK calculations. Defaults to None. @@ -87,6 +87,9 @@ def run_gizmo_robot_control_loop( sim = SimulationManager.get_instance() + if isinstance(robot, str): + robot = sim.get_robot(uid=robot) + # Enter auto-update mode. sim.set_manual_update(False) diff --git a/embodichain/lab/sim/utility/keyboard_utils.py b/embodichain/lab/sim/utility/keyboard_utils.py index 27dde849..e4e1f7f1 100644 --- a/embodichain/lab/sim/utility/keyboard_utils.py +++ b/embodichain/lab/sim/utility/keyboard_utils.py @@ -14,6 +14,11 @@ # limitations under the License. # ---------------------------------------------------------------------------- +import select +import sys +import tty +import termios +import time import cv2 import torch import numpy as np @@ -25,7 +30,7 @@ def run_keyboard_control_for_camera( - sensor: Camera, + sensor: Camera | str, trans_step: float = 0.01, rot_step: float = 1.0, vis_pose: bool = False, @@ -33,11 +38,24 @@ def run_keyboard_control_for_camera( """Run keyboard control loop for camera pose adjustment. Args: - sensor (Camera): Camera sensor to control. + sensor (Camera | str): Camera sensor or name of the camera to control. trans_step (float, optional): Translation step size. Defaults to 0.01. rot_step (float, optional): Rotation step size in degrees. Defaults to 1.0. vis_pose (bool, optional): Whether to visualize the camera pose in axis form. Defaults to False. """ + from embodichain.lab.sim import SimulationManager + + sim = SimulationManager.get_instance() + + if vis_pose and sim.is_rt_enabled: + log_warning( + "'vis_pose' is not fully supported with ray tracing enabled. Will be fixed in future updates." + ) + return + + if isinstance(sensor, str): + sensor = sim.get_sensor(uid=sensor) + if sensor.num_instances > 1: log_warning( "Multiple sensor instances detected. Keyboard control will only work for one instance." @@ -62,12 +80,10 @@ def run_keyboard_control_for_camera( marker = None if vis_pose: - from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.cfg import MarkerCfg init_axis_pose = sensor.get_arena_pose(to_matrix=True).squeeze().numpy() - sim = SimulationManager.get_instance() marker = sim.draw_marker( cfg=MarkerCfg( name="camera_axis", @@ -227,3 +243,265 @@ def run_keyboard_control_for_camera( cv2.destroyAllWindows() except Exception as e: log_warning(f"cv2.destroyAllWindows() failed: {e}") + + +def run_keyboard_control_for_light( + light: object | str, + trans_step: float = 0.01, + intensity_step: float = 1.0, + falloff_step: float = 1.0, + color_step: float = 0.05, + vis_pose: bool = False, +) -> None: + """Run keyboard control loop for light adjustment. + + Args: + light (Light | str): Light object or name of the light to control. + trans_step (float, optional): Translation step size. Defaults to 0.01. + intensity_step (float, optional): Intensity adjustment step. Defaults to 0.1. + falloff_step (float, optional): Falloff/radius adjustment step. Defaults to 0.1. + color_step (float, optional): Color channel adjustment step. Defaults to 0.05. + vis_pose (bool, optional): Whether to visualize the light position with a marker. Defaults to False. + """ + from embodichain.lab.sim.objects import Light + from embodichain.lab.sim import SimulationManager + + sim = SimulationManager.get_instance() + + if vis_pose and sim.is_rt_enabled: + log_warning( + "'vis_pose' is not fully supported with ray tracing enabled. Will be fixed in future updates." + ) + return + + if isinstance(light, str): + light: Light = sim.get_light(uid=light) + + if light.num_instances > 1: + log_warning( + "Multiple light instances detected. Keyboard control will only work for one instance." + ) + return + + log_info("\n=== Light Control ===") + log_info("Translation controls:") + log_info(" W/S: Move forward/backward (Z-axis)") + log_info(" A/D: Move left/right (Y-axis)") + log_info(" Q/E: Move up/down (X-axis)") + log_info("\nIntensity controls:") + log_info(" I/K: Increase/decrease intensity") + log_info("\nFalloff controls:") + log_info(" U/O: Increase/decrease falloff radius") + log_info("\nColor controls:") + log_info(" T/Y: Increase/decrease red channel") + log_info(" G/H: Increase/decrease green channel") + log_info(" B/N: Increase/decrease blue channel") + log_info("\nOther controls:") + log_info(" R: Reset to initial values") + log_info(" P: Print current light properties") + log_info(" ESC: Exit control mode") + + # Store initial values from config + init_pose = light.get_local_pose() + init_color = torch.as_tensor(light.cfg.color).clone() + init_intensity = float(light.cfg.intensity) + init_falloff = float(light.cfg.radius) + + # Current values + current_color = init_color.clone() + current_intensity = init_intensity + current_falloff = init_falloff + + marker = None + if vis_pose: + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.cfg import MarkerCfg + + init_marker_pose = light.get_local_pose(to_matrix=True).squeeze().numpy() + + sim = SimulationManager.get_instance() + marker = sim.draw_marker( + cfg=MarkerCfg( + name="light_marker", + marker_type="axis", + axis_xpos=[init_marker_pose], + axis_size=0.002, + axis_len=0.05, + ) + ) + + # TODO: We may add node to BatchEntity object. + marker[0].node.attach_node(light._entities[0].get_node()) + + log_info("\nLight control active. Press keys to adjust light properties...") + + # Save terminal settings + old_settings = termios.tcgetattr(sys.stdin) + tty.setcbreak(sys.stdin.fileno()) + + def get_key(): + """Non-blocking keyboard input.""" + if select.select([sys.stdin], [], [], 0)[0]: + return sys.stdin.read(1) + return None + + try: + + while True: + current_pose = light.get_local_pose().squeeze().numpy() + + # Non-blocking key input + key = get_key() + + if key is None: + continue + elif key in ["\x1b"]: # Q or ESC + if vis_pose: + sim.remove_marker("light_marker") + log_info("Exiting light control mode...") + break + + property_changed = False + new_pose = current_pose.copy() + + # Translation controls + if key in ["w", "W"]: + new_pose[2] += trans_step + property_changed = True + log_info(f"Moving forward: Z += {trans_step}") + elif key in ["s", "S"]: + new_pose[2] -= trans_step + property_changed = True + log_info(f"Moving backward: Z -= {trans_step}") + elif key in ["a", "A"]: + new_pose[1] -= trans_step + property_changed = True + log_info(f"Moving left: Y -= {trans_step}") + elif key in ["d", "D"]: + new_pose[1] += trans_step + property_changed = True + log_info(f"Moving right: Y += {trans_step}") + elif key in ["q", "Q"]: + new_pose[0] += trans_step + property_changed = True + log_info(f"Moving up: X += {trans_step}") + elif key in ["e", "E"]: + new_pose[0] -= trans_step + property_changed = True + log_info(f"Moving down: X -= {trans_step}") + + # Intensity controls + elif key in ["i", "I"]: + current_intensity += intensity_step + current_intensity = max(0.0, current_intensity) + light.set_intensity(torch.tensor(current_intensity)) + property_changed = True + log_info(f"Intensity increased to: {current_intensity:.2f}") + elif key in ["k", "K"]: + current_intensity -= intensity_step + current_intensity = max(0.0, current_intensity) + light.set_intensity(torch.tensor(current_intensity)) + property_changed = True + log_info(f"Intensity decreased to: {current_intensity:.2f}") + + # Falloff controls + elif key in ["u", "U"]: + current_falloff += falloff_step + current_falloff = max(0.0, current_falloff) + light.set_falloff(torch.tensor(current_falloff)) + property_changed = True + log_info(f"Falloff increased to: {current_falloff:.2f}") + elif key in ["o", "O"]: + current_falloff -= falloff_step + current_falloff = max(0.0, current_falloff) + light.set_falloff(torch.tensor(current_falloff)) + property_changed = True + log_info(f"Falloff decreased to: {current_falloff:.2f}") + + # Color controls - Red channel + elif key in ["t", "T"]: + current_color[0] += color_step + current_color[0] = torch.clamp(current_color[0], 0.0, 1.0) + light.set_color(current_color) + property_changed = True + log_info(f"Red channel increased to: {current_color[0]:.2f}") + elif key in ["y", "Y"]: + current_color[0] -= color_step + current_color[0] = torch.clamp(current_color[0], 0.0, 1.0) + light.set_color(current_color) + property_changed = True + log_info(f"Red channel decreased to: {current_color[0]:.2f}") + + # Color controls - Green channel + elif key in ["g", "G"]: + current_color[1] += color_step + current_color[1] = torch.clamp(current_color[1], 0.0, 1.0) + light.set_color(current_color) + property_changed = True + log_info(f"Green channel increased to: {current_color[1]:.2f}") + elif key in ["h", "H"]: + current_color[1] -= color_step + current_color[1] = torch.clamp(current_color[1], 0.0, 1.0) + light.set_color(current_color) + property_changed = True + log_info(f"Green channel decreased to: {current_color[1]:.2f}") + + # Color controls - Blue channel + elif key in ["b", "B"]: + current_color[2] += color_step + current_color[2] = torch.clamp(current_color[2], 0.0, 1.0) + light.set_color(current_color) + property_changed = True + log_info(f"Blue channel increased to: {current_color[2]:.2f}") + elif key in ["n", "N"]: + current_color[2] -= color_step + current_color[2] = torch.clamp(current_color[2], 0.0, 1.0) + light.set_color(current_color) + property_changed = True + log_info(f"Blue channel decreased to: {current_color[2]:.2f}") + + # Reset control + elif key in ["r", "R"]: + current_color = init_color.clone() + current_intensity = init_intensity + current_falloff = init_falloff + light.set_local_pose(init_pose) + light.set_color(current_color) + light.set_intensity(torch.tensor(current_intensity)) + light.set_falloff(torch.tensor(current_falloff)) + property_changed = True + log_info("Reset to initial light properties") + + # Print current properties + elif key in ["p", "P"]: + translation = current_pose[:3] + log_info("\n=== Current Light Properties ===") + log_info(f"Position: {translation}") + log_info(f"Color (RGB): {current_color.numpy()}") + log_info(f"Intensity: {current_intensity:.2f}") + log_info(f"Falloff: {current_falloff:.2f}") + + # Update pose if translation changed + if property_changed and not np.allclose(new_pose, current_pose): + light_pose = torch.as_tensor(new_pose, dtype=torch.float32).unsqueeze_( + 0 + ) + light.set_local_pose(light_pose) + + # Update simulation if any property changed + if property_changed and vis_pose: + sim.update(step=1) + + except KeyboardInterrupt: + if vis_pose: + sim.remove_marker("light_marker") + log_info("\nControl loop interrupted by user (Ctrl+C)") + except Exception as e: + log_error(f"Error in light control loop: {e}") + finally: + try: + # Restore terminal settings + termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings) + except: + pass + log_info("Light control loop terminated") From e24a36099adf75e99a9b344651c2b36cde20da02 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 14 Feb 2026 15:38:12 +0800 Subject: [PATCH 4/5] wip --- embodichain/lab/sim/sensors/camera.py | 9 +++++++++ embodichain/lab/sim/utility/keyboard_utils.py | 5 +++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/embodichain/lab/sim/sensors/camera.py b/embodichain/lab/sim/sensors/camera.py index f761fb92..c5baed17 100644 --- a/embodichain/lab/sim/sensors/camera.py +++ b/embodichain/lab/sim/sensors/camera.py @@ -263,6 +263,15 @@ def group_id(self) -> int: ) return -1 + @property + def is_attached(self) -> bool: + """Check if the camera is attached to a parent entity. + + Returns: + bool: True if the camera is attached to a parent entity, False otherwise. + """ + return self.cfg.extrinsics.parent is not None + def update(self, **kwargs) -> None: """Update the sensor data. diff --git a/embodichain/lab/sim/utility/keyboard_utils.py b/embodichain/lab/sim/utility/keyboard_utils.py index e4e1f7f1..f0646b25 100644 --- a/embodichain/lab/sim/utility/keyboard_utils.py +++ b/embodichain/lab/sim/utility/keyboard_utils.py @@ -213,8 +213,9 @@ def run_keyboard_control_for_camera( log_info("Reset to initial pose") elif key == ord("p"): new_pose_print = new_pose.copy() - new_pose_print[:3, 1] = -new_pose_print[:3, 1] - new_pose_print[:3, 2] = -new_pose_print[:3, 2] + if sensor.is_attached is False: + new_pose_print[:3, 1] = -new_pose_print[:3, 1] + new_pose_print[:3, 2] = -new_pose_print[:3, 2] translation = new_pose_print[:3, 3] rot = R.from_matrix(new_pose_print[:3, :3]) quaternion = rot.as_quat() From 48d864c6373fe86e25a857adc48bed20b1e86c2f Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 14 Feb 2026 15:51:12 +0800 Subject: [PATCH 5/5] wip --- embodichain/lab/gym/envs/embodied_env.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index f58850ec..f89621c7 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -575,7 +575,7 @@ def preview_sensor_data( cv2.imshow(window_name, cv2.cvtColor(view, cv2.COLOR_RGB2BGR)) cv2.waitKey(0) cv2.destroyWindow(window_name) - + elif method == "plt": from matplotlib import pyplot as plt