kos_sim.services

  1"""Service implementations for MuJoCo simulation."""
  2
  3import asyncio
  4import math
  5
  6import grpc
  7from google.protobuf import empty_pb2
  8from kos_protos import (
  9    actuator_pb2,
 10    actuator_pb2_grpc,
 11    common_pb2,
 12    imu_pb2,
 13    imu_pb2_grpc,
 14    process_manager_pb2,
 15    process_manager_pb2_grpc,
 16    sim_pb2,
 17    sim_pb2_grpc,
 18)
 19
 20from kos_sim import logger
 21from kos_sim.simulator import ActuatorCommand, ConfigureActuatorRequest, MujocoSimulator
 22from kos_sim.video_recorder import VideoRecorder
 23
 24
 25class SimService(sim_pb2_grpc.SimulationServiceServicer):
 26    """Implementation of SimService that wraps a MuJoCo simulation."""
 27
 28    def __init__(self, simulator: MujocoSimulator) -> None:
 29        self.simulator = simulator
 30
 31    async def Reset(  # noqa: N802
 32        self,
 33        request: sim_pb2.ResetRequest,
 34        context: grpc.ServicerContext,
 35    ) -> common_pb2.ActionResponse:
 36        """Reset the simulation to initial or specified state."""
 37        try:
 38            logger.debug("Resetting simulator")
 39            await self.simulator.reset(
 40                xyz=None if request.HasField("pos") is None else (request.pos.x, request.pos.y, request.pos.z),
 41                quat=(
 42                    None
 43                    if request.HasField("quat") is None
 44                    else (request.quat.w, request.quat.x, request.quat.y, request.quat.z)
 45                ),
 46                joint_pos=(
 47                    None
 48                    if request.joints is None
 49                    else {joint.name: joint.pos for joint in request.joints.values if joint.HasField("pos")}
 50                ),
 51                joint_vel=(
 52                    None
 53                    if request.joints is None
 54                    else {joint.name: joint.vel for joint in request.joints.values if joint.HasField("vel")}
 55                ),
 56            )
 57            return common_pb2.ActionResponse(success=True)
 58        except Exception as e:
 59            logger.error("Reset failed: %s", e)
 60            context.set_code(grpc.StatusCode.INTERNAL)
 61            context.set_details(str(e))
 62            return common_pb2.ActionResponse(success=False, error=str(e))
 63
 64    async def SetPaused(  # noqa: N802
 65        self,
 66        request: sim_pb2.SetPausedRequest,
 67        context: grpc.ServicerContext,
 68    ) -> common_pb2.ActionResponse:  # noqa: N802
 69        """Pause or unpause the simulation."""
 70        try:
 71            return common_pb2.ActionResponse(success=True)
 72        except Exception as e:
 73            logger.error("SetPaused failed: %s", e)
 74            context.set_code(grpc.StatusCode.INTERNAL)
 75            context.set_details(str(e))
 76            return common_pb2.ActionResponse(success=False, error=str(e))
 77
 78    async def Step(  # noqa: N802
 79        self,
 80        request: sim_pb2.StepRequest,
 81        context: grpc.ServicerContext,
 82    ) -> common_pb2.ActionResponse:
 83        raise NotImplementedError("Step is not implemented")
 84
 85    async def SetParameters(  # noqa: N802
 86        self,
 87        request: sim_pb2.SetParametersRequest,
 88        context: grpc.ServicerContext,
 89    ) -> common_pb2.ActionResponse:
 90        """Set simulation parameters."""
 91        try:
 92            params = request.parameters
 93            if params.HasField("time_scale"):
 94                logger.debug("Setting time scale to %f", params.time_scale)
 95                self.simulator._model.opt.timestep = self.simulator._dt / params.time_scale
 96            if params.HasField("gravity"):
 97                logger.debug("Setting gravity to %f", params.gravity)
 98                self.simulator._model.opt.gravity[2] = params.gravity
 99            return common_pb2.ActionResponse(success=True)
100        except Exception as e:
101            logger.error("SetParameters failed: %s", e)
102            context.set_code(grpc.StatusCode.INTERNAL)
103            context.set_details(str(e))
104            return common_pb2.ActionResponse(success=False, error=str(e))
105
106    async def GetParameters(  # noqa: N802
107        self,
108        request: empty_pb2.Empty,
109        context: grpc.ServicerContext,
110    ) -> sim_pb2.GetParametersResponse:
111        """Get current simulation parameters."""
112        try:
113            params = sim_pb2.SimulationParameters(
114                time_scale=self.simulator._dt / self.simulator._model.opt.timestep,
115                gravity=float(self.simulator._model.opt.gravity[2]),
116            )
117            logger.debug("Current parameters: time_scale=%f, gravity=%f", params.time_scale, params.gravity)
118            return sim_pb2.GetParametersResponse(parameters=params)
119        except Exception as e:
120            logger.error("GetParameters failed: %s", e)
121            context.set_code(grpc.StatusCode.INTERNAL)
122            context.set_details(str(e))
123            return sim_pb2.GetParametersResponse(error=common_pb2.Error(message=str(e)))
124
125
126class ActuatorService(actuator_pb2_grpc.ActuatorServiceServicer):
127    """Implementation of ActuatorService that wraps a MuJoCo simulation."""
128
129    def __init__(self, simulator: MujocoSimulator, step_lock: asyncio.Semaphore) -> None:
130        self.simulator = simulator
131        self.step_lock = step_lock
132
133    async def CommandActuators(  # noqa: N802
134        self,
135        request: actuator_pb2.CommandActuatorsRequest,
136        context: grpc.ServicerContext,
137    ) -> actuator_pb2.CommandActuatorsResponse:
138        """Implements CommandActuators by forwarding to simulator."""
139        try:
140            # Convert degrees to radians.
141            commands: dict[int, ActuatorCommand] = {
142                cmd.actuator_id: {
143                    "position": math.radians(cmd.position),
144                    "velocity": math.radians(cmd.velocity),
145                    "torque": cmd.torque,
146                }
147                for cmd in request.commands
148            }
149            async with self.step_lock:
150                await self.simulator.command_actuators(commands)
151            return actuator_pb2.CommandActuatorsResponse()
152        except Exception as e:
153            context.set_code(grpc.StatusCode.INTERNAL)
154            context.set_details(str(e))
155            return actuator_pb2.CommandActuatorsResponse()
156
157    async def GetActuatorsState(  # noqa: N802
158        self,
159        request: actuator_pb2.GetActuatorsStateRequest,
160        context: grpc.ServicerContext,
161    ) -> actuator_pb2.GetActuatorsStateResponse:
162        """Implements GetActuatorsState by reading from simulator."""
163        ids = request.actuator_ids or list(self.simulator._joint_id_to_name.keys())
164        try:
165            states = {joint_id: await self.simulator.get_actuator_state(joint_id) for joint_id in ids}
166            return actuator_pb2.GetActuatorsStateResponse(
167                states=[
168                    actuator_pb2.ActuatorStateResponse(
169                        actuator_id=joint_id,
170                        position=math.degrees(state.position),
171                        velocity=math.degrees(state.velocity),
172                        online=True,
173                    )
174                    for joint_id, state in states.items()
175                ]
176            )
177        except Exception as e:
178            logger.error("GetActuatorsState failed: %s", e)
179            context.set_code(grpc.StatusCode.INTERNAL)
180            context.set_details(str(e))
181            return actuator_pb2.GetActuatorsStateResponse()
182
183    async def ConfigureActuator(  # noqa: N802
184        self,
185        request: actuator_pb2.ConfigureActuatorRequest,
186        context: grpc.ServicerContext,
187    ) -> common_pb2.ActionResponse:
188        configuration: ConfigureActuatorRequest = {}
189        if request.HasField("torque_enabled"):
190            configuration["torque_enabled"] = request.torque_enabled
191        if request.HasField("zero_position"):
192            configuration["zero_position"] = request.zero_position
193        if request.HasField("kp"):
194            configuration["kp"] = request.kp
195        if request.HasField("kd"):
196            configuration["kd"] = request.kd
197        if request.HasField("max_torque"):
198            configuration["max_torque"] = request.max_torque
199        await self.simulator.configure_actuator(request.actuator_id, configuration)
200
201        return common_pb2.ActionResponse(success=True)
202
203
204class IMUService(imu_pb2_grpc.IMUServiceServicer):
205    """Implementation of IMUService that wraps a MuJoCo simulation."""
206
207    def __init__(
208        self,
209        simulator: MujocoSimulator,
210        acc_name: str | None = "imu_acc",
211        gyro_name: str | None = "imu_gyro",
212        mag_name: str | None = "imu_mag",
213        quat_name: str | None = "base_link_quat",
214    ) -> None:
215        self.simulator = simulator
216        self.acc_name = acc_name
217        self.gyro_name = gyro_name
218        self.mag_name = mag_name
219        self.quat_name = quat_name
220
221    async def GetValues(  # noqa: N802
222        self,
223        request: empty_pb2.Empty,
224        context: grpc.ServicerContext,
225    ) -> imu_pb2.IMUValuesResponse:
226        """Implements GetValues by reading IMU sensor data from simulator."""
227        try:
228            if self.acc_name is None or self.gyro_name is None:
229                raise ValueError("Accelerometer or gyroscope name not set")
230            acc_data = await self.simulator.get_sensor_data(self.acc_name)
231            gyro_data = await self.simulator.get_sensor_data(self.gyro_name)
232            mag_data = None if self.mag_name is None else await self.simulator.get_sensor_data(self.mag_name)
233            return imu_pb2.IMUValuesResponse(
234                accel_x=float(acc_data[0]),
235                accel_y=float(acc_data[1]),
236                accel_z=float(acc_data[2]),
237                gyro_x=float(gyro_data[0]),
238                gyro_y=float(gyro_data[1]),
239                gyro_z=float(gyro_data[2]),
240                mag_x=None if mag_data is None else float(mag_data[0]),
241                mag_y=None if mag_data is None else float(mag_data[1]),
242                mag_z=None if mag_data is None else float(mag_data[2]),
243            )
244        except Exception as e:
245            context.set_code(grpc.StatusCode.INTERNAL)
246            context.set_details(str(e))
247            return imu_pb2.IMUValuesResponse()
248
249    async def GetQuaternion(  # noqa: N802
250        self,
251        request: empty_pb2.Empty,
252        context: grpc.ServicerContext,
253    ) -> imu_pb2.QuaternionResponse:
254        """Implements GetQuaternion by reading orientation data from simulator."""
255        try:
256            if self.quat_name is None:
257                raise ValueError("Quaternion name not set")
258            quat_data = await self.simulator.get_sensor_data(self.quat_name)
259            return imu_pb2.QuaternionResponse(
260                w=float(quat_data[0]), x=float(quat_data[1]), y=float(quat_data[2]), z=float(quat_data[3])
261            )
262        except Exception as e:
263            context.set_code(grpc.StatusCode.INTERNAL)
264            context.set_details(str(e))
265            return imu_pb2.QuaternionResponse()
266
267    async def GetEuler(  # noqa: N802
268        self,
269        request: empty_pb2.Empty,
270        context: grpc.ServicerContext,
271    ) -> imu_pb2.EulerAnglesResponse:
272        """Implements GetEuler by converting orientation quaternion to Euler angles."""
273        try:
274            if self.quat_name is None:
275                raise ValueError("Quaternion name not set")
276            quat_data = await self.simulator.get_sensor_data(self.quat_name)
277
278            # Extract quaternion components
279            w, x, y, z = [float(q) for q in quat_data]
280
281            # Convert quaternion to Euler angles (roll, pitch, yaw)
282            # Roll (x-axis rotation)
283            sinr_cosp = 2 * (w * x + y * z)
284            cosr_cosp = 1 - 2 * (x * x + y * y)
285            roll = math.atan2(sinr_cosp, cosr_cosp)
286
287            # Pitch (y-axis rotation)
288            sinp = 2 * (w * y - z * x)
289            pitch = math.asin(sinp) if abs(sinp) < 1 else math.copysign(math.pi / 2, sinp)
290
291            # Yaw (z-axis rotation)
292            siny_cosp = 2 * (w * z + x * y)
293            cosy_cosp = 1 - 2 * (y * y + z * z)
294            yaw = math.atan2(siny_cosp, cosy_cosp)
295
296            # Convert to degrees
297            roll_deg = math.degrees(roll)
298            pitch_deg = math.degrees(pitch)
299            yaw_deg = math.degrees(yaw)
300
301            return imu_pb2.EulerAnglesResponse(roll=roll_deg, pitch=pitch_deg, yaw=yaw_deg)
302        except Exception as e:
303            context.set_code(grpc.StatusCode.INTERNAL)
304            context.set_details(str(e))
305            return imu_pb2.EulerAnglesResponse()
306
307
308class ProcessManagerService(process_manager_pb2_grpc.ProcessManagerServiceServicer):
309    """Implementation of ProcessManagerService that wraps a MuJoCo simulation."""
310
311    def __init__(self, simulator: MujocoSimulator, video_recorder: VideoRecorder) -> None:
312        self.simulator = simulator
313        self.video_recorder = video_recorder
314
315    async def StartKClip(  # noqa: N802
316        self, request: process_manager_pb2.KClipStartRequest, context: grpc.ServicerContext
317    ) -> process_manager_pb2.KClipStartResponse:
318        """Implements StartKClip by starting k-clip recording."""
319        if self.video_recorder is None:
320            return process_manager_pb2.KClipStartResponse(
321                error=common_pb2.Error(
322                    code=common_pb2.ErrorCode.INVALID_ARGUMENT,
323                    message="`video_recorder` is `None`. Video recording not enabled.",
324                )
325            )
326        clip_id = self.video_recorder.start_recording()
327        return process_manager_pb2.KClipStartResponse(clip_uuid=clip_id)
328
329    async def StopKClip(  # noqa: N802
330        self, request: empty_pb2.Empty, context: grpc.ServicerContext
331    ) -> process_manager_pb2.KClipStopResponse:
332        """Implements StopKClip by stopping k-clip recording."""
333        if self.video_recorder is None:
334            return process_manager_pb2.KClipStopResponse(
335                error=common_pb2.Error(
336                    code=common_pb2.ErrorCode.INVALID_ARGUMENT,
337                    message="`video_recorder` is `None`. Video recording not enabled.",
338                )
339            )
340        self.video_recorder.stop_recording()
341        return process_manager_pb2.KClipStopResponse()