Removed display_cameras from control loop args. Rename browser ui rel…

…ated code

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)

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,

lerobot/scripts/stream_video_over_http.py → 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():

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,