rlmesh 0.0.1__py3-none-any.whl

rlmesh/__init__.py

@@ -0,0 +1,4 @@
+from rlmesh.client import RLMeshClient
+
+__version__ = "0.0.1"
+__all__ = ["RLMeshClient"]

rlmesh/client/__init__.py

@@ -0,0 +1,3 @@
+from .session import RLMeshClient
+
+__all__ = ["RLMeshClient"]

rlmesh/client/print.py

@@ -0,0 +1,66 @@
+from math import ceil
+
+from rich import print
+from rich.table import Table
+from rich.panel import Panel
+from rich.console import Group, Console
+
+from rlmesh.proto import control_pb2
+
+
+def _into_columns(items, cols):
+    if cols < 1:
+        cols = 1
+    rows = ceil(len(items) / cols)
+    out = []
+    for c in range(cols):
+        start = c * rows
+        out.append(items[start : start + rows])
+    return out, rows
+
+
+def print_envs(res: control_pb2.ListEnvironmentsRes):
+    console = Console()
+    panels = []
+    for pkg, envs in res.packages.items():
+        env_list = list(envs.env_ids)
+        if not env_list:
+            continue
+
+        max_len = max(len(s) for s in env_list)
+        term_width = console.size.width or 100
+        min_col_width = max(18, min(40, max_len + 2))
+        cols = max(1, min(4, term_width // min_col_width))
+
+        cols_data, rows = _into_columns(env_list, cols)
+
+        grid = Table.grid(expand=True, pad_edge=False)
+        for _ in range(cols):
+            grid.add_column(ratio=1, justify="left", no_wrap=True)
+
+        for r in range(rows):
+            grid.add_row(
+                *[
+                    (cols_data[c][r] if r < len(cols_data[c]) else "")
+                    for c in range(cols)
+                ]
+            )
+
+        panels.append(
+            Panel(
+                grid,
+                title=f"[green]{pkg}[/green]",
+                border_style="bright_black",
+                padding=(0, 1),
+            )
+        )
+
+    group = Group(*panels)
+    print(
+        Panel(
+            group,
+            title="[bold yellow]Available Environments[/bold yellow]",
+            border_style="bold yellow",
+            expand=True,
+        )
+    )

rlmesh/client/session.py

@@ -0,0 +1,203 @@
+from typing import Iterator, Optional, Literal, Callable, Tuple
+
+import contextlib
+import grpc
+import asyncio
+import gymnasium as gym
+
+from gymnasium.wrappers import HumanRendering
+
+from rlmesh.client.print import print_envs
+from rlmesh.env import RemoteEnv
+from rlmesh.ping import compare_gcp_regions_async
+from rlmesh.proto import control_pb2_grpc, control_pb2
+from rlmesh.utils import config
+
+DEFAULT_CONTROLPLANE_ENDPOINT = "34.130.104.52:4510"
+
+
+class RLMeshClient:
+    def __init__(
+        self,
+        api_key: str | None = None,
+        *,
+        endpoint: str | None = DEFAULT_CONTROLPLANE_ENDPOINT,
+        region: str | None = None,
+        insecure: bool = True,
+    ):
+        self._addr = endpoint or config.api_url
+        self._token = api_key
+        self._region = region
+
+        self._chan = grpc.insecure_channel(self._addr) if insecure else None
+        if self._chan is None:
+            raise NotImplementedError("TLS not wired yet")
+
+        self._stub = control_pb2_grpc.SessionAPIStub(self._chan)
+        self._metadata = []
+
+        if self._token:
+            self._metadata.append(("x-rlmesh-token", self._token))
+
+        self._env = None
+
+    @contextlib.contextmanager
+    def channel(self):
+        try:
+            yield self._chan
+        finally:
+            if self._chan is not None:
+                self._chan.close()
+
+    def find_closest_region(self):
+        regions = list(self.list_regions().regions)
+        if regions:
+            return asyncio.run(compare_gcp_regions_async(regions))
+
+        return None
+
+    def list_regions(self, deadline_s: float = 5.0) -> control_pb2.ListRegionsRes:
+        req = control_pb2.ListRegionsReq()
+        return self._stub.ListRegions(req, timeout=deadline_s, metadata=self._metadata)
+
+    def list_environments(
+        self, deadline_s: float = 5.0
+    ) -> control_pb2.ListEnvironmentsRes:
+        req = control_pb2.ListEnvironmentsReq()
+        return self._stub.ListEnvironments(
+            req, timeout=deadline_s, metadata=self._metadata
+        )
+
+    def create_session(
+        self, gym_id: str, *, num_envs: int = 1, deadline_s: float = 5.0
+    ) -> control_pb2.Session:
+        if self._region is None:
+            self._region = self.find_closest_region()
+
+        req = control_pb2.InteractiveSessionSpec(
+            gym_id=gym_id, num_envs=num_envs, location=self._region
+        )
+        try:
+            res = self._stub.CreateSession(
+                req, timeout=deadline_s, metadata=self._metadata
+            )
+            return res
+        except grpc.RpcError as e:
+            if e.code() == grpc.StatusCode.NOT_FOUND:
+                try:
+                    details = e.details() if hasattr(e, "details") else str(e)
+                except Exception:
+                    details = str(e)
+                print(f"[rlmesh] CreateSession failed: {details}")
+                print("[rlmesh] Available environments:")
+                self.print_envs()
+            raise
+
+    def watch_session(
+        self, session_id: str, *, deadline_s: float = 120.0
+    ) -> Iterator[control_pb2.Session]:
+        req = control_pb2.WatchReq(session_id=session_id)
+        stream = self._stub.WatchSession(
+            req, timeout=deadline_s, metadata=self._metadata
+        )
+        for evt in stream:
+            yield evt
+
+    def wait_until_ready(
+        self, session_id: str, *, deadline_s: float = 600.0
+    ) -> control_pb2.Session:
+        ready_evt: Optional[control_pb2.Session] = None
+        for evt in self.watch_session(session_id, deadline_s=deadline_s):
+            if evt.state == control_pb2.SessionState.READY:
+                ready_evt = evt
+                break
+        if not ready_evt:
+            raise RuntimeError("session stream ended without READY")
+        return ready_evt
+
+    def describe(
+        self, session_id: str, *, deadline_s: float = 5.0
+    ) -> control_pb2.DescribeRes:
+        req = control_pb2.DescribeReq(session_id=session_id)
+        return self._stub.DescribeSession(
+            req, timeout=deadline_s, metadata=self._metadata
+        )
+
+    def get_results(
+        self, session_id: str, *, deadline_s: float = 5.0
+    ) -> control_pb2.GetResultsRes:
+        req = control_pb2.GetResultsReq(session_id=session_id)
+        return self._stub.GetResults(req, timeout=deadline_s, metadata=self._metadata)
+
+    def _connect_env(
+        self, *, session: control_pb2.Session, desc: control_pb2.DescribeRes
+    ) -> RemoteEnv:
+        node_addr = getattr(session, "address")
+        if not node_addr:
+            raise ValueError("session is missing node address")
+
+        metadata = self._metadata.copy()
+        metadata.append(("x-rlmesh-session", session.id))
+
+        return RemoteEnv(
+            desc.env,
+            transport_addr=node_addr,
+            transport_metadata=metadata,
+        )
+
+    def _create_and_connect(
+        self, gym_id: str, num_envs: int = 1
+    ) -> Tuple[RemoteEnv, control_pb2.Session]:
+        if self._env:
+            try:
+                self._env.close()
+            finally:
+                self._env = None
+
+        sess = self.create_session(gym_id, num_envs=num_envs)
+        sess = self.wait_until_ready(sess.id)
+        desc = self.describe(sess.id)
+
+        env = self._connect_env(session=sess, desc=desc)
+        env.render_mode = "rgb_array"
+        self._env = env
+        return env, sess
+
+    def make(
+        self, gym_id: str, render_mode: Literal["human", "rgb_array"] = "rgb_array"
+    ) -> gym.Env:
+        env, _ = self._create_and_connect(gym_id)
+
+        if render_mode == "human":
+            env = HumanRendering(env)
+
+        self._env = env
+        return env
+
+    def evaluate(self, gym_id: str, get_action: Callable | None = None, render=False):
+        env, sess = self._create_and_connect(gym_id)
+
+        if render:
+            env = HumanRendering(env)
+
+        obs, _ = env.reset(seed=42)
+        done = False
+
+        while not done:
+            action = (
+                env.action_space.sample() if get_action is None else get_action(obs)
+            )
+            obs, r, term, trunc, _ = env.step(action)
+            done = term or trunc
+
+        env.close()
+        results = self.get_results(sess.id)
+        return {
+            "score": results.score,
+            "duration": results.duration,
+            "timestamp": results.timestamp.ToJsonString(),
+        }
+
+    def print_envs(self):
+        res = self.list_environments()
+        return print_envs(res)

rlmesh/env/__init__.py

@@ -0,0 +1,4 @@
+from .make_env import make_env_factory, make_env
+from .remote_env import RemoteEnv
+
+__all__ = ["make_env_factory", "make_env", "RemoteEnv"]

rlmesh/env/make_env.py

@@ -0,0 +1,44 @@
+from typing import Literal, Any, Optional
+
+import gymnasium as gym
+
+
+def make_env_factory(
+    gym_id: str,
+    gym_type: Literal["gymnasium", "gym-v26", "gym-v21"] = "gymnasium",
+    gym_vars: dict[str, Any] = {},
+):
+    def env_factory(index=0, **kwargs):
+        env_vars = {
+            **(gym_vars or {}),
+            "render_mode": "rgb_array",
+            # "index": index,
+            **kwargs,
+        }
+
+        env = None
+        if gym_type == "gymnasium":
+            env = gym.make(gym_id, **env_vars)
+        elif gym_type == "gym-v26":
+            env = gym.make("GymV26Environment-v0", env_id=gym_id, **env_vars)
+        elif gym_type == "gym-v21":
+            env = gym.make("GymV21Environment-v0", env_id=gym_id, **env_vars)
+        else:
+            raise EnvironmentError(
+                f"Unsupported environment type: {gym_type}. "
+                "Please use a supported environment."
+            )
+
+        return env
+
+    return env_factory
+
+
+def make_env(
+    gym_id: str,
+    gym_type: Literal["gymnasium", "gym-v26", "gym-v21"] = "gymnasium",
+    *,
+    render_mode: Optional[str] = None,
+    **kwargs,
+):
+    return make_env_factory(gym_id, gym_type)(render_mode=render_mode, **kwargs)

rlmesh/env/remote_env.py

@@ -0,0 +1,192 @@
+from typing import Any, Optional
+
+import asyncio
+import grpc
+import gymnasium as gym
+import gymnasium.envs.registration as reg
+import numpy as np
+
+from google.protobuf.json_format import MessageToDict
+
+
+from rlmesh.gym import deserialize_space, tensor_from_ndarray, ndarray_from_tensor
+from rlmesh.proto import gym_pb2, env_pb2_grpc, env_pb2
+from rlmesh.utils import LoopRunner
+
+
+class AsyncEnvTransport:
+    def __init__(self, addr: str, metadata: list):
+        self._addr = addr
+        self._metadata = metadata
+
+        self._chan: grpc.aio.Channel | None = None
+        self._stub: env_pb2_grpc.EnvServiceStub | None = None
+        self._stream: grpc.aio.StreamStreamCall | None = None
+        self._ready = asyncio.Event()
+        self._lock = asyncio.Lock()
+
+    async def start(self):
+        self._chan = grpc.aio.insecure_channel(self._addr)
+        self._stub = env_pb2_grpc.EnvServiceStub(self._chan)
+        self._stream = self._stub.Connect(metadata=self._metadata)
+        self._ready.set()
+
+    async def _read_and_check_msg(
+        self, expected: Optional[str] = None
+    ) -> env_pb2.ServerMsg:
+        assert self._stream is not None
+        msg = await self._stream.read()
+
+        if msg is None:
+            raise RuntimeError("stream closed on reset")
+
+        assert isinstance(msg, env_pb2.ServerMsg)
+        if msg.WhichOneof("kind") == "error":
+            raise RuntimeError(
+                f"recieved an error from the server: {msg.error.message}"
+            )
+
+        if expected is not None and msg.WhichOneof("kind") != expected:
+            raise RuntimeError(f"expected {expected}, got {msg.WhichOneof('kind')}")
+
+        return msg
+
+    async def close(self):
+        try:
+            if self._stream is not None:
+                await self._stream.write(env_pb2.ClientMsg(close=env_pb2.CloseReq()))
+                await self._stream.done_writing()
+        finally:
+            if self._chan is not None:
+                await self._chan.close()
+
+    async def reset(
+        self, *, seed: int | None, options: dict | None = None
+    ) -> tuple[np.ndarray, dict]:
+        await self._ready.wait()
+        async with self._lock:
+            assert self._stream is not None
+            seeds = [] if seed is None else [int(seed)]
+            reset = env_pb2.ResetReq(seeds=seeds)
+
+            await self._stream.write(env_pb2.ClientMsg(reset=reset))
+            msg = await self._read_and_check_msg("reset_ok")
+
+            obs = ndarray_from_tensor(msg.reset_ok.obs)
+
+            info = {}
+            if msg.reset_ok.infos:
+                info = MessageToDict(msg.reset_ok.infos[0])
+
+            return obs, info
+
+    async def step(
+        self, action_msg: env_pb2.ClientMsg
+    ) -> tuple[np.ndarray, float, bool, bool, dict]:
+        await self._ready.wait()
+        async with self._lock:
+            assert self._stream is not None
+            await self._stream.write(action_msg)
+            msg = await self._read_and_check_msg()
+
+            k = msg.WhichOneof("kind")
+            if k == "step_ok":
+                s = msg.step_ok
+                obs = ndarray_from_tensor(s.obs)
+                rew = float(ndarray_from_tensor(s.rewards).reshape(-1)[0])
+                term = bool(s.terminated_mask and (s.terminated_mask[0] & 0x01))
+                trunc = bool(s.truncated_mask and (s.truncated_mask[0] & 0x01))
+                info = {}
+                if s.infos:
+                    from google.protobuf.json_format import MessageToDict
+
+                    info = MessageToDict(s.infos[0])
+                return obs, rew, term, trunc, info
+
+            if k == "reset_ok":
+                obs = ndarray_from_tensor(msg.reset_ok.obs)
+                return obs, 0.0, True, False, {}
+
+            if k == "error":
+                raise RuntimeError(f"server error: {msg.error.message}")
+
+            raise RuntimeError(f"unexpected server msg kind: {k}")
+
+    async def render(self) -> np.ndarray | None:
+        await self._ready.wait()
+        async with self._lock:
+            assert self._stream is not None
+            await self._stream.write(env_pb2.ClientMsg(render=env_pb2.RenderReq()))
+            msg = await self._read_and_check_msg("render_ok")
+            frame = ndarray_from_tensor(msg.render_ok.frame)
+            return frame
+
+
+class RemoteEnv(gym.Env):
+    metadata: dict[str, Any]
+
+    def __init__(
+        self,
+        env_spec: gym_pb2.EnvSpec,
+        *,
+        transport_addr: str,
+        transport_metadata: list,
+    ):
+        super().__init__()
+
+        self.observation_space = deserialize_space(env_spec.spaces.obs)
+        self.action_space = deserialize_space(env_spec.spaces.action)
+        self.metadata = {
+            "render_modes": list(env_spec.metadata.render_modes),
+            "render_fps": int(env_spec.metadata.render_fps or 30),
+        }
+
+        self.spec = reg.EnvSpec(id=env_spec.gym_id)
+
+        self._runner = LoopRunner()
+        self._tx = AsyncEnvTransport(transport_addr, transport_metadata)
+        self._runner.run(self._tx.start())
+        self._closed = False
+
+        self._obs = None
+        self._is_ale = env_spec.package == "ale_py"
+
+    def reset(self, *, seed: int | None = None, options: dict | None = None):
+        if self._closed:
+            raise RuntimeError("env closed")
+        obs, info = self._runner.run(self._tx.reset(seed=seed, options=options))
+        self._obs = obs
+
+        return obs, info
+
+    def step(self, action):
+        if self._closed:
+            raise RuntimeError("env closed")
+        if not self.action_space.contains(action):
+            raise ValueError("action outside action_space")
+
+        arr = np.asarray(action)
+        msg = env_pb2.ClientMsg(step=env_pb2.StepReq(actions=tensor_from_ndarray(arr)))
+        obs, r, term, trunc, info = self._runner.run(self._tx.step(msg))
+        self._obs = obs
+        return obs, r, term, trunc, info
+
+    def render(self):
+        if self._closed:
+            raise RuntimeError("env closed")
+
+        # Special Case for ALE Environment
+        # - since we know the obs is the rgb_array, we can skip sending it over twice
+        if self._is_ale and self._obs is not None:
+            return self._obs
+
+        frame = self._runner.run(self._tx.render())
+        return frame
+
+    def close(self):
+        if not self._closed:
+            try:
+                self._runner.run(self._tx.close())
+            finally:
+                self._runner.stop()
+                self._closed = True

rlmesh/env/remote_vector_env.py

@@ -0,0 +1,54 @@
+# import gymnasium as gym
+
+
+# class RemoteVectorEnv(gym.vector.VectorEnv):
+#     metadata: dict[str, Any]
+
+#     def __init__(
+#         self,
+#         env_spec: gym_pb2.EnvSpec,
+#         *,
+#         gateway_addr: str,
+#         session_id: str,
+#         bearer: str | None,
+#     ):
+#         super().__init__()
+#         # build spaces from spec
+#         self.observation_space = deserialize_space(env_spec.spaces.obs)
+#         self.action_space = deserialize_space(env_spec.spaces.action)
+#         self.metadata = {
+#             "render_modes": list(env_spec.metadata.render_modes),
+#             "render_fps": int(env_spec.metadata.render_fps or 30),
+#         }
+
+#         self._runner = _LoopRunner()
+#         self._tx = _AsyncEnvTransport(
+#             gateway_addr, session_id=session_id, bearer=bearer
+#         )
+#         self._runner.run(self._tx.start())
+#         self._closed = False
+
+#     def reset(self, *, seed: int | None = None, options: dict | None = None):
+#         if self._closed:
+#             raise RuntimeError("env closed")
+#         obs, info = self._runner.run(self._tx.reset(seed=seed, options=options))
+#         return obs, info
+
+#     def step(self, action):
+#         if self._closed:
+#             raise RuntimeError("env closed")
+#         if not self.action_space.contains(action):
+#             raise ValueError("action outside action_space")
+
+#         arr = np.asarray(action)
+#         msg = env_pb2.ClientMsg(step=env_pb2.StepReq(actions=tensor_from_ndarray(arr)))
+#         obs, r, term, trunc, info = self._runner.run(self._tx.step(msg))
+#         return obs, r, term, trunc, info
+
+#     def close(self):
+#         if not self._closed:
+#             try:
+#                 self._runner.run(self._tx.close())
+#             finally:
+#                 self._runner.stop()
+#                 self._closed = True

rlmesh/gym/__init__.py

@@ -0,0 +1,29 @@
+from .env_codec import (
+    serialize_env,
+    deserialize_env,
+    env_spec_to_json,
+    env_spec_from_json,
+)
+from .spaces_codec import (
+    serialize_space,
+    deserialize_space,
+    space_spec_to_json,
+    space_spec_from_json,
+)
+from .tensor_codec import (
+    tensor_from_ndarray,
+    ndarray_from_tensor,
+)
+
+__all__ = [
+    "serialize_env",
+    "deserialize_env",
+    "env_spec_to_json",
+    "env_spec_from_json",
+    "serialize_space",
+    "deserialize_space",
+    "space_spec_to_json",
+    "space_spec_from_json",
+    "tensor_from_ndarray",
+    "ndarray_from_tensor",
+]