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

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

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."""