add gt obs getter for benchmark metrics

kos_sim/services.py

@@ -90,6 +90,24 @@ async def Step(  # noqa: N802
     ) -> common_pb2.ActionResponse:
         raise NotImplementedError("Step is not implemented")
 
+    async def GetGroundTruthObservation(  # noqa: N802
+        self,
+        request: empty_pb2.Empty,
+        context: grpc.ServicerContext,
+    ) -> sim_pb2.GetGroundTruthObservationResponse:
+        """Get ground truth observation."""
+        try:
+            observation = await self.simulator.get_ground_truth_observation()
+            return sim_pb2.GetGroundTruthObservationResponse(
+                base_angular_velocity=observation.base_angular_velocity,
+                base_linear_velocity=observation.base_linear_velocity,
+            )
+        except Exception as e:
+            logger.error("GetGroundTruthObservation failed: %s", e)
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(str(e))
+            return sim_pb2.GetGroundTruthObservationResponse(error=common_pb2.Error(message=str(e)))
+
     async def SetParameters(  # noqa: N802
         self,
         request: sim_pb2.SetParametersRequest,

kos_sim/simulator.py

@@ -55,6 +55,12 @@ class ActuatorCommand(TypedDict):
     torque: NotRequired[float]
 
 
+@dataclass
+class GroundTruthObservation:
+    base_angular_velocity: tuple[float, float, float]
+    base_linear_velocity: tuple[float, float, float]
+
+
 def get_integrator(integrator: str) -> mujoco.mjtIntegrator:
     match integrator.lower():
         case "euler":
@@ -281,26 +287,7 @@ async def capture_frame(self, camid: int = -1, depth: bool = False) -> tuple[np.
         Returns:
             RGB image array (and optionally depth array) if depth=True
         """
-        # TODO: Use native mujoco renderder for offline and shiiiiiiiiit
         return np.zeros((480, 640, 3), dtype=np.uint8), None
-        # if self._render_mode != "offscreen" and self._render_enabled:
-        #     logger.warning("Capturing frames is more efficient in offscreen mode")
-
-        # if depth:
-        #     logger.warning("Depth is not currently supported")
-
-        # for marker in self._markers.values():
-        #     self._viewer.add_marker(**marker)
-
-        # if camid is not None:
-        #     if camid == -1:
-        #         self._viewer.handle.cam.type = mujoco.mjtCamera.mjCAMERA_FREE
-        #     else:
-        #         self._viewer.handle.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED
-        #         self._viewer.handle.cam.fixedcamid = camid
-
-        # rgb = self._viewer.read_pixels()
-        # return rgb, None
 
     async def get_sensor_data(self, name: str) -> np.ndarray:
         """Get data from a named sensor."""
@@ -417,3 +404,10 @@ async def close(self) -> None:
     @property
     def timestep(self) -> float:
         return self._model.opt.timestep
+
+    async def get_ground_truth_observation(self) -> GroundTruthObservation:
+        """Get the ground truth observation."""
+        return GroundTruthObservation(
+            base_angular_velocity=self._data.cvel[1, 0:3],
+            base_linear_velocity=self._data.cvel[1, 3:6],
+        )