luckyrobots 0.1.66__py3-none-any.whl → 0.1.72__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (60) hide show
  1. luckyrobots/__init__.py +30 -12
  2. luckyrobots/client.py +997 -0
  3. luckyrobots/config/robots.yaml +231 -71
  4. luckyrobots/engine/__init__.py +23 -0
  5. luckyrobots/{utils → engine}/check_updates.py +108 -48
  6. luckyrobots/{utils → engine}/download.py +61 -39
  7. luckyrobots/engine/manager.py +427 -0
  8. luckyrobots/grpc/__init__.py +6 -0
  9. luckyrobots/grpc/generated/__init__.py +18 -0
  10. luckyrobots/grpc/generated/agent_pb2.py +69 -0
  11. luckyrobots/grpc/generated/agent_pb2_grpc.py +283 -0
  12. luckyrobots/grpc/generated/camera_pb2.py +47 -0
  13. luckyrobots/grpc/generated/camera_pb2_grpc.py +144 -0
  14. luckyrobots/grpc/generated/common_pb2.py +43 -0
  15. luckyrobots/grpc/generated/common_pb2_grpc.py +24 -0
  16. luckyrobots/grpc/generated/hazel_rpc_pb2.py +43 -0
  17. luckyrobots/grpc/generated/hazel_rpc_pb2_grpc.py +24 -0
  18. luckyrobots/grpc/generated/media_pb2.py +39 -0
  19. luckyrobots/grpc/generated/media_pb2_grpc.py +24 -0
  20. luckyrobots/grpc/generated/mujoco_pb2.py +51 -0
  21. luckyrobots/grpc/generated/mujoco_pb2_grpc.py +230 -0
  22. luckyrobots/grpc/generated/scene_pb2.py +66 -0
  23. luckyrobots/grpc/generated/scene_pb2_grpc.py +317 -0
  24. luckyrobots/grpc/generated/telemetry_pb2.py +47 -0
  25. luckyrobots/grpc/generated/telemetry_pb2_grpc.py +143 -0
  26. luckyrobots/grpc/generated/viewport_pb2.py +50 -0
  27. luckyrobots/grpc/generated/viewport_pb2_grpc.py +144 -0
  28. luckyrobots/grpc/proto/agent.proto +213 -0
  29. luckyrobots/grpc/proto/camera.proto +41 -0
  30. luckyrobots/grpc/proto/common.proto +36 -0
  31. luckyrobots/grpc/proto/hazel_rpc.proto +32 -0
  32. luckyrobots/grpc/proto/media.proto +26 -0
  33. luckyrobots/grpc/proto/mujoco.proto +64 -0
  34. luckyrobots/grpc/proto/scene.proto +104 -0
  35. luckyrobots/grpc/proto/telemetry.proto +43 -0
  36. luckyrobots/grpc/proto/viewport.proto +45 -0
  37. luckyrobots/luckyrobots.py +252 -0
  38. luckyrobots/models/__init__.py +15 -0
  39. luckyrobots/models/camera.py +97 -0
  40. luckyrobots/models/observation.py +135 -0
  41. luckyrobots/models/randomization.py +77 -0
  42. luckyrobots/{utils/helpers.py → utils.py} +75 -40
  43. luckyrobots-0.1.72.dist-info/METADATA +262 -0
  44. luckyrobots-0.1.72.dist-info/RECORD +47 -0
  45. {luckyrobots-0.1.66.dist-info → luckyrobots-0.1.72.dist-info}/WHEEL +1 -1
  46. luckyrobots/core/luckyrobots.py +0 -624
  47. luckyrobots/core/manager.py +0 -236
  48. luckyrobots/core/models.py +0 -68
  49. luckyrobots/core/node.py +0 -273
  50. luckyrobots/message/__init__.py +0 -18
  51. luckyrobots/message/pubsub.py +0 -145
  52. luckyrobots/message/srv/client.py +0 -81
  53. luckyrobots/message/srv/service.py +0 -135
  54. luckyrobots/message/srv/types.py +0 -83
  55. luckyrobots/message/transporter.py +0 -427
  56. luckyrobots/utils/event_loop.py +0 -94
  57. luckyrobots/utils/sim_manager.py +0 -406
  58. luckyrobots-0.1.66.dist-info/METADATA +0 -253
  59. luckyrobots-0.1.66.dist-info/RECORD +0 -24
  60. {luckyrobots-0.1.66.dist-info → luckyrobots-0.1.72.dist-info}/licenses/LICENSE +0 -0
luckyrobots/client.py ADDED
@@ -0,0 +1,997 @@
1
+ """
2
+ LuckyEngine gRPC client.
3
+
4
+ Uses checked-in Python stubs generated from the `.proto` files under
5
+ `src/luckyrobots/grpc/proto/`.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import logging
11
+ import time
12
+ import statistics
13
+ from dataclasses import dataclass
14
+ from types import SimpleNamespace
15
+ from typing import Any, Optional
16
+
17
+ logger = logging.getLogger("luckyrobots.client")
18
+
19
+ try:
20
+ from .grpc.generated import agent_pb2 # type: ignore
21
+ from .grpc.generated import agent_pb2_grpc # type: ignore
22
+ from .grpc.generated import camera_pb2 # type: ignore
23
+ from .grpc.generated import camera_pb2_grpc # type: ignore
24
+ from .grpc.generated import common_pb2 # type: ignore
25
+ from .grpc.generated import media_pb2 # type: ignore
26
+ from .grpc.generated import mujoco_pb2 # type: ignore
27
+ from .grpc.generated import mujoco_pb2_grpc # type: ignore
28
+ from .grpc.generated import scene_pb2 # type: ignore
29
+ from .grpc.generated import scene_pb2_grpc # type: ignore
30
+ from .grpc.generated import telemetry_pb2 # type: ignore
31
+ from .grpc.generated import telemetry_pb2_grpc # type: ignore
32
+ from .grpc.generated import viewport_pb2 # type: ignore
33
+ from .grpc.generated import viewport_pb2_grpc # type: ignore
34
+ except Exception as e: # pragma: no cover
35
+ raise ImportError(
36
+ "Missing generated gRPC stubs. Regenerate them from the protos in "
37
+ "src/luckyrobots/grpc/proto into src/luckyrobots/grpc/generated."
38
+ ) from e
39
+
40
+ # Import Pydantic models for type-checked responses
41
+ from .models import ObservationResponse, StateSnapshot, DomainRandomizationConfig
42
+
43
+
44
+ class GrpcConnectionError(Exception):
45
+ """Raised when gRPC connection fails."""
46
+
47
+ def __init__(self, message: str):
48
+ super().__init__(message)
49
+ logger.warning("gRPC connection error: %s", message)
50
+
51
+
52
+ class LuckyEngineClient:
53
+ """
54
+ Client for connecting to the LuckyEngine gRPC server.
55
+
56
+ Provides access to all gRPC services defined by the protos under
57
+ `src/luckyrobots/rpc/proto`:
58
+ - SceneService
59
+ - MujocoService
60
+ - TelemetryService
61
+ - AgentService
62
+ - ViewportService
63
+ - CameraService
64
+
65
+ Usage:
66
+ client = LuckyEngineClient(host="127.0.0.1", port=50051)
67
+ client.connect()
68
+
69
+ # Access services
70
+ scene_info = client.scene.GetSceneInfo(client.pb.scene.GetSceneInfoRequest())
71
+ joint_state = client.mujoco.GetJointState(
72
+ client.pb.mujoco.GetJointStateRequest()
73
+ )
74
+
75
+ client.close()
76
+ """
77
+
78
+ def __init__(
79
+ self,
80
+ host: str = "127.0.0.1",
81
+ port: int = 50051,
82
+ timeout: float = 5.0,
83
+ *,
84
+ robot_name: Optional[str] = None,
85
+ ) -> None:
86
+ """
87
+ Initialize the LuckyEngine gRPC client.
88
+
89
+ Args:
90
+ host: gRPC server host address.
91
+ port: gRPC server port.
92
+ timeout: Default timeout for RPC calls in seconds.
93
+ robot_name: Default robot name for calls that require it.
94
+ """
95
+ self.host = host
96
+ self.port = port
97
+ self.timeout = timeout
98
+ self._robot_name = robot_name
99
+
100
+ self._channel = None
101
+
102
+ # Service stubs (populated after connect)
103
+ self._scene = None
104
+ self._mujoco = None
105
+ self._telemetry = None
106
+ self._agent = None
107
+ self._viewport = None
108
+ self._camera = None
109
+
110
+ # Cached agent schemas: agent_name -> (observation_names, action_names)
111
+ # Populated lazily by get_agent_schema() or fetch_schema()
112
+ self._schema_cache: dict[str, tuple[list[str], list[str]]] = {}
113
+
114
+ # Protobuf modules (for discoverability + explicit imports).
115
+ self._pb = SimpleNamespace(
116
+ common=common_pb2,
117
+ media=media_pb2,
118
+ scene=scene_pb2,
119
+ mujoco=mujoco_pb2,
120
+ telemetry=telemetry_pb2,
121
+ agent=agent_pb2,
122
+ viewport=viewport_pb2,
123
+ camera=camera_pb2,
124
+ )
125
+
126
+ def connect(self) -> None:
127
+ """
128
+ Connect to the LuckyEngine gRPC server.
129
+
130
+ Raises:
131
+ GrpcConnectionError: If connection fails.
132
+ """
133
+ try:
134
+ import grpc # type: ignore
135
+ except ImportError as e:
136
+ raise RuntimeError(
137
+ "Missing grpcio. Install with: pip install grpcio protobuf"
138
+ ) from e
139
+
140
+ target = f"{self.host}:{self.port}"
141
+ logger.info(f"Connecting to LuckyEngine gRPC server at {target}")
142
+
143
+ self._channel = grpc.insecure_channel(target)
144
+
145
+ # Create service stubs
146
+ self._scene = scene_pb2_grpc.SceneServiceStub(self._channel)
147
+ self._mujoco = mujoco_pb2_grpc.MujocoServiceStub(self._channel)
148
+ self._telemetry = telemetry_pb2_grpc.TelemetryServiceStub(self._channel)
149
+ self._agent = agent_pb2_grpc.AgentServiceStub(self._channel)
150
+ self._viewport = viewport_pb2_grpc.ViewportServiceStub(self._channel)
151
+ self._camera = camera_pb2_grpc.CameraServiceStub(self._channel)
152
+
153
+ logger.info(f"Channel opened to {target} (server not verified yet)")
154
+
155
+ def close(self) -> None:
156
+ """Close the gRPC channel."""
157
+ if self._channel is not None:
158
+ try:
159
+ self._channel.close()
160
+ except Exception as e:
161
+ logger.debug(f"Error closing gRPC channel: {e}")
162
+ self._channel = None
163
+ self._scene = None
164
+ self._mujoco = None
165
+ self._telemetry = None
166
+ self._agent = None
167
+ self._viewport = None
168
+ self._camera = None
169
+ logger.info("gRPC channel closed")
170
+
171
+ def is_connected(self) -> bool:
172
+ """Check if the client is connected."""
173
+ return self._channel is not None
174
+
175
+ def health_check(self, timeout: Optional[float] = None) -> bool:
176
+ """
177
+ Perform a health check by calling GetMujocoInfo.
178
+
179
+ Args:
180
+ timeout: Timeout in seconds (uses default if None).
181
+
182
+ Returns:
183
+ True if server responds, False otherwise.
184
+ """
185
+ if not self.is_connected():
186
+ return False
187
+
188
+ timeout = timeout or self.timeout
189
+ try:
190
+ self._mujoco.GetMujocoInfo(
191
+ self.pb.mujoco.GetMujocoInfoRequest(robot_name=self._robot_name or ""),
192
+ timeout=timeout,
193
+ )
194
+ return True
195
+ except Exception as e:
196
+ logger.debug(f"Health check failed: {e}")
197
+ return False
198
+
199
+ def wait_for_server(
200
+ self, timeout: float = 30.0, poll_interval: float = 0.5
201
+ ) -> bool:
202
+ """
203
+ Wait for the gRPC server to become available.
204
+
205
+ Args:
206
+ timeout: Maximum time to wait in seconds.
207
+ poll_interval: Time between connection attempts.
208
+
209
+ Returns:
210
+ True if server became available, False if timeout.
211
+ """
212
+ import time
213
+
214
+ start = time.perf_counter()
215
+
216
+ while time.perf_counter() - start < timeout:
217
+ if not self.is_connected():
218
+ try:
219
+ self.connect()
220
+ except Exception:
221
+ pass
222
+
223
+ if self.health_check(timeout=10.0):
224
+ logger.info(f"Connected to LuckyEngine gRPC server at {self.host}:{self.port}")
225
+ return True
226
+
227
+ time.sleep(poll_interval)
228
+
229
+ return False
230
+
231
+ # --- Protobuf modules (discoverable + explicit) ---
232
+
233
+ @property
234
+ def pb(self) -> Any:
235
+ """Access protobuf modules grouped by domain (e.g., `client.pb.scene`)."""
236
+ return self._pb
237
+
238
+ @property
239
+ def robot_name(self) -> Optional[str]:
240
+ """Default robot name used by calls that accept an optional robot_name."""
241
+ return self._robot_name
242
+
243
+ def set_robot_name(self, robot_name: str) -> None:
244
+ """Set the default robot name used by calls that accept an optional robot_name."""
245
+ self._robot_name = robot_name
246
+
247
+ # --- Service stubs ---
248
+ #
249
+ # Confirmed working:
250
+ # - mujoco: GetMujocoInfo, GetJointState, SendControl
251
+ # - agent: GetObservation, ResetAgent, GetAgentSchema, StreamAgent
252
+ #
253
+ # Placeholders (not yet confirmed working - use at your own risk):
254
+ # - scene: GetSceneInfo
255
+ # - telemetry: StreamTelemetry
256
+ # - viewport: (no methods implemented)
257
+ # - camera: ListCameras, StreamCamera
258
+
259
+ @property
260
+ def scene(self) -> Any:
261
+ """SceneService stub. [PLACEHOLDER - not confirmed working]"""
262
+ if self._scene is None:
263
+ raise GrpcConnectionError("Not connected. Call connect() first.")
264
+ return self._scene
265
+
266
+ @property
267
+ def mujoco(self) -> Any:
268
+ """MujocoService stub."""
269
+ if self._mujoco is None:
270
+ raise GrpcConnectionError("Not connected. Call connect() first.")
271
+ return self._mujoco
272
+
273
+ @property
274
+ def telemetry(self) -> Any:
275
+ """TelemetryService stub. [PLACEHOLDER - not confirmed working]"""
276
+ if self._telemetry is None:
277
+ raise GrpcConnectionError("Not connected. Call connect() first.")
278
+ return self._telemetry
279
+
280
+ @property
281
+ def agent(self) -> Any:
282
+ """AgentService stub."""
283
+ if self._agent is None:
284
+ raise GrpcConnectionError("Not connected. Call connect() first.")
285
+ return self._agent
286
+
287
+ @property
288
+ def viewport(self) -> Any:
289
+ """ViewportService stub. [PLACEHOLDER - not confirmed working]"""
290
+ if self._viewport is None:
291
+ raise GrpcConnectionError("Not connected. Call connect() first.")
292
+ return self._viewport
293
+
294
+ @property
295
+ def camera(self) -> Any:
296
+ """CameraService stub. [PLACEHOLDER - not confirmed working]"""
297
+ if self._camera is None:
298
+ raise GrpcConnectionError("Not connected. Call connect() first.")
299
+ return self._camera
300
+
301
+ # --- Convenience methods ---
302
+
303
+ def get_scene_info(self, timeout: Optional[float] = None):
304
+ """Get scene information. [PLACEHOLDER - not confirmed working]"""
305
+ raise NotImplementedError(
306
+ "get_scene_info() is not yet confirmed working. "
307
+ "Remove this check if you want to test it."
308
+ )
309
+ timeout = timeout or self.timeout
310
+ return self.scene.GetSceneInfo(
311
+ self.pb.scene.GetSceneInfoRequest(),
312
+ timeout=timeout,
313
+ )
314
+
315
+ def get_mujoco_info(self, robot_name: str = "", timeout: Optional[float] = None):
316
+ """Get MuJoCo model information."""
317
+ timeout = timeout or self.timeout
318
+ robot_name = robot_name or self._robot_name
319
+ if not robot_name:
320
+ raise ValueError(
321
+ "robot_name is required (pass `robot_name=` or set it once via "
322
+ "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
323
+ )
324
+ return self.mujoco.GetMujocoInfo(
325
+ self.pb.mujoco.GetMujocoInfoRequest(robot_name=robot_name),
326
+ timeout=timeout,
327
+ )
328
+
329
+ def get_joint_state(self, robot_name: str = "", timeout: Optional[float] = None):
330
+ """Get current joint state."""
331
+ timeout = timeout or self.timeout
332
+ robot_name = robot_name or self._robot_name
333
+ if not robot_name:
334
+ raise ValueError(
335
+ "robot_name is required (pass `robot_name=` or set it once via "
336
+ "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
337
+ )
338
+ return self.mujoco.GetJointState(
339
+ self.pb.mujoco.GetJointStateRequest(robot_name=robot_name),
340
+ timeout=timeout,
341
+ )
342
+
343
+ def send_control(
344
+ self,
345
+ controls: list[float],
346
+ robot_name: str = "",
347
+ timeout: Optional[float] = None,
348
+ ):
349
+ """Send control commands to the robot."""
350
+ timeout = timeout or self.timeout
351
+ robot_name = robot_name or self._robot_name
352
+ if not robot_name:
353
+ raise ValueError(
354
+ "robot_name is required (pass `robot_name=` or set it once via "
355
+ "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
356
+ )
357
+ return self.mujoco.SendControl(
358
+ self.pb.mujoco.SendControlRequest(robot_name=robot_name, controls=controls),
359
+ timeout=timeout,
360
+ )
361
+
362
+ def get_agent_schema(self, agent_name: str = "", timeout: Optional[float] = None):
363
+ """Get agent schema (observation/action sizes and names).
364
+
365
+ The schema is cached for subsequent get_observation() calls to enable
366
+ named access to observation values.
367
+
368
+ Args:
369
+ agent_name: Agent name (empty = default agent).
370
+ timeout: RPC timeout.
371
+
372
+ Returns:
373
+ GetAgentSchemaResponse with schema containing observation_names,
374
+ action_names, observation_size, and action_size.
375
+ """
376
+ timeout = timeout or self.timeout
377
+ resp = self.agent.GetAgentSchema(
378
+ self.pb.agent.GetAgentSchemaRequest(agent_name=agent_name),
379
+ timeout=timeout,
380
+ )
381
+
382
+ # Cache the schema for named observation access
383
+ schema = getattr(resp, "schema", None)
384
+ if schema is not None:
385
+ cache_key = agent_name or "agent_0"
386
+ obs_names = list(schema.observation_names) if schema.observation_names else []
387
+ action_names = list(schema.action_names) if schema.action_names else []
388
+ self._schema_cache[cache_key] = (obs_names, action_names)
389
+ logger.debug(
390
+ "Cached schema for %s: %d obs names, %d action names",
391
+ cache_key,
392
+ len(obs_names),
393
+ len(action_names),
394
+ )
395
+
396
+ return resp
397
+
398
+ def fetch_schema(self, agent_name: str = "", timeout: Optional[float] = None) -> None:
399
+ """Fetch and cache agent schema for named observation access.
400
+
401
+ Call this once before get_observation() to enable accessing observations
402
+ by name (e.g., obs["proj_grav_x"]).
403
+
404
+ Args:
405
+ agent_name: Agent name (empty = default agent).
406
+ timeout: RPC timeout.
407
+ """
408
+ self.get_agent_schema(agent_name=agent_name, timeout=timeout)
409
+
410
+ def get_observation(
411
+ self,
412
+ agent_name: str = "",
413
+ timeout: Optional[float] = None,
414
+ ) -> ObservationResponse:
415
+ """
416
+ Get the RL observation vector for an agent.
417
+
418
+ This returns only the flat observation vector defined by the agent's
419
+ observation spec in LuckyEngine. For sensor data (joints, telemetry,
420
+ cameras), use the dedicated methods.
421
+
422
+ Args:
423
+ agent_name: Agent name (empty = default agent).
424
+ timeout: RPC timeout.
425
+
426
+ Returns:
427
+ ObservationResponse with observation vector, actions, timestamp.
428
+ """
429
+ timeout = timeout or self.timeout
430
+
431
+ resolved_robot_name = self._robot_name
432
+ if not resolved_robot_name:
433
+ raise ValueError(
434
+ "robot_name is required (set it once via "
435
+ "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
436
+ )
437
+
438
+ # Request only the agent's RL observation (agent_frame in gRPC terms).
439
+ # Joint state and telemetry have their own dedicated methods.
440
+ resp = self.agent.GetObservation(
441
+ self.pb.agent.GetObservationRequest(
442
+ robot_name=resolved_robot_name,
443
+ agent_name=agent_name,
444
+ include_joint_state=False,
445
+ include_agent_frame=True, # The RL observation vector
446
+ include_telemetry=False,
447
+ ),
448
+ timeout=timeout,
449
+ )
450
+
451
+ # Extract observation from agent_frame.
452
+ #
453
+ # In gRPC, the "agent_frame" is the message containing the RL observation
454
+ # data from LuckyEngine's agent system. It includes:
455
+ # - observations: flat float vector matching the agent's observation spec
456
+ # - actions: the last action vector sent to the agent (echoed back)
457
+ # - timestamp_ms: wall-clock time when the observation was captured
458
+ # - frame_number: monotonic counter for ordering observations
459
+ #
460
+ # This is distinct from joint_state (raw MuJoCo qpos/qvel) and telemetry
461
+ # (debugging data like contact forces, energy, etc).
462
+ agent_frame = getattr(resp, "agent_frame", None)
463
+ observations = []
464
+ actions = []
465
+ timestamp_ms = getattr(resp, "timestamp_ms", 0)
466
+ frame_number = getattr(resp, "frame_number", 0)
467
+
468
+ if agent_frame is not None:
469
+ observations = list(agent_frame.observations) if agent_frame.observations else []
470
+ actions = list(agent_frame.actions) if agent_frame.actions else []
471
+ timestamp_ms = getattr(agent_frame, "timestamp_ms", timestamp_ms)
472
+ frame_number = getattr(agent_frame, "frame_number", frame_number)
473
+
474
+ # Look up cached schema for named access
475
+ cache_key = agent_name or "agent_0"
476
+ obs_names, action_names = self._schema_cache.get(cache_key, (None, None))
477
+
478
+ return ObservationResponse(
479
+ observation=observations,
480
+ actions=actions,
481
+ timestamp_ms=timestamp_ms,
482
+ frame_number=frame_number,
483
+ agent_name=cache_key,
484
+ observation_names=obs_names,
485
+ action_names=action_names,
486
+ )
487
+
488
+ def get_state(
489
+ self,
490
+ agent_name: str = "",
491
+ include_observation: bool = True,
492
+ include_joint_state: bool = True,
493
+ camera_names: Optional[list[str]] = None,
494
+ width: int = 0,
495
+ height: int = 0,
496
+ format: str = "raw",
497
+ timeout: Optional[float] = None,
498
+ ) -> StateSnapshot:
499
+ """
500
+ Get a bundled snapshot of multiple data sources.
501
+
502
+ Use this for efficiency when you need multiple data types in one call.
503
+ For single data types, prefer the dedicated methods:
504
+ - get_observation() for RL observation vector
505
+ - get_joint_state() for joint positions/velocities
506
+ - stream_telemetry() for telemetry (streaming only)
507
+ - stream_camera() for camera frames (streaming)
508
+
509
+ Args:
510
+ agent_name: Agent name (empty = default agent).
511
+ include_observation: Include RL observation vector.
512
+ include_joint_state: Include joint positions/velocities.
513
+ camera_names: List of camera names to include frames from.
514
+ width: Desired width for camera frames (0 = native).
515
+ height: Desired height for camera frames (0 = native).
516
+ format: Image format ("raw" or "jpeg").
517
+ timeout: RPC timeout.
518
+
519
+ Returns:
520
+ StateSnapshot with requested data.
521
+ """
522
+ timeout = timeout or self.timeout
523
+
524
+ resolved_robot_name = self._robot_name
525
+ if not resolved_robot_name:
526
+ raise ValueError(
527
+ "robot_name is required (set it once via "
528
+ "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
529
+ )
530
+
531
+ cameras = []
532
+ if camera_names:
533
+ for name in camera_names:
534
+ cameras.append(
535
+ self.pb.agent.GetCameraFrameRequest(
536
+ name=name,
537
+ width=width,
538
+ height=height,
539
+ format=format,
540
+ )
541
+ )
542
+
543
+ resp = self.agent.GetObservation(
544
+ self.pb.agent.GetObservationRequest(
545
+ robot_name=resolved_robot_name,
546
+ agent_name=agent_name,
547
+ include_joint_state=include_joint_state,
548
+ include_agent_frame=include_observation,
549
+ include_telemetry=False, # Telemetry is streaming-only
550
+ cameras=cameras,
551
+ ),
552
+ timeout=timeout,
553
+ )
554
+
555
+ # Build ObservationResponse if requested
556
+ obs_response = None
557
+ if include_observation:
558
+ agent_frame = getattr(resp, "agent_frame", None)
559
+ observations = []
560
+ actions = []
561
+ if agent_frame is not None:
562
+ observations = list(agent_frame.observations) if agent_frame.observations else []
563
+ actions = list(agent_frame.actions) if agent_frame.actions else []
564
+
565
+ # Look up cached schema for named access
566
+ cache_key = agent_name or "agent_0"
567
+ obs_names, action_names = self._schema_cache.get(cache_key, (None, None))
568
+
569
+ obs_response = ObservationResponse(
570
+ observation=observations,
571
+ actions=actions,
572
+ timestamp_ms=getattr(resp, "timestamp_ms", 0),
573
+ frame_number=getattr(resp, "frame_number", 0),
574
+ agent_name=cache_key,
575
+ observation_names=obs_names,
576
+ action_names=action_names,
577
+ )
578
+
579
+ return StateSnapshot(
580
+ observation=obs_response,
581
+ joint_state=getattr(resp, "joint_state", None) if include_joint_state else None,
582
+ camera_frames=list(getattr(resp, "camera_frames", [])) if camera_names else None,
583
+ timestamp_ms=getattr(resp, "timestamp_ms", 0),
584
+ frame_number=getattr(resp, "frame_number", 0),
585
+ )
586
+
587
+ def stream_agent(self, agent_name: str = "", target_fps: int = 30):
588
+ """
589
+ Start streaming agent observations.
590
+
591
+ Returns an iterator of AgentFrame messages.
592
+ """
593
+ return self.agent.StreamAgent(
594
+ self.pb.agent.StreamAgentRequest(
595
+ agent_name=agent_name, target_fps=target_fps
596
+ ),
597
+ )
598
+
599
+ def reset_agent(
600
+ self,
601
+ agent_name: str = "",
602
+ randomization_cfg: Optional[Any] = None,
603
+ timeout: Optional[float] = None,
604
+ ):
605
+ """
606
+ Reset a specific agent (full reset: clear buffers, reset state, resample commands, apply MuJoCo state).
607
+
608
+ Useful for multi-env RL where individual agents need to be reset without resetting the entire scene.
609
+
610
+ Args:
611
+ agent_name: Agent logical name. Convention is `agent_0`, `agent_1`, ...
612
+ Empty string means "default agent" (agent_0).
613
+ randomization_cfg: Optional domain randomization config for this reset.
614
+ timeout: Timeout in seconds (uses default if None).
615
+
616
+ Returns:
617
+ ResetAgentResponse with success and message fields.
618
+ """
619
+ timeout = timeout or self.timeout
620
+
621
+ # Build request with optional randomization config
622
+ request_kwargs = {"agent_name": agent_name}
623
+
624
+ if randomization_cfg is not None:
625
+ randomization_proto = self._randomization_to_proto(randomization_cfg)
626
+ request_kwargs["dr_config"] = randomization_proto
627
+
628
+ return self.agent.ResetAgent(
629
+ self.pb.agent.ResetAgentRequest(**request_kwargs),
630
+ timeout=timeout,
631
+ )
632
+
633
+ def step(
634
+ self,
635
+ actions: list[float],
636
+ agent_name: str = "",
637
+ timeout_ms: int = 0,
638
+ timeout: Optional[float] = None,
639
+ ) -> ObservationResponse:
640
+ """
641
+ Synchronous RL step: apply action, wait for physics, return observation.
642
+
643
+ This is the recommended interface for RL training as it eliminates
644
+ one network round-trip compared to separate SendControl + GetObservation.
645
+
646
+ Args:
647
+ actions: Action vector to apply for this step.
648
+ agent_name: Agent name (empty = default agent).
649
+ timeout_ms: Server-side timeout for the step in milliseconds (0 = use default).
650
+ timeout: RPC timeout in seconds.
651
+
652
+ Returns:
653
+ ObservationResponse with observation after physics step.
654
+ """
655
+ timeout = timeout or self.timeout
656
+
657
+ resp = self.agent.Step(
658
+ self.pb.agent.StepRequest(
659
+ agent_name=agent_name,
660
+ actions=actions,
661
+ timeout_ms=timeout_ms,
662
+ ),
663
+ timeout=timeout,
664
+ )
665
+
666
+ if not resp.success:
667
+ raise RuntimeError(f"Step failed: {resp.message}")
668
+
669
+ # Extract observation from AgentFrame
670
+ agent_frame = resp.observation
671
+ observations = list(agent_frame.observations) if agent_frame.observations else []
672
+ actions_out = list(agent_frame.actions) if agent_frame.actions else []
673
+ timestamp_ms = getattr(agent_frame, "timestamp_ms", 0)
674
+ frame_number = getattr(agent_frame, "frame_number", 0)
675
+
676
+ # Look up cached schema for named access
677
+ cache_key = agent_name or "agent_0"
678
+ obs_names, action_names = self._schema_cache.get(cache_key, (None, None))
679
+
680
+ return ObservationResponse(
681
+ observation=observations,
682
+ actions=actions_out,
683
+ timestamp_ms=timestamp_ms,
684
+ frame_number=frame_number,
685
+ agent_name=cache_key,
686
+ observation_names=obs_names,
687
+ action_names=action_names,
688
+ )
689
+
690
+ def set_simulation_mode(
691
+ self,
692
+ mode: str = "fast",
693
+ timeout: Optional[float] = None,
694
+ ):
695
+ """
696
+ Set simulation timing mode.
697
+
698
+ Args:
699
+ mode: "realtime", "deterministic", or "fast"
700
+ - realtime: Physics runs at 1x wall-clock speed (for visualization)
701
+ - deterministic: Physics runs at fixed rate (for reproducibility)
702
+ - fast: Physics runs as fast as possible (for RL training)
703
+ timeout: RPC timeout in seconds.
704
+
705
+ Returns:
706
+ SetSimulationModeResponse with success and current mode.
707
+ """
708
+ timeout = timeout or self.timeout
709
+
710
+ mode_map = {
711
+ "realtime": 0,
712
+ "deterministic": 1,
713
+ "fast": 2,
714
+ }
715
+ mode_value = mode_map.get(mode.lower(), 2) # Default to fast
716
+
717
+ return self.scene.SetSimulationMode(
718
+ self.pb.scene.SetSimulationModeRequest(mode=mode_value),
719
+ timeout=timeout,
720
+ )
721
+
722
+ def get_simulation_mode(self, timeout: Optional[float] = None):
723
+ """
724
+ Get current simulation timing mode.
725
+
726
+ Returns:
727
+ GetSimulationModeResponse with mode field (0=realtime, 1=deterministic, 2=fast).
728
+ """
729
+ timeout = timeout or self.timeout
730
+
731
+ return self.scene.GetSimulationMode(
732
+ self.pb.scene.GetSimulationModeRequest(),
733
+ timeout=timeout,
734
+ )
735
+
736
+ def _randomization_to_proto(self, randomization_cfg: Any):
737
+ """Convert domain randomization config to proto message.
738
+
739
+ Accepts any object with randomization config fields (DomainRandomizationConfig,
740
+ PhysicsDRCfg from luckylab, or similar).
741
+ """
742
+ proto_kwargs = {}
743
+
744
+ # Helper to get attribute value, checking for None
745
+ def get_val(name: str, default=None):
746
+ val = getattr(randomization_cfg, name, default)
747
+ # Handle both None and empty tuples/lists
748
+ if val is None or (isinstance(val, (tuple, list)) and len(val) == 0):
749
+ return None
750
+ return val
751
+
752
+ # Initial state randomization
753
+ pose_pos = get_val("pose_position_noise")
754
+ if pose_pos is not None:
755
+ proto_kwargs["pose_position_noise"] = list(pose_pos)
756
+
757
+ pose_ori = get_val("pose_orientation_noise")
758
+ if pose_ori is not None and pose_ori != 0.0:
759
+ proto_kwargs["pose_orientation_noise"] = pose_ori
760
+
761
+ joint_pos = get_val("joint_position_noise")
762
+ if joint_pos is not None and joint_pos != 0.0:
763
+ proto_kwargs["joint_position_noise"] = joint_pos
764
+
765
+ joint_vel = get_val("joint_velocity_noise")
766
+ if joint_vel is not None and joint_vel != 0.0:
767
+ proto_kwargs["joint_velocity_noise"] = joint_vel
768
+
769
+ # Physics parameters (ranges)
770
+ friction = get_val("friction_range")
771
+ if friction is not None:
772
+ proto_kwargs["friction_range"] = list(friction)
773
+
774
+ restitution = get_val("restitution_range")
775
+ if restitution is not None:
776
+ proto_kwargs["restitution_range"] = list(restitution)
777
+
778
+ mass_scale = get_val("mass_scale_range")
779
+ if mass_scale is not None:
780
+ proto_kwargs["mass_scale_range"] = list(mass_scale)
781
+
782
+ com_offset = get_val("com_offset_range")
783
+ if com_offset is not None:
784
+ proto_kwargs["com_offset_range"] = list(com_offset)
785
+
786
+ # Motor/actuator
787
+ motor_strength = get_val("motor_strength_range")
788
+ if motor_strength is not None:
789
+ proto_kwargs["motor_strength_range"] = list(motor_strength)
790
+
791
+ motor_offset = get_val("motor_offset_range")
792
+ if motor_offset is not None:
793
+ proto_kwargs["motor_offset_range"] = list(motor_offset)
794
+
795
+ # External disturbances
796
+ push_interval = get_val("push_interval_range")
797
+ if push_interval is not None:
798
+ proto_kwargs["push_interval_range"] = list(push_interval)
799
+
800
+ push_velocity = get_val("push_velocity_range")
801
+ if push_velocity is not None:
802
+ proto_kwargs["push_velocity_range"] = list(push_velocity)
803
+
804
+ # Terrain
805
+ terrain_type = get_val("terrain_type")
806
+ if terrain_type is not None and terrain_type != "":
807
+ proto_kwargs["terrain_type"] = terrain_type
808
+
809
+ terrain_diff = get_val("terrain_difficulty")
810
+ if terrain_diff is not None and terrain_diff != 0.0:
811
+ proto_kwargs["terrain_difficulty"] = terrain_diff
812
+
813
+ return self.pb.agent.DomainRandomizationConfig(**proto_kwargs)
814
+
815
+ def stream_telemetry(self, target_fps: int = 30):
816
+ """
817
+ Start streaming telemetry data. [PLACEHOLDER - not confirmed working]
818
+
819
+ Returns an iterator of TelemetryFrame messages.
820
+ """
821
+ raise NotImplementedError(
822
+ "stream_telemetry() is not yet confirmed working. "
823
+ "Remove this check if you want to test it."
824
+ )
825
+ return self.telemetry.StreamTelemetry(
826
+ self.pb.telemetry.StreamTelemetryRequest(target_fps=target_fps),
827
+ )
828
+
829
+ def list_cameras(self, timeout: Optional[float] = None):
830
+ """List available cameras. [PLACEHOLDER - not confirmed working]"""
831
+ raise NotImplementedError(
832
+ "list_cameras() is not yet confirmed working. "
833
+ "Remove this check if you want to test it."
834
+ )
835
+ timeout = timeout or self.timeout
836
+ return self.camera.ListCameras(
837
+ self.pb.camera.ListCamerasRequest(),
838
+ timeout=timeout,
839
+ )
840
+
841
+ def stream_camera(
842
+ self,
843
+ camera_id: Optional[int] = None,
844
+ camera_name: Optional[str] = None,
845
+ target_fps: int = 30,
846
+ width: int = 640,
847
+ height: int = 480,
848
+ format: str = "raw",
849
+ ):
850
+ """
851
+ Start streaming camera frames. [PLACEHOLDER - not confirmed working]
852
+
853
+ Args:
854
+ camera_id: Camera entity ID (use either id or name).
855
+ camera_name: Camera name (use either id or name).
856
+ target_fps: Desired frames per second.
857
+ width: Desired width (0 = native).
858
+ height: Desired height (0 = native).
859
+ format: Image format ("raw" or "jpeg").
860
+
861
+ Returns an iterator of ImageFrame messages.
862
+ """
863
+ raise NotImplementedError(
864
+ "stream_camera() is not yet confirmed working. "
865
+ "Remove this check if you want to test it."
866
+ )
867
+ if camera_id is not None:
868
+ request = self.pb.camera.StreamCameraRequest(
869
+ id=self.pb.common.EntityId(id=camera_id),
870
+ target_fps=target_fps,
871
+ width=width,
872
+ height=height,
873
+ format=format,
874
+ )
875
+ elif camera_name is not None:
876
+ request = self.pb.camera.StreamCameraRequest(
877
+ name=camera_name,
878
+ target_fps=target_fps,
879
+ width=width,
880
+ height=height,
881
+ format=format,
882
+ )
883
+ else:
884
+ raise ValueError("Either camera_id or camera_name must be provided")
885
+
886
+ return self.camera.StreamCamera(request)
887
+
888
+ def benchmark(
889
+ self,
890
+ duration_seconds: float = 5.0,
891
+ method: str = "get_observation",
892
+ print_results: bool = True,
893
+ ) -> "BenchmarkResult":
894
+ """
895
+ Benchmark gRPC performance.
896
+
897
+ Measures actual FPS and latency for observation calls.
898
+
899
+ Args:
900
+ duration_seconds: How long to run the benchmark.
901
+ method: Which method to benchmark ("get_observation" or "stream_agent").
902
+ print_results: Whether to print results to console.
903
+
904
+ Returns:
905
+ BenchmarkResult with FPS and latency statistics.
906
+ """
907
+ if method not in ("get_observation", "stream_agent"):
908
+ raise ValueError(f"Unknown method: {method}. Use 'get_observation' or 'stream_agent'.")
909
+
910
+ latencies: list[float] = []
911
+ frame_times: list[float] = []
912
+ frame_count = 0
913
+ start_time = time.perf_counter()
914
+ last_frame_time = start_time
915
+
916
+ if method == "get_observation":
917
+ while time.perf_counter() - start_time < duration_seconds:
918
+ call_start = time.perf_counter()
919
+ self.get_observation()
920
+ call_end = time.perf_counter()
921
+
922
+ latencies.append((call_end - call_start) * 1000) # ms
923
+ frame_times.append(call_end - last_frame_time)
924
+ last_frame_time = call_end
925
+ frame_count += 1
926
+
927
+ elif method == "stream_agent":
928
+ stream = self.stream_agent(target_fps=1000) # Request max FPS
929
+ for frame in stream:
930
+ now = time.perf_counter()
931
+ if now - start_time >= duration_seconds:
932
+ break
933
+
934
+ frame_times.append(now - last_frame_time)
935
+ last_frame_time = now
936
+ frame_count += 1
937
+
938
+ total_time = time.perf_counter() - start_time
939
+
940
+ # Calculate statistics
941
+ actual_fps = frame_count / total_time if total_time > 0 else 0
942
+
943
+ if len(frame_times) > 1:
944
+ # Remove first frame (warmup)
945
+ frame_times = frame_times[1:]
946
+ fps_from_frames = 1.0 / statistics.mean(frame_times) if frame_times else 0
947
+ else:
948
+ fps_from_frames = actual_fps
949
+
950
+ result = BenchmarkResult(
951
+ method=method,
952
+ duration_seconds=total_time,
953
+ frame_count=frame_count,
954
+ actual_fps=actual_fps,
955
+ avg_latency_ms=statistics.mean(latencies) if latencies else 0,
956
+ min_latency_ms=min(latencies) if latencies else 0,
957
+ max_latency_ms=max(latencies) if latencies else 0,
958
+ std_latency_ms=statistics.stdev(latencies) if len(latencies) > 1 else 0,
959
+ p50_latency_ms=statistics.median(latencies) if latencies else 0,
960
+ p99_latency_ms=sorted(latencies)[int(len(latencies) * 0.99)] if latencies else 0,
961
+ )
962
+
963
+ if print_results:
964
+ print(f"\n{'=' * 50}")
965
+ print(f"Benchmark Results ({method})")
966
+ print(f"{'=' * 50}")
967
+ print(f"Duration: {result.duration_seconds:.2f}s")
968
+ print(f"Frames: {result.frame_count}")
969
+ print(f"Actual FPS: {result.actual_fps:.1f}")
970
+ if latencies:
971
+ print(f"{'─' * 50}")
972
+ print(f"Latency (ms):")
973
+ print(f" avg: {result.avg_latency_ms:.2f}")
974
+ print(f" min: {result.min_latency_ms:.2f}")
975
+ print(f" max: {result.max_latency_ms:.2f}")
976
+ print(f" std: {result.std_latency_ms:.2f}")
977
+ print(f" p50: {result.p50_latency_ms:.2f}")
978
+ print(f" p99: {result.p99_latency_ms:.2f}")
979
+ print(f"{'=' * 50}\n")
980
+
981
+ return result
982
+
983
+
984
+ @dataclass
985
+ class BenchmarkResult:
986
+ """Results from a benchmark run."""
987
+
988
+ method: str
989
+ duration_seconds: float
990
+ frame_count: int
991
+ actual_fps: float
992
+ avg_latency_ms: float
993
+ min_latency_ms: float
994
+ max_latency_ms: float
995
+ std_latency_ms: float
996
+ p50_latency_ms: float
997
+ p99_latency_ms: float