Merge pull request #140 from utn-mi/juelg/refactor-ik

refactor(ik): moved rl ik code into IK class in common module

Author: juelg

python/rcsss/_core/common.pyi

@@ -14,6 +14,7 @@ __all__ = [
     "GConfig",
     "GState",
     "Gripper",
+    "IK",
     "IdentityRotMatrix",
     "IdentityRotQuartVec",
     "IdentityTranslation",
@@ -43,6 +44,12 @@ class Gripper:
     def set_normalized_width(self, width: float, force: float = 0) -> None: ...
     def shut(self) -> None: ...
 
+class IK:
+    def __init__(self, urdf_path: str, max_duration_ms: int = 300) -> None: ...
+    def ik(
+        self, pose: Pose, q0: numpy.ndarray[typing.Literal[7], numpy.dtype[numpy.float64]], tcp_offset: Pose = ...
+    ) -> numpy.ndarray[typing.Literal[7], numpy.dtype[numpy.float64]] | None: ...
+
 class NRobotsWithGripper:
     def __init__(self, robots_with_gripper: list[RobotWithGripper]) -> None: ...
     def get_cartesian_position(self, idxs: list[int]) -> list[Pose]: ...
@@ -143,6 +150,7 @@ class RState:
 class Robot:
     def get_base_pose_in_world_coordinates(self) -> Pose: ...
     def get_cartesian_position(self) -> Pose: ...
+    def get_ik(self) -> IK | None: ...
     def get_joint_position(self) -> numpy.ndarray[typing.Literal[7], numpy.dtype[numpy.float64]]: ...
     def get_parameters(self) -> RConfig: ...
     def get_state(self) -> RState: ...

python/rcsss/_core/hw/__init__.pyi

@@ -47,15 +47,15 @@ class FHState(rcsss._core.common.GState):
     def width(self) -> float: ...
 
 class FR3(rcsss._core.common.Robot):
-    def __init__(self, ip: str, filename: str | None = None) -> None: ...
+    def __init__(self, ip: str, ik: rcsss._core.common.IK | None = None) -> None: ...
     def automatic_error_recovery(self) -> None: ...
     def double_tap_robot_to_continue(self) -> None: ...
     def get_parameters(self) -> FR3Config: ...
     def get_state(self) -> FR3State: ...
-    def set_cartesian_position_internal(self, pose: rcsss._core.common.Pose) -> None: ...
-    def set_cartesian_position_rl(
+    def set_cartesian_position_ik(
         self, pose: rcsss._core.common.Pose, max_time: float, elbow: float | None, max_force: float | None = 5
     ) -> None: ...
+    def set_cartesian_position_internal(self, pose: rcsss._core.common.Pose) -> None: ...
     def set_default_robot_behavior(self) -> None: ...
     def set_guiding_mode(self, enabled: bool) -> None: ...
     def set_parameters(self, cfg: FR3Config) -> bool: ...

python/rcsss/_core/sim.pyi

@@ -86,13 +86,12 @@ class FHState(rcsss._core.common.GState):
     def max_unnormalized_width(self) -> float: ...
 
 class FR3(rcsss._core.common.Robot):
-    def __init__(self, sim: Sim, id: str, rlmdl: str) -> None: ...
+    def __init__(self, sim: Sim, id: str, ik: rcsss._core.common.IK) -> None: ...
     def get_parameters(self) -> FR3Config: ...
     def get_state(self) -> FR3State: ...
     def set_parameters(self, cfg: FR3Config) -> bool: ...
 
 class FR3Config(rcsss._core.common.RConfig):
-    ik_duration_in_milliseconds: int
     joint_rotational_tolerance: float
     realtime: bool
     seconds_between_callbacks: float

python/rcsss/control/record.py

@@ -4,7 +4,7 @@
 from typing import cast
 
 import numpy as np
-from rcsss import hw
+from rcsss import common, hw
 
 np.set_printoptions(precision=27)
 
@@ -194,7 +194,8 @@ def __init__(
         poses: list[Pose] | None = None,
         urdf_path: str | None = None,
     ):
-        self.r: dict[str, hw.FR3] = {key: hw.FR3(ip, urdf_path) for key, ip in name2ip.items()}
+        self.ik = common.IK(urdf_path) if urdf_path else None
+        self.r: dict[str, hw.FR3] = {key: hw.FR3(ip, self.ik) for key, ip in name2ip.items()}
         # TODO: this config should be given to the constructor
         cfg = hw.FHConfig()
         cfg.epsilon_inner = 0.1

python/rcsss/control/vive.py

@@ -1,30 +1,28 @@
 import logging
 import threading
+import typing
 from enum import IntFlag, auto
 from socket import AF_INET, SOCK_DGRAM, socket
 from struct import pack, unpack
 
-import gymnasium as gym
 import numpy as np
-import rcsss
 from rcsss._core.common import RPY, Pose
-from rcsss._core.sim import CameraType
-from rcsss.camera.sim import SimCameraConfig, SimCameraSet, SimCameraSetConfig
-from rcsss.config import read_config_yaml
-from rcsss.control.fr3_desk import FCI, Desk
 from rcsss.envs.base import (
-    CameraSetWrapper,
     ControlMode,
-    FR3Env,
     GripperDictType,
-    GripperWrapper,
     LimitedTQuartRelDictType,
     RelativeActionSpace,
-    RelativeTo,
+    RobotInstance,
+)
+from rcsss.envs.factories import (
+    default_fr3_hw_gripper_cfg,
+    default_fr3_hw_robot_cfg,
+    default_fr3_sim_gripper_cfg,
+    default_fr3_sim_robot_cfg,
+    default_mujoco_cameraset_cfg,
+    fr3_hw_env,
+    fr3_sim_env,
 )
-from rcsss.envs.hw import FR3HW
-from rcsss.envs.sim import CollisionGuard, FR3Sim
-from rcsss.sim import FR3, FR3Config, Sim
 
 # import matplotlib.pyplot as plt
 
@@ -34,16 +32,18 @@
 EGO_LOCK = False
 VIVE_HOST = "192.168.100.1"
 VIVE_PORT = 54321
-USE_REAL_ROBOT = True
 INCLUDE_ROTATION = True
 ROBOT_IP = "192.168.101.1"
+ROBOT_INSTANCE = RobotInstance.HARDWARE
 
 
 class Button(IntFlag):
     L_TRIGGER = auto()
     L_SQUEEZE = auto()
+    LT_CLICK = auto()
     R_TRIGGER = auto()
     R_SQUEEZE = auto()
+    RT_CLICK = auto()
 
 
 class UDPViveActionServer(threading.Thread):
@@ -200,91 +200,38 @@ def environment_step_loop(self):
                 self._env.step(action)
 
 
-def hw():
-
-    cfg = read_config_yaml("config.yaml")
-    d = Desk(ROBOT_IP, cfg.hw.username, cfg.hw.password)
-    with FCI(d, unlock=False, lock_when_done=False):
-
-        robot = rcsss.hw.FR3(ROBOT_IP, str(rcsss.scenes["lab"].parent / "fr3.urdf"))
-        rcfg = rcsss.hw.FR3Config()
-        rcfg.tcp_offset = rcsss.common.Pose(rcsss.common.FrankaHandTCPOffset())
-        rcfg.speed_factor = 0.2
-        rcfg.controller = rcsss.hw.IKController.robotics_library
-        robot.set_parameters(rcfg)
-
-        # env = FR3Env(robot, ControlMode.CARTESIAN_TQuart)
-        env = FR3Env(robot, ControlMode.JOINTS)
-        env_hw: gym.Env = FR3HW(env)
-        gripper_cfg = rcsss.hw.FHConfig()
-        gripper_cfg.epsilon_inner = gripper_cfg.epsilon_outer = 0.5
-        gripper_cfg.speed = 0.1
-        gripper_cfg.force = 30
-        gripper = rcsss.hw.FrankaHand(ROBOT_IP, gripper_cfg)
-        # gripper.homing()
-        env_hw = GripperWrapper(env_hw, gripper, binary=True)
-
-        # TODO: camera
-        env_hw = CollisionGuard.env_from_xml_paths(
-            env_hw,
-            str(rcsss.scenes["fr3_empty_world"]),
-            str(rcsss.scenes["lab"].parent / "fr3.urdf"),
-            gripper=True,
-            check_home_collision=False,
-            camera=True,
+def main():
+    # if ROBOT_INSTANCE == RobotInstance.HARDWARE:
+    #     user, pw = load_creds_fr3_desk()
+    #     resource_manger = FCI(Desk(ROBOT_IP, user, pw), unlock=False, lock_when_done=False)
+    # else:
+    #     resource_manger = DummyResourceManager()
+
+    # with resource_manger:
+
+    if ROBOT_INSTANCE == RobotInstance.HARDWARE:
+        env_rel = fr3_hw_env(
+            ip=ROBOT_IP,
+            control_mode=ControlMode.CARTESIAN_TQuart,
+            robot_cfg=default_fr3_hw_robot_cfg(),
+            collision_guard=False,
+            gripper_cfg=default_fr3_hw_gripper_cfg(),
+            max_relative_movement=0.5,
+        )
+    else:
+        env_rel = fr3_sim_env(
             control_mode=ControlMode.CARTESIAN_TQuart,
-            tcp_offset=rcsss.common.Pose(rcsss.common.FrankaHandTCPOffset()),
+            # control_mode=ControlMode.JOINTS,
+            robot_cfg=default_fr3_sim_robot_cfg(),
+            gripper_cfg=default_fr3_sim_gripper_cfg(),
+            camera_set_cfg=default_mujoco_cameraset_cfg(),
+            max_relative_movement=0.5,
         )
 
-        env_rel = RelativeActionSpace(env_hw, relative_to=RelativeTo.CONFIGURED_ORIGIN)
-        env_rel.reset()
-        with UDPViveActionServer(VIVE_HOST, VIVE_PORT, env_rel) as action_server:
-            action_server.environment_step_loop()
-
-
-def sim():
-    simulation = Sim(rcsss.scenes["fr3_empty_world"])
-    robot = FR3(simulation, "0", str(rcsss.scenes["lab"].parent / "fr3.urdf"))
-    fr3_config = FR3Config()
-    fr3_config.realtime = False
-    # TODO: We might need a TCP offset with only translation here
-    env_sim = FR3Sim(FR3Env(robot, ControlMode.JOINTS), simulation)
-
-    cameras = {
-        "wrist": SimCameraConfig(identifier="eye-in-hand_0", type=int(CameraType.fixed)),
-        "default_free": SimCameraConfig(identifier="", type=int(CameraType.default_free)),
-    }
-    cam_cfg = SimCameraSetConfig(cameras=cameras, resolution_width=640, resolution_height=480, frame_rate=10)
-    camera_set = SimCameraSet(simulation, cam_cfg)
-    env_cam: gym.Env = CameraSetWrapper(env_sim, camera_set)
-
-    gripper_cfg = rcsss.sim.FHConfig()
-    gripper = rcsss.sim.FrankaHand(simulation, "0", gripper_cfg)
-    env_cam = GripperWrapper(env_cam, gripper)
-
-    env_cam = CollisionGuard.env_from_xml_paths(
-        env_cam,
-        str(rcsss.scenes["fr3_empty_world"]),
-        str(rcsss.scenes["lab"].parent / "fr3.urdf"),
-        gripper=True,
-        camera=False,
-        check_home_collision=False,
-        control_mode=ControlMode.CARTESIAN_TQuart,
-    )
-    env_rel = RelativeActionSpace(env_cam, relative_to=RelativeTo.CONFIGURED_ORIGIN)
     env_rel.reset()
-    with UDPViveActionServer(VIVE_HOST, VIVE_PORT, env_rel) as action_server:
-        action_server.environment_step_loop()
-
 
-def main():
-    if "lab" not in rcsss.scenes:
-        logger.error("This pip package was not built with the UTN lab models, aborting.")
-        return
-    if USE_REAL_ROBOT:
-        hw()
-    else:
-        sim()
+    with UDPViveActionServer(VIVE_HOST, VIVE_PORT, typing.cast(RelativeActionSpace, env_rel)) as action_server:
+        action_server.environment_step_loop()
 
 
 if __name__ == "__main__":

python/rcsss/envs/base.py

@@ -457,11 +457,11 @@ def check_depth(depth):
             camera_name: (
                 {
                     self.RGB_KEY: frame.camera.color.data,
+                    self.DEPTH_KEY: frame.camera.depth.data,  # type: ignore
                 }
                 if check_depth(frame.camera.depth)
                 else {
                     self.RGB_KEY: frame.camera.color.data,
-                    self.DEPTH_KEY: frame.camera.depth.data,  # type: ignore
                 }
             )
             for camera_name, frame in frameset.frames.items()

python/rcsss/envs/factories.py

@@ -1,5 +1,6 @@
 import logging
-import sys
+import os
+from os import PathLike
 
 import gymnasium as gym
 import rcsss
@@ -63,12 +64,17 @@ def fr3_hw_env(
     gripper_cfg: rcsss.hw.FHConfig | None = None,
     camera_set: BaseHardwareCameraSet | None = None,
     max_relative_movement: float | None = None,
+    urdf_path: str | None = None,
 ) -> gym.Env:
-    if "lab" not in rcsss.scenes:
-        # TODO: mujoco xml and urdf as arguments
-        logger.error("This pip package was not built with the UTN lab models, aborting.")
-        sys.exit()
-    robot = rcsss.hw.FR3(ip, str(rcsss.scenes["lab"].parent / "fr3.urdf"))
+    if urdf_path is None and "lab" in rcsss.scenes:
+        urdf_path = str(rcsss.scenes["lab"].parent / "fr3.urdf")
+        assert os.path.exists(urdf_path), "Automatic deduced urdf path does not exist. Corrupted models directory."
+        logger.info("Using automatic found urdf.")
+    elif urdf_path is None:
+        logger.warning("No urdf path was found. Proceeding, but set_cartesian methods will result in errors.")
+
+    ik = rcsss.common.IK(urdf_path) if urdf_path is not None else None
+    robot = rcsss.hw.FR3(ip, ik)
     robot.set_parameters(robot_cfg)
 
     env: gym.Env = FR3Env(robot, ControlMode.JOINTS if collision_guard else control_mode)
@@ -102,7 +108,6 @@ def fr3_hw_env(
 def default_fr3_sim_robot_cfg():
     cfg = sim.FR3Config()
     cfg.tcp_offset = rcsss.common.Pose(rcsss.common.FrankaHandTCPOffset())
-    cfg.ik_duration_in_milliseconds = 300
     cfg.realtime = False
     return cfg
 
@@ -126,12 +131,24 @@ def fr3_sim_env(
     gripper_cfg: rcsss.sim.FHConfig | None = None,
     camera_set_cfg: SimCameraSetConfig | None = None,
     max_relative_movement: float | None = None,
+    urdf_path: str | None = None,
+    mjcf: str | PathLike = "fr3_empty_world",
 ) -> gym.Env[ObsArmsGrCam, LimitedJointsRelDictType]:
-    if "lab" not in rcsss.scenes:
-        logger.error("This pip package was not built with the UTN lab models, aborting.")
-        sys.exit()
-    simulation = sim.Sim(rcsss.scenes["fr3_empty_world"])
-    robot = rcsss.sim.FR3(simulation, "0", str(rcsss.scenes["lab"].parent / "fr3.urdf"))
+
+    if urdf_path is None and "lab" in rcsss.scenes:
+        urdf_path = str(rcsss.scenes["lab"].parent / "fr3.urdf")
+        assert os.path.exists(urdf_path), "Automatic deduced urdf path does not exist. Corrupted models directory."
+        logger.info("Using automatic found urdf.")
+    elif urdf_path is None:
+        msg = "This pip package was not built with the UTN lab models, please pass the urdf and mjcf path."
+        raise ValueError(msg)
+
+    if mjcf not in rcsss.scenes:
+        logger.warning("mjcf not found as key in scenes, interpreting mjcf as path the mujoco scene xml")
+
+    simulation = sim.Sim(rcsss.scenes.get(mjcf, mjcf))  # type: ignore
+    ik = rcsss.common.IK(urdf_path)
+    robot = rcsss.sim.FR3(simulation, "0", ik)
     robot.set_parameters(robot_cfg)
     env: gym.Env = FR3Env(robot, control_mode)
     env = FR3Sim(env, simulation)

python/rcsss/envs/sim.py

@@ -108,7 +108,7 @@ def env_from_xml_paths(
         cls,
         env: gym.Env,
         mjmld: str,
-        rlmdl: str,
+        urdf: str,
         id="0",
         gripper=False,
         check_home_collision=True,
@@ -118,9 +118,9 @@ def env_from_xml_paths(
     ) -> "CollisionGuard":
         assert isinstance(env.unwrapped, FR3Env)
         simulation = sim.Sim(mjmld)
-        robot = rcsss.sim.FR3(simulation, id, rlmdl)
+        ik = rcsss.common.IK(urdf, max_duration_ms=300)
+        robot = rcsss.sim.FR3(simulation, id, ik)
         cfg = sim.FR3Config()
-        cfg.ik_duration_in_milliseconds = 300
         cfg.realtime = False
         if tcp_offset is not None:
             cfg.tcp_offset = tcp_offset

src/common/CMakeLists.txt

@@ -1,3 +1,3 @@
 add_library(common)
-target_sources(common PRIVATE Pose.cpp NRobotsWithGripper.cpp Robot.cpp)
-target_link_libraries(common PUBLIC Eigen3::Eigen)
+target_sources(common PRIVATE Pose.cpp NRobotsWithGripper.cpp Robot.cpp IK.cpp)
+target_link_libraries(common PUBLIC Eigen3::Eigen mdl)