refactor(wrappers): base environment is now platform based

- base environment is no longer the robot but for hardware an empty environment and for simulation the simulator
- greenlet implementation: MultiRobotWrapper steps the simulation only once for environment wraps
- adapted initialization of the environments wraps
- added tests to test sim greenlet implementation

Author: juelg

README.md

@@ -41,7 +41,6 @@ Flexibly compose your Gymnasium environment to fit your exact training needs. *F
 ```python
 from time import sleep
 
-import gymnasium as gym
 import numpy as np
 from rcs._core.sim import SimConfig
 from rcs.camera.sim import SimCameraSet
@@ -51,7 +50,8 @@ from rcs.envs.base import (
     GripperWrapper,
     RelativeActionSpace,
     RelativeTo,
-    RobotEnv,
+    RobotWrapper,
+    SimEnv,
 )
 from rcs.envs.sim import GripperWrapperSim, RobotSimWrapper
 from rcs.envs.utils import (
@@ -82,18 +82,19 @@ if __name__ == "__main__":
 
     # base env
     robot = rcs.sim.SimRobot(simulation, ik, robot_cfg)
-    env: gym.Env = RobotEnv(robot, ControlMode.CARTESIAN_TQuat)
+    env = SimEnv(simulation)
+    env = RobotWrapper(env, robot, ControlMode.CARTESIAN_TQuat)
 
     # gripper
     gripper = sim.SimGripper(simulation, gripper_cfg)
     env = GripperWrapper(env, gripper, binary=True)
 
-    env = RobotSimWrapper(env, simulation)
-    env = GripperWrapperSim(env, gripper)
+    env = RobotSimWrapper(env)
+    env = GripperWrapperSim(env)
 
     # camera
     camera_set = SimCameraSet(simulation, cameras, physical_units=True, render_on_demand=True)
-    env = CameraSetWrapper(env, camera_set, include_depth=True)
+    env = CameraSetWrapper(env, camera_set, include_depth=True)  # type: ignore
 
     # relative actions bounded by 10cm translation and 10 degree rotation
     env = RelativeActionSpace(env, max_mov=(0.1, np.deg2rad(10)), relative_to=RelativeTo.LAST_STEP)
@@ -104,14 +105,13 @@ if __name__ == "__main__":
     env.reset()
 
     # access low level robot api to get current cartesian position
-    print(env.unwrapped.robot.get_cartesian_position())
+    print(env.get_wrapper_attr("robot").get_cartesian_position())
 
     for _ in range(10):
         # move 1cm in x direction (forward) and close gripper
         act = {"tquat": [0.01, 0, 0, 0, 0, 0, 1], "gripper": [0]}
         obs, reward, terminated, truncated, info = env.step(act)
         print(obs)
-
 ```
 
 > **Note:** This and other examples can be found in the [`examples/`]() folder.

examples/fr3/fr3_env_cartesian_control.py

@@ -56,7 +56,7 @@ def main():
     env_rel.reset()
 
     # access low level robot api to get current cartesian position
-    print(env_rel.unwrapped.robot.get_cartesian_position())  # type: ignore
+    print(env_rel.get_wrapper_attr("robot").get_cartesian_position())  # type: ignore
 
     for _ in range(100):
         for _ in range(10):

examples/fr3/fr3_env_joint_control.py

@@ -56,7 +56,7 @@ def main():
         input("the robot is going to move, press enter whenever you are ready")
 
     # access low level robot api to get current cartesian position
-    print(env_rel.unwrapped.robot.get_joint_position())  # type: ignore
+    print(env_rel.get_wrapper_attr("robot").get_joint_position())  # type: ignore
 
     for _ in range(100):
         obs, info = env_rel.reset()

examples/fr3/fr3_readme.py

@@ -1,6 +1,5 @@
 from time import sleep
 
-import gymnasium as gym
 import numpy as np
 from rcs._core.sim import SimConfig
 from rcs.camera.sim import SimCameraSet
@@ -10,7 +9,8 @@
     GripperWrapper,
     RelativeActionSpace,
     RelativeTo,
-    RobotEnv,
+    RobotWrapper,
+    SimEnv,
 )
 from rcs.envs.sim import GripperWrapperSim, RobotSimWrapper
 from rcs.envs.utils import (
@@ -41,14 +41,15 @@
 
     # base env
     robot = rcs.sim.SimRobot(simulation, ik, robot_cfg)
-    env: gym.Env = RobotEnv(robot, ControlMode.CARTESIAN_TQuat)
+    env = SimEnv(simulation)
+    env = RobotWrapper(env, robot, ControlMode.CARTESIAN_TQuat)
 
     # gripper
     gripper = sim.SimGripper(simulation, gripper_cfg)
     env = GripperWrapper(env, gripper, binary=True)
 
-    env = RobotSimWrapper(env, simulation)
-    env = GripperWrapperSim(env, gripper)
+    env = RobotSimWrapper(env)
+    env = GripperWrapperSim(env)
 
     # camera
     camera_set = SimCameraSet(simulation, cameras, physical_units=True, render_on_demand=True)
@@ -63,7 +64,7 @@
     env.reset()
 
     # access low level robot api to get current cartesian position
-    print(env.unwrapped.robot.get_cartesian_position())
+    print(env.get_wrapper_attr("robot").get_cartesian_position())
 
     for _ in range(10):
         # move 1cm in x direction (forward) and close gripper

examples/fr3/grasp_demo.py

@@ -6,7 +6,8 @@
 import mujoco
 import numpy as np
 from rcs._core.common import Pose
-from rcs.envs.base import GripperWrapper, RobotEnv
+from rcs._core.sim import SimRobot
+from rcs.envs.base import GripperWrapper
 from rcs.envs.creators import FR3SimplePickUpSimEnvCreator
 
 logger = logging.getLogger(__name__)
@@ -16,8 +17,8 @@
 class PickUpDemo:
     def __init__(self, env: gym.Env):
         self.env = env
-        self.unwrapped: RobotEnv = cast(RobotEnv, self.env.unwrapped)
-        self.home_pose = self.unwrapped.robot.get_cartesian_position()
+        self._robot = cast(SimRobot, self.env.get_wrapper_attr("robot"))
+        self.home_pose = self._robot.get_cartesian_position()
 
     def _action(self, pose: Pose, gripper: list[float]) -> dict[str, Any]:
         return {"xyzrpy": pose.xyzrpy(), "gripper": [gripper]}
@@ -32,7 +33,7 @@ def get_object_pose(self, geom_name) -> Pose:
         ) * Pose(
             rpy_vector=np.array([0, 0, np.pi]), translation=[0, 0, 0]  # type: ignore
         )
-        return self.unwrapped.robot.to_pose_in_robot_coordinates(obj_pose_world_coordinates)
+        return self._robot.to_pose_in_robot_coordinates(obj_pose_world_coordinates)
 
     def generate_waypoints(self, start_pose: Pose, end_pose: Pose, num_waypoints: int) -> list[Pose]:
         waypoints = []
@@ -45,12 +46,13 @@ def step(self, action: dict) -> dict:
         return self.env.step(action)[0]
 
     def plan_linear_motion(self, geom_name: str, delta_up: float, num_waypoints: int = 20) -> list[Pose]:
-        end_eff_pose = self.unwrapped.robot.get_cartesian_position()
+        end_eff_pose = self._robot.get_cartesian_position()
         goal_pose = self.get_object_pose(geom_name=geom_name)
         goal_pose *= Pose(translation=np.array([0, 0, delta_up]), quaternion=np.array([1, 0, 0, 0]))  # type: ignore
         return self.generate_waypoints(end_eff_pose, goal_pose, num_waypoints=num_waypoints)
 
     def execute_motion(self, waypoints: list[Pose], gripper: list[float] = GripperWrapper.BINARY_GRIPPER_OPEN) -> dict:
+        obs = {}
         for i in range(len(waypoints)):
             obs = self.step(self._action(waypoints[i], gripper))
         return obs
@@ -65,13 +67,13 @@ def grasp(self, geom_name: str):
         self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_OPEN)
 
         for _ in range(4):
-            self.step(self._action(self.unwrapped.robot.get_cartesian_position(), GripperWrapper.BINARY_GRIPPER_CLOSED))
+            self.step(self._action(self._robot.get_cartesian_position(), GripperWrapper.BINARY_GRIPPER_CLOSED))
 
         waypoints = self.plan_linear_motion(geom_name=geom_name, delta_up=0.2, num_waypoints=60)
         self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_CLOSED)
 
     def move_home(self):
-        end_eff_pose = self.unwrapped.robot.get_cartesian_position()
+        end_eff_pose = self._robot.get_cartesian_position()
         waypoints = self.generate_waypoints(end_eff_pose, self.home_pose, num_waypoints=60)
         self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_CLOSED)
 
@@ -90,7 +92,7 @@ def main():
     sleep(3)
     for _ in range(100):
         env.reset()
-        print(env.unwrapped.robot.get_cartesian_position().translation())  # type: ignore
+        print(env.get_wrapper_attr("robot").get_cartesian_position().translation())  # type: ignore
         controller = PickUpDemo(env)
         controller.pickup("box_geom")
 

examples/fr3/grasp_digit_demo.py

@@ -5,7 +5,8 @@
 import mujoco
 import numpy as np
 from rcs._core.common import Pose
-from rcs.envs.base import GripperWrapper, RobotEnv
+from rcs._core.sim import SimRobot
+from rcs.envs.base import GripperWrapper
 from rcs_tacto.creators import FR3TactoSimplePickUpSimEnvCreator
 from tqdm import tqdm
 
@@ -16,8 +17,8 @@
 class PickUpDemo:
     def __init__(self, env: gym.Env):
         self.env = env
-        self.unwrapped: RobotEnv = cast(RobotEnv, self.env.unwrapped)
-        self.home_pose = self.unwrapped.robot.get_cartesian_position()
+        self._robot = cast(SimRobot, self.env.get_wrapper_attr("robot"))
+        self.home_pose = self._robot.get_cartesian_position()
 
     def _action(self, pose: Pose, gripper: list[float]) -> dict[str, Any]:
         return {"xyzrpy": pose.xyzrpy(), "gripper": gripper}
@@ -30,7 +31,7 @@ def get_object_pose(self, geom_name) -> Pose:
         obj_pose_world_coordinates = Pose(
             translation=data.geom_xpos[geom_id], rotation=data.geom_xmat[geom_id].reshape(3, 3)
         )
-        return self.unwrapped.robot.to_pose_in_robot_coordinates(obj_pose_world_coordinates)
+        return self._robot.to_pose_in_robot_coordinates(obj_pose_world_coordinates)
 
     def generate_waypoints(self, start_pose: Pose, end_pose: Pose, num_waypoints: int) -> list[Pose]:
         waypoints = []
@@ -43,7 +44,7 @@ def step(self, action: dict) -> dict:
         return self.env.step(action)[0]
 
     def plan_linear_motion(self, geom_name: str, delta_up: float, num_waypoints: int = 200) -> list[Pose]:
-        end_eff_pose = self.unwrapped.robot.get_cartesian_position()
+        end_eff_pose = self._robot.get_cartesian_position()
         goal_pose = self.get_object_pose(geom_name=geom_name)
         goal_pose *= Pose(translation=np.array([0, 0, delta_up]), quaternion=np.array([1, 0, 0, 0]))  # type: ignore
         return self.generate_waypoints(end_eff_pose, goal_pose, num_waypoints=num_waypoints)
@@ -68,7 +69,7 @@ def grasp(self, geom_name: str):
         self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_CLOSED)
 
     def move_home(self):
-        end_eff_pose = self.unwrapped.robot.get_cartesian_position()
+        end_eff_pose = self._robot.get_cartesian_position()
         waypoints = self.generate_waypoints(end_eff_pose, self.home_pose, num_waypoints=10)
         self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_CLOSED)
 

examples/fr3/grasp_ompl_demo.py

@@ -6,7 +6,8 @@
 import mujoco
 import numpy as np
 from rcs._core.common import Pose
-from rcs.envs.base import ControlMode, GripperWrapper, RobotEnv
+from rcs._core.sim import SimRobot
+from rcs.envs.base import ControlMode, GripperWrapper
 from rcs.envs.creators import FR3SimplePickUpSimEnvCreator
 from rcs.ompl.mj_ompl import MjOMPL
 
@@ -38,9 +39,9 @@
 class OmplTrajectoryDemo:
     def __init__(self, env: gym.Env, planner: MjOMPL):
         self.env = env
-        self.unwrapped: RobotEnv = cast(RobotEnv, self.env.unwrapped)
-        self.home_pose: Pose = self.unwrapped.robot.get_cartesian_position()
-        self.home_qpos: np.ndarray = self.unwrapped.robot.get_joint_position()
+        self._robot = cast(SimRobot, self.env.get_wrapper_attr("robot"))
+        self.home_pose: Pose = self._robot.get_cartesian_position()
+        self.home_qpos: np.ndarray = self._robot.get_joint_position()
         self.sol_path = None
         self.planner = planner
 
@@ -60,7 +61,7 @@ def get_object_pose(self, geom_name) -> Pose:
         ) * Pose(
             rpy_vector=np.array([0, 0, np.pi]), translation=[0, 0, 0]  # type: ignore
         )
-        return self.unwrapped.robot.to_pose_in_robot_coordinates(obj_pose_world_coordinates)
+        return self._robot.to_pose_in_robot_coordinates(obj_pose_world_coordinates)
 
     def plan_path_to_object(self, obj_name: str, delta_up):
         self.move_home()
@@ -83,7 +84,7 @@ def approach_and_grasp(self, obj_name: str, delta_up: float = 0.2):
 
         obj_pose_grasp = obj_pose_og * Pose(translation=np.array([0, 0, delta_up]), quaternion=np.array([1, 0, 0, 0]))  # type: ignore
         waypoints = self.generate_waypoints(
-            start_pose=self.unwrapped.robot.get_cartesian_position(), end_pose=obj_pose_grasp, num_waypoints=5
+            start_pose=self._robot.get_cartesian_position(), end_pose=obj_pose_grasp, num_waypoints=5
         )
         for waypoint in waypoints:
             self.step(self._jaction(waypoint, GripperWrapper.BINARY_GRIPPER_OPEN))  # type: ignore
@@ -108,7 +109,7 @@ def execute_motion(self, waypoints: list[Pose], gripper: list[float] = GripperWr
         return obs
 
     def move_home(self):
-        end_eff_pose = self.unwrapped.robot.get_cartesian_position()
+        end_eff_pose = self._robot.get_cartesian_position()
         waypoints = self.generate_waypoints(end_eff_pose, self.home_pose, num_waypoints=15)
         self.execute_motion(waypoints=waypoints, gripper=GripperWrapper.BINARY_GRIPPER_CLOSED)
 

extensions/rcs_fr3/src/rcs_fr3/creators.py

@@ -13,10 +13,11 @@
     ControlMode,
     GripperWrapper,
     HandWrapper,
+    HardwareEnv,
     MultiRobotWrapper,
     RelativeActionSpace,
     RelativeTo,
-    RobotEnv,
+    RobotWrapper,
 )
 from rcs.envs.creators import RCSHardwareEnvCreator
 from rcs.hand.tilburg_hand import TilburgHand
@@ -91,7 +92,8 @@ def __call__(  # type: ignore
         robot = hw.Franka(ip, ik)
         robot.set_config(robot_cfg)
 
-        env: gym.Env = RobotEnv(robot, ControlMode.JOINTS if collision_guard is not None else control_mode)
+        env = HardwareEnv()
+        env = RobotWrapper(env, robot, ControlMode.JOINTS if collision_guard is not None else control_mode)
 
         env = FR3HW(env)
         if isinstance(gripper_cfg, hw.FHConfig):
@@ -154,7 +156,8 @@ def __call__(  # type: ignore
 
         envs = {}
         for key, ip in name2ip.items():
-            env: gym.Env = RobotEnv(robots[key], control_mode)
+            env = HardwareEnv()
+            env = RobotWrapper(env, robots[key], control_mode)
             env = FR3HW(env)
             if gripper_cfg is not None:
                 gripper = hw.FrankaHand(ip, gripper_cfg)

extensions/rcs_fr3/src/rcs_fr3/envs.py

@@ -2,7 +2,7 @@
 from typing import Any, SupportsFloat, cast
 
 import gymnasium as gym
-from rcs.envs.base import RobotEnv
+from rcs._core.common import RobotPlatform
 from rcs_fr3._core import hw
 
 _logger = logging.getLogger(__name__)
@@ -11,9 +11,9 @@
 class FR3HW(gym.Wrapper):
     def __init__(self, env):
         super().__init__(env)
-        self.unwrapped: RobotEnv
-        assert isinstance(self.unwrapped.robot, hw.Franka), "Robot must be a hw.Franka instance."
-        self.hw_robot = cast(hw.Franka, self.unwrapped.robot)
+        assert self.env.get_wrapper_attr("PLATFORM") == RobotPlatform.HARDWARE, "Base environment must be hardware."
+        assert isinstance(self.get_wrapper_attr("robot"), hw.Franka), "Robot must be a hw.Franka instance."
+        self.hw_robot = cast(hw.Franka, self.get_wrapper_attr("robot"))
         self._robot_state_keys: list[str] | None = None
 
     def step(self, action: Any) -> tuple[dict[str, Any], SupportsFloat, bool, bool, dict]:
@@ -24,14 +24,12 @@ def step(self, action: Any) -> tuple[dict[str, Any], SupportsFloat, bool, bool,
         except hw.exceptions.FrankaControlException as e:
             _logger.error("FrankaControlException: %s", e)
             self.hw_robot.automatic_error_recovery()
-            # TODO: this does not work if some wrappers are in between
-            # FR3HW and RobotEnv
             return self.get_obs(), 0, False, True, {}
 
     def get_obs(self, obs: dict | None = None) -> dict[str, Any]:
         if obs is None:
-            obs = dict(self.unwrapped.get_obs())
-        robot_state = cast(hw.FrankaState, self.unwrapped.robot.get_state())
+            obs = dict(self.get_wrapper_attr("get_robot_obs")())
+        robot_state = cast(hw.FrankaState, self.hw_robot.get_state())
         obs["robot_state"] = self._rs2dict(robot_state.robot_state)
         return obs
 
@@ -44,7 +42,7 @@ def _rs2dict(self, state: hw.RobotState):
         return {key: getattr(state, key) for key in self._robot_state_keys}
 
     def reset(
-        self, seed: int | None = None, options: dict[str, Any] | None = None
+        self, *, seed: int | None = None, options: dict[str, Any] | None = None
     ) -> tuple[dict[str, Any], dict[str, Any]]:
         return super().reset(seed=seed, options=options)