luckyrobots 0.1.68__py3-none-any.whl → 0.1.70__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 (57) hide show
  1. luckyrobots/__init__.py +23 -12
  2. luckyrobots/client.py +800 -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 +61 -0
  11. luckyrobots/grpc/generated/agent_pb2_grpc.py +255 -0
  12. luckyrobots/grpc/generated/camera_pb2.py +45 -0
  13. luckyrobots/grpc/generated/camera_pb2_grpc.py +155 -0
  14. luckyrobots/grpc/generated/common_pb2.py +39 -0
  15. luckyrobots/grpc/generated/common_pb2_grpc.py +27 -0
  16. luckyrobots/grpc/generated/media_pb2.py +35 -0
  17. luckyrobots/grpc/generated/media_pb2_grpc.py +27 -0
  18. luckyrobots/grpc/generated/mujoco_pb2.py +47 -0
  19. luckyrobots/grpc/generated/mujoco_pb2_grpc.py +248 -0
  20. luckyrobots/grpc/generated/scene_pb2.py +54 -0
  21. luckyrobots/grpc/generated/scene_pb2_grpc.py +248 -0
  22. luckyrobots/grpc/generated/telemetry_pb2.py +43 -0
  23. luckyrobots/grpc/generated/telemetry_pb2_grpc.py +154 -0
  24. luckyrobots/grpc/generated/viewport_pb2.py +48 -0
  25. luckyrobots/grpc/generated/viewport_pb2_grpc.py +155 -0
  26. luckyrobots/grpc/proto/agent.proto +152 -0
  27. luckyrobots/grpc/proto/camera.proto +41 -0
  28. luckyrobots/grpc/proto/common.proto +36 -0
  29. luckyrobots/grpc/proto/hazel_rpc.proto +32 -0
  30. luckyrobots/grpc/proto/media.proto +26 -0
  31. luckyrobots/grpc/proto/mujoco.proto +64 -0
  32. luckyrobots/grpc/proto/scene.proto +70 -0
  33. luckyrobots/grpc/proto/telemetry.proto +43 -0
  34. luckyrobots/grpc/proto/viewport.proto +45 -0
  35. luckyrobots/luckyrobots.py +212 -0
  36. luckyrobots/models/__init__.py +13 -0
  37. luckyrobots/models/camera.py +97 -0
  38. luckyrobots/models/observation.py +135 -0
  39. luckyrobots/{utils/helpers.py → utils.py} +75 -40
  40. luckyrobots-0.1.70.dist-info/METADATA +262 -0
  41. luckyrobots-0.1.70.dist-info/RECORD +44 -0
  42. {luckyrobots-0.1.68.dist-info → luckyrobots-0.1.70.dist-info}/WHEEL +1 -1
  43. luckyrobots/core/luckyrobots.py +0 -628
  44. luckyrobots/core/manager.py +0 -236
  45. luckyrobots/core/models.py +0 -68
  46. luckyrobots/core/node.py +0 -273
  47. luckyrobots/message/__init__.py +0 -18
  48. luckyrobots/message/pubsub.py +0 -145
  49. luckyrobots/message/srv/client.py +0 -81
  50. luckyrobots/message/srv/service.py +0 -135
  51. luckyrobots/message/srv/types.py +0 -83
  52. luckyrobots/message/transporter.py +0 -427
  53. luckyrobots/utils/event_loop.py +0 -94
  54. luckyrobots/utils/sim_manager.py +0 -413
  55. luckyrobots-0.1.68.dist-info/METADATA +0 -253
  56. luckyrobots-0.1.68.dist-info/RECORD +0 -24
  57. {luckyrobots-0.1.68.dist-info → luckyrobots-0.1.70.dist-info}/licenses/LICENSE +0 -0
luckyrobots/client.py ADDED
@@ -0,0 +1,800 @@
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
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"Connected to LuckyEngine gRPC server at {target}")
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
+ return True
225
+
226
+ time.sleep(poll_interval)
227
+
228
+ return False
229
+
230
+ # --- Protobuf modules (discoverable + explicit) ---
231
+
232
+ @property
233
+ def pb(self) -> Any:
234
+ """Access protobuf modules grouped by domain (e.g., `client.pb.scene`)."""
235
+ return self._pb
236
+
237
+ @property
238
+ def robot_name(self) -> Optional[str]:
239
+ """Default robot name used by calls that accept an optional robot_name."""
240
+ return self._robot_name
241
+
242
+ def set_robot_name(self, robot_name: str) -> None:
243
+ """Set the default robot name used by calls that accept an optional robot_name."""
244
+ self._robot_name = robot_name
245
+
246
+ # --- Service stubs ---
247
+ #
248
+ # Confirmed working:
249
+ # - mujoco: GetMujocoInfo, GetJointState, SendControl
250
+ # - agent: GetObservation, ResetAgent, GetAgentSchema, StreamAgent
251
+ #
252
+ # Placeholders (not yet confirmed working - use at your own risk):
253
+ # - scene: GetSceneInfo
254
+ # - telemetry: StreamTelemetry
255
+ # - viewport: (no methods implemented)
256
+ # - camera: ListCameras, StreamCamera
257
+
258
+ @property
259
+ def scene(self) -> Any:
260
+ """SceneService stub. [PLACEHOLDER - not confirmed working]"""
261
+ if self._scene is None:
262
+ raise GrpcConnectionError("Not connected. Call connect() first.")
263
+ return self._scene
264
+
265
+ @property
266
+ def mujoco(self) -> Any:
267
+ """MujocoService stub."""
268
+ if self._mujoco is None:
269
+ raise GrpcConnectionError("Not connected. Call connect() first.")
270
+ return self._mujoco
271
+
272
+ @property
273
+ def telemetry(self) -> Any:
274
+ """TelemetryService stub. [PLACEHOLDER - not confirmed working]"""
275
+ if self._telemetry is None:
276
+ raise GrpcConnectionError("Not connected. Call connect() first.")
277
+ return self._telemetry
278
+
279
+ @property
280
+ def agent(self) -> Any:
281
+ """AgentService stub."""
282
+ if self._agent is None:
283
+ raise GrpcConnectionError("Not connected. Call connect() first.")
284
+ return self._agent
285
+
286
+ @property
287
+ def viewport(self) -> Any:
288
+ """ViewportService stub. [PLACEHOLDER - not confirmed working]"""
289
+ if self._viewport is None:
290
+ raise GrpcConnectionError("Not connected. Call connect() first.")
291
+ return self._viewport
292
+
293
+ @property
294
+ def camera(self) -> Any:
295
+ """CameraService stub. [PLACEHOLDER - not confirmed working]"""
296
+ if self._camera is None:
297
+ raise GrpcConnectionError("Not connected. Call connect() first.")
298
+ return self._camera
299
+
300
+ # --- Convenience methods ---
301
+
302
+ def get_scene_info(self, timeout: Optional[float] = None):
303
+ """Get scene information. [PLACEHOLDER - not confirmed working]"""
304
+ raise NotImplementedError(
305
+ "get_scene_info() is not yet confirmed working. "
306
+ "Remove this check if you want to test it."
307
+ )
308
+ timeout = timeout or self.timeout
309
+ return self.scene.GetSceneInfo(
310
+ self.pb.scene.GetSceneInfoRequest(),
311
+ timeout=timeout,
312
+ )
313
+
314
+ def get_mujoco_info(self, robot_name: str = "", timeout: Optional[float] = None):
315
+ """Get MuJoCo model information."""
316
+ timeout = timeout or self.timeout
317
+ robot_name = robot_name or self._robot_name
318
+ if not robot_name:
319
+ raise ValueError(
320
+ "robot_name is required (pass `robot_name=` or set it once via "
321
+ "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
322
+ )
323
+ return self.mujoco.GetMujocoInfo(
324
+ self.pb.mujoco.GetMujocoInfoRequest(robot_name=robot_name),
325
+ timeout=timeout,
326
+ )
327
+
328
+ def get_joint_state(self, robot_name: str = "", timeout: Optional[float] = None):
329
+ """Get current joint state."""
330
+ timeout = timeout or self.timeout
331
+ robot_name = robot_name or self._robot_name
332
+ if not robot_name:
333
+ raise ValueError(
334
+ "robot_name is required (pass `robot_name=` or set it once via "
335
+ "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
336
+ )
337
+ return self.mujoco.GetJointState(
338
+ self.pb.mujoco.GetJointStateRequest(robot_name=robot_name),
339
+ timeout=timeout,
340
+ )
341
+
342
+ def send_control(
343
+ self,
344
+ controls: list[float],
345
+ robot_name: str = "",
346
+ timeout: Optional[float] = None,
347
+ ):
348
+ """Send control commands to the robot."""
349
+ timeout = timeout or self.timeout
350
+ robot_name = robot_name or self._robot_name
351
+ if not robot_name:
352
+ raise ValueError(
353
+ "robot_name is required (pass `robot_name=` or set it once via "
354
+ "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
355
+ )
356
+ return self.mujoco.SendControl(
357
+ self.pb.mujoco.SendControlRequest(robot_name=robot_name, controls=controls),
358
+ timeout=timeout,
359
+ )
360
+
361
+ def get_agent_schema(self, agent_name: str = "", timeout: Optional[float] = None):
362
+ """Get agent schema (observation/action sizes and names).
363
+
364
+ The schema is cached for subsequent get_observation() calls to enable
365
+ named access to observation values.
366
+
367
+ Args:
368
+ agent_name: Agent name (empty = default agent).
369
+ timeout: RPC timeout.
370
+
371
+ Returns:
372
+ GetAgentSchemaResponse with schema containing observation_names,
373
+ action_names, observation_size, and action_size.
374
+ """
375
+ timeout = timeout or self.timeout
376
+ resp = self.agent.GetAgentSchema(
377
+ self.pb.agent.GetAgentSchemaRequest(agent_name=agent_name),
378
+ timeout=timeout,
379
+ )
380
+
381
+ # Cache the schema for named observation access
382
+ schema = getattr(resp, "schema", None)
383
+ if schema is not None:
384
+ cache_key = agent_name or "agent_0"
385
+ obs_names = list(schema.observation_names) if schema.observation_names else []
386
+ action_names = list(schema.action_names) if schema.action_names else []
387
+ self._schema_cache[cache_key] = (obs_names, action_names)
388
+ logger.debug(
389
+ "Cached schema for %s: %d obs names, %d action names",
390
+ cache_key,
391
+ len(obs_names),
392
+ len(action_names),
393
+ )
394
+
395
+ return resp
396
+
397
+ def fetch_schema(self, agent_name: str = "", timeout: Optional[float] = None) -> None:
398
+ """Fetch and cache agent schema for named observation access.
399
+
400
+ Call this once before get_observation() to enable accessing observations
401
+ by name (e.g., obs["proj_grav_x"]).
402
+
403
+ Args:
404
+ agent_name: Agent name (empty = default agent).
405
+ timeout: RPC timeout.
406
+ """
407
+ self.get_agent_schema(agent_name=agent_name, timeout=timeout)
408
+
409
+ def get_observation(
410
+ self,
411
+ agent_name: str = "",
412
+ timeout: Optional[float] = None,
413
+ ) -> ObservationResponse:
414
+ """
415
+ Get the RL observation vector for an agent.
416
+
417
+ This returns only the flat observation vector defined by the agent's
418
+ observation spec in LuckyEngine. For sensor data (joints, telemetry,
419
+ cameras), use the dedicated methods.
420
+
421
+ Args:
422
+ agent_name: Agent name (empty = default agent).
423
+ timeout: RPC timeout.
424
+
425
+ Returns:
426
+ ObservationResponse with observation vector, actions, timestamp.
427
+ """
428
+ timeout = timeout or self.timeout
429
+
430
+ resolved_robot_name = self._robot_name
431
+ if not resolved_robot_name:
432
+ raise ValueError(
433
+ "robot_name is required (set it once via "
434
+ "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
435
+ )
436
+
437
+ # Request only the agent's RL observation (agent_frame in gRPC terms).
438
+ # Joint state and telemetry have their own dedicated methods.
439
+ resp = self.agent.GetObservation(
440
+ self.pb.agent.GetObservationRequest(
441
+ robot_name=resolved_robot_name,
442
+ agent_name=agent_name,
443
+ include_joint_state=False,
444
+ include_agent_frame=True, # The RL observation vector
445
+ include_telemetry=False,
446
+ ),
447
+ timeout=timeout,
448
+ )
449
+
450
+ # Extract observation from agent_frame.
451
+ #
452
+ # In gRPC, the "agent_frame" is the message containing the RL observation
453
+ # data from LuckyEngine's agent system. It includes:
454
+ # - observations: flat float vector matching the agent's observation spec
455
+ # - actions: the last action vector sent to the agent (echoed back)
456
+ # - timestamp_ms: wall-clock time when the observation was captured
457
+ # - frame_number: monotonic counter for ordering observations
458
+ #
459
+ # This is distinct from joint_state (raw MuJoCo qpos/qvel) and telemetry
460
+ # (debugging data like contact forces, energy, etc).
461
+ agent_frame = getattr(resp, "agent_frame", None)
462
+ observations = []
463
+ actions = []
464
+ timestamp_ms = getattr(resp, "timestamp_ms", 0)
465
+ frame_number = getattr(resp, "frame_number", 0)
466
+
467
+ if agent_frame is not None:
468
+ observations = list(agent_frame.observations) if agent_frame.observations else []
469
+ actions = list(agent_frame.actions) if agent_frame.actions else []
470
+ timestamp_ms = getattr(agent_frame, "timestamp_ms", timestamp_ms)
471
+ frame_number = getattr(agent_frame, "frame_number", frame_number)
472
+
473
+ # Look up cached schema for named access
474
+ cache_key = agent_name or "agent_0"
475
+ obs_names, action_names = self._schema_cache.get(cache_key, (None, None))
476
+
477
+ return ObservationResponse(
478
+ observation=observations,
479
+ actions=actions,
480
+ timestamp_ms=timestamp_ms,
481
+ frame_number=frame_number,
482
+ agent_name=cache_key,
483
+ observation_names=obs_names,
484
+ action_names=action_names,
485
+ )
486
+
487
+ def get_state(
488
+ self,
489
+ agent_name: str = "",
490
+ include_observation: bool = True,
491
+ include_joint_state: bool = True,
492
+ camera_names: Optional[list[str]] = None,
493
+ width: int = 0,
494
+ height: int = 0,
495
+ format: str = "raw",
496
+ timeout: Optional[float] = None,
497
+ ) -> StateSnapshot:
498
+ """
499
+ Get a bundled snapshot of multiple data sources.
500
+
501
+ Use this for efficiency when you need multiple data types in one call.
502
+ For single data types, prefer the dedicated methods:
503
+ - get_observation() for RL observation vector
504
+ - get_joint_state() for joint positions/velocities
505
+ - stream_telemetry() for telemetry (streaming only)
506
+ - stream_camera() for camera frames (streaming)
507
+
508
+ Args:
509
+ agent_name: Agent name (empty = default agent).
510
+ include_observation: Include RL observation vector.
511
+ include_joint_state: Include joint positions/velocities.
512
+ camera_names: List of camera names to include frames from.
513
+ width: Desired width for camera frames (0 = native).
514
+ height: Desired height for camera frames (0 = native).
515
+ format: Image format ("raw" or "jpeg").
516
+ timeout: RPC timeout.
517
+
518
+ Returns:
519
+ StateSnapshot with requested data.
520
+ """
521
+ timeout = timeout or self.timeout
522
+
523
+ resolved_robot_name = self._robot_name
524
+ if not resolved_robot_name:
525
+ raise ValueError(
526
+ "robot_name is required (set it once via "
527
+ "LuckyEngineClient(robot_name=...) / client.set_robot_name(...))."
528
+ )
529
+
530
+ cameras = []
531
+ if camera_names:
532
+ for name in camera_names:
533
+ cameras.append(
534
+ self.pb.agent.GetCameraFrameRequest(
535
+ name=name,
536
+ width=width,
537
+ height=height,
538
+ format=format,
539
+ )
540
+ )
541
+
542
+ resp = self.agent.GetObservation(
543
+ self.pb.agent.GetObservationRequest(
544
+ robot_name=resolved_robot_name,
545
+ agent_name=agent_name,
546
+ include_joint_state=include_joint_state,
547
+ include_agent_frame=include_observation,
548
+ include_telemetry=False, # Telemetry is streaming-only
549
+ cameras=cameras,
550
+ ),
551
+ timeout=timeout,
552
+ )
553
+
554
+ # Build ObservationResponse if requested
555
+ obs_response = None
556
+ if include_observation:
557
+ agent_frame = getattr(resp, "agent_frame", None)
558
+ observations = []
559
+ actions = []
560
+ if agent_frame is not None:
561
+ observations = list(agent_frame.observations) if agent_frame.observations else []
562
+ actions = list(agent_frame.actions) if agent_frame.actions else []
563
+
564
+ # Look up cached schema for named access
565
+ cache_key = agent_name or "agent_0"
566
+ obs_names, action_names = self._schema_cache.get(cache_key, (None, None))
567
+
568
+ obs_response = ObservationResponse(
569
+ observation=observations,
570
+ actions=actions,
571
+ timestamp_ms=getattr(resp, "timestamp_ms", 0),
572
+ frame_number=getattr(resp, "frame_number", 0),
573
+ agent_name=cache_key,
574
+ observation_names=obs_names,
575
+ action_names=action_names,
576
+ )
577
+
578
+ return StateSnapshot(
579
+ observation=obs_response,
580
+ joint_state=getattr(resp, "joint_state", None) if include_joint_state else None,
581
+ camera_frames=list(getattr(resp, "camera_frames", [])) if camera_names else None,
582
+ timestamp_ms=getattr(resp, "timestamp_ms", 0),
583
+ frame_number=getattr(resp, "frame_number", 0),
584
+ )
585
+
586
+ def stream_agent(self, agent_name: str = "", target_fps: int = 30):
587
+ """
588
+ Start streaming agent observations.
589
+
590
+ Returns an iterator of AgentFrame messages.
591
+ """
592
+ return self.agent.StreamAgent(
593
+ self.pb.agent.StreamAgentRequest(
594
+ agent_name=agent_name, target_fps=target_fps
595
+ ),
596
+ )
597
+
598
+ def reset_agent(self, agent_name: str = "", timeout: Optional[float] = None):
599
+ """
600
+ Reset a specific agent (full reset: clear buffers, reset state, resample commands, apply MuJoCo state).
601
+
602
+ Useful for multi-env RL where individual agents need to be reset without resetting the entire scene.
603
+
604
+ Args:
605
+ agent_name: Agent logical name. Convention is `agent_0`, `agent_1`, ...
606
+ Empty string means "default agent" (agent_0).
607
+ timeout: Timeout in seconds (uses default if None).
608
+
609
+ Returns:
610
+ ResetAgentResponse with success and message fields.
611
+ """
612
+ timeout = timeout or self.timeout
613
+ return self.agent.ResetAgent(
614
+ self.pb.agent.ResetAgentRequest(agent_name=agent_name),
615
+ timeout=timeout,
616
+ )
617
+
618
+ def stream_telemetry(self, target_fps: int = 30):
619
+ """
620
+ Start streaming telemetry data. [PLACEHOLDER - not confirmed working]
621
+
622
+ Returns an iterator of TelemetryFrame messages.
623
+ """
624
+ raise NotImplementedError(
625
+ "stream_telemetry() is not yet confirmed working. "
626
+ "Remove this check if you want to test it."
627
+ )
628
+ return self.telemetry.StreamTelemetry(
629
+ self.pb.telemetry.StreamTelemetryRequest(target_fps=target_fps),
630
+ )
631
+
632
+ def list_cameras(self, timeout: Optional[float] = None):
633
+ """List available cameras. [PLACEHOLDER - not confirmed working]"""
634
+ raise NotImplementedError(
635
+ "list_cameras() is not yet confirmed working. "
636
+ "Remove this check if you want to test it."
637
+ )
638
+ timeout = timeout or self.timeout
639
+ return self.camera.ListCameras(
640
+ self.pb.camera.ListCamerasRequest(),
641
+ timeout=timeout,
642
+ )
643
+
644
+ def stream_camera(
645
+ self,
646
+ camera_id: Optional[int] = None,
647
+ camera_name: Optional[str] = None,
648
+ target_fps: int = 30,
649
+ width: int = 640,
650
+ height: int = 480,
651
+ format: str = "raw",
652
+ ):
653
+ """
654
+ Start streaming camera frames. [PLACEHOLDER - not confirmed working]
655
+
656
+ Args:
657
+ camera_id: Camera entity ID (use either id or name).
658
+ camera_name: Camera name (use either id or name).
659
+ target_fps: Desired frames per second.
660
+ width: Desired width (0 = native).
661
+ height: Desired height (0 = native).
662
+ format: Image format ("raw" or "jpeg").
663
+
664
+ Returns an iterator of ImageFrame messages.
665
+ """
666
+ raise NotImplementedError(
667
+ "stream_camera() is not yet confirmed working. "
668
+ "Remove this check if you want to test it."
669
+ )
670
+ if camera_id is not None:
671
+ request = self.pb.camera.StreamCameraRequest(
672
+ id=self.pb.common.EntityId(id=camera_id),
673
+ target_fps=target_fps,
674
+ width=width,
675
+ height=height,
676
+ format=format,
677
+ )
678
+ elif camera_name is not None:
679
+ request = self.pb.camera.StreamCameraRequest(
680
+ name=camera_name,
681
+ target_fps=target_fps,
682
+ width=width,
683
+ height=height,
684
+ format=format,
685
+ )
686
+ else:
687
+ raise ValueError("Either camera_id or camera_name must be provided")
688
+
689
+ return self.camera.StreamCamera(request)
690
+
691
+ def benchmark(
692
+ self,
693
+ duration_seconds: float = 5.0,
694
+ method: str = "get_observation",
695
+ print_results: bool = True,
696
+ ) -> "BenchmarkResult":
697
+ """
698
+ Benchmark gRPC performance.
699
+
700
+ Measures actual FPS and latency for observation calls.
701
+
702
+ Args:
703
+ duration_seconds: How long to run the benchmark.
704
+ method: Which method to benchmark ("get_observation" or "stream_agent").
705
+ print_results: Whether to print results to console.
706
+
707
+ Returns:
708
+ BenchmarkResult with FPS and latency statistics.
709
+ """
710
+ if method not in ("get_observation", "stream_agent"):
711
+ raise ValueError(f"Unknown method: {method}. Use 'get_observation' or 'stream_agent'.")
712
+
713
+ latencies: list[float] = []
714
+ frame_times: list[float] = []
715
+ frame_count = 0
716
+ start_time = time.perf_counter()
717
+ last_frame_time = start_time
718
+
719
+ if method == "get_observation":
720
+ while time.perf_counter() - start_time < duration_seconds:
721
+ call_start = time.perf_counter()
722
+ self.get_observation()
723
+ call_end = time.perf_counter()
724
+
725
+ latencies.append((call_end - call_start) * 1000) # ms
726
+ frame_times.append(call_end - last_frame_time)
727
+ last_frame_time = call_end
728
+ frame_count += 1
729
+
730
+ elif method == "stream_agent":
731
+ stream = self.stream_agent(target_fps=1000) # Request max FPS
732
+ for frame in stream:
733
+ now = time.perf_counter()
734
+ if now - start_time >= duration_seconds:
735
+ break
736
+
737
+ frame_times.append(now - last_frame_time)
738
+ last_frame_time = now
739
+ frame_count += 1
740
+
741
+ total_time = time.perf_counter() - start_time
742
+
743
+ # Calculate statistics
744
+ actual_fps = frame_count / total_time if total_time > 0 else 0
745
+
746
+ if len(frame_times) > 1:
747
+ # Remove first frame (warmup)
748
+ frame_times = frame_times[1:]
749
+ fps_from_frames = 1.0 / statistics.mean(frame_times) if frame_times else 0
750
+ else:
751
+ fps_from_frames = actual_fps
752
+
753
+ result = BenchmarkResult(
754
+ method=method,
755
+ duration_seconds=total_time,
756
+ frame_count=frame_count,
757
+ actual_fps=actual_fps,
758
+ avg_latency_ms=statistics.mean(latencies) if latencies else 0,
759
+ min_latency_ms=min(latencies) if latencies else 0,
760
+ max_latency_ms=max(latencies) if latencies else 0,
761
+ std_latency_ms=statistics.stdev(latencies) if len(latencies) > 1 else 0,
762
+ p50_latency_ms=statistics.median(latencies) if latencies else 0,
763
+ p99_latency_ms=sorted(latencies)[int(len(latencies) * 0.99)] if latencies else 0,
764
+ )
765
+
766
+ if print_results:
767
+ print(f"\n{'=' * 50}")
768
+ print(f"Benchmark Results ({method})")
769
+ print(f"{'=' * 50}")
770
+ print(f"Duration: {result.duration_seconds:.2f}s")
771
+ print(f"Frames: {result.frame_count}")
772
+ print(f"Actual FPS: {result.actual_fps:.1f}")
773
+ if latencies:
774
+ print(f"{'─' * 50}")
775
+ print(f"Latency (ms):")
776
+ print(f" avg: {result.avg_latency_ms:.2f}")
777
+ print(f" min: {result.min_latency_ms:.2f}")
778
+ print(f" max: {result.max_latency_ms:.2f}")
779
+ print(f" std: {result.std_latency_ms:.2f}")
780
+ print(f" p50: {result.p50_latency_ms:.2f}")
781
+ print(f" p99: {result.p99_latency_ms:.2f}")
782
+ print(f"{'=' * 50}\n")
783
+
784
+ return result
785
+
786
+
787
+ @dataclass
788
+ class BenchmarkResult:
789
+ """Results from a benchmark run."""
790
+
791
+ method: str
792
+ duration_seconds: float
793
+ frame_count: int
794
+ actual_fps: float
795
+ avg_latency_ms: float
796
+ min_latency_ms: float
797
+ max_latency_ms: float
798
+ std_latency_ms: float
799
+ p50_latency_ms: float
800
+ p99_latency_ms: float