Introduce message and endpoint abstractions for MuJoCo Studio communication.
- `messages.py`: Defines base classes for `Message`, `Snapshot`, and `Event`, along with `SnapshotChannel` and `EventChannel` protocols. It also includes concrete message types for state, reset, model, perturbation, pause state, and exit events. - `endpoints.py`: Provides `ViewerEndpoint` and `SimEndpoint` classes to manage message routing between the simulation and the viewer via the defined channels. A `make_endpoints` function is included to create both endpoints. PiperOrigin-RevId: 939659912 Change-Id: I61a52a7d25431bc4ed38a184813882145eb421d1
This commit is contained in:
Matija Kecman
committed by
Copybara-Service
Copybara-Service
parent
f594778813
commit
e0bb2b0665
@@ -0,0 +1,132 @@
|
||||
# Copyright 2026 DeepMind Technologies Limited
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# https://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Endpoints route messages to the right channel."""
|
||||
|
||||
from mujoco.experimental.studio import messages as vp
|
||||
|
||||
|
||||
class ViewerEndpoint:
|
||||
"""Viewer endpoint to route Messages to the Sim using the right Channel.
|
||||
|
||||
The viewer has an outgoing EventChannel but no outgoing SnapshotChannel.
|
||||
The viewer has an incoming EventChannel and SnapshotChannel.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
s2v_snapshot: vp.SnapshotChannel,
|
||||
s2v_events: vp.EventChannel,
|
||||
v2s_events: vp.EventChannel,
|
||||
):
|
||||
self._s2v_snapshot = s2v_snapshot
|
||||
self._s2v_events = s2v_events
|
||||
self._v2s_events = v2s_events
|
||||
self._is_closed = False
|
||||
|
||||
def send_to_sim(self, message: vp.Message) -> None:
|
||||
if self._is_closed:
|
||||
raise RuntimeError('ViewerEndpoint is closed')
|
||||
if isinstance(message, vp.Event):
|
||||
self._v2s_events.put(message)
|
||||
elif isinstance(message, vp.Snapshot):
|
||||
raise TypeError('viewer cannot produce snapshots')
|
||||
else:
|
||||
raise TypeError(f'expected Event, got {type(message).__name__}')
|
||||
|
||||
def get_sim_events(self) -> list[vp.Event]:
|
||||
"""Returns all pending events from the simulation."""
|
||||
if self._is_closed:
|
||||
raise RuntimeError('ViewerEndpoint is closed')
|
||||
return self._s2v_events.get()
|
||||
|
||||
def get_sim_snapshots(self) -> list[vp.Snapshot]:
|
||||
"""Returns all pending latest snapshots from the simulation, one per type."""
|
||||
if self._is_closed:
|
||||
raise RuntimeError('ViewerEndpoint is closed')
|
||||
return self._s2v_snapshot.get()
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close all channels owned by this endpoint (idempotent)."""
|
||||
if not self._is_closed:
|
||||
self._is_closed = True
|
||||
self._v2s_events.close()
|
||||
self._s2v_events.close()
|
||||
self._s2v_snapshot.close()
|
||||
|
||||
|
||||
class SimEndpoint:
|
||||
"""Sim endpoint to route Messages to the Viewer using the right Channel.
|
||||
|
||||
The simulation has an outgoing SnapshotChannel and EventChannel.
|
||||
The simulation has an incoming EventChannel but no incoming SnapshotChannel.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
s2v_snapshot: vp.SnapshotChannel,
|
||||
s2v_events: vp.EventChannel,
|
||||
v2s_events: vp.EventChannel,
|
||||
):
|
||||
self._s2v_snapshot = s2v_snapshot
|
||||
self._s2v_events = s2v_events
|
||||
self._v2s_events = v2s_events
|
||||
self._is_closed = False
|
||||
|
||||
def send_to_viewer(self, message: vp.Message) -> None:
|
||||
if self._is_closed:
|
||||
raise RuntimeError('SimEndpoint is closed')
|
||||
if isinstance(message, vp.Event):
|
||||
self._s2v_events.put(message)
|
||||
elif isinstance(message, vp.Snapshot):
|
||||
self._s2v_snapshot.put(message)
|
||||
else:
|
||||
raise TypeError(
|
||||
f'expected Snapshot or Event, got {type(message).__name__}'
|
||||
)
|
||||
|
||||
def get_viewer_events(self) -> list[vp.Event]:
|
||||
"""Returns all pending events from the viewer."""
|
||||
if self._is_closed:
|
||||
raise RuntimeError('SimEndpoint is closed')
|
||||
return self._v2s_events.get()
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close all channels owned by this endpoint."""
|
||||
if not self._is_closed:
|
||||
self._is_closed = True
|
||||
self._s2v_snapshot.close()
|
||||
self._s2v_events.close()
|
||||
self._v2s_events.close()
|
||||
|
||||
|
||||
def make_endpoints(
|
||||
*,
|
||||
s2v_snapshot: vp.SnapshotChannel,
|
||||
s2v_events: vp.EventChannel,
|
||||
v2s_events: vp.EventChannel,
|
||||
) -> tuple[ViewerEndpoint, SimEndpoint]:
|
||||
"""Returns viewer and simulation endpoints."""
|
||||
viewer_endpoint = ViewerEndpoint(
|
||||
s2v_snapshot=s2v_snapshot,
|
||||
s2v_events=s2v_events,
|
||||
v2s_events=v2s_events,
|
||||
)
|
||||
sim_endpoint = SimEndpoint(
|
||||
s2v_snapshot=s2v_snapshot,
|
||||
s2v_events=s2v_events,
|
||||
v2s_events=v2s_events,
|
||||
)
|
||||
return viewer_endpoint, sim_endpoint
|
||||
@@ -0,0 +1,140 @@
|
||||
# Copyright 2026 DeepMind Technologies Limited
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# https://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Messages and Channels for Studio."""
|
||||
|
||||
|
||||
import dataclasses
|
||||
from typing import Protocol
|
||||
from typing import runtime_checkable
|
||||
import mujoco
|
||||
from mujoco.experimental.studio import sim
|
||||
import numpy as np
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Message and Channel types for passing data between the sim and the viewer.
|
||||
# Concrete messages are defined later in this file, there is one per Event;
|
||||
# add more if needed in your script.
|
||||
# Concrete Channel implementations live in the launch_* modules.
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
class Message:
|
||||
"""Base class for data sent between the sim and the viewer.
|
||||
|
||||
Derive from Snapshot or Event, depending on the desired delivery semantics.
|
||||
"""
|
||||
|
||||
|
||||
class Snapshot(Message):
|
||||
"""Base class for latest-wins, droppable messages.
|
||||
|
||||
Each new snapshot replaces the previous one, so the receiver always sees the
|
||||
most recent value. Intermediate values may be silently dropped, making
|
||||
snapshots suitable for frequently updated state where only the latest value
|
||||
matters.
|
||||
"""
|
||||
|
||||
|
||||
class Event(Message):
|
||||
"""Base class for reliable, ordered, never dropped messages.
|
||||
|
||||
Events are queued in order and every event is delivered exactly once. Use
|
||||
events for discrete actions that must not be lost.
|
||||
"""
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SnapshotChannel(Protocol):
|
||||
"""Channel to transport Snapshot messages with latest-wins semantics."""
|
||||
|
||||
def put(self, value: Snapshot) -> None:
|
||||
"""Overwrites the latest snapshot of the same type."""
|
||||
...
|
||||
|
||||
def get(self) -> list[Snapshot]:
|
||||
"""Returns the latest pending snapshots of each type."""
|
||||
...
|
||||
|
||||
def close(self) -> None:
|
||||
"""Releases resources held by the channel."""
|
||||
...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class EventChannel(Protocol):
|
||||
"""Channel to transport Event messages in an ordered stream, nothing dropped."""
|
||||
|
||||
def put(self, value: Event) -> None:
|
||||
"""Puts an event into the channel, appending to the stream."""
|
||||
...
|
||||
|
||||
def get(self) -> list[Event]:
|
||||
"""Returns all pending events, in order."""
|
||||
...
|
||||
|
||||
def close(self) -> None:
|
||||
"""Releases resources held by the channel."""
|
||||
...
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Concrete Event and Snapshot types.
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class StateSnapshot(Snapshot):
|
||||
"""A snapshot message that transports a MuJoCo state."""
|
||||
|
||||
state: np.ndarray
|
||||
state_sig: int
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class ResetEvent(Event):
|
||||
"""An event requesting to reset the simulation."""
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class ModelEvent(Event):
|
||||
"""An event that transports a MuJoCo model."""
|
||||
|
||||
model: mujoco.MjModel
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class OptionEvent(Event):
|
||||
"""An event that transports updated MuJoCo model options."""
|
||||
|
||||
opt: mujoco.MjOption
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class PerturbEvent(Event):
|
||||
"""Carries perturbation forces from the viewer to the simulation."""
|
||||
|
||||
state: np.ndarray
|
||||
state_sig: int
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class SetPauseStateEvent(Event):
|
||||
"""An event requesting to set the pause state."""
|
||||
|
||||
pause_state: sim.PauseState
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class ExitEvent(Event):
|
||||
"""An event requesting to exit."""
|
||||
Reference in New Issue
Block a user