Typed Factory for Gym Envs

.github/workflows/py.yaml

@@ -27,24 +27,31 @@ jobs:
     env:
       CC: clang-15
       CXX: clang++-15
+      GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
     steps:
       - uses: actions/checkout@v4
       - name: Install CPP dependencies
         run: sudo apt install $(cat debian_deps.txt)
+      - name: Install virtual frame buffer tool
+        run: sudo apt install xvfb
       - name: Set up Python 3.11
         uses: actions/setup-python@v5
         with:
           python-version: "3.11"
           cache: "pip"
       - name: Install building dependencies
         run: python -m pip install -r requirements_dev.txt
+      - name: Download the models using cmake
+        env:
+          MODEL_TOKEN: ${{ secrets.MODEL_TOKEN }}
+        run: cmake -G Ninja -B build -DINCLUDE_UTN_MODELS=ON -DGITLAB_MODELS_TOKEN="$MODEL_TOKEN" -DMUJOCO_VERSION=3.1.5
       - name: Create python wheel
         run: python -m build --no-isolation --wheel --outdir dist/
       - name: Install the wheel package
-        run: python -m pip install --no-build-isolation dist/*.whl
+        run: python -m pip install -v --no-build-isolation dist/*.whl
       - name: Code linting
         run: make pylint
       - name: Check that stub files are up-to-date
         run: make stubgen && git diff --exit-code
       - name: Ensure all unittests(pytest) are passing
-        run: make pytest
+        run: xvfb-run make pytest

README.md

@@ -12,7 +12,7 @@ sudo apt install $(cat debian_deps.txt)
 python3 -m venv .venv
 source .venv/bin/activate
 pip config --site set global.no-build-isolation false
-pip install --requirements requirements_dev.txt
+pip install -r requirements_dev.txt
 ```
 3. Build and install RCS:
 ```shell

python/rcsss/envs/base.py

@@ -111,14 +111,26 @@ class CameraDictType(RCSpaceType):
     ]
 
 
-# joining works with inhertiance but need to inherit from protocol again
-class ArmObsType(TQuartDictType, JointsDictType): ...
+# joining works with inheritance but need to inherit from protocol again
+class ArmObsType(TQuartDictType, JointsDictType, TRPYDictType): ...
 
 
 CartOrJointContType: TypeAlias = TQuartDictType | JointsDictType | TRPYDictType
 LimitedCartOrJointContType: TypeAlias = LimitedTQuartRelDictType | LimitedJointsRelDictType | LimitedTRPYRelDictType
 
 
+class ObsArmsGr(ArmObsType, GripperDictType):
+    pass
+
+
+class ObsArmsGrCam(ArmObsType, GripperDictType, CameraDictType):
+    pass
+
+
+class ObsArmsGrCamCG(ArmObsType, GripperDictType, CameraDictType):
+    pass
+
+
 class ControlMode(Enum):
     JOINTS = 1
     CARTESIAN_TRPY = 2
@@ -131,7 +143,8 @@ class FR3Env(gym.Env):
     Top view of on the robot. Robot faces into x direction.
     z direction faces upwards. (Right handed coordinate axis)
         ^ x
-    <-  RobotBase
+        |
+    <-- RobotBase
     y
     """
 
@@ -178,6 +191,7 @@ def get_obs(self) -> ArmObsType:
                 [self.robot.get_cartesian_position().translation(), self.robot.get_cartesian_position().rotation_q()]
             ),
             joints=self.robot.get_joint_position(),
+            xyzrpy=self.robot.get_cartesian_position().xyzrpy(),
         )
 
     def step(self, action: CartOrJointContType) -> tuple[ArmObsType, float, bool, bool, dict]:

python/rcsss/envs/hw.py

@@ -19,15 +19,15 @@ def __init__(self, env):
         assert isinstance(self.unwrapped.robot, hw.FR3), "Robot must be a hw.FR3 instance."
         self.hw_robot = cast(hw.FR3, self.unwrapped.robot)
 
-    def step(self, action: Any) -> tuple[Any, SupportsFloat, bool, bool, dict]:
+    def step(self, action: Any) -> tuple[dict[str, Any], SupportsFloat, bool, bool, dict]:
         try:
             return super().step(action)
         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 FR3Env
-            return self.unwrapped.get_obs(), 0, False, True, {}
+            return dict(self.unwrapped.get_obs()), 0, False, True, {}
 
     def reset(
         self, seed: int | None = None, options: dict[str, Any] | None = None

python/rcsss/envs/sim.py

@@ -25,14 +25,15 @@ def __init__(self, env, simulation: sim.Sim):
         self.sim = simulation
 
     def step(self, action: dict[str, Any]) -> tuple[dict[str, Any], float, bool, bool, dict]:
-        obs, _, _, _, info = super().step(action)
+        _, _, _, _, info = super().step(action)
         self.sim.step_until_convergence()
         state = self.sim_robot.get_state()
         info["collision"] = state.collision
         info["ik_success"] = state.ik_success
         info["is_sim_converged"] = self.sim.is_converged()
         # truncate episode if collision
-        return obs, 0, False, state.collision or not state.ik_success, info
+
+        return dict(self.unwrapped.get_obs()), 0, False, state.collision or not state.ik_success, info
 
     def reset(
         self, seed: int | None = None, options: dict[str, Any] | None = None