refactor: rename FR3Env to RobotEnv

Author: juelg

python/examples/grasp_demo.py

@@ -5,7 +5,7 @@
 import mujoco
 import numpy as np
 from rcsss._core.common import Pose
-from rcsss.envs.base import FR3Env, GripperWrapper
+from rcsss.envs.base import GripperWrapper, RobotEnv
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -14,7 +14,7 @@
 class PickUpDemo:
     def __init__(self, env: gym.Env):
         self.env = env
-        self.unwrapped: FR3Env = cast(FR3Env, self.env.unwrapped)
+        self.unwrapped: RobotEnv = cast(RobotEnv, self.env.unwrapped)
         self.home_pose = self.unwrapped.robot.get_cartesian_position()
 
     def _action(self, pose: Pose, gripper: float) -> dict[str, Any]:

python/examples/grasp_demo_lab.py

@@ -5,7 +5,7 @@
 import mujoco
 import numpy as np
 from rcsss._core.common import Pose
-from rcsss.envs.base import FR3Env, GripperWrapper
+from rcsss.envs.base import GripperWrapper, RobotEnv
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -14,7 +14,7 @@
 class PickUpDemo:
     def __init__(self, env: gym.Env):
         self.env = env
-        self.unwrapped: FR3Env = cast(FR3Env, self.env.unwrapped)
+        self.unwrapped: RobotEnv = cast(RobotEnv, self.env.unwrapped)
         self.home_pose = self.unwrapped.robot.get_cartesian_position()
 
     def _action(self, pose: Pose, gripper: float) -> dict[str, Any]:

python/rcsss/envs/base.py

@@ -148,8 +148,8 @@ class RobotInstance(Enum):
     SIMULATION = auto()
 
 
-class FR3Env(gym.Env):
-    """Joint Gym Environment for Franka Research 3.
+class RobotEnv(gym.Env):
+    """Joint Gym Environment for a single robot arm.
 
     Top view of on the robot. Robot faces into x direction.
     z direction faces upwards. (Right handed coordinate axis)
@@ -274,7 +274,7 @@ def __init__(
         max_mov: float | tuple[float, float] | None = None,
     ):
         super().__init__(env)
-        self.unwrapped: FR3Env
+        self.unwrapped: RobotEnv
         self.action_space: gym.spaces.Dict
         self.relative_to = relative_to
         if (
@@ -483,7 +483,7 @@ class CameraSetWrapper(ActObsInfoWrapper):
 
     def __init__(self, env, camera_set: BaseCameraSet, include_depth: bool = False):
         super().__init__(env)
-        self.unwrapped: FR3Env
+        self.unwrapped: RobotEnv
         self.camera_set = camera_set
         self.include_depth = include_depth
 
@@ -574,7 +574,7 @@ class GripperWrapper(ActObsInfoWrapper):
 
     def __init__(self, env, gripper: common.Gripper, binary: bool = True):
         super().__init__(env)
-        self.unwrapped: FR3Env
+        self.unwrapped: RobotEnv
         self.observation_space: gym.spaces.Dict
         self.observation_space.spaces.update(get_space(GripperDictType).spaces)
         self.action_space: gym.spaces.Dict

python/rcsss/envs/creators.py

@@ -13,10 +13,10 @@
 from rcsss.envs.base import (
     CameraSetWrapper,
     ControlMode,
-    FR3Env,
     GripperWrapper,
     RelativeActionSpace,
     RelativeTo,
+    RobotEnv,
 )
 from rcsss.envs.hw import FR3HW
 from rcsss.envs.sim import (
@@ -86,7 +86,7 @@ def __call__(  # type: ignore
         robot = rcsss.hw.FR3(ip, ik)
         robot.set_parameters(robot_cfg)
 
-        env: gym.Env = FR3Env(robot, ControlMode.JOINTS if collision_guard is not None else control_mode)
+        env: gym.Env = RobotEnv(robot, ControlMode.JOINTS if collision_guard is not None else control_mode)
 
         env = FR3HW(env)
         if gripper_cfg is not None:
@@ -186,7 +186,7 @@ def __call__(  # type: ignore
         ik = rcsss.common.IK(urdf_path)
         robot = rcsss.sim.SimRobot(simulation, "0", ik)
         robot.set_parameters(robot_cfg)
-        env: gym.Env = FR3Env(robot, control_mode)
+        env: gym.Env = RobotEnv(robot, control_mode)
         env = FR3Sim(env, simulation, sim_wrapper)
 
         if camera_set_cfg is not None:

python/rcsss/envs/hw.py

@@ -3,15 +3,15 @@
 
 import gymnasium as gym
 from rcsss import hw
-from rcsss.envs.base import FR3Env
+from rcsss.envs.base import RobotEnv
 
 _logger = logging.getLogger(__name__)
 
 
 class FR3HW(gym.Wrapper):
     def __init__(self, env):
         super().__init__(env)
-        self.unwrapped: FR3Env
+        self.unwrapped: RobotEnv
         assert isinstance(self.unwrapped.robot, hw.FR3), "Robot must be a hw.FR3 instance."
         self.hw_robot = cast(hw.FR3, self.unwrapped.robot)
 
@@ -22,7 +22,7 @@ def step(self, action: Any) -> tuple[dict[str, Any], SupportsFloat, bool, bool,
             _logger.error("FrankaControlException: %s", e)
             self.hw_robot.automatic_error_recovery()
             # TODO: this does not work if some wrappers are in between
-            # FR3HW and FR3Env
+            # FR3HW and RobotEnv
             return dict(self.unwrapped.get_obs()), 0, False, True, {}
 
     def reset(

python/rcsss/envs/sim.py

@@ -5,7 +5,7 @@
 import numpy as np
 import rcsss
 from rcsss import sim
-from rcsss.envs.base import ControlMode, FR3Env, GripperWrapper
+from rcsss.envs.base import ControlMode, GripperWrapper, RobotEnv
 from rcsss.envs.space_utils import ActObsInfoWrapper, Vec7Type
 from rcsss.envs.utils import default_fr3_sim_robot_cfg
 
@@ -20,7 +20,7 @@ class SimWrapper(gym.Wrapper):
 
     def __init__(self, env: gym.Env, simulation: sim.Sim):
         super().__init__(env)
-        self.unwrapped: FR3Env
+        self.unwrapped: RobotEnv
         assert isinstance(self.unwrapped.robot, sim.SimRobot), "Robot must be a sim.SimRobot instance."
         self.sim = simulation
 
@@ -31,7 +31,7 @@ def __init__(self, env, simulation: sim.Sim, sim_wrapper: Type[SimWrapper] | Non
         if sim_wrapper is not None:
             env = sim_wrapper(env, simulation)
         super().__init__(env)
-        self.unwrapped: FR3Env
+        self.unwrapped: RobotEnv
         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
@@ -87,7 +87,7 @@ def __init__(
         truncate_on_collision: bool = True,
     ):
         super().__init__(env)
-        self.unwrapped: FR3Env
+        self.unwrapped: RobotEnv
         self.collision_env = collision_env
         self.sim = simulation
         self.last_obs: tuple[dict[str, Any], dict[str, Any]] | None = None
@@ -111,7 +111,7 @@ def step(self, action: dict[str, Any]) -> tuple[dict[str, Any], SupportsFloat, b
 
         if self.to_joint_control:
             fr3_env = self.collision_env.unwrapped
-            assert isinstance(fr3_env, FR3Env), "Collision env must be an FR3Env instance."
+            assert isinstance(fr3_env, RobotEnv), "Collision env must be an RobotEnv instance."
             action[self.unwrapped.joints_key] = fr3_env.robot.get_joint_position()
 
         if info["collision"]:
@@ -158,7 +158,7 @@ def env_from_xml_paths(
         sim_gui: bool = True,
         truncate_on_collision: bool = True,
     ) -> "CollisionGuard":
-        assert isinstance(env.unwrapped, FR3Env)
+        assert isinstance(env.unwrapped, RobotEnv)
         simulation = sim.Sim(mjmld)
         ik = rcsss.common.IK(urdf, max_duration_ms=300)
         robot = rcsss.sim.SimRobot(simulation, id, ik)
@@ -177,7 +177,7 @@ def env_from_xml_paths(
                 to_joint_control = True
         else:
             control_mode = env.unwrapped.get_control_mode()
-        c_env: gym.Env = FR3Env(robot, control_mode)
+        c_env: gym.Env = RobotEnv(robot, control_mode)
         c_env = FR3Sim(c_env, simulation)
         if gripper:
             gripper_cfg = sim.SimGripperConfig()
@@ -223,7 +223,7 @@ class PickCubeSuccessWrapper(gym.Wrapper):
 
     def __init__(self, env):
         super().__init__(env)
-        self.unwrapped: FR3Env
+        self.unwrapped: RobotEnv
         assert isinstance(self.unwrapped.robot, sim.SimRobot), "Robot must be a sim.SimRobot instance."
         self.sim = env.get_wrapper_attr("sim")
 
@@ -247,7 +247,7 @@ class CamRobot(gym.Wrapper):
 
     def __init__(self, env, cam_robot_joints: Vec7Type, scene: str = "lab_simple_pick_up_digit_hand"):
         super().__init__(env)
-        self.unwrapped: FR3Env
+        self.unwrapped: RobotEnv
         assert isinstance(self.unwrapped.robot, sim.SimRobot), "Robot must be a sim.SimRobot instance."
         self.sim = env.get_wrapper_attr("sim")
         self.cam_robot = rcsss.sim.SimRobot(

python/tests/test_sim_envs.py

@@ -5,9 +5,9 @@
 import rcsss
 from rcsss.envs.base import (
     ControlMode,
-    FR3Env,
     GripperDictType,
     JointsDictType,
+    RobotEnv,
     TQuatDictType,
     TRPYDictType,
 )
@@ -166,7 +166,7 @@ def test_collision_guard_trpy(self, cfg, gripper_cfg):
             max_relative_movement=None,
         )
         obs, _ = env.reset()
-        unwrapped = cast(FR3Env, env.unwrapped)
+        unwrapped = cast(RobotEnv, env.unwrapped)
         p1 = unwrapped.robot.get_joint_position()
         # an obvious below ground collision action
         obs["xyzrpy"][0] = 0.4
@@ -286,7 +286,7 @@ def test_collision_guard_tquat(self, cfg, gripper_cfg):
             max_relative_movement=None,
         )
         obs, _ = env.reset()
-        unwrapped = cast(FR3Env, env.unwrapped)
+        unwrapped = cast(RobotEnv, env.unwrapped)
         p1 = unwrapped.robot.get_joint_position()
         # an obvious below ground collision action
         obs["tquat"][0] = 0.4
@@ -379,7 +379,7 @@ def test_collision_guard_joints(self, cfg, gripper_cfg):
             max_relative_movement=None,
         )
         env.reset()
-        unwrapped = cast(FR3Env, env.unwrapped)
+        unwrapped = cast(RobotEnv, env.unwrapped)
         p1 = unwrapped.robot.get_joint_position()
         # the below action is a test_case where there is an obvious collision regardless of the gripper action
         collision_act = JointsDictType(joints=np.array([0, 1.78, 0, -1.45, 0, 0, 0], dtype=np.float32))