Merge pull request #340 from RobotControlStack/fix/sleep-bottom-of-stack

fix(hw): move rate limit to bottom of stack

Author: juelg

examples/inference/franka.py

@@ -15,7 +15,6 @@
 from rcs.envs.base import ControlMode, RelativeTo
 from rcs.envs.configs import EmptyWorldFR3Duo
 from rcs.envs.storage_wrapper import StorageWrapper
-from rcs.utils import SimpleFrameRate
 
 # from rcs_duobench.tasks.bin_sort import BinSortEnvConfig
 from vlagents.client import RemoteAgent
@@ -133,7 +132,6 @@ def __init__(self, env: gym.Env, cfg: InferenceConfig):
             cfg.jpeg_encoding,
             cfg.image_size,
         )
-        self.frame_rate = SimpleFrameRate(self._cfg.fps)
         self._action_buffer = []
 
     def submit_command(self, command: str) -> None:
@@ -298,9 +296,6 @@ def loop(self):
 
             obs_dict = self.obs_rcs2agents(obs)
 
-            if ROBOT_INSTANCE == RobotPlatform.HARDWARE:
-                self.frame_rate()
-
 
 def command_loop(controller: ModelInference) -> None:
     prompt = "Command [e=start, r=record, s=success/reset, q=stop/reset, o=reload, x=exit]: "
@@ -393,6 +388,7 @@ def get_env(cfg: InferenceConfig) -> gym.Env:
         # interpolation window of the controllers must match the rate at which we stream actions
         hw_cfg.robot_cfgs["left"].policy_rate = cfg.fps
         hw_cfg.robot_cfgs["right"].policy_rate = cfg.fps
+        hw_cfg.frequency = cfg.fps
         hw_cfg.gripper_cfgs["left"].serial_number = ROBOTIQ_SERIAL["left"]
         hw_cfg.gripper_cfgs["right"].serial_number = ROBOTIQ_SERIAL["right"]
         env_rel = env_creator.create_env(hw_cfg)

examples/teleop/franka.py

@@ -159,6 +159,7 @@ def get_env():
         # interpolation window of the controllers must match the rate at which we stream actions
         hw_cfg.robot_cfgs["left"].policy_rate = RECORD_FPS
         hw_cfg.robot_cfgs["right"].policy_rate = RECORD_FPS
+        hw_cfg.frequency = RECORD_FPS
         hw_cfg.gripper_cfgs["left"].serial_number = ROBOTIQ_SERIAL["left"]  # type: ignore
         hw_cfg.gripper_cfgs["right"].serial_number = ROBOTIQ_SERIAL["right"]  # type: ignore
         env_rel = env_creator.create_env(hw_cfg)

extensions/rcs_fr3/src/rcs_fr3/creators.py

@@ -161,6 +161,8 @@ class FR3HardwareEnvCreatorConfig:
     camera_cfgs: dict[str, HardwareCameraCreatorConfig] | None = None
     max_relative_movement: float | tuple[float, float] | None = None
     relative_to: RelativeTo = RelativeTo.LAST_STEP
+    frequency: float | None = None
+    """Control frequency in Hz, rate limits env.step(). None disables rate limiting."""
     wrapper_cfg: WrapperConfig = field(default_factory=WrapperConfig)
 
 
@@ -173,6 +175,8 @@ class FR3MultiHardwareEnvCreatorConfig:
     max_relative_movement: float | tuple[float, float] | None = None
     relative_to: RelativeTo = RelativeTo.LAST_STEP
     robot_to_shared_base_frame: dict[str, rcs.common.Pose] | None = None
+    frequency: float | None = None
+    """Control frequency in Hz, rate limits env.step(). None disables rate limiting."""
     wrapper_cfg: WrapperConfig = field(default_factory=WrapperConfig)
 
 
@@ -185,7 +189,7 @@ def create_env(self, cfg: FR3HardwareEnvCreatorConfig) -> gym.Env:
         )
         robot = hw.Franka(cfg.robot_cfg, ik)
 
-        env: gym.Env = HardwareEnv()
+        env: gym.Env = HardwareEnv(frequency=cfg.frequency)
         env = RobotWrapper(env, robot, cfg.control_mode, home_on_reset=cfg.wrapper_cfg.home_on_reset)
         env = FR3HW(env)
         if isinstance(cfg.gripper_cfg, rcs.hand.tilburg_hand.THConfig):
@@ -229,6 +233,7 @@ def create_env(self, cfg: FR3MultiHardwareEnvCreatorConfig) -> gym.Env:
                     camera_cfgs=None,
                     max_relative_movement=cfg.max_relative_movement,
                     relative_to=cfg.relative_to,
+                    frequency=cfg.frequency,
                     wrapper_cfg=cfg.wrapper_cfg,
                 )
             )

extensions/rcs_panda/src/rcs_panda/creators.py

@@ -119,6 +119,8 @@ class PandaHardwareEnvCreatorConfig:
     camera_cfgs: dict[str, HardwareCameraCreatorConfig] | None = None
     max_relative_movement: float | tuple[float, float] | None = None
     relative_to: RelativeTo = RelativeTo.LAST_STEP
+    frequency: float | None = None
+    """Control frequency in Hz, rate limits env.step(). None disables rate limiting."""
     wrapper_cfg: WrapperConfig = field(default_factory=WrapperConfig)
 
 
@@ -131,6 +133,8 @@ class PandaMultiHardwareEnvCreatorConfig:
     max_relative_movement: float | tuple[float, float] | None = None
     relative_to: RelativeTo = RelativeTo.LAST_STEP
     robot_to_shared_base_frame: dict[str, rcs.common.Pose] | None = None
+    frequency: float | None = None
+    """Control frequency in Hz, rate limits env.step(). None disables rate limiting."""
     wrapper_cfg: WrapperConfig = field(default_factory=WrapperConfig)
 
 
@@ -143,7 +147,7 @@ def create_env(self, cfg: PandaHardwareEnvCreatorConfig) -> gym.Env:
         )
         robot = hw.Franka(cfg.robot_cfg, ik)
 
-        env: gym.Env = HardwareEnv()
+        env: gym.Env = HardwareEnv(frequency=cfg.frequency)
         env = RobotWrapper(env, robot, cfg.control_mode, home_on_reset=cfg.wrapper_cfg.home_on_reset)
         env = PandaHW(env)
         if isinstance(cfg.gripper_cfg, rcs.hand.tilburg_hand.THConfig):
@@ -185,6 +189,7 @@ def create_env(self, cfg: PandaMultiHardwareEnvCreatorConfig) -> gym.Env:
                     camera_cfgs=None,
                     max_relative_movement=cfg.max_relative_movement,
                     relative_to=cfg.relative_to,
+                    frequency=cfg.frequency,
                     wrapper_cfg=cfg.wrapper_cfg,
                 )
             )

extensions/rcs_so101/src/rcs_so101/creators.py

@@ -85,6 +85,8 @@ class SO101HardwareEnvCreatorConfig:
     camera_cfgs: dict[str, HardwareCameraCreatorConfig] | None = None
     max_relative_movement: float | tuple[float, float] | None = None
     relative_to: RelativeTo = RelativeTo.LAST_STEP
+    frequency: float | None = None
+    """Control frequency in Hz, rate limits env.step(). None disables rate limiting."""
     wrapper_cfg: WrapperConfig = field(default_factory=WrapperConfig)
 
 
@@ -96,7 +98,7 @@ def create_env(self, cfg: SO101HardwareEnvCreatorConfig) -> gym.Env:
             urdf=cfg.robot_cfg.kinematic_model_path.endswith(".urdf"),
         )
         robot = SO101(cfg=cfg.robot_cfg, ik=ik)
-        env: gym.Env = HardwareEnv()
+        env: gym.Env = HardwareEnv(frequency=cfg.frequency)
         env = RobotWrapper(env, robot, cfg.control_mode, home_on_reset=cfg.wrapper_cfg.home_on_reset)
 
         gripper = SO101Gripper(robot._hf_robot, robot)

extensions/rcs_ur5e/src/rcs_ur5e/creators.py

@@ -99,6 +99,8 @@ class UR5eHardwareEnvCreatorConfig:
     camera_cfgs: dict[str, HardwareCameraCreatorConfig] | None = None
     max_relative_movement: float | tuple[float, float] | None = None
     relative_to: RelativeTo = RelativeTo.LAST_STEP
+    frequency: float | None = None
+    """Control frequency in Hz, rate limits env.step(). None disables rate limiting."""
     wrapper_cfg: WrapperConfig = field(default_factory=WrapperConfig)
 
 
@@ -110,7 +112,7 @@ def create_env(self, cfg: UR5eHardwareEnvCreatorConfig) -> gym.Env:
             urdf=cfg.robot_cfg.kinematic_model_path.endswith(".urdf"),
         )
         robot = UR5e(cfg.robot_cfg, ik)
-        env: gym.Env = HardwareEnv()
+        env: gym.Env = HardwareEnv(frequency=cfg.frequency)
         env = RobotWrapper(env, robot, cfg.control_mode, home_on_reset=cfg.wrapper_cfg.home_on_reset)
 
         if cfg.gripper_cfg is not None:

extensions/rcs_xarm7/src/rcs_xarm7/creators.py

@@ -91,6 +91,8 @@ class XArm7HardwareEnvCreatorConfig:
     hand_cfg: THConfig | None = None
     max_relative_movement: float | tuple[float, float] | None = None
     relative_to: RelativeTo = RelativeTo.LAST_STEP
+    frequency: float | None = None
+    """Control frequency in Hz, rate limits env.step(). None disables rate limiting."""
     wrapper_cfg: WrapperConfig = field(default_factory=WrapperConfig)
 
 
@@ -105,7 +107,7 @@ def create_env(self, cfg: XArm7HardwareEnvCreatorConfig) -> gym.Env:
             urdf=cfg.robot_cfg.kinematic_model_path.endswith(".urdf"),
         )
         robot = XArm7(cfg=cfg.robot_cfg, ik=ik)
-        env: gym.Env = HardwareEnv()
+        env: gym.Env = HardwareEnv(frequency=cfg.frequency)
         env = RobotWrapper(env, robot, cfg.control_mode, home_on_reset=cfg.wrapper_cfg.home_on_reset)
 
         camera_set = _create_hardware_camera_set(cfg.camera_cfgs)

extensions/rcs_yam/src/rcs_yam/creators.py

@@ -88,6 +88,8 @@ class YamHardwareEnvCreatorConfig:
     camera_cfgs: dict[str, HardwareCameraCreatorConfig] | None = None
     max_relative_movement: float | tuple[float, float] | None = None
     relative_to: RelativeTo = RelativeTo.LAST_STEP
+    frequency: float | None = None
+    """Control frequency in Hz, rate limits env.step(). None disables rate limiting."""
     wrapper_cfg: WrapperConfig = field(default_factory=WrapperConfig)
 
 
@@ -99,7 +101,7 @@ def create_env(self, cfg: YamHardwareEnvCreatorConfig) -> gym.Env:
             urdf=cfg.robot_cfg.kinematic_model_path.endswith(".urdf"),
         )
         robot = Yam(cfg.robot_cfg, ik)
-        env: gym.Env = HardwareEnv()
+        env: gym.Env = HardwareEnv(frequency=cfg.frequency)
         env = RobotWrapper(env, robot, cfg.control_mode, home_on_reset=cfg.wrapper_cfg.home_on_reset)
 
         if cfg.gripper_cfg is not None:
@@ -127,6 +129,8 @@ class YamMultiHardwareEnvCreatorConfig:
     max_relative_movement: float | tuple[float, float] | None = None
     relative_to: RelativeTo = RelativeTo.LAST_STEP
     robot_to_shared_base_frame: dict[str, rcs.common.Pose] | None = None
+    frequency: float | None = None
+    """Control frequency in Hz, rate limits env.step(). None disables rate limiting."""
     wrapper_cfg: WrapperConfig = field(default_factory=WrapperConfig)
 
 
@@ -144,6 +148,7 @@ def create_env(self, cfg: YamMultiHardwareEnvCreatorConfig) -> gym.Env:
                     camera_cfgs=None,
                     max_relative_movement=cfg.max_relative_movement,
                     relative_to=cfg.relative_to,
+                    frequency=cfg.frequency,
                     wrapper_cfg=cfg.wrapper_cfg,
                 )
             )

python/rcs/envs/base.py

@@ -211,6 +211,28 @@ def reset(
 class HardwareEnv(BaseEnv):
     PLATFORM = RobotPlatform.HARDWARE
 
+    def __init__(self, frequency: float | None = None) -> None:
+        """
+        Args:
+            frequency: Control frequency in Hz. Each env step is rate limited such that
+                the loop runs at this frequency, analogous to `SimConfig.frequency` in realtime sim.
+                None disables rate limiting.
+        """
+        super().__init__()
+        assert frequency is not None and frequency > 0, "frequency must be set to a positive value"
+        self.frame_rate = SimpleFrameRate(frequency, "Hardware Loop")
+
+    def step(self, action: dict[str, Any]) -> tuple[dict[str, Any], float, bool, bool, dict]:
+        ret = super().step(action)
+        self.frame_rate()
+        return ret
+
+    def reset(
+        self, *, seed: int | None = None, options: dict[str, Any] | None = None
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        self.frame_rate.reset()
+        return super().reset(seed=seed, options=options)
+
 
 class SimEnv(BaseEnv):
     PLATFORM = RobotPlatform.SIMULATION

python/rcs/operator/interface.py

@@ -11,7 +11,6 @@
 from rcs._core.common import RobotPlatform
 from rcs.envs.base import ArmWithGripper, ControlMode, RelativeTo
 from rcs.sim.sim import Sim
-from rcs.utils import SimpleFrameRate
 
 logger = logging.getLogger(__name__)
 
@@ -121,9 +120,6 @@ def _translate_keys(self, actions):
         return translated
 
     def environment_step_loop(self):
-        rate_limiter = SimpleFrameRate(
-            self.env_frequency if self.robot_platform == RobotPlatform.HARDWARE else None, "env loop"
-        )
 
         # 0. Initial Reset to get current positions for untracked robots
         self._last_obs, _ = self.env.reset()
@@ -180,7 +176,6 @@ def environment_step_loop(self):
 
                 self._last_obs, _, _, _, _ = self.env.step(hold_actions)
                 self.operator.set_camera(self._last_obs)
-                rate_limiter()
                 continue
 
             for controller in cmds.reset_origin_to_current:
@@ -202,11 +197,8 @@ def environment_step_loop(self):
             self._last_obs, _, _, _, _ = self.env.step(actions)
             self.operator.set_camera(self._last_obs)
 
-            rate_limiter()
-
     def sync_robot_to_operator(self, duration: float = 3.0):
         print(f"Command: Syncing robot to operator (duration: {duration}s)...")
-        rate_limiter = SimpleFrameRate(self.env_frequency, "sync loop")
         num_steps = int(duration * self.env_frequency)
 
         # 1. Capture the initial state for interpolation
@@ -239,6 +231,5 @@ def sync_robot_to_operator(self, duration: float = 3.0):
 
             self._last_obs, _, _, _, _ = self.env.step(interp_actions)
             self.operator.set_camera(self._last_obs)
-            rate_limiter()
 
         print("Sync Complete.")