Merge pull request #183 from utn-mi/juelg/refactor-sim-attribute-name-handling

refactor(sim): mujoco ids into config

Author: juelg

python/examples/fr3.py

@@ -69,14 +69,14 @@ def main():
             urdf_path = get_urdf_path(URDF_PATH, allow_none_if_not_found=False)
             assert urdf_path is not None
             ik = rcs.common.IK(urdf_path)
-            robot = rcs.sim.SimRobot(simulation, "0", ik)
             cfg = sim.SimRobotConfig()
+            cfg.add_id("0")
             cfg.tcp_offset = rcs.common.Pose(rcs.common.FrankaHandTCPOffset())
-            cfg.realtime = False
-            robot.set_parameters(cfg)
+            robot = rcs.sim.SimRobot(simulation, ik, cfg)
 
             gripper_cfg_sim = sim.SimGripperConfig()
-            gripper = sim.SimGripper(simulation, "0", gripper_cfg_sim)
+            gripper_cfg_sim.add_id("0")
+            gripper = sim.SimGripper(simulation, gripper_cfg_sim)
 
             # add camera to have a rendering gui
             cameras = {

python/rcs/_core/sim.pyi

@@ -113,17 +113,23 @@ class SimCameraSetConfig:
     def frame_rate(self, arg0: int) -> None: ...
 
 class SimGripper(rcs._core.common.Gripper):
-    def __init__(self, sim: Sim, id: str, cfg: SimGripperConfig) -> None: ...
+    def __init__(self, sim: Sim, cfg: SimGripperConfig) -> None: ...
     def get_parameters(self) -> SimGripperConfig: ...
     def get_state(self) -> SimGripperState: ...
     def set_parameters(self, cfg: SimGripperConfig) -> bool: ...
 
 class SimGripperConfig(rcs._core.common.GripperConfig):
+    actuator: str
+    collision_geoms: list[str]
+    collision_geoms_fingers: list[str]
     epsilon_inner: float
     epsilon_outer: float
     ignored_collision_geoms: list[str]
+    joint1: str
+    joint2: str
     seconds_between_callbacks: float
     def __init__(self) -> None: ...
+    def add_id(self, id: str) -> None: ...
 
 class SimGripperState(rcs._core.common.GripperState):
     def __init__(self) -> None: ...
@@ -140,20 +146,27 @@ class SimGripperState(rcs._core.common.GripperState):
 
 class SimRobot(rcs._core.common.Robot):
     def __init__(
-        self, sim: Sim, id: str, ik: rcs._core.common.IK, register_convergence_callback: bool = True
+        self, sim: Sim, ik: rcs._core.common.IK, cfg: SimRobotConfig, register_convergence_callback: bool = True
     ) -> None: ...
     def get_parameters(self) -> SimRobotConfig: ...
     def get_state(self) -> SimRobotState: ...
     def set_joints_hard(self, q: numpy.ndarray[M, numpy.dtype[numpy.float64]]) -> None: ...
     def set_parameters(self, cfg: SimRobotConfig) -> bool: ...
 
 class SimRobotConfig(rcs._core.common.RobotConfig):
+    actuators: list[str]
+    arm_collision_geoms: list[str]
+    attachment_site: str
+    base: str
     joint_rotational_tolerance: float
+    joints: list[str]
     realtime: bool
+    robot_type: rcs._core.common.RobotType
     seconds_between_callbacks: float
     tcp_offset: rcs._core.common.Pose
     trajectory_trace: bool
     def __init__(self) -> None: ...
+    def add_id(self, id: str) -> None: ...
 
 class SimRobotState(rcs._core.common.RobotState):
     def __init__(self) -> None: ...

python/rcs/envs/creators.py

@@ -190,8 +190,7 @@ def __call__(  # type: ignore
         simulation = sim.Sim(mjb_file)
 
         ik = rcs.common.IK(urdf_path)
-        robot = rcs.sim.SimRobot(simulation, "0", ik)
-        robot.set_parameters(robot_cfg)
+        robot = rcs.sim.SimRobot(simulation, ik, robot_cfg)
         env: gym.Env = RobotEnv(robot, control_mode)
         env = FR3Sim(env, simulation, sim_wrapper)
 
@@ -200,7 +199,7 @@ def __call__(  # type: ignore
             env = CameraSetWrapper(env, camera_set, include_depth=True)
 
         if gripper_cfg is not None and isinstance(gripper_cfg, rcs.sim.SimGripperConfig):
-            gripper = sim.SimGripper(simulation, "0", gripper_cfg)
+            gripper = sim.SimGripper(simulation, gripper_cfg)
             env = GripperWrapper(env, gripper, binary=True)
             env = GripperWrapperSim(env, gripper)
 

python/rcs/envs/sim.py

@@ -158,15 +158,15 @@ def env_from_xml_paths(
         sim_gui: bool = True,
         truncate_on_collision: bool = True,
     ) -> "CollisionGuard":
+        # TODO: this needs to support non FR3 robots
         assert isinstance(env.unwrapped, RobotEnv)
         simulation = sim.Sim(mjmld)
         ik = rcs.common.IK(urdf, max_duration_ms=300)
-        robot = rcs.sim.SimRobot(simulation, id, ik)
-        cfg = sim.SimRobotConfig()
+        cfg = default_fr3_sim_robot_cfg(mjmld, id)
         cfg.realtime = False
         if tcp_offset is not None:
             cfg.tcp_offset = tcp_offset
-        robot.set_parameters(cfg)
+        robot = rcs.sim.SimRobot(simulation, ik, cfg)
         to_joint_control = False
         if control_mode is not None:
             if control_mode != env.unwrapped.get_control_mode():
@@ -181,7 +181,8 @@ def env_from_xml_paths(
         c_env = FR3Sim(c_env, simulation)
         if gripper:
             gripper_cfg = sim.SimGripperConfig()
-            fh = sim.SimGripper(simulation, id, gripper_cfg)
+            gripper_cfg.add_id(id)
+            fh = sim.SimGripper(simulation, gripper_cfg)
             c_env = GripperWrapper(c_env, fh)
             c_env = GripperWrapperSim(c_env, fh)
         return cls(
@@ -250,8 +251,9 @@ def __init__(self, env, cam_robot_joints: VecType, scene: str = "lab_simple_pick
         self.unwrapped: RobotEnv
         assert isinstance(self.unwrapped.robot, sim.SimRobot), "Robot must be a sim.SimRobot instance."
         self.sim = env.get_wrapper_attr("sim")
+        cfg = default_fr3_sim_robot_cfg(scene, "1")
         self.cam_robot = rcs.sim.SimRobot(
-            self.sim, "1", env.unwrapped.robot.get_ik(), register_convergence_callback=False
+            self.sim, env.unwrapped.robot.get_ik(), cfg, register_convergence_callback=False
         )
         self.cam_robot.set_parameters(default_fr3_sim_robot_cfg(scene))
         self.cam_robot_joints = cam_robot_joints

python/rcs/envs/utils.py

@@ -17,14 +17,16 @@
 logger.setLevel(logging.INFO)
 
 
-def default_fr3_sim_robot_cfg(mjcf: str | Path = "fr3_empty_world") -> sim.SimRobotConfig:
+def default_fr3_sim_robot_cfg(mjcf: str | Path = "fr3_empty_world", idx: str = "0") -> sim.SimRobotConfig:
     cfg = sim.SimRobotConfig()
     cfg.tcp_offset = get_tcp_offset(mjcf)
     cfg.realtime = False
+    cfg.robot_type = rcs.common.RobotType.FR3
+    cfg.add_id(idx)
     return cfg
 
 
-def default_fr3_hw_robot_cfg(async_control: bool = False):
+def default_fr3_hw_robot_cfg(async_control: bool = False) -> FR3Config:
     robot_cfg = FR3Config()
     robot_cfg.tcp_offset = rcs.common.Pose(rcs.common.FrankaHandTCPOffset())
     robot_cfg.speed_factor = 0.1
@@ -33,7 +35,7 @@ def default_fr3_hw_robot_cfg(async_control: bool = False):
     return robot_cfg
 
 
-def default_fr3_hw_gripper_cfg(async_control: bool = False):
+def default_fr3_hw_gripper_cfg(async_control: bool = False) -> rcs.hw.FHConfig:
     gripper_cfg = rcs.hw.FHConfig()
     gripper_cfg.epsilon_inner = gripper_cfg.epsilon_outer = 0.1
     gripper_cfg.speed = 0.1
@@ -65,11 +67,13 @@ def default_realsense(name2id: dict[str, str] | None) -> RealSenseCameraSet | No
     return RealSenseCameraSet(cam_cfg)
 
 
-def default_fr3_sim_gripper_cfg():
-    return sim.SimGripperConfig()
+def default_fr3_sim_gripper_cfg(idx: str = "0") -> sim.SimGripperConfig:
+    cfg = sim.SimGripperConfig()
+    cfg.add_id(idx)
+    return cfg
 
 
-def default_mujoco_cameraset_cfg():
+def default_mujoco_cameraset_cfg() -> SimCameraSetConfig:
     cameras = {
         "wrist": SimCameraConfig(identifier="wrist_0", type=int(CameraType.fixed)),
         "default_free": SimCameraConfig(identifier="", type=int(CameraType.default_free)),
@@ -81,7 +85,7 @@ def default_mujoco_cameraset_cfg():
     )
 
 
-def get_tcp_offset(mjcf: str | Path):
+def get_tcp_offset(mjcf: str | Path) -> rcs.common.Pose:
     """Reads out tcp offset set in mjcf file.
 
     Convention: The tcp offset is stored in the model as a numeric attribute named "tcp_offset".

src/pybind/rcs.cpp

@@ -413,7 +413,16 @@ PYBIND11_MODULE(_core, m) {
                      &rcs::sim::SimRobotConfig::seconds_between_callbacks)
       .def_readwrite("realtime", &rcs::sim::SimRobotConfig::realtime)
       .def_readwrite("trajectory_trace",
-                     &rcs::sim::SimRobotConfig::trajectory_trace);
+                     &rcs::sim::SimRobotConfig::trajectory_trace)
+      .def_readwrite("robot_type", &rcs::sim::SimRobotConfig::robot_type)
+      .def_readwrite("arm_collision_geoms",
+                     &rcs::sim::SimRobotConfig::arm_collision_geoms)
+      .def_readwrite("joints", &rcs::sim::SimRobotConfig::joints)
+      .def_readwrite("attachment_site",
+                     &rcs::sim::SimRobotConfig::attachment_site)
+      .def_readwrite("actuators", &rcs::sim::SimRobotConfig::actuators)
+      .def_readwrite("base", &rcs::sim::SimRobotConfig::base)
+      .def("add_id", &rcs::sim::SimRobotConfig::add_id, py::arg("id"));
   py::class_<rcs::sim::SimRobotState, rcs::common::RobotState>(sim,
                                                                "SimRobotState")
       .def(py::init<>())
@@ -436,7 +445,15 @@ PYBIND11_MODULE(_core, m) {
       .def_readwrite("seconds_between_callbacks",
                      &rcs::sim::SimGripperConfig::seconds_between_callbacks)
       .def_readwrite("ignored_collision_geoms",
-                     &rcs::sim::SimGripperConfig::ignored_collision_geoms);
+                     &rcs::sim::SimGripperConfig::ignored_collision_geoms)
+      .def_readwrite("collision_geoms",
+                     &rcs::sim::SimGripperConfig::collision_geoms)
+      .def_readwrite("collision_geoms_fingers",
+                     &rcs::sim::SimGripperConfig::collision_geoms_fingers)
+      .def_readwrite("joint1", &rcs::sim::SimGripperConfig::joint1)
+      .def_readwrite("joint2", &rcs::sim::SimGripperConfig::joint2)
+      .def_readwrite("actuator", &rcs::sim::SimGripperConfig::actuator)
+      .def("add_id", &rcs::sim::SimGripperConfig::add_id, py::arg("id"));
   py::class_<rcs::sim::SimGripperState, rcs::common::GripperState>(
       sim, "SimGripperState")
       .def(py::init<>())
@@ -462,18 +479,19 @@ PYBIND11_MODULE(_core, m) {
       .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>, const std::string &,
+      .def(py::init<std::shared_ptr<rcs::sim::Sim>,
                     const rcs::sim::SimGripperConfig &>(),
-           py::arg("sim"), py::arg("id"), py::arg("cfg"))
+           py::arg("sim"), py::arg("cfg"))
       .def("get_parameters", &rcs::sim::SimGripper::get_parameters)
       .def("get_state", &rcs::sim::SimGripper::get_state)
       .def("set_parameters", &rcs::sim::SimGripper::set_parameters,
            py::arg("cfg"));
   py::class_<rcs::sim::SimRobot, rcs::common::Robot,
              std::shared_ptr<rcs::sim::SimRobot>>(sim, "SimRobot")
-      .def(py::init<std::shared_ptr<rcs::sim::Sim>, const std::string &,
-                    std::shared_ptr<rcs::common::IK>, bool>(),
-           py::arg("sim"), py::arg("id"), py::arg("ik"),
+      .def(py::init<std::shared_ptr<rcs::sim::Sim>,
+                    std::shared_ptr<rcs::common::IK>, rcs::sim::SimRobotConfig,
+                    bool>(),
+           py::arg("sim"), py::arg("ik"), py::arg("cfg"),
            py::arg("register_convergence_callback") = true)
       .def("get_parameters", &rcs::sim::SimRobot::get_parameters)
       .def("set_parameters", &rcs::sim::SimRobot::set_parameters,

src/sim/SimGripper.cpp

@@ -7,49 +7,33 @@
 #include <string>
 #include <tuple>
 
-struct {
-  // TODO(juelg): add camera mount collision (first a name needs to be
-  // added to the model)
-  std::vector<std::string> collision_geoms{"hand_c", "d435i_collision",
-                                           "finger_0_left", "finger_0_right"};
-
-  std::vector<std::string> collision_geoms_fingers{"finger_0_left",
-                                                   "finger_0_right"};
-  std::string joint1 = "finger_joint1";
-  std::string joint2 = "finger_joint2";
-  std::string actuator = "actuator8";
-} const gripper_names;
-
 namespace rcs {
 namespace sim {
 
-SimGripper::SimGripper(std::shared_ptr<Sim> sim, const std::string &id,
-                       const SimGripperConfig &cfg)
-    : sim{sim}, cfg{cfg}, id{id} {
+SimGripper::SimGripper(std::shared_ptr<Sim> sim, const SimGripperConfig &cfg)
+    : sim{sim}, cfg{cfg} {
   this->state = SimGripperState();
-  this->actuator_id = mj_name2id(this->sim->m, mjOBJ_ACTUATOR,
-                                 (gripper_names.actuator + "_" + id).c_str());
+  this->actuator_id =
+      mj_name2id(this->sim->m, mjOBJ_ACTUATOR, this->cfg.actuator.c_str());
   if (this->actuator_id == -1) {
     throw std::runtime_error(
-        std::string("No actuator named " + gripper_names.actuator));
+        std::string("No actuator named " + this->cfg.actuator));
   }
   // this->tendon_id =
   //     mj_name2id(this->sim->m, mjOBJ_TENDON, ("split_" + id).c_str());
   this->joint_id_1 = this->sim->m->jnt_qposadr[mj_name2id(
-      this->sim->m, mjOBJ_JOINT, (gripper_names.joint1 + "_" + id).c_str())];
+      this->sim->m, mjOBJ_JOINT, this->cfg.joint1.c_str())];
   if (this->joint_id_1 == -1) {
-    throw std::runtime_error(
-        std::string("No joint named " + gripper_names.joint1));
+    throw std::runtime_error(std::string("No joint named " + this->cfg.joint1));
   }
   this->joint_id_2 = this->sim->m->jnt_qposadr[mj_name2id(
-      this->sim->m, mjOBJ_JOINT, (gripper_names.joint2 + "_" + id).c_str())];
+      this->sim->m, mjOBJ_JOINT, this->cfg.joint2.c_str())];
   if (this->joint_id_2 == -1) {
-    throw std::runtime_error(
-        std::string("No joint named " + gripper_names.joint2));
+    throw std::runtime_error(std::string("No joint named " + this->cfg.joint2));
   }
   // Collision geoms
-  this->add_collision_geoms(gripper_names.collision_geoms, this->cgeom, false);
-  this->add_collision_geoms(gripper_names.collision_geoms_fingers, this->cfgeom,
+  this->add_collision_geoms(this->cfg.collision_geoms, this->cgeom, false);
+  this->add_collision_geoms(this->cfg.collision_geoms_fingers, this->cfgeom,
                             false);
 
   this->sim->register_all_cb(std::bind(&SimGripper::convergence_callback, this),
@@ -68,7 +52,7 @@ void SimGripper::add_collision_geoms(const std::vector<std::string> &cgeoms_str,
     cgeoms_set.clear();
   }
   for (size_t i = 0; i < std::size(cgeoms_str); ++i) {
-    std::string name = cgeoms_str[i] + "_" + id;
+    std::string name = cgeoms_str[i];
 
     int coll_id = mj_name2id(this->sim->m, mjOBJ_GEOM, name.c_str());
     if (coll_id == -1) {

src/sim/SimGripper.h

@@ -18,6 +18,29 @@ struct SimGripperConfig : common::GripperConfig {
   double epsilon_outer = 0.005;
   double seconds_between_callbacks = 0.05;  // 20 Hz
   std::vector<std::string> ignored_collision_geoms = {};
+  std::vector<std::string> collision_geoms{"hand_c", "d435i_collision",
+                                           "finger_0_left", "finger_0_right"};
+
+  std::vector<std::string> collision_geoms_fingers{"finger_0_left",
+                                                   "finger_0_right"};
+  std::string joint1 = "finger_joint1";
+  std::string joint2 = "finger_joint2";
+  std::string actuator = "actuator8";
+
+  void add_id(const std::string &id) {
+    for (auto &s : this->collision_geoms) {
+      s = s + "_" + id;
+    }
+    for (auto &s : this->collision_geoms_fingers) {
+      s = s + "_" + id;
+    }
+    for (auto &s : this->ignored_collision_geoms) {
+      s = s + "_" + id;
+    }
+    this->joint1 = this->joint1 + "_" + id;
+    this->joint2 = this->joint2 + "_" + id;
+    this->actuator = this->actuator + "_" + id;
+  }
 };
 
 struct SimGripperState : common::GripperState {
@@ -38,7 +61,6 @@ class SimGripper : public common::Gripper {
   int joint_id_1;
   int joint_id_2;
   SimGripperState state;
-  std::string id;
   bool convergence_callback();
   bool collision_callback();
   std::set<size_t> cgeom;
@@ -49,8 +71,7 @@ class SimGripper : public common::Gripper {
   void m_reset();
 
  public:
-  SimGripper(std::shared_ptr<Sim> sim, const std::string &id,
-             const SimGripperConfig &cfg);
+  SimGripper(std::shared_ptr<Sim> sim, const SimGripperConfig &cfg);
   ~SimGripper() override;
 
   bool set_parameters(const SimGripperConfig &cfg);