kos_sim.server

  1"""Server and simulation loop for KOS."""
  2
  3import argparse
  4import asyncio
  5import itertools
  6import logging
  7import os
  8import time
  9import traceback
 10from concurrent import futures
 11from dataclasses import dataclass
 12from pathlib import Path
 13
 14import colorlogging
 15import grpc
 16from kos_protos import actuator_pb2_grpc, imu_pb2_grpc, process_manager_pb2_grpc, sim_pb2_grpc
 17from kscale import K
 18from kscale.web.gen.api import RobotURDFMetadataOutput
 19from mujoco_scenes.mjcf import list_scenes
 20
 21from kos_sim import logger
 22from kos_sim.services import ActuatorService, IMUService, ProcessManagerService, SimService
 23from kos_sim.simulator import MujocoSimulator
 24from kos_sim.utils import get_sim_artifacts_path
 25from kos_sim.video_recorder import VideoRecorder
 26
 27
 28@dataclass
 29class SimulationRandomizationConfig:
 30    command_delay_min: float = 0.0
 31    command_delay_max: float = 0.0
 32    joint_pos_delta_noise: float = 0.0
 33    joint_pos_noise: float = 0.0
 34    joint_vel_noise: float = 0.0
 35
 36
 37@dataclass
 38class RenderingConfig:
 39    render: bool = True
 40    render_frequency: float = 1
 41    video_output_dir: str | Path | None = None
 42    frame_width: int = 640
 43    frame_height: int = 480
 44    camera: str | None = None
 45
 46
 47@dataclass
 48class PhysicsConfig:
 49    gravity: bool = True
 50    suspended: bool = False
 51    start_height: float = 1.5
 52    dt: float = 0.0001
 53
 54
 55@dataclass
 56class SimulationServerConfig:
 57    randomization: SimulationRandomizationConfig
 58    rendering: RenderingConfig
 59    physics: PhysicsConfig
 60    model_path: str | Path
 61    model_metadata: RobotURDFMetadataOutput
 62    mujoco_scene: str = "smooth"
 63    host: str = "localhost"
 64    port: int = 50051
 65    sleep_time: float = 1e-6
 66
 67
 68class SimulationServer:
 69    def __init__(
 70        self,
 71        config: SimulationServerConfig,
 72    ) -> None:
 73        self.simulator = MujocoSimulator(
 74            model_path=config.model_path,
 75            model_metadata=config.model_metadata,
 76            dt=config.physics.dt,
 77            gravity=config.physics.gravity,
 78            render_mode="window" if config.rendering.render else "offscreen",
 79            suspended=config.physics.suspended,
 80            start_height=config.physics.start_height,
 81            command_delay_min=config.randomization.command_delay_min,
 82            command_delay_max=config.randomization.command_delay_max,
 83            joint_pos_delta_noise=config.randomization.joint_pos_delta_noise,
 84            joint_pos_noise=config.randomization.joint_pos_noise,
 85            joint_vel_noise=config.randomization.joint_vel_noise,
 86            mujoco_scene=config.mujoco_scene,
 87            camera=config.rendering.camera,
 88            frame_width=config.rendering.frame_width,
 89            frame_height=config.rendering.frame_height,
 90        )
 91        self.host = config.host
 92        self.port = config.port
 93        self._sleep_time = config.sleep_time
 94        self._stop_event = asyncio.Event()
 95        self._server = None
 96        self._step_lock = asyncio.Semaphore(1)
 97        self._render_decimation = int(1.0 / config.rendering.render_frequency)
 98
 99        # Initialize video recorder if needed
100        self.video_recorder = None
101        if config.rendering.video_output_dir is not None:
102            self.video_recorder = VideoRecorder(
103                simulator=self.simulator,
104                output_dir=config.rendering.video_output_dir,
105                fps=int(self.simulator._control_frequency),
106                frame_width=config.rendering.frame_width,
107                frame_height=config.rendering.frame_height,
108            )
109
110    async def _grpc_server_loop(self) -> None:
111        """Run the async gRPC server."""
112        # Create async gRPC server
113        self._server = grpc.aio.server(futures.ThreadPoolExecutor(max_workers=10))
114
115        assert self._server is not None
116
117        # Add our services (these need to be modified to be async as well)
118        actuator_service = ActuatorService(self.simulator, self._step_lock)
119        imu_service = IMUService(self.simulator)
120        sim_service = SimService(self.simulator)
121        process_manager_service = ProcessManagerService(self.simulator, self.video_recorder)
122
123        actuator_pb2_grpc.add_ActuatorServiceServicer_to_server(actuator_service, self._server)
124        imu_pb2_grpc.add_IMUServiceServicer_to_server(imu_service, self._server)
125        sim_pb2_grpc.add_SimulationServiceServicer_to_server(sim_service, self._server)
126        process_manager_pb2_grpc.add_ProcessManagerServiceServicer_to_server(process_manager_service, self._server)
127
128        # Start the server
129        self._server.add_insecure_port(f"{self.host}:{self.port}")
130        await self._server.start()
131        logger.info("Server started on %s:%d", self.host, self.port)
132        await self._server.wait_for_termination()
133
134    async def simulation_loop(self) -> None:
135        """Run the simulation loop asynchronously."""
136        start_time = time.time()
137        num_renders = 0
138        num_steps = 0
139
140        try:
141            while not self._stop_event.is_set():
142                while self.simulator._sim_time < time.time():
143                    # Run one control loop.
144                    async with self._step_lock:
145                        for _ in range(self.simulator._sim_decimation):
146                            await self.simulator.step()
147                    await asyncio.sleep(self._sleep_time)
148
149                if num_steps % self._render_decimation == 0:
150                    await self.simulator.render()
151                    num_renders += 1
152
153                if self.video_recorder is not None and self.video_recorder.is_recording:
154                    await self.video_recorder.capture_frame()
155
156                # Sleep until the next control update.
157                current_time = time.time()
158                if current_time < self.simulator._sim_time:
159                    await asyncio.sleep(self.simulator._sim_time - current_time)
160                num_steps += 1
161                logger.debug(
162                    "Simulation time: %f, rendering frequency: %f",
163                    self.simulator._sim_time,
164                    num_renders / (time.time() - start_time),
165                )
166
167        except Exception as e:
168            logger.error("Simulation loop failed: %s", e)
169            logger.error("Traceback: %s", traceback.format_exc())
170
171        finally:
172            await self.stop()
173
174    async def start(self) -> None:
175        """Start both the gRPC server and simulation loop asynchronously."""
176        grpc_task = asyncio.create_task(self._grpc_server_loop())
177        sim_task = asyncio.create_task(self.simulation_loop())
178
179        try:
180            await asyncio.gather(grpc_task, sim_task)
181        except asyncio.CancelledError:
182            await self.stop()
183
184    async def stop(self) -> None:
185        """Stop the simulation and cleanup resources asynchronously."""
186        logger.info("Shutting down simulation...")
187        self._stop_event.set()
188        if self.video_recorder is not None and self.video_recorder.is_recording:
189            self.video_recorder.stop_recording()
190        if self._server is not None:
191            await self._server.stop(0)
192        await self.simulator.close()
193
194
195async def get_model_metadata(api: K, model_name: str) -> RobotURDFMetadataOutput:
196    model_path = get_sim_artifacts_path() / model_name / "metadata.json"
197    if model_path.exists():
198        return RobotURDFMetadataOutput.model_validate_json(model_path.read_text())
199    model_path.parent.mkdir(parents=True, exist_ok=True)
200    robot_class = await api.get_robot_class(model_name)
201    metadata = robot_class.metadata
202    if metadata is None:
203        raise ValueError(f"No metadata found for model {model_name}")
204    model_path.write_text(metadata.model_dump_json())
205    return metadata
206
207
208async def serve(
209    model_name: str,
210    host: str = "localhost",
211    port: int = 50051,
212    rendering: RenderingConfig = RenderingConfig(),
213    physics: PhysicsConfig = PhysicsConfig(),
214    randomization: SimulationRandomizationConfig = SimulationRandomizationConfig(),
215    mujoco_scene: str = "smooth",
216) -> None:
217    async with K() as api:
218        model_dir, model_metadata = await asyncio.gather(
219            api.download_and_extract_urdf(model_name),
220            get_model_metadata(api, model_name),
221        )
222
223    model_path = next(
224        itertools.chain(
225            model_dir.glob("*.mjcf"),
226            model_dir.glob("*.xml"),
227        )
228    )
229
230    config = SimulationServerConfig(
231        model_path=model_path,
232        model_metadata=model_metadata,
233        physics=physics,
234        rendering=rendering,
235        randomization=randomization,
236    )
237
238    server = SimulationServer(config)
239    await server.start()
240
241
242async def run_server() -> None:
243    parser = argparse.ArgumentParser(description="Start the simulation gRPC server.")
244    parser.add_argument("model_name", type=str, help="Name of the model to simulate")
245    parser.add_argument("--host", type=str, default="localhost", help="Host to listen on")
246    parser.add_argument("--port", type=int, default=50051, help="Port to listen on")
247    parser.add_argument("--dt", type=float, default=0.001, help="Simulation timestep")
248    parser.add_argument("--no-gravity", action="store_true", help="Disable gravity")
249    parser.add_argument("--no-render", action="store_true", help="Disable rendering")
250    parser.add_argument("--render-frequency", type=float, default=1, help="Render frequency (Hz)")
251    parser.add_argument("--suspended", action="store_true", help="Suspended simulation")
252    parser.add_argument("--command-delay-min", type=float, default=0.0, help="Minimum command delay")
253    parser.add_argument("--command-delay-max", type=float, default=0.0, help="Maximum command delay")
254    parser.add_argument("--start-height", type=float, default=1.5, help="Start height")
255    parser.add_argument("--joint-pos-delta-noise", type=float, default=0.0, help="Joint position delta noise (degrees)")
256    parser.add_argument("--joint-pos-noise", type=float, default=0.0, help="Joint position noise (degrees)")
257    parser.add_argument("--joint-vel-noise", type=float, default=0.0, help="Joint velocity noise (degrees/second)")
258    parser.add_argument("--scene", choices=list_scenes(), default="smooth", help="Mujoco scene to use")
259    parser.add_argument("--camera", type=str, default=None, help="Camera to use")
260    parser.add_argument("--debug", action="store_true", help="Enable debug logging")
261    parser.add_argument("--video-output-dir", type=str, default="videos", help="Directory to save videos")
262    parser.add_argument("--frame-width", type=int, default=640, help="Frame width")
263    parser.add_argument("--frame-height", type=int, default=480, help="Frame height")
264
265    args = parser.parse_args()
266
267    colorlogging.configure(level=logging.DEBUG if args.debug else logging.INFO)
268
269    model_name = args.model_name
270    host = args.host
271    port = args.port
272    dt = args.dt
273    gravity = not args.no_gravity
274    render = not args.no_render
275    render_frequency = args.render_frequency
276    suspended = args.suspended
277    start_height = args.start_height
278    command_delay_min = args.command_delay_min
279    command_delay_max = args.command_delay_max
280    joint_pos_delta_noise = args.joint_pos_delta_noise
281    joint_pos_noise = args.joint_pos_noise
282    joint_vel_noise = args.joint_vel_noise
283    mujoco_scene = args.scene
284    camera = args.camera
285    frame_width = args.frame_width
286    frame_height = args.frame_height
287
288    video_output_dir = args.video_output_dir if not render else None
289
290    logger.info("Model name: %s", model_name)
291    logger.info("Port: %d", port)
292    logger.info("DT: %f", dt)
293    logger.info("Gravity: %s", gravity)
294    logger.info("Render: %s", render)
295    logger.info("Render frequency: %f", render_frequency)
296    logger.info("Suspended: %s", suspended)
297    logger.info("Start height: %f", start_height)
298    logger.info("Command delay min: %f", command_delay_min)
299    logger.info("Command delay max: %f", command_delay_max)
300    logger.info("Joint pos delta noise: %f", joint_pos_delta_noise)
301    logger.info("Joint pos noise: %f", joint_pos_noise)
302    logger.info("Joint vel noise: %f", joint_vel_noise)
303    logger.info("Mujoco scene: %s", mujoco_scene)
304    logger.info("Camera: %s", camera)
305    logger.info("Video output dir: %s", video_output_dir)
306    logger.info("Frame width: %d", frame_width)
307    logger.info("Frame height: %d", frame_height)
308
309    if video_output_dir is not None:
310        os.makedirs(video_output_dir, exist_ok=True)
311
312    rendering = RenderingConfig(
313        render=render,
314        render_frequency=render_frequency,
315        video_output_dir=video_output_dir,
316        frame_width=frame_width,
317        frame_height=frame_height,
318        camera=camera,
319    )
320
321    physics = PhysicsConfig(
322        dt=dt,
323        gravity=gravity,
324        suspended=suspended,
325        start_height=start_height,
326    )
327
328    randomization = SimulationRandomizationConfig(
329        command_delay_min=command_delay_min,
330        command_delay_max=command_delay_max,
331        joint_pos_delta_noise=joint_pos_delta_noise,
332        joint_pos_noise=joint_pos_noise,
333        joint_vel_noise=joint_vel_noise,
334    )
335
336    await serve(
337        model_name=model_name,
338        host=host,
339        port=port,
340        rendering=rendering,
341        physics=physics,
342        randomization=randomization,
343        mujoco_scene=mujoco_scene,
344    )
345
346
347def main() -> None:
348    asyncio.run(run_server())
349
350
351if __name__ == "__main__":
352    # python -m kos_sim.server
353    main()