From e0bb2b0665a5e8cb67f8778ab0b23562761c4834 Mon Sep 17 00:00:00 2001 From: Matija Kecman Date: Mon, 29 Jun 2026 00:23:17 -0700 Subject: [PATCH] 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 --- .../mujoco/experimental/studio/endpoints.py | 132 +++++++++++++++++ python/mujoco/experimental/studio/messages.py | 140 ++++++++++++++++++ 2 files changed, 272 insertions(+) create mode 100644 python/mujoco/experimental/studio/endpoints.py create mode 100644 python/mujoco/experimental/studio/messages.py diff --git a/python/mujoco/experimental/studio/endpoints.py b/python/mujoco/experimental/studio/endpoints.py new file mode 100644 index 00000000..edcbe2ce --- /dev/null +++ b/python/mujoco/experimental/studio/endpoints.py @@ -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 diff --git a/python/mujoco/experimental/studio/messages.py b/python/mujoco/experimental/studio/messages.py new file mode 100644 index 00000000..655fbb85 --- /dev/null +++ b/python/mujoco/experimental/studio/messages.py @@ -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."""