diff --git a/kos_sim/services.py b/kos_sim/services.py index c52fb56..3b1c925 100644 --- a/kos_sim/services.py +++ b/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, diff --git a/kos_sim/simulator.py b/kos_sim/simulator.py index c07e325..2ccaf31 100644 --- a/kos_sim/simulator.py +++ b/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], + )