diff --git a/docs/source-en/rst_source/development/interfaces.rst b/docs/source-en/rst_source/development/interfaces.rst index 2295585b9..7e7ebf836 100644 --- a/docs/source-en/rst_source/development/interfaces.rst +++ b/docs/source-en/rst_source/development/interfaces.rst @@ -127,7 +127,9 @@ normal-CLI or Dashboard runtime hook: Common flags: ``--env-endpoint``, ``--vla-endpoint``. The default ``http`` sends JSON over ``POST /call``, encoding NumPy arrays as -``{"__ndarray__": , "dtype": ..., "shape": ...}``; switch to ``socket`` +``{"__ndarray__": , "dtype": ..., "shape": ...}`` and NumPy scalars as +``{"__npscalar__": , "dtype": ...}`` so dtypes survive the round trip; +switch to ``socket`` for large or history-stacked nested-NumPy observations to move length-prefixed pickle frames and skip repeated JSON encoding. Pickle is unsafe on untrusted input, so only point ``socket`` at trusted endpoints. diff --git a/docs/source-en/rst_source/usage/advanced_deployment.rst b/docs/source-en/rst_source/usage/advanced_deployment.rst index 0b0c363ee..12406d678 100644 --- a/docs/source-en/rst_source/usage/advanced_deployment.rst +++ b/docs/source-en/rst_source/usage/advanced_deployment.rst @@ -84,3 +84,20 @@ that service, and each ``*_PORT`` with the free port you chose at startup. Any of the three endpoint flags can be omitted; when one is unset, RPent spawns that service locally on a free port. All three default to HTTP when the protocol is omitted, and all three accept ``socket://HOST:PORT``. + +Custom RLinf checkout +--------------------- + +The environment and VLA services import ``rlinf``. When starting them +manually against a development RLinf checkout instead of the installed +package, point ``PYTHONPATH`` at the checkout first: + +.. code-block:: bash + + export PYTHONPATH=/path/to/rlinf:$PYTHONPATH + +Servers spawned by RPent get this automatically: the checkout is resolved +from ``RPENT_RLINF_ROOT`` (or ``RLINF_REPO_PATH``), falling back to the +``rlinf`` directory next to the RPent checkout. A resolved path that does +not exist is harmless — Python ignores invalid ``PYTHONPATH`` entries — +so the servers import the installed ``rlinf`` package. diff --git a/docs/source-zh/rst_source/development/interfaces.rst b/docs/source-zh/rst_source/development/interfaces.rst index 9a116b3a7..855066d7c 100644 --- a/docs/source-zh/rst_source/development/interfaces.rst +++ b/docs/source-zh/rst_source/development/interfaces.rst @@ -120,6 +120,7 @@ runtime 钩子中解析）： 常见：``--env-endpoint``、``--vla-endpoint``。``http`` 为默认，走 ``POST /call`` 传 JSON，其中 NumPy 数组编码为 ``{"__ndarray__": , "dtype": ..., "shape": ...}``； +NumPy 标量编码为 ``{"__npscalar__": , "dtype": ...}`` 以保留精确 dtype； 观测数据很大、或是多帧堆叠的嵌套 NumPy 字典时可改 ``socket``，用带长度前缀的 pickle 数据帧传输，省掉反复的 JSON 编解码。pickle 不适合不可信输入，socket 只应连接可信端点。 diff --git a/docs/source-zh/rst_source/usage/advanced_deployment.rst b/docs/source-zh/rst_source/usage/advanced_deployment.rst index 27937f448..96fa73006 100644 --- a/docs/source-zh/rst_source/usage/advanced_deployment.rst +++ b/docs/source-zh/rst_source/usage/advanced_deployment.rst @@ -81,3 +81,18 @@ SAM3 服务只加载一次模型，可以由多个 RPent 运行复用。 某项未指定时，RPent 会在当前机器上启动对应服务并自动选择空闲端口。三个服务省略 protocol 时都默认使用 HTTP，也都可以通过 ``socket://HOST:PORT`` 改用 socket RPC。 + +自定义 RLinf 源码 +----------------- + +环境服务和 VLA 服务需要导入 ``rlinf``。手动启动时如果使用开发中的 RLinf +源码而非已安装的包，请先把 ``PYTHONPATH`` 指向该源码目录： + +.. code-block:: bash + + export PYTHONPATH=/path/to/rlinf:$PYTHONPATH + +由 RPent 自动拉起的 server 无需手动设置：RLinf 源码路径从 +``RPENT_RLINF_ROOT`` （或 ``RLINF_REPO_PATH`` ）解析，默认回退到 RPent +仓库旁边的 ``rlinf`` 目录。解析出的路径不存在也无妨：Python 会忽略无效的 +``PYTHONPATH`` 条目，server 将导入已安装的 ``rlinf`` 包。 diff --git a/robots/dual_franka/env_server.py b/robots/dual_franka/env_server.py index aba0fa014..02fbfece0 100644 --- a/robots/dual_franka/env_server.py +++ b/robots/dual_franka/env_server.py @@ -17,22 +17,15 @@ from __future__ import annotations import queue -import sys import time from typing import Any import numpy as np from robots.dual_franka.runtime_config import load_runtime_config -from robots.franka.env_server import _to_numpy_tree, main -from rpent.utils.config import get_repo_root, get_rlinf_repo_path +from robots.franka.env_server import main from rpent.utils.logging import get_logger - -# Resolve the RLinf checkout before the deferred ``import rlinf`` executes. -RPENT_ROOT = get_repo_root() -RLINF_REPO_PATH = get_rlinf_repo_path() or (RPENT_ROOT.parent / "rlinf").resolve() -if str(RLINF_REPO_PATH) not in sys.path: - sys.path.insert(0, str(RLINF_REPO_PATH)) +from rpent.utils.serialization import to_numpy_tree logger = get_logger("dual_franka_env_server") @@ -145,7 +138,7 @@ def reset(self) -> dict[str, Any]: observation, info = self.env.reset() return { "ok": True, - "info": _to_numpy_tree(info), + "info": to_numpy_tree(info), "robot_state": self.get_robot_state(), "states": self._strip_batch(observation.get("states")), } @@ -154,7 +147,7 @@ def reset(self) -> dict[str, Any]: @staticmethod def _strip_batch(value: Any) -> Any: - array = _to_numpy_tree(value) + array = to_numpy_tree(value) if isinstance(array, np.ndarray) and array.ndim > 0 and array.shape[0] == 1: return array[0] if isinstance(array, list) and len(array) == 1: @@ -195,7 +188,7 @@ def _attach_camera_snapshots(self, output: dict[str, Any]) -> dict[str, Any]: snapshot_getter = self.env.env.call( "get_wrapper_attr", "get_raw_camera_snapshot" )[0] - snapshot = _to_numpy_tree(snapshot_getter()) + snapshot = to_numpy_tree(snapshot_getter()) output["raw_camera_frames"] = snapshot.get("raw_frames", {}) output["raw_camera_depths"] = snapshot.get("raw_depths", {}) perception = self._capture_perception_camera_snapshot() @@ -268,8 +261,8 @@ def _wrapped_states(self) -> np.ndarray: def get_robot_state(self) -> dict[str, Any]: left, right = self._arm_states() return { - "left_arm": _to_numpy_tree(left), - "right_arm": _to_numpy_tree(right), + "left_arm": to_numpy_tree(left), + "right_arm": to_numpy_tree(right), "action_dim": self.action_dim, "per_arm_dim": self.per_arm_dim, "action_scale": self.action_scale.tolist(), @@ -302,7 +295,7 @@ def get_camera_meta(self) -> dict[str, Any] | None: metadata_getter = self.env.env.call( "get_wrapper_attr", "get_raw_camera_metadata" )[0] - metadata.update(_to_numpy_tree(metadata_getter())) + metadata.update(to_numpy_tree(metadata_getter())) metadata.update(self._perception_camera_meta) return { "cameras": cameras, @@ -530,9 +523,9 @@ def chunk_step( for action in np.asarray(actions, dtype=np.float32): observation, _reward, term, trunc, info = self.env.step(action[None, :]) observations.append(self._decorate_step_observation(observation)) - terminated = terminated or bool(np.asarray(_to_numpy_tree(term)).any()) - truncated = truncated or bool(np.asarray(_to_numpy_tree(trunc)).any()) - last_info = _to_numpy_tree(info) + terminated = terminated or bool(np.asarray(to_numpy_tree(term)).any()) + truncated = truncated or bool(np.asarray(to_numpy_tree(trunc)).any()) + last_info = to_numpy_tree(info) if terminated or truncated: break return { diff --git a/robots/dual_franka/vla_server.py b/robots/dual_franka/vla_server.py index 7f5ad1496..9c9080aa7 100644 --- a/robots/dual_franka/vla_server.py +++ b/robots/dual_franka/vla_server.py @@ -18,7 +18,6 @@ import argparse import os -import sys import time from typing import Any @@ -27,17 +26,10 @@ from omegaconf import OmegaConf from rpent.robots.components.vla_facade_base import BaseVLAFacade -from rpent.utils.config import get_repo_root, get_rlinf_repo_path from rpent.utils.logging import get_logger logger = get_logger("dual_franka_vla_server") -# Resolve the RLinf checkout before the deferred ``import rlinf`` executes. -RPENT_ROOT = get_repo_root() -RLINF_REPO_PATH = get_rlinf_repo_path() or (RPENT_ROOT.parent / "rlinf").resolve() -if str(RLINF_REPO_PATH) not in sys.path: - sys.path.insert(0, str(RLINF_REPO_PATH)) - def build_model_cfg(model_path: str, repo_id: str) -> Any: """Build the RLinf OpenPI inference config for dual Franka.""" diff --git a/robots/franka/env_server.py b/robots/franka/env_server.py index 2036a6e71..91d3b62da 100644 --- a/robots/franka/env_server.py +++ b/robots/franka/env_server.py @@ -17,8 +17,6 @@ from __future__ import annotations import argparse -import dataclasses -import sys import time from collections.abc import Callable from typing import Any @@ -27,34 +25,11 @@ from robots.franka.runtime_config import load_runtime_config from rpent.robots.components.env_facade_base import BaseEnvFacade -from rpent.utils.config import get_repo_root, get_rlinf_repo_path from rpent.utils.logging import get_logger +from rpent.utils.serialization import to_numpy_tree logger = get_logger("franka_env_server") -# Resolve the RLinf checkout before the deferred ``import rlinf`` executes. -RPENT_ROOT = get_repo_root() -RLINF_REPO_PATH = get_rlinf_repo_path() or (RPENT_ROOT.parent / "rlinf").resolve() -if str(RLINF_REPO_PATH) not in sys.path: - sys.path.insert(0, str(RLINF_REPO_PATH)) - - -def _to_numpy_tree(value: Any) -> Any: - """Convert tensors and nested values into pickle-safe CPU/numpy data.""" - if hasattr(value, "detach") and hasattr(value, "cpu"): - return value.detach().cpu().numpy() - if dataclasses.is_dataclass(value): - return _to_numpy_tree(dataclasses.asdict(value)) - if isinstance(value, dict): - return {key: _to_numpy_tree(item) for key, item in value.items()} - if isinstance(value, list): - return [_to_numpy_tree(item) for item in value] - if isinstance(value, tuple): - return tuple(_to_numpy_tree(item) for item in value) - if isinstance(value, np.generic): - return value.item() - return value - class FrankaEnvFacade(BaseEnvFacade): """Expose the single-Franka ``env.*`` protocol from a Ray-backed worker.""" @@ -151,14 +126,14 @@ def reset(self) -> dict[str, Any]: observation, info = self.env.reset() return { "ok": True, - "info": _to_numpy_tree(info), + "info": to_numpy_tree(info), "robot_state": self.get_robot_state(), "states": self._strip_batch(observation.get("states")), } @staticmethod def _strip_batch(value: Any) -> Any: - array = _to_numpy_tree(value) + array = to_numpy_tree(value) if isinstance(array, np.ndarray) and array.ndim > 0 and array.shape[0] == 1: return array[0] if isinstance(array, list) and len(array) == 1: @@ -223,7 +198,7 @@ def _raw_tcp_pose(self) -> np.ndarray: def get_robot_state(self) -> dict[str, Any]: return { - "raw_base_state": _to_numpy_tree(self._raw_state()), + "raw_base_state": to_numpy_tree(self._raw_state()), "action_dim": self.action_dim, "action_scale": self.action_scale.tolist(), "use_relative_frame": self.use_relative_frame, @@ -234,7 +209,7 @@ def get_camera_meta(self) -> dict[str, Any] | None: metadata = self.env.env.call("get_camera_metadata")[0] except Exception as exc: return {"error": str(exc), "error_type": type(exc).__name__} - metadata = _to_numpy_tree(metadata) + metadata = to_numpy_tree(metadata) main_key = self.cfg.env.eval.get("main_image_key") names = sorted(metadata.get("cameras", {})) extras = [name for name in names if name != main_key] @@ -371,7 +346,7 @@ def set_gripper(self, *, open: bool) -> dict[str, Any]: gripper=command, ) time.sleep(self.controller["gripper_settle_s"]) - state = _to_numpy_tree(self._raw_state()) + state = to_numpy_tree(self._raw_state()) reached = bool(state.get("gripper_open")) == bool(open) iterations += 1 if reached: @@ -397,9 +372,9 @@ def chunk_step( for action in np.asarray(actions, dtype=np.float32): observation, _reward, term, trunc, info = self.env.step(action[None, :]) observations.append(self._strip_observation(observation)) - terminated = terminated or bool(np.asarray(_to_numpy_tree(term)).any()) - truncated = truncated or bool(np.asarray(_to_numpy_tree(trunc)).any()) - last_info = _to_numpy_tree(info) + terminated = terminated or bool(np.asarray(to_numpy_tree(term)).any()) + truncated = truncated or bool(np.asarray(to_numpy_tree(trunc)).any()) + last_info = to_numpy_tree(info) if terminated or truncated: break return { diff --git a/robots/libero/env_server.py b/robots/libero/env_server.py index 6553db8a8..86e57d9d5 100644 --- a/robots/libero/env_server.py +++ b/robots/libero/env_server.py @@ -24,11 +24,8 @@ import numpy as np from rpent.robots.components.env_facade_base import BaseEnvFacade -from rpent.utils.config import ( - get_repo_root, - get_rlinf_repo_path, -) from rpent.utils.logging import get_logger +from rpent.utils.serialization import to_numpy_tree # MuJoCo env vars must be set BEFORE importing anything that touches MuJoCo. os.environ.setdefault("MUJOCO_GL", "egl") @@ -39,16 +36,12 @@ logger = get_logger("env_server") -RPENT_ROOT = get_repo_root() -RLINF_REPO_PATH = get_rlinf_repo_path() or (RPENT_ROOT.parent / "rlinf").resolve() -if str(RLINF_REPO_PATH) not in sys.path: - sys.path.insert(0, str(RLINF_REPO_PATH)) os.environ.setdefault("ROBOT_PLATFORM", "LIBERO") # torch and LiberoEnv are only imported at call time (after --cuda-device # sets CUDA_VISIBLE_DEVICES in main()); LiberoEnv transitively imports torch. if TYPE_CHECKING: - import torch # noqa: F401 (referenced at runtime in _to_numpy_tree) + import torch # noqa: F401 (transitive dep of LiberoEnv; type-check only) from rlinf.envs.libero.libero_env import LiberoEnv @@ -137,20 +130,6 @@ def make_env( # --------------------------------------------------------------------------- -def _to_numpy_tree(x): - """Recursively convert torch tensors to CPU numpy arrays so the result - pickles cleanly across the agent/env_server wire.""" - if hasattr(x, "detach") and hasattr(x, "cpu") and hasattr(x, "numpy"): - return x.detach().cpu().numpy() - if isinstance(x, dict): - return {k: _to_numpy_tree(v) for k, v in x.items()} - if isinstance(x, list): - return [_to_numpy_tree(v) for v in x] - if isinstance(x, tuple): - return tuple(_to_numpy_tree(v) for v in x) - return x - - class LiberoEnvFacade(BaseEnvFacade): """Implements :class:`robots.libero.env_client.LiberoEnvClient` over :class:`rlinf.envs.libero.libero_env.LiberoEnv`. @@ -210,20 +189,20 @@ def _expand_chunk(self, actions) -> np.ndarray: def reset(self): obs, info = self._env.reset() - obs = self._strip_obs(_to_numpy_tree(obs)) - return obs, _to_numpy_tree(info) + obs = self._strip_obs(to_numpy_tree(obs)) + return obs, to_numpy_tree(info) def step(self, action): obs, rew, term, trunc, info = self._env.step(self._expand_action(action)) - obs = self._strip_obs(_to_numpy_tree(obs)) - term = self._strip(_to_numpy_tree(term)) - trunc = self._strip(_to_numpy_tree(trunc)) + obs = self._strip_obs(to_numpy_tree(obs)) + term = self._strip(to_numpy_tree(term)) + trunc = self._strip(to_numpy_tree(trunc)) return ( obs, - self._strip(_to_numpy_tree(rew)), + self._strip(to_numpy_tree(rew)), term, trunc, - _to_numpy_tree(info), + to_numpy_tree(info), ) def chunk_step(self, actions, *, return_all_frames: bool = False): @@ -241,20 +220,20 @@ def chunk_step(self, actions, *, return_all_frames: bool = False): obs_list, rew, term, trunc, info = self._env.chunk_step( self._expand_chunk(actions) ) - obs_list = [self._strip_obs(_to_numpy_tree(o)) for o in obs_list] - term = self._strip(_to_numpy_tree(term)) - trunc = self._strip(_to_numpy_tree(trunc)) + obs_list = [self._strip_obs(to_numpy_tree(o)) for o in obs_list] + term = self._strip(to_numpy_tree(term)) + trunc = self._strip(to_numpy_tree(trunc)) obs_field = obs_list if return_all_frames else obs_list[-1] return ( obs_field, - self._strip(_to_numpy_tree(rew)), + self._strip(to_numpy_tree(rew)), term, trunc, - _to_numpy_tree(info), + to_numpy_tree(info), ) def raw_obs(self) -> dict: - return _to_numpy_tree(self._env.current_raw_obs[self._env_idx]) + return to_numpy_tree(self._env.current_raw_obs[self._env_idx]) def get_env_meta(self) -> dict: """Return the meta info this server was launched with.""" @@ -274,7 +253,7 @@ def render_camera( width: int = 1024, depth: bool = False, ): - return _to_numpy_tree( + return to_numpy_tree( self._env.render_camera( camera_name=camera_name, height=height, @@ -289,7 +268,7 @@ def get_camera_meta( height: int = 256, width: int = 256, ) -> dict | None: - return _to_numpy_tree( + return to_numpy_tree( self._env.get_camera_meta( camera_name=camera_name, height=height, width=width ) diff --git a/robots/robotwin/env_server.py b/robots/robotwin/env_server.py index a8633be50..fae1876a7 100644 --- a/robots/robotwin/env_server.py +++ b/robots/robotwin/env_server.py @@ -31,6 +31,7 @@ from rpent.robots.components.env_facade_base import BaseEnvFacade from rpent.utils.logging import get_logger +from rpent.utils.serialization import to_numpy_tree logger = get_logger("robotwin_env_server") @@ -62,21 +63,6 @@ def _teardown_env(env: Any) -> None: from robots.robotwin.robot_spec import RoboTwinActionType # noqa: E402 -def _to_numpy_tree(value: Any) -> Any: - """Convert RPC results to numpy arrays and plain Python values.""" - if hasattr(value, "detach") and hasattr(value, "cpu") and hasattr(value, "numpy"): - return value.detach().cpu().numpy() - if isinstance(value, dict): - return {key: _to_numpy_tree(item) for key, item in value.items()} - if isinstance(value, list): - return [_to_numpy_tree(item) for item in value] - if isinstance(value, tuple): - return tuple(_to_numpy_tree(item) for item in value) - if isinstance(value, np.generic): - return value.item() - return value - - class RoboTwinEnvFacade(BaseEnvFacade): """Expose common and typed RoboTwin environment RPC contracts.""" @@ -229,7 +215,7 @@ def plan_arm_path(self, arm: str, target_pose) -> dict[str, Any]: return self._env.plan_arm_path(0, arm, target_pose) def _dispatch(self, method: str, args: tuple, kwargs: dict) -> Any: - return _to_numpy_tree(super()._dispatch(method, args, kwargs)) + return to_numpy_tree(super()._dispatch(method, args, kwargs)) def build_env_cfg( diff --git a/rpent/robots/components/pi05_vla_server.py b/rpent/robots/components/pi05_vla_server.py index d0df87e81..8750c36eb 100644 --- a/rpent/robots/components/pi05_vla_server.py +++ b/rpent/robots/components/pi05_vla_server.py @@ -23,7 +23,6 @@ import argparse import os -import sys import time from typing import Any @@ -32,20 +31,11 @@ from omegaconf import OmegaConf from rpent.robots.components.vla_facade_base import BaseVLAFacade -from rpent.utils.config import ( - get_pi05_checkpoint_path, - get_repo_root, - get_rlinf_repo_path, -) +from rpent.utils.config import get_pi05_checkpoint_path from rpent.utils.logging import get_logger logger = get_logger("vla_server") -RPENT_ROOT = get_repo_root() -RLINF_REPO_PATH = get_rlinf_repo_path() or (RPENT_ROOT.parent / "rlinf").resolve() -if str(RLINF_REPO_PATH) not in sys.path: - sys.path.insert(0, str(RLINF_REPO_PATH)) - # --------------------------------------------------------------------------- # Embodiment registry # --------------------------------------------------------------------------- diff --git a/rpent/utils/__init__.py b/rpent/utils/__init__.py index f11cea45a..30629712c 100644 --- a/rpent/utils/__init__.py +++ b/rpent/utils/__init__.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Utility helpers: config, logging, path resolution, templates.""" +"""Utility helpers: config, logging, path resolution, serialization, templates.""" from rpent.utils.logging import get_logger, get_output_dir, init_output_dir from rpent.utils.rpc import RpcClient, RpcError, make_rpc_client, parse_endpoint @@ -20,6 +20,7 @@ SocketRpcClient, SocketRpcServer, ) +from rpent.utils.serialization import to_numpy_tree from rpent.utils.templates import ( default_variables, substitute, @@ -39,4 +40,5 @@ "init_output_dir", "substitute", "substitute_text", + "to_numpy_tree", ] diff --git a/rpent/utils/config.py b/rpent/utils/config.py index d1e908480..e8e73bbf7 100644 --- a/rpent/utils/config.py +++ b/rpent/utils/config.py @@ -55,9 +55,9 @@ def get_libero_type() -> str: return os.environ.get("LIBERO_TYPE", "pro") -def get_rlinf_repo_path() -> Path | None: - """Return the configured RLinf checkout path, or *None*.""" +def get_rlinf_repo_path() -> Path: + """Return the RLinf path that RPC servers import.""" env = os.environ.get("RPENT_RLINF_ROOT") or os.environ.get("RLINF_REPO_PATH") if env: return Path(env).expanduser().resolve() - return None + return (get_repo_root().parent / "rlinf").resolve() diff --git a/rpent/utils/daemon.py b/rpent/utils/daemon.py index da7e53a70..c82fe8c49 100644 --- a/rpent/utils/daemon.py +++ b/rpent/utils/daemon.py @@ -23,7 +23,7 @@ import threading from typing import Callable -from rpent.utils.config import get_repo_root +from rpent.utils.config import get_repo_root, get_rlinf_repo_path from rpent.utils.logging import get_logger logger = get_logger("daemon") @@ -85,6 +85,12 @@ def __init__( self.cmd = cmd self.subprocess_env = os.environ.copy() self.subprocess_env.update(env_overrides or {}) + # Put the RLinf checkout first on the child's PYTHONPATH + rlinf_path = str(get_rlinf_repo_path()) + pythonpath = self.subprocess_env.get("PYTHONPATH", "") + self.subprocess_env["PYTHONPATH"] = ( + f"{rlinf_path}{os.pathsep}{pythonpath}" if pythonpath else rlinf_path + ) self.log_path = log_path self.cwd = cwd self._proc: subprocess.Popen | None = None diff --git a/rpent/utils/rpc/http_rpc.py b/rpent/utils/rpc/http_rpc.py index 07bf2ad8d..f71f69960 100644 --- a/rpent/utils/rpc/http_rpc.py +++ b/rpent/utils/rpc/http_rpc.py @@ -51,7 +51,8 @@ def _is_direct_url(url: str) -> bool: def _from_json(obj: Any) -> Any: """Rehydrate ``{"__ndarray__": , "dtype": ..., "shape": [...]}`` - back into ndarrays. Everything else is passed through unchanged. + back into ndarrays and ``{"__npscalar__": , "dtype": ...}`` + back into numpy scalars. Everything else is passed through unchanged. """ if isinstance(obj, dict): if "__ndarray__" in obj and set(obj) <= {"__ndarray__", "dtype", "shape"}: @@ -61,6 +62,8 @@ def _from_json(obj: Any) -> Any: # so callers can mutate the returned array like they would with # a pickle round-tripped one. return arr.reshape(obj.get("shape", (-1,))).copy() + if "__npscalar__" in obj and set(obj) <= {"__npscalar__", "dtype"}: + return np.dtype(obj["dtype"]).type(obj["__npscalar__"]) return {k: _from_json(v) for k, v in obj.items()} if isinstance(obj, list): return [_from_json(v) for v in obj] @@ -143,7 +146,7 @@ def call( class _NumpyEncoder(json.JSONEncoder): - """JSON encoder that tags numpy arrays and normalizes numpy scalars.""" + """JSON encoder that tags numpy arrays and scalars for faithful decode.""" def default(self, obj: Any) -> Any: if isinstance(obj, np.ndarray): @@ -152,12 +155,11 @@ def default(self, obj: Any) -> Any: "dtype": str(obj.dtype), "shape": list(obj.shape), } - if isinstance(obj, (np.integer,)): - return int(obj) - if isinstance(obj, (np.floating,)): - return float(obj) - if isinstance(obj, (np.bool_,)): - return bool(obj) + if isinstance(obj, np.generic): + item = obj.item() + # boolean, int, unsigned int, float + if obj.dtype.kind in "biuf" and isinstance(item, (bool, int, float)): + return {"__npscalar__": item, "dtype": str(obj.dtype)} return super().default(obj) diff --git a/rpent/utils/serialization.py b/rpent/utils/serialization.py new file mode 100644 index 000000000..bc70f110f --- /dev/null +++ b/rpent/utils/serialization.py @@ -0,0 +1,35 @@ +# Copyright 2026 The RPent Authors. +# +# 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. + +"""Helpers for converting values into pickle-safe wire formats.""" + +from __future__ import annotations + +import dataclasses +from typing import Any + + +def to_numpy_tree(value: Any) -> Any: + """Recursively convert tensors and nested values into pickle-safe data.""" + if hasattr(value, "detach") and hasattr(value, "cpu") and hasattr(value, "numpy"): + return value.detach().cpu().numpy() + if dataclasses.is_dataclass(value): + return to_numpy_tree(dataclasses.asdict(value)) + if isinstance(value, dict): + return {key: to_numpy_tree(item) for key, item in value.items()} + if isinstance(value, list): + return [to_numpy_tree(item) for item in value] + if isinstance(value, tuple): + return tuple(to_numpy_tree(item) for item in value) + return value diff --git a/tests/unit_tests/rpent/utils/rpc/test_transport_loopback.py b/tests/unit_tests/rpent/utils/rpc/test_transport_loopback.py index 6cf99dfa4..44efa7278 100644 --- a/tests/unit_tests/rpent/utils/rpc/test_transport_loopback.py +++ b/tests/unit_tests/rpent/utils/rpc/test_transport_loopback.py @@ -166,17 +166,24 @@ def test_transport_round_trips_nested_numpy_payloads(transport: Transport) -> No "metadata": { "count": np.int64(6), "valid": np.bool_(True), + "score": np.float32(1.5), "labels": ["left", "right"], }, }, ) np.testing.assert_array_equal(result["values"], original * 2.5) - assert result["metadata"] == { - "count": 6, - "valid": True, - "labels": ["left", "right"], - } + metadata = result["metadata"] + # Numpy scalars keep their exact dtype on both transports: the + # socket transport natively via pickle, the HTTP transport via + # the ``__npscalar__`` tag. + assert isinstance(metadata["count"], np.int64) + assert metadata["count"] == 6 + assert isinstance(metadata["valid"], np.bool_) + assert metadata["valid"] == np.bool_(True) + assert isinstance(metadata["score"], np.float32) + assert metadata["score"] == np.float32(1.5) + assert metadata["labels"] == ["left", "right"] result["values"][0, 0] = -1 assert original[0, 0] == 0 diff --git a/tests/unit_tests/rpent/utils/test_daemon_lifecycle.py b/tests/unit_tests/rpent/utils/test_daemon_lifecycle.py index 2d95f8080..b7e8d93de 100644 --- a/tests/unit_tests/rpent/utils/test_daemon_lifecycle.py +++ b/tests/unit_tests/rpent/utils/test_daemon_lifecycle.py @@ -14,6 +14,7 @@ from __future__ import annotations +import os import sys import time from collections.abc import Callable @@ -122,3 +123,93 @@ def test_daemon_force_kills_after_terminate_timeout(tmp_path: Path) -> None: _wait_until(lambda: daemon.poll() is not None) finally: daemon.stop(timeout=1.0) + + +def test_daemon_prepends_checkout_already_on_pythonpath( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + checkout = tmp_path / "rlinf" + checkout.mkdir() + resolved = str(checkout.resolve()) + monkeypatch.setenv("RPENT_RLINF_ROOT", str(checkout)) + # The checkout is already inherited on the path: it is prepended again + # anyway — duplicate entries are harmless (Python ignores them). + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(("/other", resolved))) + daemon = ProcessDaemon("env-probe", [sys.executable, "-c", "pass"]) + assert daemon.subprocess_env["PYTHONPATH"] == os.pathsep.join( + (resolved, "/other", resolved) + ) + + +def test_daemon_sets_pythonpath_when_parent_has_none( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + checkout = tmp_path / "rlinf" + checkout.mkdir() + monkeypatch.setenv("RPENT_RLINF_ROOT", str(checkout)) + monkeypatch.delenv("PYTHONPATH", raising=False) + daemon = ProcessDaemon("env-probe", [sys.executable, "-c", "pass"]) + assert daemon.subprocess_env["PYTHONPATH"] == str(checkout.resolve()) + + +def test_daemon_prepends_rlinf_checkout_to_pythonpath( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + checkout = tmp_path / "rlinf" + checkout.mkdir() + monkeypatch.setenv("RPENT_RLINF_ROOT", str(checkout)) + monkeypatch.setenv("PYTHONPATH", "/other") + daemon = ProcessDaemon("env-probe", [sys.executable, "-c", "pass"]) + assert daemon.subprocess_env["PYTHONPATH"] == os.pathsep.join( + (str(checkout.resolve()), "/other") + ) + + +def test_daemon_prepends_missing_checkout_harmlessly( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + missing = tmp_path / "missing" + monkeypatch.setenv("RPENT_RLINF_ROOT", str(missing)) + monkeypatch.setenv("PYTHONPATH", "/other") + daemon = ProcessDaemon("env-probe", [sys.executable, "-c", "pass"]) + # Python ignores non-existent PYTHONPATH entries, so the child then + # imports the installed rlinf package; the path is prepended anyway. + assert daemon.subprocess_env["PYTHONPATH"] == os.pathsep.join( + (str(missing.resolve()), "/other") + ) + + +def test_spawned_daemon_inherits_rlinf_pythonpath( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + checkout = tmp_path / "rlinf" + checkout.mkdir() + monkeypatch.setenv("RPENT_RLINF_ROOT", str(checkout)) + log_path = tmp_path / "pythonpath-probe.log" + daemon = ProcessDaemon( + "pythonpath-probe", + [ + sys.executable, + "-c", + "import os; print('PYTHONPATH=' + os.environ.get('PYTHONPATH', ''))", + ], + log_path=str(log_path), + cwd=str(tmp_path), + ) + daemon.start() + + def _printed() -> bool: + return log_path.exists() and "PYTHONPATH=" in log_path.read_text() + + try: + _wait_until(_printed) + finally: + daemon.stop(timeout=1.0) + + value = log_path.read_text().strip().partition("PYTHONPATH=")[2] + assert value.split(os.pathsep)[0] == str(checkout.resolve()) diff --git a/tests/unit_tests/rpent/utils/test_serialization.py b/tests/unit_tests/rpent/utils/test_serialization.py new file mode 100644 index 000000000..cdd6039ba --- /dev/null +++ b/tests/unit_tests/rpent/utils/test_serialization.py @@ -0,0 +1,91 @@ +# Copyright 2026 The RPent Authors. +# +# 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. + +"""Tests for :func:`rpent.utils.serialization.to_numpy_tree`.""" + +from __future__ import annotations + +import dataclasses +from typing import Any + +import numpy as np + +from rpent.utils.serialization import to_numpy_tree + + +class _FakeTensor: + """Duck-typed tensor stand-in (torch is not a unit-test dependency).""" + + def __init__(self, array: np.ndarray) -> None: + self._array = array + + def detach(self) -> _FakeTensor: + return self + + def cpu(self) -> _FakeTensor: + return self + + def numpy(self) -> np.ndarray: + return self._array + + +@dataclasses.dataclass +class _FakeObservation: + state: np.ndarray + reward: np.float32 + + +def test_tensor_like_values_become_numpy_arrays() -> None: + array = np.arange(3, dtype=np.float32) + # CPU tensors convert zero-copy: the array shares memory with the source. + assert to_numpy_tree(_FakeTensor(array)) is array + + +def test_numpy_arrays_pass_through_unchanged() -> None: + array = np.arange(3) + assert to_numpy_tree(array) is array + + +def test_numpy_scalars_pass_through_unchanged() -> None: + # np.generic values are pickle-safe as-is and the RPC transports + # round-trip them with dtype fidelity, so they need no unwrapping. + for value in (np.float32(1.5), np.bool_(True), np.int64(7)): + assert to_numpy_tree(value) is value + + +def test_nested_containers_recurse() -> None: + value: Any = { + "states": [_FakeTensor(np.zeros(2))], + "info": (np.float32(0.25), {"done": np.bool_(False)}), + } + result = to_numpy_tree(value) + assert result["states"][0].tolist() == [0.0, 0.0] + assert result["info"][0] == np.float32(0.25) + assert isinstance(result["info"][0], np.float32) + assert result["info"][1] == {"done": np.bool_(False)} + assert isinstance(result["info"][1]["done"], np.bool_) + + +def test_dataclasses_recurse() -> None: + result = to_numpy_tree(_FakeObservation(np.ones(2), np.float32(3.5))) + assert isinstance(result, dict) + assert result["state"].tolist() == [1.0, 1.0] + assert result["reward"] == np.float32(3.5) + assert isinstance(result["reward"], np.float32) + + +def test_plain_values_pass_through() -> None: + assert to_numpy_tree(None) is None + assert to_numpy_tree("x") == "x" + assert to_numpy_tree(3) == 3