From e6fda954065226649c487bcd986a02e63ce5f601 Mon Sep 17 00:00:00 2001 From: Jack Vial Date: Tue, 24 Dec 2024 14:55:21 -0500 Subject: [PATCH] Removed display_cameras from control loop args. Rename browser ui related code --- .../common/robot_devices/control_context.py | 33 ++++++++++--------- lerobot/common/robot_devices/control_utils.py | 5 --- ...ideo_over_http.py => browser_ui_server.py} | 2 +- lerobot/scripts/control_robot.py | 11 ++----- ...am_video_template.html => browser_ui.html} | 0 5 files changed, 21 insertions(+), 30 deletions(-) rename lerobot/scripts/{stream_video_over_http.py => browser_ui_server.py} (98%) rename lerobot/templates/{stream_video_template.html => browser_ui.html} (100%) diff --git a/lerobot/common/robot_devices/control_context.py b/lerobot/common/robot_devices/control_context.py index 5b018f1e1..60abe0d79 100644 --- a/lerobot/common/robot_devices/control_context.py +++ b/lerobot/common/robot_devices/control_context.py @@ -9,9 +9,7 @@ import time from lerobot.common.robot_devices.robots.utils import Robot from lerobot.common.robot_devices.utils import busy_wait -from lerobot.common.robot_devices.control_utils import ( - log_control_info -) +from lerobot.common.robot_devices.control_utils import log_control_info class ControlPhase: @@ -85,6 +83,9 @@ def __init__(self, config: ControlContextConfig): # Subscribe to all messages self.command_sub_socket.setsockopt_string(zmq.SUBSCRIBE, "") + def update_config(self, config: ControlContextConfig): + self.config = config + def calculate_window_size(self, images: Dict[str, np.ndarray]): """Calculate required window size based on images""" max_width = 0 @@ -147,7 +148,7 @@ def handle_events(self): self.pressed_keys.remove(key_name) return self.events - + def handle_browser_events(self): """ Translate a key pressed in the web UI to the same event logic used in Pygame. @@ -189,13 +190,13 @@ def publish_observations(self, observation: Dict[str, np.ndarray], log_items: li """ Encode and publish the full observation object via ZeroMQ PUB socket. Includes observation data, events, and config information. - + Args: observation (Dict[str, np.ndarray]): Dictionary containing observation data, including images and state information """ processed_data = {} - + # Process observation data for key, value in observation.items(): if "image" in key: @@ -211,17 +212,17 @@ def publish_observations(self, observation: Dict[str, np.ndarray], log_items: li "type": "image", "encoding": "jpeg_base64", "data": b64_jpeg, - "shape": image.shape + "shape": image.shape, } else: tensor_data = value.detach().cpu().numpy() if torch.is_tensor(value) else value - + processed_data[key] = { "type": "tensor", "data": tensor_data.tolist(), - "shape": tensor_data.shape + "shape": tensor_data.shape, } - + # Add events and config information events_data = self.get_events() config_data = { @@ -230,18 +231,18 @@ def publish_observations(self, observation: Dict[str, np.ndarray], log_items: li "assign_rewards": self.config.assign_rewards, "control_phase": self.config.control_phase, "num_episodes": self.config.num_episodes, - "current_episode": self.current_episode_index + "current_episode": self.current_episode_index, } - + message = { "type": "observation_update", "timestamp": time.time(), "data": processed_data, "events": events_data, "config": config_data, - "log_items": log_items + "log_items": log_items, } - + # Send JSON over ZeroMQ self.publisher_socket.send_json(message) @@ -293,7 +294,7 @@ def update_current_episode(self, episode_index): def get_events(self): return self.events.copy() - + def log_control_info(self, start_loop_t): log_items = [] fps = self.config.fps @@ -333,7 +334,7 @@ def read_image_from_camera(cap): assign_rewards=True, control_phase=ControlPhase.RECORD, num_episodes=200, - fps=30 + fps=30, ) context = ControlContext(config) context.update_current_episode(199) diff --git a/lerobot/common/robot_devices/control_utils.py b/lerobot/common/robot_devices/control_utils.py index 4a7e3248a..fec237f68 100644 --- a/lerobot/common/robot_devices/control_utils.py +++ b/lerobot/common/robot_devices/control_utils.py @@ -185,14 +185,12 @@ def warmup_record( robot, enable_teleoperation, warmup_time_s, - display_cameras, fps, control_context ): control_loop( robot=robot, control_time_s=warmup_time_s, - display_cameras=display_cameras, fps=fps, teleoperate=enable_teleoperation, control_context=control_context, @@ -203,7 +201,6 @@ def record_episode( robot, dataset, episode_time_s, - display_cameras, policy, device, use_amp, @@ -213,7 +210,6 @@ def record_episode( control_loop( robot=robot, control_time_s=episode_time_s, - display_cameras=display_cameras, dataset=dataset, policy=policy, device=device, @@ -229,7 +225,6 @@ def control_loop( robot, control_time_s=None, teleoperate=False, - display_cameras=False, # TODO - remove this dataset: LeRobotDataset | None = None, policy=None, device=None, diff --git a/lerobot/scripts/stream_video_over_http.py b/lerobot/scripts/browser_ui_server.py similarity index 98% rename from lerobot/scripts/stream_video_over_http.py rename to lerobot/scripts/browser_ui_server.py index f411a2e6e..ee3a6f581 100644 --- a/lerobot/scripts/stream_video_over_http.py +++ b/lerobot/scripts/browser_ui_server.py @@ -89,7 +89,7 @@ def handle_keydown_event(data): @app.route("/") def index(): """Render the main page.""" - return render_template("stream_video_template.html") + return render_template("browser_ui.html") @socketio.on("connect") def handle_connect(): diff --git a/lerobot/scripts/control_robot.py b/lerobot/scripts/control_robot.py index a2dd9597e..fc6fc9ef8 100644 --- a/lerobot/scripts/control_robot.py +++ b/lerobot/scripts/control_robot.py @@ -192,7 +192,6 @@ def teleoperate( control_time_s=teleop_time_s, fps=fps, teleoperate=True, - display_cameras=display_cameras, control_context=control_context, ) @@ -291,18 +290,15 @@ def record( ) ) - warmup_record(robot, enable_teleoperation, warmup_time_s, display_cameras, fps, control_context) + warmup_record(robot, enable_teleoperation, warmup_time_s, fps, control_context) if has_method(robot, "teleop_safety_stop"): robot.teleop_safety_stop() - # @TODO(jackvial): Maybe add an update_config method to ControlContext so we don't - # initialize a new class - control_context = ControlContext( - config=ControlContextConfig( + control_context.update_config( + ControlContextConfig( robot=robot, control_phase=ControlPhase.RECORD, - display_cameras=display_cameras, play_sounds=play_sounds, assign_rewards=False, num_episodes=num_episodes, @@ -325,7 +321,6 @@ def record( dataset=dataset, robot=robot, episode_time_s=episode_time_s, - display_cameras=display_cameras, policy=policy, device=device, use_amp=use_amp, diff --git a/lerobot/templates/stream_video_template.html b/lerobot/templates/browser_ui.html similarity index 100% rename from lerobot/templates/stream_video_template.html rename to lerobot/templates/browser_ui.html