feat(sim): sim config support for async and realtime mode

- added frequency flag to sim config
- pybind for sim config
- support for async sim
- fixed realtime mode in sim

Author: juelg

python/rcs/_core/sim.pyi

@@ -16,6 +16,7 @@ __all__ = [
     "Sim",
     "SimCameraConfig",
     "SimCameraSet",
+    "SimConfig",
     "SimGripper",
     "SimGripperConfig",
     "SimGripperState",
@@ -86,8 +87,10 @@ class Sim:
     def __init__(self, mjmdl: int, mjdata: int) -> None: ...
     def _start_gui_server(self, id: str) -> None: ...
     def _stop_gui_server(self) -> None: ...
+    def get_config(self) -> SimConfig: ...
     def is_converged(self) -> bool: ...
     def reset(self) -> None: ...
+    def set_config(self, cfg: SimConfig) -> bool: ...
     def step(self, k: int) -> None: ...
     def step_until_convergence(self) -> None: ...
 
@@ -106,6 +109,13 @@ class SimCameraSet:
     @property
     def _sim(self) -> Sim: ...
 
+class SimConfig:
+    async_control: bool
+    frequency: int
+    max_convergence_steps: int
+    realtime: bool
+    def __init__(self) -> None: ...
+
 class SimGripper(rcs._core.common.Gripper):
     def __init__(self, sim: Sim, cfg: SimGripperConfig) -> None: ...
     def get_parameters(self) -> SimGripperConfig: ...

python/rcs/envs/creators.py

@@ -47,6 +47,7 @@ def __call__(  # type: ignore
         robot_cfg: rcs.sim.SimRobotConfig,
         collision_guard: bool = False,
         gripper_cfg: rcs.sim.SimGripperConfig | None = None,
+        sim_cfg: rcs.sim.SimConfig | None = None,
         hand_cfg: rcs.sim.SimTilburgHandConfig | None = None,
         cameras: dict[str, SimCameraConfig] | None = None,
         max_relative_movement: float | tuple[float, float] | None = None,
@@ -75,7 +76,7 @@ def __call__(  # type: ignore
         Returns:
             gym.Env: The configured simulation environment for the FR3 robot.
         """
-        simulation = sim.Sim(robot_cfg.mjcf_scene_path)
+        simulation = sim.Sim(robot_cfg.mjcf_scene_path, sim_cfg)
 
         ik = rcs.common.Pin(
             robot_cfg.kinematic_model_path,
@@ -137,6 +138,7 @@ def __call__(  # type: ignore
         cameras: dict[str, SimCameraConfig] | None = None,
         hand_cfg: rcs.sim.SimTilburgHandConfig | None = None,
         gripper_cfg: rcs.sim.SimGripperConfig | None = None,
+        sim_cfg: rcs.sim.SimConfig | None = None,
         random_pos_args: dict | None = None,
     ) -> gym.Env:
         mode = "gripper"
@@ -170,6 +172,7 @@ def __call__(  # type: ignore
             collision_guard=False,
             gripper_cfg=_gripper_cfg,
             hand_cfg=_hand_cfg,
+            sim_cfg=sim_cfg,
             cameras=cameras,
             max_relative_movement=(0.2, np.deg2rad(45)) if delta_actions else None,
             relative_to=RelativeTo.LAST_STEP,
@@ -213,8 +216,12 @@ def __call__(  # type: ignore
             translation=np.array([0.0, 0.0, 0.1034]),  # type: ignore
             rotation=np.array([[0.707, 0.707, 0], [-0.707, 0.707, 0], [0, 0, 1]]),  # type: ignore
         )
+        sim_cfg = sim.SimConfig()
+        sim_cfg.realtime = False
+        sim_cfg.async_control = True
+        sim_cfg.frequency = 30  # in Hz
 
-        return SimTaskEnvCreator()(robot_cfg, render_mode, control_mode, delta_actions, cameras)
+        return SimTaskEnvCreator()(robot_cfg, render_mode, control_mode, delta_actions, cameras, sim_cfg=sim_cfg)
 
 
 class FR3LabDigitGripperPickUpSimEnvCreator(EnvCreator):

python/rcs/envs/sim.py

@@ -14,7 +14,9 @@
 from rcs.envs.utils import default_sim_robot_cfg, default_sim_tilburg_hand_cfg
 
 import rcs
+import time
 from rcs import sim
+from rcs.utils import SimpleFrameRate
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -33,7 +35,7 @@ def __init__(self, env: gym.Env, simulation: sim.Sim):
 
 
 class RobotSimWrapper(gym.Wrapper):
-    def __init__(self, env, simulation: sim.Sim, sim_wrapper: Type[SimWrapper] | None = None):
+    def __init__(self, env, simulation: sim.Sim, sim_wrapper: SimWrapper | None = None):
         self.sim_wrapper = sim_wrapper
         if sim_wrapper is not None:
             env = sim_wrapper(env, simulation)
@@ -42,10 +44,20 @@ def __init__(self, env, simulation: sim.Sim, sim_wrapper: Type[SimWrapper] | Non
         assert isinstance(self.unwrapped.robot, sim.SimRobot), "Robot must be a sim.SimRobot instance."
         self.sim_robot = cast(sim.SimRobot, self.unwrapped.robot)
         self.sim = simulation
+        cfg = self.sim.get_config()
+        self.frame_rate = SimpleFrameRate(1/cfg.frequency, "RobotSimWrapper")
 
     def step(self, action: dict[str, Any]) -> tuple[dict[str, Any], float, bool, bool, dict]:
         _, _, _, _, info = super().step(action)
-        self.sim.step_until_convergence()
+        cfg = self.sim.get_config()
+        if cfg.async_control:
+            self.sim.step(round(1/cfg.frequency / self.sim.model.opt.timestep))
+            if cfg.realtime:
+                self.frame_rate.frame_rate = 1/cfg.frequency
+                self.frame_rate()
+
+        else:
+            self.sim.step_until_convergence()
         state = self.sim_robot.get_state()
         info["collision"] = state.collision
         info["ik_success"] = state.ik_success
@@ -348,7 +360,7 @@ def reset(
 class RandomCubePos(SimWrapper):
     """Wrapper to randomly place cube in the lab environments."""
 
-    def __init__(self, env: gym.Env, simulation: sim.Sim, include_rotation: bool = False):
+    def __init__(self, env: gym.Env, simulation: sim.Sim, include_rotation: bool = True):
         super().__init__(env, simulation)
         self.include_rotation = include_rotation
 
@@ -361,7 +373,7 @@ def reset(
         iso_cube = np.array([0.498, 0.0, 0.226])
         iso_cube_pose = rcs.common.Pose(translation=np.array(iso_cube), rpy_vector=np.array([0, 0, 0]))  # type: ignore
         iso_cube = self.unwrapped.robot.to_pose_in_world_coordinates(iso_cube_pose).translation()
-        pos_z = 0.826
+        pos_z = 0.0288 / 2
         pos_x = iso_cube[0] + np.random.random() * 0.2 - 0.1
         pos_y = iso_cube[1] + np.random.random() * 0.2 - 0.1
 

python/rcs/sim/__init__.py

@@ -9,6 +9,7 @@
     SimTilburgHand,
     SimTilburgHandConfig,
     SimTilburgHandState,
+    SimConfig,
 )
 from rcs.sim.sim import Sim, gui_loop
 
@@ -25,4 +26,5 @@
     "SimTilburgHandState",
     "gui_loop",
     "SimCameraConfig",
+    "SimConfig",
 ]

python/rcs/sim/sim.py

@@ -12,7 +12,7 @@
 import mujoco.viewer
 from rcs._core.sim import GuiClient as _GuiClient
 from rcs._core.sim import Sim as _Sim
-from rcs.sim import egl_bootstrap
+from rcs.sim import egl_bootstrap, SimConfig
 from rcs.utils import SimpleFrameRate
 
 egl_bootstrap.bootstrap()
@@ -42,7 +42,7 @@ def gui_loop(gui_uuid: str, close_event):
 
 
 class Sim(_Sim):
-    def __init__(self, mjmdl: str | PathLike):
+    def __init__(self, mjmdl: str | PathLike, cfg: SimConfig = None):
         mjmdl = Path(mjmdl)
         if mjmdl.suffix == ".xml":
             self.model = mj.MjModel.from_xml_path(str(mjmdl))
@@ -58,6 +58,8 @@ def __init__(self, mjmdl: str | PathLike):
         self._gui_client: Optional[_GuiClient] = None
         self._gui_process: Optional[mp.context.SpawnProcess] = None
         self._stop_event: Optional[EventClass] = None
+        if cfg is not None:
+            self.set_config(cfg)
 
     def close_gui(self):
         if self._stop_event is not None:

python/rcs/utils.py

@@ -36,7 +36,7 @@ def __call__(self):
             sleep(sleep_time)
         if self._last_print is None or perf_counter() - self._last_print > 10:
             self._last_print = perf_counter()
-            logger.info(f"FPS {self.loop_name}: {1 / (perf_counter() - self.t)}")
+            logger.debug(f"FPS {self.loop_name}: {1 / (perf_counter() - self.t)}")
 
         self.t = perf_counter()
 

src/pybind/rcs.cpp

@@ -487,6 +487,18 @@ PYBIND11_MODULE(_core, m) {
       .def_readonly("last_width", &rcs::sim::SimGripperState::last_width)
       .def_readonly("collision", &rcs::sim::SimGripperState::collision);
 
+  py::class_<rcs::sim::SimConfig>(
+      sim, "SimConfig")
+      .def(py::init<>())
+      .def_readwrite("async_control",
+                     &rcs::sim::SimConfig::async_control)
+      .def_readwrite("realtime",
+                     &rcs::sim::SimConfig::realtime)
+      .def_readwrite("frequency",
+                     &rcs::sim::SimConfig::frequency)
+      .def_readwrite("max_convergence_steps",
+                     &rcs::sim::SimConfig::max_convergence_steps);
+
   py::class_<rcs::sim::Sim, std::shared_ptr<rcs::sim::Sim>>(sim, "Sim")
       .def(py::init([](long m, long d) {
              return std::make_shared<rcs::sim::Sim>((mjModel *)m, (mjData *)d);
@@ -495,10 +507,14 @@ PYBIND11_MODULE(_core, m) {
       .def("step_until_convergence", &rcs::sim::Sim::step_until_convergence,
            py::call_guard<py::gil_scoped_release>())
       .def("is_converged", &rcs::sim::Sim::is_converged)
+      .def("set_config", &rcs::sim::Sim::set_config, py::arg("cfg"))
+      .def("get_config", &rcs::sim::Sim::get_config)
       .def("step", &rcs::sim::Sim::step, py::arg("k"))
       .def("reset", &rcs::sim::Sim::reset)
       .def("_start_gui_server", &rcs::sim::Sim::start_gui_server, py::arg("id"))
       .def("_stop_gui_server", &rcs::sim::Sim::stop_gui_server);
+
+
   py::class_<rcs::sim::SimGripper, rcs::common::Gripper,
              std::shared_ptr<rcs::sim::SimGripper>>(sim, "SimGripper")
       .def(py::init<std::shared_ptr<rcs::sim::Sim>,

src/sim/sim.cpp

@@ -28,16 +28,12 @@ bool get_last_return_value(ConditionCallback cb) {
 
 Sim::Sim(mjModel* m, mjData* d) : m(m), d(d), renderer(m){};
 
-bool Sim::set_config(const Config& cfg) {
-  if (cfg.async) {
-    /* Not implemented :) */
-    return false;
-  }
+bool Sim::set_config(const SimConfig& cfg) {
   this->cfg = cfg;
   return true;
 }
 
-Config Sim::get_config() { return this->cfg; }
+SimConfig Sim::get_config() { return this->cfg; }
 
 void Sim::invoke_callbacks() {
   for (int i = 0; i < std::size(this->callbacks); ++i) {
@@ -115,9 +111,6 @@ void Sim::step(size_t k) {
     this->invoke_callbacks();
     mj_step2(this->m, this->d);
     this->invoke_rendering_callbacks();
-    if (this->cfg.realtime) {
-      std::this_thread::sleep_for(std::chrono::milliseconds(10));
-    }
   }
 }
 

src/sim/sim.h

@@ -26,10 +26,11 @@ class Renderer {
   std::unordered_map<std::string, mjrContext*> ctxs;
 };
 
-struct Config {
-  bool async = false;
+struct SimConfig {
+  bool async_control = false;
   bool realtime = false;
-  int max_convergence_steps = 5000;
+  int frequency = 30; // in Hz
+  int max_convergence_steps = 500;
 };
 
 struct Callback {
@@ -56,7 +57,7 @@ struct RenderingCallback {
 
 class Sim {
  private:
-  Config cfg;
+  SimConfig cfg;
   std::vector<Callback> callbacks;
   std::vector<ConditionCallback> any_callbacks;
   std::vector<ConditionCallback> all_callbacks;
@@ -74,8 +75,8 @@ class Sim {
   mjModel* m;
   mjData* d;
   Sim(mjModel* m, mjData* d);
-  bool set_config(const Config& cfg);
-  Config get_config();
+  bool set_config(const SimConfig& cfg);
+  SimConfig get_config();
   bool is_converged();
   void step_until_convergence();
   void step(size_t k);