diff --git a/README.md b/README.md index daaebecce..7d940421b 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,7 @@ RPent is built for four kinds of users: - [2026/08] 🔥 RPent supports RoboCasa with RLDX-1 as manipulation model. See the [RoboCasa setup and Target50 guide](robots/robocasa/README.md) and [full documentation](https://rpent.readthedocs.io/en/latest/rst_source/usage/robocasa.html). - [2026/08] 🔥 RPent supports the non-reasoning mode, which reduces average execution time by ~40%. +- [2026/08] 🔥 RPent supports BEHAVIOR with Pi0.5 for long-horizon household tasks. Doc: [BEHAVIOR](https://rpent.readthedocs.io/en/latest/rst_source/usage/behavior.html). - [2026/08] 🔥 RPent supports exploration mode for LIBERO. Doc: [LIBERO exploration mode](https://rpent.readthedocs.io/en/latest/rst_source/usage/libero.html#exploration-and-local-memory-evaluation). - [2026/08] 🔥 RPent supports RoboTwin with LingBot-VLA for dual-arm manipulation tasks. Doc: [RoboTwin](https://rpent.readthedocs.io/en/latest/rst_source/usage/robotwin.html). - [2026/07] 🔥 Our first RPent publication, [Harness VLA: Steering Frozen VLAs into Reliable Manipulation Primitives via Memory-Guided Agents](https://arxiv.org/abs/2607.08448), is released. @@ -81,6 +82,7 @@ RPent is built for four kinds of users: diff --git a/README.zh-CN.md b/README.zh-CN.md index e2cfc69d5..462b27444 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -40,6 +40,7 @@ RPent 面向以下四类用户： - [2026/08] 🔥 支持 RoboCasa，使用 RLDX-1 作为操作模型。参见 [RoboCasa 安装与 Target50 指南](robots/robocasa/README.md)和 [完整中文文档](https://rpent.readthedocs.io/zh-cn/latest/rst_source/usage/robocasa.html)。 - [2026/08] 🔥 新增非推理（non-reasoning）模式，平均执行时间降低约 40%。 +- [2026/08] 🔥 支持 BEHAVIOR，使用 Pi0.5 处理家庭长程任务。文档：[BEHAVIOR](https://rpent.readthedocs.io/zh-cn/latest/rst_source/usage/behavior.html)。 - [2026/08] 🔥 支持 LIBERO 探索模式。文档：[LIBERO 探索模式](https://rpent.readthedocs.io/zh-cn/latest/rst_source/usage/libero.html#memory)。 - [2026/08] 🔥 支持 RoboTwin，使用 LingBot-VLA 处理双臂操作任务。文档：[RoboTwin](https://rpent.readthedocs.io/zh-cn/latest/rst_source/usage/robotwin.html)。 - [2026/07] 🔥 RPent 首篇论文 [Harness VLA: Steering Frozen VLAs into Reliable Manipulation Primitives via Memory-Guided Agents](https://arxiv.org/abs/2607.08448) 发布。 @@ -81,6 +82,7 @@ RPent 面向以下四类用户： diff --git a/docs/source-en/index.rst b/docs/source-en/index.rst index 2d1a9f835..bba527528 100644 --- a/docs/source-en/index.rst +++ b/docs/source-en/index.rst @@ -54,8 +54,8 @@ Welcome to RPent :link-type: doc :text-align: center - Drive the LIBERO / RoboCasa simulators or a Franka / SO-101 arm, - switch planners, and pick action primitives. + Drive the LIBERO / BEHAVIOR / RoboCasa simulators or a Franka / + SO-101 arm, switch planners, and pick action primitives. .. grid-item-card:: Development Tutorial :link: rst_source/development/architecture @@ -87,6 +87,7 @@ Welcome to RPent LIBERO RoboCasa RoboTwin + BEHAVIOR Franka SO-101 Advanced Deployment diff --git a/docs/source-en/rst_source/development/add_robot.rst b/docs/source-en/rst_source/development/add_robot.rst index 5b6cbfd9f..3a69863e5 100644 --- a/docs/source-en/rst_source/development/add_robot.rst +++ b/docs/source-en/rst_source/development/add_robot.rst @@ -131,9 +131,9 @@ template, and output slug; ``runtime_components`` and ``frame_channels`` describe the robot-specific rows and camera views rendered by the frontend. See ``robots/libero/robot_spec.py`` for the reference shape. -That's the entire registration step — ``_resolve_robot(name)`` does an -``importlib.import_module(f"robots.{name}")``, so dropping the package under -``robots/`` on disk is enough. No central list to update. +The standard source tree currently includes ``libero``, ``robocasa``, +``robotwin``, and ``behavior`` robot packages. New robot packages should follow +the same entry-point contract before being wired into a release. The sections below describe what each referenced module must contain. ``_add_cli_args`` / ``_parse_config`` are covered in §4 and the runtime hook diff --git a/docs/source-en/rst_source/development/architecture.rst b/docs/source-en/rst_source/development/architecture.rst index 2da933700..ebcace1c4 100644 --- a/docs/source-en/rst_source/development/architecture.rst +++ b/docs/source-en/rst_source/development/architecture.rst @@ -158,10 +158,10 @@ Dashboard description, and three runner hooks (``add_cli_args`` / ``parse_config`` / ``init_runtime``). See :doc:`interfaces` for what each field must provide. -The loader itself does not maintain a list of robot names. The -current CLI restricts ``--robot`` to ``libero`` and ``robocasa``; adding a -new name therefore also requires updating the CLI choices. See -:doc:`add_robot` for the complete procedure. +The standard source tree currently includes ``libero``, ``robocasa``, +``robotwin``, and ``behavior`` robot packages. New integrations should follow +the same ``RobotSpec`` / ``get_toolkit`` package shape described in +:doc:`add_robot`. Planner, Toolkit, and RPC transports ------------------------------------- diff --git a/docs/source-en/rst_source/installation.rst b/docs/source-en/rst_source/installation.rst index 14e76a494..5d9ddc56c 100644 --- a/docs/source-en/rst_source/installation.rst +++ b/docs/source-en/rst_source/installation.rst @@ -35,6 +35,7 @@ Other environment configurations are available when needed: pip install -e ".[robocasa]" # RoboCasa pip install -e ".[robotwin]" # RoboTwin + pip install -e ".[behavior]" # BEHAVIOR ``.[libero-pro]`` is the recommended default. @@ -51,6 +52,10 @@ Available extras: - LIBERO-PRO + openpi Pi0.5 VLA + SAM 3.0 + RLinf runtime * - ``.[libero-plus]`` - LIBERO-plus + openpi Pi0.5 VLA + SAM 3.0 + RLinf runtime + * - ``.[behavior]`` + - RPent-side BEHAVIOR dependencies; the full OmniGibson/Isaac Sim stack + uses the source-editable dual-venv workflow and licensed assets documented + in :doc:`usage/behavior` * - ``.[robocasa]`` - RoboCasa365 simulator + the RLDX-1 VLA; see :doc:`usage/robocasa` * - ``.[robotwin]`` diff --git a/docs/source-en/rst_source/usage/behavior.rst b/docs/source-en/rst_source/usage/behavior.rst new file mode 100644 index 000000000..d5aa58ace --- /dev/null +++ b/docs/source-en/rst_source/usage/behavior.rst @@ -0,0 +1,296 @@ +BEHAVIOR +======== + +`BEHAVIOR-1K `_ provides long-horizon household +tasks in OmniGibson. RPent exposes ``turning_on_radio`` and +``picking_up_trash`` as a standard sibling robot plugin under +``robots/behavior``. + +The integration follows the same lightweight contract as LIBERO, RoboCasa, and +RoboTwin: ``get_robot_spec()`` supplies CLI/config/runtime hooks and +``get_toolkit()`` supplies the public tools. BEHAVIOR-specific lifecycle code +stays inside ``robots/behavior``. The common ``--explore`` entry point supports +BEHAVIOR while preserving its one-attempt-per-session environment lifecycle. + +Installation status +------------------- + +BEHAVIOR is source-editable and uses two independent Python 3.10 environments: + +- the **RPent venv** runs the CLI, planner, Dashboard, and MemoryManager; +- the **BEHAVIOR venv** runs RLinf, OmniGibson, Isaac Sim, and Pi0.5. + +The ``.[behavior]`` extra installs only RPent-side dependencies. It does not +install the complete simulator, assets, or checkpoint, and a normal wheel does +not promise a directly runnable BEHAVIOR stack. From a source checkout, run: + +.. code-block:: bash + + python -m pip install -e ".[behavior]" + export RPENT_REPRO_ROOT="$PWD/.behavior-runtime" + export UV_CACHE_DIR="$RPENT_REPRO_ROOT/uv-cache" + behavior-install-runtime + +The installer keeps RPent editable in both venvs, clones the reviewed RLinf +revision, invokes the official RLinf BEHAVIOR installer, applies the reviewed +CUDA/OpenPI compatibility pins, verifies critical imports and CUDA, and writes +freezes plus source identities under ``$RPENT_REPRO_ROOT/manifests``. Use a new +``RPENT_REPRO_ROOT`` for a fresh install; the script refuses to overwrite a +wrong or dirty RLinf checkout. + +Motion planning uses NVlabs/cuRobo v0.8.0 at commit +``4ea77366ca48ee453e7df139e39fa6532af49f3b`` in the BEHAVIOR venv. +The installer applies constraints before its final repin, retaining NumPy +1.26.4, Torch 2.5.1+cu124 and Isaac Sim 4.5.0.0. Do not install cuRobo with +an unconstrained resolver: a NumPy 2 upgrade is incompatible with this stack. + +Simulator assets +---------------- + +Accept the BEHAVIOR/OmniGibson licences, choose a dedicated data root, and use +the standard asset command. It invokes the three official OmniGibson download +functions in the BEHAVIOR venv rather than importing OmniGibson into the RPent +environment: + +.. code-block:: bash + + export OMNIGIBSON_DATA_PATH=/path/to/BEHAVIOR-1K-datasets + export BEHAVIOR_PYTHON="$RPENT_REPRO_ROOT/venvs/behavior/bin/python" + behavior-download-assets --accept-license --skip-existing + +Omit ``--accept-license`` to let the official downloader display its +interactive licence prompt. The flag is an explicit non-interactive +confirmation; do not use it unless you accept the licence terms. + +The final data root must contain: + +.. code-block:: text + + BEHAVIOR-1K-datasets/ + 2025-challenge-task-instances/ + behavior-1k-assets/ + scenes/ + omnigibson-robot-assets/ + omnigibson.key + +Pi0.5 checkpoint +---------------- + +Download the reviewed checkpoint into a directory outside the source tree: + +.. code-block:: bash + + export PI05_CHECKPOINT_PATH=/path/to/RLinf-Pi05-BEHAVIOR-1K-PT50-CS32 + "$RPENT_REPRO_ROOT/venvs/rpent/bin/hf" download \ + RLinf/RLinf-Pi05-BEHAVIOR-1K-PT50-CS32 \ + --local-dir "$PI05_CHECKPOINT_PATH" + +``behavior-download-assets --verify`` verifies the required OmniGibson layout +and the source-controlled checkpoint size/SHA binding. The shared #136 Pi0.5 +component receives head, left-wrist, right-wrist, and raw R1Pro proprio data. +The raw RPC result is ``[1, 32, 23]``; the common client returns ``[32, 23]``. + +.. code-block:: bash + + behavior-download-assets --verify + +DINOv2 configuration +-------------------- + +BEHAVIOR keeps a reviewed `DINOv2 `_ +ViT-S/14 deployment for whole-image embeddings and episode-memory retrieval. +Provide a DINOv2 source archive and the ``dinov2_vits14_pretrain.pth`` weights: + +.. code-block:: bash + + export DINOV2_SOURCE_ARCHIVE=/path/to/dinov2-source.tar.gz + export DINOV2_WEIGHTS=/path/to/dinov2_vits14_pretrain.pth + + curl -L \ + https://github.com/facebookresearch/dinov2/archive/7764ea0f912e53c92e82eb78a2a1631e92725fc8.tar.gz \ + -o "$DINOV2_SOURCE_ARCHIVE" + curl -L \ + https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_pretrain.pth \ + -o "$DINOV2_WEIGHTS" + +DINOv2 is the shared visual-memory component. It is not a segmentation model, +does not replace SAM3 masks, and does not replace current public observations +or MemoryManager Markdown/YAML material. The accepted DINOv2 source revision +and both asset SHA-256 identities are pinned in +``robots/behavior/dino_v2/encoder.py``; the runtime rejects mismatched assets. + +The DINO episode-memory catalog is a rebuildable derivative of official +BEHAVIOR demonstration data. Build it with the standard command below; do not +edit the catalog by hand. + +.. code-block:: bash + + behavior-build-memory \ + --selection-manifest /path/to/selection_manifest.json \ + --video-root /path/to/2025-challenge-demos \ + --rollups-dir /path/to/episode_rollups \ + --source-archive "$DINOV2_SOURCE_ARCHIVE" \ + --weights "$DINOV2_WEIGHTS" \ + --cuda-device 2 \ + --output-dir /path/to/behavior-dino-catalog + +Task identity +------------- + +Use ``--task-name`` and ``--public-seed``. Public seeds map to fixed official +activity instances through ``robots/behavior/task_specs.py``. + +.. list-table:: + :header-rows: 1 + :widths: 24 42 16 18 + + * - Task + - Instruction + - Explore seeds + - Eval seeds + * - ``turning_on_radio`` + - Turn on the radio receiver on the living-room table. + - ``0`` + - ``1``-``9`` + * - ``picking_up_trash`` + - Put the three living-room soda cans into the kitchen trash can. + - ``0``-``9`` + - ``10``-``19`` + +One evaluation run +------------------ + +Bind each CUDA child to one physical GPU explicitly: + +.. code-block:: bash + + "$RPENT_REPRO_ROOT/venvs/rpent/bin/rpent" --robot behavior \ + --task-name turning_on_radio --public-seed 1 \ + --behavior-mode eval \ + --planner codex --model gpt-5.5 \ + --behavior-repo "$RPENT_REPRO_ROOT/RLinf" \ + --behavior-python "$RPENT_REPRO_ROOT/venvs/behavior/bin/python" \ + --activity-instance-dir \ + "$OMNIGIBSON_DATA_PATH/2025-challenge-task-instances" \ + --policy-checkpoint "$PI05_CHECKPOINT_PATH" \ + --behavior-env-cuda-device 0 \ + --behavior-model-cuda-device 1 \ + --dino-source-archive "$DINOV2_SOURCE_ARCHIVE" \ + --dino-weights "$DINOV2_WEIGHTS" \ + --memory-profile local \ + --memory-dir /path/to/behavior-memory \ + --behavior-memory-dir /path/to/reviewed-behavior-episode-memory + +The first environment load can take several minutes. The environment, VLA, and +DINO are separate processes, and each receives only its explicitly selected GPU. + +Official MemoryManager +---------------------- + +BEHAVIOR uses the same Markdown/YAML ``MemoryManager`` format and common memory +tools as the other robots. The DINO episode-memory catalog is a separate visual +experience retrieval source; when configured, its advisory is attached to +public tool receipts and remains historical guidance only. + +- Eval creates one ``MemoryManager`` with ``read_only`` access. +- Explore creates one ``MemoryManager`` with ``inbox_write`` access scoped to + ``/_internal/inbox/``. +- ``MEMORY.md``, ``global/``, ``suite/``, and ``task_only/`` hold the published + corpus. Solved audit/recipe pairs are copied to ``task_only/``. +- When merge processes a valid root-level draft, the cell inbox is archived to + ``_internal/merged/``. An inbox containing only invalid drafts + stays in place; conflicting prose is archived under ``_internal/conflicts/``. + +An absent or empty corpus is valid, but it contains no advice. Pass the same +explicit ``--memory-dir`` to runs that should share reviewed memory. + +Use ``--behavior-memory-dir`` only for the reviewed DINO episode-memory catalog. +Omitting it selects a legal empty episode catalog and does not download or +silently substitute task-specific memory. + +Use the standard RPent Explore entry point for a bounded sequence of sessions: + +.. code-block:: bash + + "$RPENT_REPRO_ROOT/venvs/rpent/bin/rpent" --robot behavior \ + --behavior-mode explore \ + --explore \ + --explore-sessions 3 \ + --task-name picking_up_trash --public-seed 0 \ + --output-dir /path/to/behavior-explore \ + --memory-dir /path/to/behavior-memory \ + --planner codex --model gpt-5.5 \ + --behavior-repo "$RPENT_REPRO_ROOT/RLinf" \ + --behavior-python "$RPENT_REPRO_ROOT/venvs/behavior/bin/python" \ + --activity-instance-dir \ + "$OMNIGIBSON_DATA_PATH/2025-challenge-task-instances" \ + --policy-checkpoint "$PI05_CHECKPOINT_PATH" \ + --behavior-env-cuda-device 0 \ + --behavior-model-cuda-device 1 \ + --dino-source-archive "$DINOV2_SOURCE_ARCHIVE" \ + --dino-weights "$DINOV2_WEIGHTS" + +For BEHAVIOR, one session is exactly one attempt. Each session starts a fresh +environment sidecar, episode, and ``sessions/session_NNN`` output directory; +the VLA and DINO sidecars stay shared across sessions. The planner cannot reset +inside an invocation, and ``--explore-attempts-per-session`` values above zero +are rejected. + +Runtime and Dashboard +--------------------- + +The runtime has four component roles: + +- ``env``: the task-scoped official BEHAVIOR/OmniGibson environment; +- ``vla``: the shared ``rpent/robots/components/pi05_vla_server.py`` service; +- ``dino``: the shared ``robots/behavior/dino_v2/server.py`` episode-memory + embedding service; +- ``memory``: the task-scoped official MemoryManager. + +Start a Dashboard Session with: + +.. code-block:: bash + + export RPENT_BEHAVIOR_PYTHON="$RPENT_REPRO_ROOT/venvs/behavior/bin/python" + "$RPENT_REPRO_ROOT/venvs/rpent/bin/rpent" \ + --robot behavior --dashboard \ + --task-name turning_on_radio --public-seed 1 \ + --behavior-mode eval \ + --behavior-repo "$RPENT_REPRO_ROOT/RLinf" \ + --behavior-python "$RPENT_BEHAVIOR_PYTHON" \ + --activity-instance-dir \ + "$OMNIGIBSON_DATA_PATH/2025-challenge-task-instances" \ + --policy-checkpoint "$PI05_CHECKPOINT_PATH" \ + --dino-source-archive "$DINOV2_SOURCE_ARCHIVE" \ + --dino-weights "$DINOV2_WEIGHTS" \ + --memory-profile local --memory-dir /path/to/behavior-memory \ + --output-dir /path/to/behavior-dashboard-run + +The Dashboard uses the common Start Session flow and head/left-wrist/right- +wrist camera views. BEHAVIOR does not add robot-local manual buttons, a manual +control backend, or ``env.dashboard_*`` RPC methods. The public contract +registers nine executable planner primitives: ``pi0_nav_pick``, ``observe``, +``pixel_to_world``, ``navigate_to``, ``move_to``, ``rotate_wrist``, ``close``, +``open``, and ``press``. ``move_to(hand=both)`` coordinates both arms through +cuRobo collision-checked trajectories; wrist rotation uses the same planner. +Navigation executes a bounded straight base segment or rotation and rejects +obstructed paths. RGB-D projections use the current physical camera frame; +R1Pro has no movable head camera, so non-center ``head_view`` presets are rejected. +``press`` advances an already aligned hand +at most 2 cm for at most 10 seconds, stopping on external contact or episode end. +Contact is not verified button contact; visual hand checks remain unverified. +Planning, collision, tracking and duration failures are reported explicitly. +Only raw ``info["done"]["success"]`` establishes task success. Motion primitives +return their final observation for the next policy call and streaming video; +VLA chunks additionally record each returned environment frame. + +The main logs are: + +.. code-block:: text + + /run.log + /behavior_vla_server.log + /behavior_dino_server.log + /tasks//behavior_env_server.log + /tasks//episode.mp4 + /tasks//terminal_receipt.json diff --git a/docs/source-zh/index.rst b/docs/source-zh/index.rst index 8740cfdd3..302258908 100644 --- a/docs/source-zh/index.rst +++ b/docs/source-zh/index.rst @@ -46,7 +46,7 @@ :link-type: doc :text-align: center - 使用 LIBERO / RoboCasa 仿真环境或 Franka / SO-101 机械臂， + 使用 LIBERO / BEHAVIOR / RoboCasa 仿真环境或 Franka / SO-101 机械臂， 切换 planner 并选择动作原语。 .. grid-item-card:: 开发教程 @@ -79,6 +79,7 @@ LIBERO RoboCasa RoboTwin + BEHAVIOR Franka SO-101 高级部署 diff --git a/docs/source-zh/rst_source/development/add_robot.rst b/docs/source-zh/rst_source/development/add_robot.rst index 2ecc2dcf1..5b236ac02 100644 --- a/docs/source-zh/rst_source/development/add_robot.rst +++ b/docs/source-zh/rst_source/development/add_robot.rst @@ -120,8 +120,8 @@ RPent 的整体进程划分、服务职责和通信方式见 :doc:`系统设计 slug，``runtime_components`` 与 ``frame_channels`` 描述前端展示的环境专用服务行 和相机视图。完整结构参考 ``robots/libero/robot_spec.py``。 -``_resolve_robot(name)`` 通过 ``importlib.import_module(f"robots.{name}")`` -动态加载机器人包。因此，只需将机器人包放在 ``robots/`` 下，无需维护中央注册列表。 +标准源码树当前包含 ``libero``、``robocasa``、``robotwin`` 和 ``behavior`` +robot package。新的 robot package 应先沿用相同入口契约，再接入发布版本。 下文依次说明这些模块需要实现的内容。``_add_cli_args`` 和 ``_parse_config`` 见第 4 节，runtime 钩子见第 5 节。Dashboard spec 只由 Dashboard runner 使用。 diff --git a/docs/source-zh/rst_source/development/architecture.rst b/docs/source-zh/rst_source/development/architecture.rst index 1b909410a..a1c63d73c 100644 --- a/docs/source-zh/rst_source/development/architecture.rst +++ b/docs/source-zh/rst_source/development/architecture.rst @@ -138,9 +138,9 @@ planner 后端集中在 ``rpent/planner/``， 钩子（``add_cli_args`` / ``parse_config`` / ``init_runtime``）。各字段要填什么见 :doc:`interfaces`。 -加载器本身不维护机器人名称列表。当前 CLI 将 ``--robot`` 限定为 ``libero`` -和 ``robocasa``；接入新的机器人名称时，还需要同步更新 CLI 的可选值。完整步骤见 -:doc:`add_robot`。 +标准源码树当前包含 ``libero``、``robocasa``、``robotwin`` 和 ``behavior`` +robot package。新的集成应沿用 :doc:`add_robot` 中说明的 ``RobotSpec`` / +``get_toolkit`` package 结构。 Planner、Toolkit 与 RPC 传输层 ------------------------------ diff --git a/docs/source-zh/rst_source/installation.rst b/docs/source-zh/rst_source/installation.rst index 786300518..279a6c50f 100644 --- a/docs/source-zh/rst_source/installation.rst +++ b/docs/source-zh/rst_source/installation.rst @@ -42,6 +42,7 @@ RPent 可以通过一条 ``pip install`` 命令完成安装，并提供多种可 pip install -e ".[robocasa]" # RoboCasa pip install -e ".[robotwin]" # RoboTwin + pip install -e ".[behavior]" # BEHAVIOR ``.[libero-pro]`` 是默认推荐的依赖组合。 @@ -58,6 +59,10 @@ RPent 可以通过一条 ``pip install`` 命令完成安装，并提供多种可 - LIBERO-PRO + openpi Pi0.5 VLA + SAM 3.0 + RLinf 运行时 * - ``.[libero-plus]`` - LIBERO-plus + openpi Pi0.5 VLA + SAM 3.0 + RLinf 运行时 + * - ``.[behavior]`` + - BEHAVIOR 的 RPent 侧依赖；完整 OmniGibson/Isaac Sim 运行环境需使用 + 源码 editable 双 venv 流程及受许可约束的仿真资产，详见 + :doc:`usage/behavior` * - ``.[robocasa]`` - RoboCasa365 仿真器 + RLDX-1 VLA，详见 :doc:`usage/robocasa` * - ``.[robotwin]`` diff --git a/docs/source-zh/rst_source/usage/behavior.rst b/docs/source-zh/rst_source/usage/behavior.rst new file mode 100644 index 000000000..71126ebb6 --- /dev/null +++ b/docs/source-zh/rst_source/usage/behavior.rst @@ -0,0 +1,279 @@ +BEHAVIOR +======== + +`BEHAVIOR-1K `_ 基于 OmniGibson 提供长程家庭任务。 +RPent 当前把 ``turning_on_radio`` 和 ``picking_up_trash`` 作为标准 sibling +robot plugin 接入，源码位于 ``robots/behavior``。 + +接入合同与 LIBERO、RoboCasa、RoboTwin 相同：``get_robot_spec()`` 提供 +CLI/config/runtime hooks，``get_toolkit()`` 提供公开工具。BEHAVIOR 生命周期逻辑 +留在 ``robots/behavior`` 内。公共 ``--explore`` 入口现已支持 BEHAVIOR，同时保留 +每个 session 只运行一个 attempt 的环境生命周期。 + +安装状态 +-------- + +BEHAVIOR 以源码 editable 方式运行，并使用两个相互独立的 Python 3.10 环境： + +- **RPent venv**：运行 CLI、planner、Dashboard 和 MemoryManager； +- **BEHAVIOR venv**：运行 RLinf、OmniGibson、Isaac Sim 和 Pi0.5。 + +``.[behavior]`` 只安装 RPent 侧依赖，不包含完整模拟器、资产或 checkpoint； +普通 wheel 不承诺可直接运行 BEHAVIOR。请在源码 checkout 中执行： + +.. code-block:: bash + + python -m pip install -e ".[behavior]" + export RPENT_REPRO_ROOT="$PWD/.behavior-runtime" + export UV_CACHE_DIR="$RPENT_REPRO_ROOT/uv-cache" + behavior-install-runtime + +安装器会在两个 venv 中保持 RPent editable，克隆已审查的 RLinf revision，调用 +官方 RLinf BEHAVIOR 安装器，应用已审查的 CUDA/OpenPI 兼容性 pin，验证关键 import +和 CUDA，并在 ``$RPENT_REPRO_ROOT/manifests`` 写入 freeze 与源码身份。fresh install +应使用新的 ``RPENT_REPRO_ROOT``；脚本不会覆盖 revision 错误或 dirty 的 RLinf +checkout。 + +运动规划在 BEHAVIOR venv 中使用 NVlabs/cuRobo v0.8.0，固定 commit +``4ea77366ca48ee453e7df139e39fa6532af49f3b``。安装器在最终 repin 前使用 +constraints，保留 NumPy 1.26.4、Torch 2.5.1+cu124 和 Isaac Sim 4.5.0.0。 +不要无约束安装 cuRobo：resolver 升级到 NumPy 2 会破坏此环境的兼容性。 + +仿真资产 +-------- + +接受 BEHAVIOR/OmniGibson 许可后，选择独立数据根并使用标准资产命令。该命令会在 +BEHAVIOR venv 中调用三个 OmniGibson 官方下载函数，不会把 OmniGibson import 到 +RPent 环境： + +.. code-block:: bash + + export OMNIGIBSON_DATA_PATH=/path/to/BEHAVIOR-1K-datasets + export BEHAVIOR_PYTHON="$RPENT_REPRO_ROOT/venvs/behavior/bin/python" + behavior-download-assets --accept-license --skip-existing + +不传 ``--accept-license`` 时，官方下载器会显示交互式许可确认。该参数代表明确的 +非交互许可确认；仅在接受许可条款后使用。 + +最终数据根必须包含： + +.. code-block:: text + + BEHAVIOR-1K-datasets/ + 2025-challenge-task-instances/ + behavior-1k-assets/ + scenes/ + omnigibson-robot-assets/ + omnigibson.key + +Pi0.5 checkpoint +---------------- + +将已审查 checkpoint 下载到源码树之外： + +.. code-block:: bash + + export PI05_CHECKPOINT_PATH=/path/to/RLinf-Pi05-BEHAVIOR-1K-PT50-CS32 + "$RPENT_REPRO_ROOT/venvs/rpent/bin/hf" download \ + RLinf/RLinf-Pi05-BEHAVIOR-1K-PT50-CS32 \ + --local-dir "$PI05_CHECKPOINT_PATH" + +``behavior-download-assets --verify`` 会检查 OmniGibson 必需目录，以及源码中固定的 +checkpoint size/SHA binding。#136 的共享 Pi0.5 component 接收 head、left wrist、 +right wrist 和 raw R1Pro proprio；原始 RPC 输出为 ``[1, 32, 23]``，公共 client +返回 ``[32, 23]``。 + +.. code-block:: bash + + behavior-download-assets --verify + +DINOv2 配置 +------------- + +BEHAVIOR 保留经审查的 `DINOv2 `_ +ViT-S/14 部署，用于整图 embedding 与 episode memory 检索。运行前提供 DINOv2 +源码归档和 ``dinov2_vits14_pretrain.pth`` 权重： + +.. code-block:: bash + + export DINOV2_SOURCE_ARCHIVE=/path/to/dinov2-source.tar.gz + export DINOV2_WEIGHTS=/path/to/dinov2_vits14_pretrain.pth + + curl -L \ + https://github.com/facebookresearch/dinov2/archive/7764ea0f912e53c92e82eb78a2a1631e92725fc8.tar.gz \ + -o "$DINOV2_SOURCE_ARCHIVE" + curl -L \ + https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_pretrain.pth \ + -o "$DINOV2_WEIGHTS" + +DINOv2 是共享视觉 memory component；它不是分割模型，不替代 SAM3 mask、当前 +公开观察或 MemoryManager 的 Markdown/YAML 语料。接受的源码 revision 和两个资产 +SHA-256 固定在 ``robots/behavior/dino_v2/encoder.py``；runtime 会拒绝不匹配的资产。 + +DINO episode-memory catalog 是官方 BEHAVIOR demonstration 数据的可重建派生物。 +使用下面的标准命令生成；不要手工编辑 catalog。 + +.. code-block:: bash + + behavior-build-memory \ + --selection-manifest /path/to/selection_manifest.json \ + --video-root /path/to/2025-challenge-demos \ + --rollups-dir /path/to/episode_rollups \ + --source-archive "$DINOV2_SOURCE_ARCHIVE" \ + --weights "$DINOV2_WEIGHTS" \ + --cuda-device 2 \ + --output-dir /path/to/behavior-dino-catalog + +任务身份 +-------- + +使用 ``--task-name`` 和 ``--public-seed``。public seed 通过 +``robots/behavior/task_specs.py`` 固定映射到官方 activity instance。 + +.. list-table:: + :header-rows: 1 + :widths: 24 42 16 18 + + * - 任务 + - 指令 + - Explore seeds + - Eval seeds + * - ``turning_on_radio`` + - 打开客厅桌上的收音机。 + - ``0`` + - ``1``-``9`` + * - ``picking_up_trash`` + - 把客厅的三个汽水罐放进厨房垃圾桶。 + - ``0``-``9`` + - ``10``-``19`` + +运行一次 Eval +------------- + +每个 CUDA 子进程必须显式绑定一个物理 GPU： + +.. code-block:: bash + + "$RPENT_REPRO_ROOT/venvs/rpent/bin/rpent" --robot behavior \ + --task-name turning_on_radio --public-seed 1 \ + --behavior-mode eval \ + --planner codex --model gpt-5.5 \ + --behavior-repo "$RPENT_REPRO_ROOT/RLinf" \ + --behavior-python "$RPENT_REPRO_ROOT/venvs/behavior/bin/python" \ + --activity-instance-dir \ + "$OMNIGIBSON_DATA_PATH/2025-challenge-task-instances" \ + --policy-checkpoint "$PI05_CHECKPOINT_PATH" \ + --behavior-env-cuda-device 0 \ + --behavior-model-cuda-device 1 \ + --dino-source-archive "$DINOV2_SOURCE_ARCHIVE" \ + --dino-weights "$DINOV2_WEIGHTS" \ + --memory-profile local \ + --memory-dir /path/to/behavior-memory \ + --behavior-memory-dir /path/to/reviewed-behavior-episode-memory + +首次加载环境通常需要数分钟。env、VLA 和 DINO 是不同进程，每个进程只接收自己 +显式选择的 GPU。 + +官方 MemoryManager +------------------ + +BEHAVIOR 使用和其他机器人相同的 Markdown/YAML ``MemoryManager`` 格式与公共 memory +工具。DINO episode-memory catalog 是独立的视觉经验检索源；配置后，它的 advisory +会附加到公开 tool receipt，并始终只作为历史建议。 + +- Eval 只构造一个 ``read_only`` MemoryManager； +- Explore 只构造一个 ``inbox_write`` MemoryManager，写入范围限定为 + ``/_internal/inbox/``； +- ``MEMORY.md``、``global/``、``suite/`` 和 ``task_only/`` 保存已发布语料， + 成功的 audit/recipe 对复制到 ``task_only/``； +- merge 处理有效的根级草稿后，将该 cell 的 inbox 归档到 + ``_internal/merged/``。只有无效草稿的 inbox 保留原位； + 冲突文本归档到 ``_internal/conflicts/``。 + +缺失或空 corpus 是合法状态，但不会提供任何建议。需要共享已审查 memory 的运行应 +显式传入同一个 ``--memory-dir``。 + +``--behavior-memory-dir`` 只用于经审查的 DINO episode-memory catalog。省略该参数 +会选择合法的空 episode catalog，不会下载或静默替换为特定任务 memory。 + +使用标准 RPent Explore 入口运行一组有界 session： + +.. code-block:: bash + + "$RPENT_REPRO_ROOT/venvs/rpent/bin/rpent" --robot behavior \ + --behavior-mode explore \ + --explore \ + --explore-sessions 3 \ + --task-name picking_up_trash --public-seed 0 \ + --output-dir /path/to/behavior-explore \ + --memory-dir /path/to/behavior-memory \ + --planner codex --model gpt-5.5 \ + --behavior-repo "$RPENT_REPRO_ROOT/RLinf" \ + --behavior-python "$RPENT_REPRO_ROOT/venvs/behavior/bin/python" \ + --activity-instance-dir \ + "$OMNIGIBSON_DATA_PATH/2025-challenge-task-instances" \ + --policy-checkpoint "$PI05_CHECKPOINT_PATH" \ + --behavior-env-cuda-device 0 \ + --behavior-model-cuda-device 1 \ + --dino-source-archive "$DINOV2_SOURCE_ARCHIVE" \ + --dino-weights "$DINOV2_WEIGHTS" + +对 BEHAVIOR 而言，一个 session 就是一个 attempt。每个 session 都启动 fresh env +sidecar、new episode，并写入独立的 ``sessions/session_NNN`` 目录；VLA 与 DINO +sidecar 在各 session 之间共享。planner 无权在单次 invocation 内 reset，且 +``--explore-attempts-per-session`` 大于零会被拒绝。 + +Runtime 与 Dashboard +-------------------- + +runtime 有四个 component role： + +- ``env``：task-scoped 官方 BEHAVIOR/OmniGibson 环境； +- ``vla``：共享 ``rpent/robots/components/pi05_vla_server.py`` 服务； +- ``dino``：共享 ``robots/behavior/dino_v2/server.py`` episode-memory + embedding 服务； +- ``memory``：task-scoped 官方 MemoryManager。 + +启动 Dashboard Session： + +.. code-block:: bash + + export RPENT_BEHAVIOR_PYTHON="$RPENT_REPRO_ROOT/venvs/behavior/bin/python" + "$RPENT_REPRO_ROOT/venvs/rpent/bin/rpent" \ + --robot behavior --dashboard \ + --task-name turning_on_radio --public-seed 1 \ + --behavior-mode eval \ + --behavior-repo "$RPENT_REPRO_ROOT/RLinf" \ + --behavior-python "$RPENT_BEHAVIOR_PYTHON" \ + --activity-instance-dir \ + "$OMNIGIBSON_DATA_PATH/2025-challenge-task-instances" \ + --policy-checkpoint "$PI05_CHECKPOINT_PATH" \ + --dino-source-archive "$DINOV2_SOURCE_ARCHIVE" \ + --dino-weights "$DINOV2_WEIGHTS" \ + --memory-profile local --memory-dir /path/to/behavior-memory \ + --output-dir /path/to/behavior-dashboard-run + +Dashboard 使用公共 Start Session 流程与 head/left-wrist/right-wrist 相机视图。 +BEHAVIOR 不增加 robot-local 手动按钮、手动控制 backend 或 +``env.dashboard_*`` RPC。公开合同注册 9 个可执行 planner primitive： +``pi0_nav_pick``、``observe``、``pixel_to_world``、``navigate_to``、``move_to``、 +``rotate_wrist``、``close``、``open`` 和 ``press``。 +``move_to(hand=both)`` 通过 cuRobo 碰撞检查轨迹协调双臂，腕旋转复用同一规划器。 +导航执行有界直线底盘运动或转向，路径被阻挡时拒绝执行。 +RGB-D 投影使用当前物理相机帧；R1Pro 没有可动头部相机，因此拒绝非 center 的 +``head_view`` 预设。``press`` 沿已对准的手部方向推进，最多 2 cm、10 秒， +遇外部接触或 episode 结束即停；接触不等于已验证按钮接触，视觉手部检查仍未验证。 +规划、碰撞、跟踪和时长限制导致的失败均明确返回。 +任务成功仅认原始 ``info["done"]["success"]``。运动原语返回最终观测供后续策略调用 +和流式视频使用；VLA chunk 另外记录每个实际返回的环境帧。 + +主要日志： + +.. code-block:: text + + /run.log + /behavior_vla_server.log + /behavior_dino_server.log + /tasks//behavior_env_server.log + /tasks//episode.mp4 + /tasks//terminal_receipt.json diff --git a/pyproject.toml b/pyproject.toml index c46a82160..63afad07b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,9 @@ dependencies = [ [project.scripts] rpent = "rpent.cli.main:main" rpent-memory = "rpent.cli.memory:main" +behavior-download-assets = "robots.behavior.assets_cli:main" +behavior-build-memory = "robots.behavior.build_memory_cli:main" +behavior-install-runtime = "robots.behavior.install_runtime:main" [project.urls] Homepage = "https://github.com/RLinf/RPent" @@ -56,6 +59,7 @@ Documentation = "https://rpent.readthedocs.io/en/latest/" test = [ "pytest", "pytest-timeout", + "scipy>=1.10,<2", ] rlinf = [ "rpent-rlinf", @@ -68,6 +72,17 @@ sam3 = [ "torchvision", "pillow>=10", ] +behavior = [ + "scipy>=1.10,<2", + "rpent[rlinf]", + "rpent-openpi @ git+https://github.com/RLinf/openpi.git@rpent", + # OpenPI's tool.uv source is not propagated through package metadata. + "lerobot @ git+https://github.com/huggingface/lerobot.git@0cf864870cf29f4738d3ade893e6fd13fbd7cdb5", + "torch", + "torchvision", + "pillow>=10", + "imageio-ffmpeg>=0.5", +] libero = [ "rpent[rlinf]", "rpent[sam3]", @@ -115,7 +130,7 @@ include-package-data = true [tool.setuptools.packages.find] where = ["."] -include = ["rpent*"] +include = ["rpent*", "robots", "robots.behavior*"] exclude = ["tests*", "docs*", "examples*", "docker*", "toolkits*"] [tool.setuptools.package-data] diff --git a/robots/behavior/__init__.py b/robots/behavior/__init__.py new file mode 100644 index 000000000..db7e30c76 --- /dev/null +++ b/robots/behavior/__init__.py @@ -0,0 +1,19 @@ +# 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. + +"""BEHAVIOR robot extension.""" + +from robots.behavior.robot_spec import get_robot_spec, get_toolkit + +__all__ = ["get_robot_spec", "get_toolkit"] diff --git a/robots/behavior/assets_cli.py b/robots/behavior/assets_cli.py new file mode 100644 index 000000000..7e5a568e0 --- /dev/null +++ b/robots/behavior/assets_cli.py @@ -0,0 +1,278 @@ +# 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. + +"""Download and verify BEHAVIOR assets. + +This command requires an RPent source checkout with ``robots.behavior`` +available through editable install. +""" + +from __future__ import annotations + +import argparse +import hashlib +import os +import subprocess +from collections.abc import Sequence +from pathlib import Path + +from robots.behavior.dino_v2.encoder import ( + EXPECTED_SOURCE_ARCHIVE_SHA256, + EXPECTED_WEIGHTS_SHA256, +) +from robots.behavior.policy_checkpoint import validate_policy_checkpoint + +_DOWNLOAD_SNIPPET = """ +import sys + +from omnigibson.utils.asset_utils import ( + download_2025_challenge_task_instances, + download_behavior_1k_assets, + download_omnigibson_robot_assets, +) + +actions = set(filter(None, sys.argv[1].split(","))) +accept_license = sys.argv[2] == "1" +if "robot" in actions: + download_omnigibson_robot_assets() +if "behavior" in actions: + download_behavior_1k_assets(accept_license=accept_license) +if "challenge" in actions: + download_2025_challenge_task_instances() +""" + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="behavior-download-assets", + description=( + "Download BEHAVIOR assets through the official OmniGibson API, " + "or verify an existing installation." + ), + ) + parser.add_argument( + "--behavior-python", + type=Path, + default=None, + help=( + "Python from the BEHAVIOR venv. Defaults to BEHAVIOR_PYTHON, " + "RPENT_BEHAVIOR_PYTHON, or $RPENT_REPRO_ROOT/venvs/behavior/bin/python." + ), + ) + parser.add_argument( + "--data-path", + type=Path, + default=None, + help="BEHAVIOR data root (defaults to OMNIGIBSON_DATA_PATH).", + ) + parser.add_argument( + "--checkpoint", + type=Path, + default=None, + help="Pi0.5 checkpoint directory (defaults to PI05_CHECKPOINT_PATH).", + ) + parser.add_argument( + "--dino-source-archive", + type=Path, + default=None, + help="Pinned DINOv2 source archive (defaults to DINOV2_SOURCE_ARCHIVE).", + ) + parser.add_argument( + "--dino-weights", + type=Path, + default=None, + help="DINOv2 ViT-S/14 weights (defaults to DINOV2_WEIGHTS).", + ) + parser.add_argument( + "--accept-license", + action="store_true", + help=( + "Confirm acceptance of the BEHAVIOR data licence non-interactively. " + "Without this flag, the official downloader prompts when needed." + ), + ) + parser.add_argument( + "--skip-existing", + action="store_true", + help="Do not call a downloader for an asset directory already present.", + ) + parser.add_argument( + "--verify", + action="store_true", + help="Verify existing simulator, checkpoint, and DINO assets without downloading.", + ) + return parser + + +def _env_path(value: Path | None, *names: str) -> Path | None: + if value is not None: + return value.expanduser().resolve() + for name in names: + configured = os.environ.get(name) + if configured: + return Path(configured).expanduser().resolve() + return None + + +def _behavior_python(value: Path | None) -> Path: + resolved = _env_path(value, "BEHAVIOR_PYTHON", "RPENT_BEHAVIOR_PYTHON") + if resolved is None: + repro_root = os.environ.get("RPENT_REPRO_ROOT") + if repro_root: + resolved = ( + Path(repro_root).expanduser().resolve() + / "venvs" + / "behavior" + / "bin" + / "python" + ) + if resolved is None or not resolved.is_file(): + raise ValueError( + "BEHAVIOR Python is unavailable; pass --behavior-python or set " + "BEHAVIOR_PYTHON" + ) + return resolved + + +def _require_data_path(value: Path | None) -> Path: + resolved = _env_path(value, "OMNIGIBSON_DATA_PATH") + if resolved is None: + raise ValueError( + "BEHAVIOR data root is unavailable; pass --data-path or set " + "OMNIGIBSON_DATA_PATH" + ) + return resolved + + +def _require_file(path: Path | None, *, label: str) -> Path: + if path is None or not path.is_file(): + raise ValueError(f"missing required {label}: {path}") + return path + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def verify_assets( + *, + data_path: Path, + checkpoint: Path | None, + dino_source_archive: Path | None, + dino_weights: Path | None, +) -> None: + required_directories = ( + data_path / "behavior-1k-assets" / "scenes", + data_path / "omnigibson-robot-assets", + data_path / "2025-challenge-task-instances", + ) + for path in required_directories: + if not path.is_dir(): + raise ValueError(f"missing required directory: {path}") + _require_file(data_path / "omnigibson.key", label="OmniGibson licence key") + + if checkpoint is None: + raise ValueError( + "Pi0.5 checkpoint is unavailable; pass --checkpoint or set " + "PI05_CHECKPOINT_PATH" + ) + validate_policy_checkpoint(checkpoint) + + source = _require_file(dino_source_archive, label="DINOv2 source archive") + weights = _require_file(dino_weights, label="DINOv2 weights") + source_sha256 = _sha256_file(source) + if source_sha256 != EXPECTED_SOURCE_ARCHIVE_SHA256: + raise ValueError( + "DINOv2 source archive SHA-256 mismatch: " + f"expected {EXPECTED_SOURCE_ARCHIVE_SHA256}, got {source_sha256}" + ) + weights_sha256 = _sha256_file(weights) + if weights_sha256 != EXPECTED_WEIGHTS_SHA256: + raise ValueError( + "DINOv2 weights SHA-256 mismatch: " + f"expected {EXPECTED_WEIGHTS_SHA256}, got {weights_sha256}" + ) + print("BEHAVIOR assets: OK") + + +def download_assets( + *, + behavior_python: Path, + data_path: Path, + accept_license: bool, + skip_existing: bool, +) -> None: + actions = ["robot", "behavior", "challenge"] + if skip_existing: + present = { + "robot": (data_path / "omnigibson-robot-assets").is_dir(), + "behavior": ( + (data_path / "behavior-1k-assets").is_dir() + and (data_path / "omnigibson.key").is_file() + ), + "challenge": (data_path / "2025-challenge-task-instances").is_dir(), + } + actions = [name for name in actions if not present[name]] + if not actions: + print("BEHAVIOR assets already exist; nothing to download.") + return + + data_path.mkdir(parents=True, exist_ok=True) + env = os.environ.copy() + env["OMNIGIBSON_DATA_PATH"] = str(data_path) + subprocess.run( + [ + str(behavior_python), + "-c", + _DOWNLOAD_SNIPPET, + ",".join(actions), + "1" if accept_license else "0", + ], + check=True, + env=env, + ) + + +def main(argv: Sequence[str] | None = None) -> int: + parser = _build_parser() + args = parser.parse_args(argv) + try: + data_path = _require_data_path(args.data_path) + if args.verify: + verify_assets( + data_path=data_path, + checkpoint=_env_path(args.checkpoint, "PI05_CHECKPOINT_PATH"), + dino_source_archive=_env_path( + args.dino_source_archive, "DINOV2_SOURCE_ARCHIVE" + ), + dino_weights=_env_path(args.dino_weights, "DINOV2_WEIGHTS"), + ) + else: + download_assets( + behavior_python=_behavior_python(args.behavior_python), + data_path=data_path, + accept_license=bool(args.accept_license), + skip_existing=bool(args.skip_existing), + ) + except ValueError as error: + parser.error(str(error)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/robots/behavior/build_memory_cli.py b/robots/behavior/build_memory_cli.py new file mode 100644 index 000000000..837e6c7dc --- /dev/null +++ b/robots/behavior/build_memory_cli.py @@ -0,0 +1,83 @@ +# 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. + +"""Build the BEHAVIOR DINO episode catalog from official demonstration data.""" + +from __future__ import annotations + +import argparse +import json +import os +from collections.abc import Sequence +from pathlib import Path + +from robots.behavior.sft_offline_converter import compile_runtime_catalog + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="behavior-build-memory", + description=( + "Build the derived BEHAVIOR DINO episode-memory catalog from " + "official demonstration media and reviewed episode rollups." + ), + ) + parser.add_argument("--selection-manifest", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument("--video-root", required=True, type=Path, action="append") + parser.add_argument("--rollups-dir", required=True, type=Path) + parser.add_argument("--source-archive", required=True, type=Path) + parser.add_argument("--weights", required=True, type=Path) + parser.add_argument("--cache-dir", type=Path, default=None) + parser.add_argument( + "--cuda-device", + required=True, + help="Single CUDA device id made visible to the DINO compiler.", + ) + parser.add_argument("--batch-size", type=int, default=32) + return parser + + +def _single_cuda_device(value: str) -> str: + device = str(value).strip() + if not device or "," in device or not device.isdigit(): + raise ValueError("--cuda-device must be one numeric device id") + return device + + +def main(argv: Sequence[str] | None = None) -> int: + parser = _build_parser() + args = parser.parse_args(argv) + try: + os.environ["CUDA_VISIBLE_DEVICES"] = _single_cuda_device(args.cuda_device) + result = compile_runtime_catalog( + selection_manifest=args.selection_manifest.expanduser().resolve(), + output_dir=args.output_dir.expanduser().resolve(), + video_roots=tuple(path.expanduser().resolve() for path in args.video_root), + rollups_dir=args.rollups_dir.expanduser().resolve(), + source_archive=args.source_archive.expanduser().resolve(), + weights=args.weights.expanduser().resolve(), + cache_dir=None + if args.cache_dir is None + else args.cache_dir.expanduser().resolve(), + batch_size=int(args.batch_size), + ) + except ValueError as error: + parser.error(str(error)) + print(json.dumps(dict(result), sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/robots/behavior/dino_v2/__init__.py b/robots/behavior/dino_v2/__init__.py new file mode 100644 index 000000000..4934bc53d --- /dev/null +++ b/robots/behavior/dino_v2/__init__.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. + +"""BEHAVIOR DINOv2 encoder, RPC client, and server.""" + +from robots.behavior.dino_v2.client import BehaviorDinoClient +from robots.behavior.dino_v2.encoder import ( + DINOV2_DIMENSION, + DISTANCE_METRIC, + Dinov2DeploymentPaths, + Dinov2Engine, + Dinov2RevisionIdentity, +) +from robots.behavior.dino_v2.server import BehaviorDinoFacade + +__all__ = [ + "DINOV2_DIMENSION", + "DISTANCE_METRIC", + "BehaviorDinoFacade", + "BehaviorDinoClient", + "Dinov2DeploymentPaths", + "Dinov2Engine", + "Dinov2RevisionIdentity", +] diff --git a/robots/behavior/dino_v2/client.py b/robots/behavior/dino_v2/client.py new file mode 100644 index 000000000..5f381d1d3 --- /dev/null +++ b/robots/behavior/dino_v2/client.py @@ -0,0 +1,87 @@ +# 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. + +"""RPC client for the optional BEHAVIOR DINOv2 component.""" + +from __future__ import annotations + +import threading +from collections.abc import Mapping +from typing import Any + +import numpy as np + +from robots.behavior.dino_v2.encoder import DINOV2_DIMENSION, l2_normalize_row +from rpent.utils.rpc import RpcClient + + +class BehaviorDinoClient: + """Small checked RPC wrapper around a DINOv2 encoder service.""" + + def __init__( + self, + client: RpcClient, + *, + expected_meta: dict[str, Any] | None = None, + ) -> None: + self._client = client + self._close_lock = threading.Lock() + self._transport_closed = False + meta = self.get_meta() + if expected_meta: + mismatches = { + key: {"expected": expected, "actual": meta.get(key)} + for key, expected in expected_meta.items() + if meta.get(key) != expected + } + if mismatches: + raise RuntimeError(f"dino_meta mismatch: {mismatches!r}") + self.server_meta = dict(meta) + + def _call(self, method: str, **kwargs: Any) -> Any: + return self._client.call(method, kwargs=kwargs, timeout_s=120.0) + + def get_meta(self) -> dict[str, Any]: + payload = self._client.call("dino.get_meta", timeout_s=5.0) + if not isinstance(payload, Mapping): + raise TypeError("dino.get_meta must return a mapping") + if payload.get("dimension") != DINOV2_DIMENSION: + raise RuntimeError("DINO service dimension does not match CLS384") + return dict(payload) + + def encode_batch( + self, images: list[np.ndarray | None] + ) -> tuple[np.ndarray | None, ...]: + payload = self._call("dino.encode_batch", images=images) + if not isinstance(payload, list): + raise TypeError("dino.encode_batch must return a list") + result: list[np.ndarray | None] = [] + for index, item in enumerate(payload): + if item is None: + result.append(None) + continue + result.append(l2_normalize_row(item, path=f"dino.output[{index}]")) + return tuple(result) + + def close_transport(self) -> None: + with self._close_lock: + if self._transport_closed: + return + close = getattr(self._client, "close", None) + if callable(close): + close() + self._transport_closed = True + + +__all__ = ["BehaviorDinoClient"] diff --git a/robots/behavior/dino_v2/encoder.py b/robots/behavior/dino_v2/encoder.py new file mode 100644 index 000000000..147bbdafd --- /dev/null +++ b/robots/behavior/dino_v2/encoder.py @@ -0,0 +1,512 @@ +# 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. + +"""Pinned DINOv2 ViT-S/14 RGB224 CLS384 embedding contract. + +The encoder identity is portable and path-free. Deployment paths are checked +only when an actual backend is materialized. Tests and offline builders may +inject a backend; the default backend imports torch lazily after verifying both +frozen assets, so importing this module itself remains lightweight. +""" + +from __future__ import annotations + +import hashlib +import importlib +import os +import tarfile +import tempfile +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from typing import Any, Protocol + +import numpy as np + +from robots.behavior.memory.schema import MemoryValidationError, fail, require_sha256 + +MODEL_ID = "facebookresearch/dinov2_vits14" +MODEL_REVISION = "facebookresearch/dinov2@7764ea0f912e53c92e82eb78a2a1631e92725fc8" +EXPECTED_SOURCE_COMMIT = "7764ea0f912e53c92e82eb78a2a1631e92725fc8" +EXPECTED_SOURCE_ARCHIVE_SHA256 = ( + "c27dcdaf50e9fb5bbdf2bb529da357716372e19c6afab17d5350f3f0094aed4b" +) +EXPECTED_WEIGHTS_SHA256 = ( + "b938bf1bc15cd2ec0feacfe3a1bb553fe8ea9ca46a7e1d8d00217f29aef60cd9" +) +PREPROCESS_ID = "rpent_dinov2_vits14_rgb224_bicubic_antialias_v1" +EXTRACTOR_ID = "dinov2_vits14_cls_token_v1" +DINOV2_DIMENSION = 384 +MAX_BATCH_SIZE = 32 +DISTANCE_METRIC = "one_minus_cosine_on_l2_cls384" + + +class Dinov2Backend(Protocol): + torch_version: str + torchvision_version: str + device: str + eval_mode: bool + parameters_frozen: bool + inference_only: bool + + def encode_batch(self, images: Sequence[np.ndarray]) -> np.ndarray: ... + def close(self) -> None: ... + + +BackendLoader = Callable[ + ["Dinov2RevisionIdentity", "Dinov2DeploymentPaths"], Dinov2Backend +] + + +@dataclass(frozen=True, slots=True) +class Dinov2RevisionIdentity: + model_id: str + model_revision: str + source_commit: str + source_archive_sha256: str + weights_sha256: str + torch_version: str + torchvision_version: str + device: str + preprocess_id: str = PREPROCESS_ID + extractor_id: str = EXTRACTOR_ID + dimension: int = DINOV2_DIMENSION + + def __post_init__(self) -> None: + expected = { + "model_id": MODEL_ID, + "model_revision": MODEL_REVISION, + "source_commit": EXPECTED_SOURCE_COMMIT, + "source_archive_sha256": EXPECTED_SOURCE_ARCHIVE_SHA256, + "weights_sha256": EXPECTED_WEIGHTS_SHA256, + "device": "cuda", + "preprocess_id": PREPROCESS_ID, + "extractor_id": EXTRACTOR_ID, + } + for field, value in expected.items(): + if getattr(self, field) != value: + fail( + "MEMORY_DINOV2_IDENTITY_MISMATCH", + f"embedding.{field}", + f"expected {value!r}", + ) + require_sha256( + self.source_archive_sha256, path="embedding.source_archive_sha256" + ) + require_sha256(self.weights_sha256, path="embedding.weights_sha256") + if self.dimension != DINOV2_DIMENSION: + fail( + "MEMORY_DINOV2_DIMENSION_INVALID", "embedding.dimension", "expected 384" + ) + for field in ("torch_version", "torchvision_version"): + value = getattr(self, field) + if not isinstance(value, str) or not value or value.strip() != value: + fail( + "MEMORY_DINOV2_IDENTITY_INVALID", + f"embedding.{field}", + "must be exact non-empty version", + ) + + def as_dict(self) -> dict[str, Any]: + return { + "model_id": self.model_id, + "model_revision": self.model_revision, + "source_commit": self.source_commit, + "source_archive_sha256": self.source_archive_sha256, + "weights_sha256": self.weights_sha256, + "torch_version": self.torch_version, + "torchvision_version": self.torchvision_version, + "device": self.device, + "preprocess_id": self.preprocess_id, + "extractor_id": self.extractor_id, + "dimension": self.dimension, + } + + +@dataclass(frozen=True, slots=True) +class Dinov2DeploymentPaths: + source_archive_path: Path + weights_path: Path + cache_dir: Path | None = None + + def __post_init__(self) -> None: + for field in ("source_archive_path", "weights_path"): + value = getattr(self, field) + if not isinstance(value, Path) or not value.is_absolute(): + fail("MEMORY_DINOV2_DEPLOYMENT_INVALID", field, "must be absolute Path") + if self.cache_dir is not None and ( + not isinstance(self.cache_dir, Path) or not self.cache_dir.is_absolute() + ): + fail( + "MEMORY_DINOV2_DEPLOYMENT_INVALID", + "cache_dir", + "must be absolute Path or None", + ) + + +def l2_normalize_row(value: Any, *, path: str) -> np.ndarray: + row = np.asarray(value, dtype=np.float64) + if row.shape != (DINOV2_DIMENSION,) or not np.isfinite(row).all(): + fail("MEMORY_DINOV2_VECTOR_INVALID", path, "expected finite vector[384]") + norm = float(np.linalg.norm(row)) + if norm <= 0.0: + fail("MEMORY_DINOV2_VECTOR_INVALID", path, "cannot normalize zero vector") + result = np.asarray(row / norm, dtype=np.float32) + second = float(np.linalg.norm(result.astype(np.float64))) + if second <= 0.0: + fail("MEMORY_DINOV2_VECTOR_INVALID", path, "float32 normalization collapsed") + result = result / np.float32(second) + return np.ascontiguousarray(result, dtype=np.float32) + + +def l2_matrix(values: Any, *, path: str) -> np.ndarray: + matrix = np.asarray(values, dtype=np.float32) + if ( + matrix.ndim != 2 + or matrix.shape[1] != DINOV2_DIMENSION + or not np.isfinite(matrix).all() + ): + fail("MEMORY_DINOV2_MATRIX_INVALID", path, "expected finite matrix[N,384]") + return ( + np.stack( + [ + l2_normalize_row(row, path=f"{path}[{index}]") + for index, row in enumerate(matrix) + ], + axis=0, + ).astype(np.float32, copy=False) + if matrix.shape[0] + else np.zeros((0, DINOV2_DIMENSION), dtype=np.float32) + ) + + +def _sha256_file(path: Path, *, label: str) -> str: + if not path.is_file(): + fail("MEMORY_DINOV2_ASSET_MISSING", label, f"missing file: {path}") + digest = hashlib.sha256() + try: + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + except OSError as exc: + fail("MEMORY_DINOV2_ASSET_UNREADABLE", label, f"{type(exc).__name__}: {exc}") + return digest.hexdigest() + + +def _safe_extract_source(source_archive: Path, destination: Path) -> Path: + try: + with tarfile.open(source_archive, mode="r:*") as archive: + members = archive.getmembers() + if not members: + fail( + "MEMORY_DINOV2_SOURCE_ARCHIVE_INVALID", + "source_archive", + "archive is empty", + ) + for member in members: + portable = PurePosixPath(member.name) + if ( + portable.is_absolute() + or ".." in portable.parts + or member.issym() + or member.islnk() + or not (member.isfile() or member.isdir()) + ): + fail( + "MEMORY_DINOV2_SOURCE_ARCHIVE_INVALID", + "source_archive", + f"unsafe archive member {member.name!r}", + ) + archive.extractall(destination) + except MemoryValidationError: + raise + except (OSError, tarfile.TarError) as exc: + fail( + "MEMORY_DINOV2_SOURCE_ARCHIVE_INVALID", + "source_archive", + f"{type(exc).__name__}: {exc}", + ) + hubconf_paths = tuple(destination.rglob("hubconf.py")) + if len(hubconf_paths) != 1: + fail( + "MEMORY_DINOV2_SOURCE_ARCHIVE_INVALID", + "source_archive", + f"expected exactly one hubconf.py, found {len(hubconf_paths)}", + ) + return hubconf_paths[0].parent + + +class _TorchDinov2Backend: + def __init__( + self, + identity: Dinov2RevisionIdentity, + deployment: Dinov2DeploymentPaths, + ) -> None: + source_sha = _sha256_file( + deployment.source_archive_path, label="source_archive" + ) + weights_sha = _sha256_file(deployment.weights_path, label="weights") + if source_sha != identity.source_archive_sha256: + fail( + "MEMORY_DINOV2_ASSET_SHA256_MISMATCH", + "source_archive", + f"expected {identity.source_archive_sha256}, actual {source_sha}", + ) + if weights_sha != identity.weights_sha256: + fail( + "MEMORY_DINOV2_ASSET_SHA256_MISMATCH", + "weights", + f"expected {identity.weights_sha256}, actual {weights_sha}", + ) + + # Heavy imports remain after complete asset validation and after the + # service entry point has set CUDA_VISIBLE_DEVICES. + torch = importlib.import_module("torch") + torchvision = importlib.import_module("torchvision") + if str(torch.__version__) != identity.torch_version: + fail( + "MEMORY_DINOV2_BACKEND_IDENTITY_MISMATCH", + "torch_version", + f"expected {identity.torch_version!r}, actual {torch.__version__!r}", + ) + if str(torchvision.__version__) != identity.torchvision_version: + fail( + "MEMORY_DINOV2_BACKEND_IDENTITY_MISMATCH", + "torchvision_version", + f"expected {identity.torchvision_version!r}, actual {torchvision.__version__!r}", + ) + if identity.device != "cuda" or not torch.cuda.is_available(): + fail( + "MEMORY_DINOV2_CUDA_UNAVAILABLE", + "device", + "the frozen encoder requires a visible CUDA device", + ) + + temporary_parent = deployment.cache_dir + if temporary_parent is None and Path("/dev/shm").is_dir(): + temporary_parent = Path("/dev/shm") + if temporary_parent is not None: + try: + temporary_parent.mkdir(parents=True, exist_ok=True) + except OSError as exc: + fail( + "MEMORY_DINOV2_CACHE_INVALID", + "cache_dir", + f"{type(exc).__name__}: {exc}", + ) + self._temporary = tempfile.TemporaryDirectory( + prefix="rpent-dinov2-source-", + dir=os.fspath(temporary_parent) if temporary_parent is not None else None, + ) + source_root = _safe_extract_source( + deployment.source_archive_path, + Path(self._temporary.name), + ) + try: + model = torch.hub.load( + os.fspath(source_root), + "dinov2_vits14", + source="local", + pretrained=False, + ) + state = torch.load( + deployment.weights_path, + map_location="cpu", + weights_only=True, + ) + model.load_state_dict(state, strict=True) + model.requires_grad_(False) + model.eval() + model.to(device="cuda") + except Exception as exc: + self._temporary.cleanup() + fail( + "MEMORY_DINOV2_MODEL_LOAD_FAILED", + "encoder.backend", + f"{type(exc).__name__}: {exc}", + ) + if model.training or any( + parameter.requires_grad for parameter in model.parameters() + ): + self._temporary.cleanup() + fail( + "MEMORY_DINOV2_MODEL_NOT_FROZEN", + "encoder.backend", + "model must be eval-only and frozen", + ) + self._torch = torch + self._functional = importlib.import_module("torchvision.transforms.functional") + transforms = importlib.import_module("torchvision.transforms") + self._bicubic = transforms.InterpolationMode.BICUBIC + self._model = model + self.torch_version = str(torch.__version__) + self.torchvision_version = str(torchvision.__version__) + self.device = "cuda" + self.eval_mode = True + self.parameters_frozen = True + self.inference_only = True + + def _preprocess(self, image: np.ndarray) -> Any: + torch = self._torch + tensor = torch.from_numpy(image).permute(2, 0, 1) + height, width = image.shape[:2] + if height <= width: + resized_height = 256 + resized_width = int(round(width * 256.0 / height)) + else: + resized_width = 256 + resized_height = int(round(height * 256.0 / width)) + tensor = self._functional.resize( + tensor, + [resized_height, resized_width], + interpolation=self._bicubic, + antialias=True, + ) + tensor = self._functional.center_crop(tensor, [224, 224]) + tensor = tensor.to(dtype=torch.float32).div_(255.0) + return self._functional.normalize( + tensor, + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225], + ) + + def encode_batch(self, images: Sequence[np.ndarray]) -> np.ndarray: + if self._model.training or any( + parameter.requires_grad for parameter in self._model.parameters() + ): + fail( + "MEMORY_DINOV2_MODEL_NOT_FROZEN", + "encoder.backend", + "model state changed after admission", + ) + batch = self._torch.stack([self._preprocess(image) for image in images]) + batch = batch.to(device="cuda", non_blocking=False) + with self._torch.inference_mode(): + output = self._model(batch) + if not isinstance(output, self._torch.Tensor): + fail( + "MEMORY_DINOV2_OUTPUT_INVALID", + "encoder.output", + f"expected Tensor, got {type(output).__name__}", + ) + return output.detach().to(device="cpu", dtype=self._torch.float32).numpy() + + def close(self) -> None: + self._model = None + self._temporary.cleanup() + + +def _default_backend_loader( + identity: Dinov2RevisionIdentity, + deployment: Dinov2DeploymentPaths, +) -> Dinov2Backend: + return _TorchDinov2Backend(identity, deployment) + + +class Dinov2Engine: + def __init__( + self, + identity: Dinov2RevisionIdentity, + deployment: Dinov2DeploymentPaths, + *, + backend_loader: BackendLoader | None = None, + ) -> None: + self._identity = identity + self._deployment = deployment + self._loader = backend_loader or _default_backend_loader + self._backend: Dinov2Backend | None = None + self._closed = False + + def revision_metadata(self) -> dict[str, Any]: + return self._identity.as_dict() + + def _backend_instance(self) -> Dinov2Backend: + if self._closed: + fail("MEMORY_DINOV2_ENCODER_CLOSED", "encoder", "encoder is closed") + if self._backend is None: + backend = self._loader(self._identity, self._deployment) + expected = { + "torch_version": self._identity.torch_version, + "torchvision_version": self._identity.torchvision_version, + "device": self._identity.device, + "eval_mode": True, + "parameters_frozen": True, + "inference_only": True, + } + for field, wanted in expected.items(): + actual = getattr(backend, field, None) + if actual != wanted: + fail( + "MEMORY_DINOV2_BACKEND_IDENTITY_MISMATCH", + field, + f"expected {wanted!r}, actual {actual!r}", + ) + self._backend = backend + return self._backend + + def encode_batch( + self, values: Sequence[np.ndarray | None] + ) -> tuple[np.ndarray | None, ...]: + if len(values) > MAX_BATCH_SIZE: + fail( + "MEMORY_DINOV2_BATCH_TOO_LARGE", + "embedding_input", + "max batch size is 32", + ) + result: list[np.ndarray | None] = [None] * len(values) + positions: list[int] = [] + images: list[np.ndarray] = [] + for index, value in enumerate(values): + if value is None: + continue + image = np.asarray(value) + if image.dtype != np.uint8 or image.ndim != 3 or image.shape[2] != 3: + fail( + "MEMORY_DINOV2_INPUT_INVALID", + f"embedding_input[{index}]", + "expected RGB8 [H,W,3]", + ) + positions.append(index) + images.append(np.ascontiguousarray(image)) + if not images: + if self._closed: + fail("MEMORY_DINOV2_ENCODER_CLOSED", "encoder", "encoder is closed") + return tuple(result) + raw = np.asarray(self._backend_instance().encode_batch(tuple(images))) + if raw.shape != (len(images), DINOV2_DIMENSION): + fail("MEMORY_DINOV2_OUTPUT_INVALID", "encoder.output", "expected [N,384]") + for row, position in enumerate(positions): + result[position] = l2_normalize_row(raw[row], path=f"encoder.output[{row}]") + return tuple(result) + + def close(self) -> None: + if self._closed: + return + self._closed = True + backend, self._backend = self._backend, None + if backend is not None: + backend.close() + + +__all__ = [ + "DINOV2_DIMENSION", + "DISTANCE_METRIC", + "EXPECTED_SOURCE_ARCHIVE_SHA256", + "Dinov2DeploymentPaths", + "Dinov2Engine", + "Dinov2RevisionIdentity", + "MemoryValidationError", + "l2_matrix", + "l2_normalize_row", +] diff --git a/robots/behavior/dino_v2/server.py b/robots/behavior/dino_v2/server.py new file mode 100644 index 000000000..09582bf25 --- /dev/null +++ b/robots/behavior/dino_v2/server.py @@ -0,0 +1,195 @@ +# 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. + +"""DINOv2 encoder RPC server for BEHAVIOR memory retrieval.""" + +from __future__ import annotations + +import argparse +import hashlib +import os +import re +import sys +import threading +from pathlib import Path +from typing import Any + +import numpy as np + +# Support direct execution from an RPent checkout before package imports. +if __package__ in (None, ""): + sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +from rpent.utils.rpc import RpcFacade # noqa: E402 + + +def _single_cuda_device(value: Any) -> str | None: + if value in (None, ""): + return None + device = str(value) + if re.fullmatch(r"[0-9]+", device) is None: + raise ValueError("--cuda-device must be one physical GPU ordinal") + return device + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _resolve_required_path(value: str | None, *, env_name: str, label: str) -> Path: + raw = value or os.environ.get(env_name) + if not raw: + raise RuntimeError( + f"DINO {label} path is required; set --{label.replace('_', '-')} " + f"or {env_name}" + ) + path = Path(raw).expanduser().resolve() + if not path.is_file(): + raise RuntimeError(f"DINO {label} path is missing: {path}") + return path + + +class BehaviorDinoFacade(RpcFacade): + """Expose the BEHAVIOR DINOv2 engine through the shared RPC facade.""" + + def __init__(self, encoder: Any, meta: dict[str, Any]) -> None: + super().__init__() + self._encoder = encoder + self._meta = dict(meta) + self._close_lock = threading.Lock() + self._closed = False + self._register_rpc() + + def _register_rpc(self) -> None: + self._rpc["dino.encode_batch"] = self.encode_batch + self._rpc["dino.get_meta"] = self.get_meta + self._readonly_methods.add("dino.get_meta") + + def get_meta(self) -> dict[str, Any]: + return {**self._meta, "pid": os.getpid()} + + def encode_batch(self, *, images: list[Any]) -> list[Any]: + result = self._encoder.encode_batch( + [ + None if image is None else np.asarray(image, dtype=np.uint8) + for image in images + ] + ) + return [ + None if item is None else np.asarray(item, dtype=np.float32) + for item in result + ] + + def close(self) -> None: + with self._close_lock: + if self._closed: + return + self._encoder.close() + self._closed = True + + +def _materialize_encoder(args: argparse.Namespace) -> tuple[Any, dict[str, Any]]: + # Heavy imports begin only after main() has applied CUDA_VISIBLE_DEVICES. + import torch + import torchvision + + from robots.behavior.dino_v2.encoder import ( + DINOV2_DIMENSION, + MODEL_ID, + MODEL_REVISION, + Dinov2DeploymentPaths, + Dinov2Engine, + Dinov2RevisionIdentity, + ) + + source_archive = _resolve_required_path( + args.source_archive, + env_name="RPENT_BEHAVIOR_DINOV2_SOURCE_ARCHIVE", + label="source_archive", + ) + weights = _resolve_required_path( + args.weights, + env_name="RPENT_BEHAVIOR_DINOV2_WEIGHTS", + label="weights", + ) + device = "cuda" if torch.cuda.is_available() else "cpu" + if device != "cuda": + raise RuntimeError( + "DINO service requires CUDA; CPU fallback is not a BEHAVIOR runtime component" + ) + identity = Dinov2RevisionIdentity( + model_id=MODEL_ID, + model_revision=MODEL_REVISION, + source_commit=MODEL_REVISION.rsplit("@", 1)[-1], + source_archive_sha256=_sha256_file(source_archive), + weights_sha256=_sha256_file(weights), + torch_version=str(torch.__version__), + torchvision_version=str(torchvision.__version__), + device=device, + ) + deployment = Dinov2DeploymentPaths( + source_archive_path=source_archive, + weights_path=weights, + cache_dir=Path(args.cache_dir).expanduser().resolve() + if args.cache_dir + else None, + ) + encoder = Dinov2Engine(identity, deployment) + # Force backend construction before get_meta advertises the deployment. + blank = np.zeros((224, 224, 3), dtype=np.uint8) + encoder.encode_batch([blank]) + return encoder, { + "status": "ok", + "runtime": "behavior_dino", + "model_id": MODEL_ID, + "model_revision": MODEL_REVISION, + "dimension": DINOV2_DIMENSION, + "device": device, + "checkpoint_binding": identity.as_dict(), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, required=True) + parser.add_argument("--cuda-device", default=None) + parser.add_argument("--source-archive", default=None) + parser.add_argument("--weights", default=None) + parser.add_argument("--cache-dir", default=None) + parser.add_argument("--parent-watch", action="store_true") + args = parser.parse_args() + cuda_device = _single_cuda_device(args.cuda_device) + if cuda_device is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = cuda_device + + encoder, meta = _materialize_encoder(args) + facade = BehaviorDinoFacade(encoder, meta) + facade.serve( + transport="http", + host=args.host, + port=args.port, + parent_watch=args.parent_watch, + ) + + +if __name__ == "__main__": + main() + + +__all__ = ["BehaviorDinoFacade", "main"] diff --git a/robots/behavior/env_client.py b/robots/behavior/env_client.py new file mode 100644 index 000000000..5c32c161c --- /dev/null +++ b/robots/behavior/env_client.py @@ -0,0 +1,332 @@ +# 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. + +"""RPC client for one BEHAVIOR environment.""" + +from __future__ import annotations + +import base64 +import copy +from typing import Any + +import numpy as np + +from robots.behavior.schemas import ( + ACTION_DIM, + validate_action_chunk, + validate_move_both_targets, + validate_move_both_visual_hand_checks, + validate_observe_request, + validate_relative_navigation_motion, +) +from robots.behavior.terminal_success import validate_official_success_receipt +from rpent.robots.components.env_client_base import BaseEnvClient +from rpent.utils.rpc import RpcClient + +_POST_SUCCESS_ALLOWED = frozenset( + { + "env.get_env_meta", + "env.current_observation", + } +) +_IMAGE_BYTE_FIELDS = frozenset( + { + "_depth_image_bytes", + "_image_bytes", + "_image_left_wrist_bytes", + "_depth_left_wrist_bytes", + "_image_right_wrist_bytes", + "_depth_right_wrist_bytes", + } +) + + +def _info_from_rpc_result(ret: Any) -> Any: + if isinstance(ret, (tuple, list)): + if len(ret) == 5: + return ret[4] + if len(ret) == 2: + return ret[1] + if isinstance(ret, dict): + return ret.get("info", ret) + return None + + +def _decode_observe_images(result: Any) -> Any: + """Decode only public image fields returned by ``env.observe``.""" + + if not isinstance(result, dict): + return result + decoded = dict(result) + for field in _IMAGE_BYTE_FIELDS: + payload = decoded.get(field) + if ( + isinstance(payload, dict) + and set(payload) == {"encoding", "data"} + and payload.get("encoding") == "base64" + and isinstance(payload.get("data"), str) + ): + decoded[field] = base64.b64decode(payload["data"], validate=True) + return decoded + + +class BehaviorEnvClient(BaseEnvClient): + """Remote implementation of the BEHAVIOR single-env protocol. + + Construction verifies only the requested immutable metadata and deliberately + does not reset the simulator. Runtime initialization owns the first reset. + """ + + _TIMEOUT_S = { + **BaseEnvClient._TIMEOUT_S, + "env.reset": 300.0, + "env.step": 60.0, + "env.chunk_step": 600.0, + "env.current_observation": 120.0, + "env.observe": 120.0, + "env.pixel_to_world": 120.0, + "env.move_to": 600.0, + "env.navigate_to": 600.0, + "env.rotate_wrist": 600.0, + "env.close_gripper": 120.0, + "env.open_gripper": 120.0, + "env.press": 300.0, + } + + def __init__(self, client: RpcClient, *, expected_meta: dict[str, Any]) -> None: + # BaseEnvClient.__init__ performs an automatic reset and exact metadata + # equality. BEHAVIOR reset is explicitly owned by runtime initialization. + self._client = client + self.last_obs: dict[str, Any] | None = None + self.last_info: dict[str, Any] = {} + self.episode_done = False + self.total_env_steps = 0 + self.vla_endpoint: str | None = None + self._official_success_latched = False + self._official_success_receipt: dict[str, Any] | None = None + server_meta = self._client.call( + "env.get_env_meta", + timeout_s=self._TIMEOUT_S["default"], + ) + if not isinstance(server_meta, dict): + raise RuntimeError(f"env_meta must be a mapping, got {type(server_meta)!r}") + mismatches = { + key: {"expected": expected, "actual": server_meta.get(key)} + for key, expected in expected_meta.items() + if server_meta.get(key) != expected + } + if mismatches: + raise RuntimeError(f"env_meta mismatch: {mismatches!r}") + self.server_meta = dict(server_meta) + + @staticmethod + def _raw_success(info: Any) -> bool: + done = info.get("done") if isinstance(info, dict) else None + value = done.get("success") if isinstance(done, dict) else None + return isinstance(value, (bool, np.bool_)) and bool(value) + + @staticmethod + def _receipt_from_info(info: Any) -> dict[str, Any] | None: + runtime = info.get("_rpent") if isinstance(info, dict) else None + if not isinstance(runtime, dict): + return None + receipt = runtime.get("official_success_receipt") + return copy.deepcopy(receipt) if isinstance(receipt, dict) else None + + def _latch_success_response(self, ret: Any) -> None: + info = _info_from_rpc_result(ret) + if not isinstance(info, dict): + return + runtime = info.get("_rpent") + if isinstance(runtime, dict): + steps = runtime.get("total_env_steps", runtime.get("global_env_steps")) + if isinstance(steps, (int, np.integer)) and not isinstance(steps, bool): + self.total_env_steps = max(self.total_env_steps, int(steps)) + if self._raw_success(info): + self.episode_done = True + self._official_success_latched = True + self._official_success_receipt = validate_official_success_receipt( + self._receipt_from_info(info) + ) + + def _rpc_call( + self, + method: str, + *, + args: tuple = (), + kwargs: dict[str, Any] | None = None, + timeout_s: float | None = None, + ) -> Any: + if self._official_success_latched and method not in _POST_SUCCESS_ALLOWED: + raise RuntimeError( + "raw task success is terminal; no further RPC is allowed" + ) + ret = self._client.call( + method, + args=args, + kwargs=kwargs or {}, + timeout_s=( + timeout_s + if timeout_s is not None + else self._TIMEOUT_S.get(method, self._TIMEOUT_S["default"]) + ), + ) + self._latch_success_response(ret) + return ret + + @property + def official_success_latched(self) -> bool: + return self._official_success_latched + + @property + def official_success_receipt(self) -> dict[str, Any] | None: + return copy.deepcopy(self._official_success_receipt) + + def reset(self) -> tuple[dict[str, Any], dict[str, Any]]: + ret = self._rpc_call("env.reset", timeout_s=self._TIMEOUT_S["env.reset"]) + if not isinstance(ret, (tuple, list)) or len(ret) != 2: + raise TypeError("env.reset must return (observation, info)") + obs, info = ret + if not isinstance(obs, dict) or not isinstance(info, dict): + raise TypeError("env.reset must return observation/info mappings") + self.total_env_steps = 0 + self.episode_done = False + self.last_obs = obs + self.last_info = info + return obs, info + + def step(self, action: Any) -> tuple[Any, Any, Any, Any, dict[str, Any]]: + array = np.asarray(action, dtype=np.float32) + if array.shape != (ACTION_DIM,) or not np.isfinite(array).all(): + raise ValueError(f"BEHAVIOR action must be finite [{ACTION_DIM}]") + ret = self._rpc_call("env.step", args=(np.ascontiguousarray(array),)) + return self._note_gym_result(ret, "env.step") + + def chunk_step( + self, + actions: Any, + *, + return_all_frames: bool = False, + ) -> tuple[Any, Any, Any, Any, dict[str, Any]]: + action_array = validate_action_chunk(actions) + ret = self._rpc_call( + "env.chunk_step", + args=(action_array,), + kwargs={"return_all_frames": bool(return_all_frames)}, + ) + return self._note_gym_result(ret, "env.chunk_step") + + def _note_gym_result(self, ret: Any, method: str) -> tuple: + if not isinstance(ret, (tuple, list)) or len(ret) != 5: + raise TypeError(f"{method} must return a gym 5-tuple") + obs, _reward, _terminated, _truncated, info = ret + if not isinstance(info, dict): + raise TypeError(f"{method} info must be a mapping") + if isinstance(obs, list): + if not obs: + raise ValueError(f"{method} returned an empty observation list") + self.last_obs = obs[-1] + elif isinstance(obs, dict): + self.last_obs = obs + else: + raise TypeError(f"{method} observation must be a mapping or list") + self.last_info = info + return tuple(ret) + + def current_observation(self) -> tuple[dict[str, Any], dict[str, Any]]: + ret = self._rpc_call("env.current_observation") + if not isinstance(ret, (tuple, list)) or len(ret) != 2: + raise TypeError("env.current_observation must return (observation, info)") + obs, info = ret + if not isinstance(obs, dict) or not isinstance(info, dict): + raise TypeError("env.current_observation returned invalid payload") + self.last_obs = obs + self.last_info = info + return obs, info + + def get_camera_meta(self, camera_name: str, **kwargs: Any) -> dict[str, Any]: + value = self._rpc_call( + "env.get_camera_meta", + kwargs={"camera_name": camera_name, **kwargs}, + ) + if not isinstance(value, dict): + raise TypeError("env.get_camera_meta must return a mapping") + return value + + def render_camera(self, camera_name: str, **kwargs: Any) -> np.ndarray: + value = self._rpc_call( + "env.render_camera", + kwargs={"camera_name": camera_name, **kwargs}, + ) + image = np.asarray(value) + if image.dtype != np.uint8 or image.ndim != 3 or image.shape[-1] != 3: + raise TypeError("env.render_camera must return uint8[H,W,3]") + return image + + def get_task_language(self) -> str: + value = self._rpc_call("env.get_task_language") + if not isinstance(value, str): + raise TypeError("env.get_task_language must return a string") + return value + + def observe(self, **kwargs: Any) -> dict[str, Any]: + request = validate_observe_request(**kwargs) + result = _decode_observe_images(self._rpc_call("env.observe", kwargs=request)) + if not isinstance(result, dict): + raise TypeError("env.observe must return a mapping") + return result + + def pixel_to_world(self, **kwargs: Any) -> dict[str, Any]: + return self._rpc_call("env.pixel_to_world", kwargs=kwargs) + + def navigate_to(self, **kwargs: Any) -> dict[str, Any]: + if "relative_motion" in kwargs and kwargs["relative_motion"] is not None: + kwargs = { + **kwargs, + "relative_motion": validate_relative_navigation_motion( + kwargs["relative_motion"] + ), + } + return self._rpc_call("env.navigate_to", kwargs=kwargs) + + def move_to(self, **kwargs: Any) -> dict[str, Any]: + if kwargs.get("hand") == "both": + kwargs = { + **kwargs, + "targets": validate_move_both_targets(kwargs.get("targets")), + "visual_hand_checks": validate_move_both_visual_hand_checks( + kwargs.get("visual_hand_checks") + ), + } + return self._rpc_call("env.move_to", kwargs=kwargs) + + def rotate_wrist(self, **kwargs: Any) -> dict[str, Any]: + return self._rpc_call("env.rotate_wrist", kwargs=kwargs) + + def close_gripper(self, **kwargs: Any) -> dict[str, Any]: + return self._rpc_call("env.close_gripper", kwargs=kwargs) + + def open_gripper(self, **kwargs: Any) -> dict[str, Any]: + return self._rpc_call("env.open_gripper", kwargs=kwargs) + + def press(self, **kwargs: Any) -> dict[str, Any]: + return self._rpc_call("env.press", kwargs=kwargs) + + def close_transport(self) -> None: + close = getattr(self._client, "close", None) + if callable(close): + close() + + +__all__ = ["BehaviorEnvClient"] diff --git a/robots/behavior/env_server.py b/robots/behavior/env_server.py new file mode 100644 index 000000000..1e5142bb1 --- /dev/null +++ b/robots/behavior/env_server.py @@ -0,0 +1,320 @@ +# 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. + +"""BEHAVIOR environment RPC adapter. + +OmniGibson scene operations stay on the process main thread. The facade uses +the common main-thread RPC serve mixin so transport handling can run on a +daemon thread while BEHAVIOR dispatch remains serialized on the calling thread. +""" + +from __future__ import annotations + +import argparse +import base64 +import os +import re +import signal +import sys +from pathlib import Path +from typing import Any + +import numpy as np + +# Support direct execution from an RPent checkout before package imports. +if __package__ in (None, ""): + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from robots.behavior.schemas import ( # noqa: E402 + ACTION_DIM, + DEFAULT_ACTION_CHUNK, + validate_action_chunk, +) +from robots.behavior.task_specs import get_task_spec # noqa: E402 +from rpent.robots.components.env_facade_base import BaseEnvFacade # noqa: E402 +from rpent.utils.rpc.main_thread_serve import MainThreadServeMixin # noqa: E402 + +_IMAGE_BYTE_FIELDS = frozenset( + { + "_depth_image_bytes", + "_image_bytes", + "_image_left_wrist_bytes", + "_depth_left_wrist_bytes", + "_image_right_wrist_bytes", + "_depth_right_wrist_bytes", + } +) + + +def _single_cuda_device(value: Any) -> str | None: + if value in (None, ""): + return None + device = str(value) + if re.fullmatch(r"[0-9]+", device) is None: + raise ValueError("--cuda-device must be one physical GPU ordinal") + return device + + +def _encode_observe_images(result: Any) -> Any: + """Encode only public image byte fields at the ``env.observe`` boundary.""" + + if not isinstance(result, dict): + return result + encoded = dict(result) + for field in _IMAGE_BYTE_FIELDS: + payload = encoded.get(field) + if isinstance(payload, bytes): + encoded[field] = { + "encoding": "base64", + "data": base64.b64encode(payload).decode("ascii"), + } + return encoded + + +class BehaviorEnvFacade(MainThreadServeMixin, BaseEnvFacade): + """Expose one official RLinf BEHAVIOR backend through common ENV RPC.""" + + def __init__(self, *, backend: Any, meta: dict[str, Any]) -> None: + self._backend = backend + self._meta = dict(meta) + self._closed = False + super().__init__() + + def _register_rpc(self) -> None: + super()._register_rpc() + self._rpc.update( + { + "env.current_observation": self.current_observation, + "env.observe": self.observe, + "env.pixel_to_world": self.pixel_to_world, + "env.navigate_to": self.navigate_to, + "env.move_to": self.move_to, + "env.rotate_wrist": self.rotate_wrist, + "env.close_gripper": self.close_gripper, + "env.open_gripper": self.open_gripper, + "env.press": self.press, + } + ) + self._readonly_methods.update( + { + "env.current_observation", + } + ) + + def _call_backend(self, name: str, *args: Any, **kwargs: Any) -> Any: + method = getattr(self._backend, name, None) + if not callable(method): + raise RuntimeError(f"backend does not expose {name}()") + return method(*args, **kwargs) + + @property + def total_env_steps(self) -> int: + value = getattr(self._backend, "total_env_steps", 0) + if isinstance(value, (int, np.integer)) and not isinstance( + value, (bool, np.bool_) + ): + return max(0, int(value)) + return 0 + + def get_env_meta(self) -> dict[str, Any]: + return dict(self._meta) + + def reset(self) -> tuple[dict[str, Any], dict[str, Any]]: + result = self._call_backend("reset") + if not isinstance(result, (tuple, list)) or len(result) != 2: + raise TypeError("backend reset must return (observation, info)") + observation, info = result + if not isinstance(observation, dict) or not isinstance(info, dict): + raise TypeError("backend reset must return observation/info mappings") + observation.setdefault("task_descriptions", self._meta["task_language"]) + return observation, info + + def step(self, action: Any) -> tuple[Any, Any, Any, Any, dict[str, Any]]: + result = self._call_backend("step", action) + return self._require_gym_result(result, "step") + + def chunk_step( + self, + actions: Any, + *, + return_all_frames: bool = False, + ) -> tuple[Any, Any, Any, Any, dict[str, Any]]: + action_array = validate_action_chunk(actions) + result = self._call_backend( + "chunk_step", + action_array, + return_all_frames=bool(return_all_frames), + ) + return self._require_gym_result(result, "chunk_step") + + @staticmethod + def _require_gym_result(result: Any, method: str) -> tuple: + if not isinstance(result, (tuple, list)) or len(result) != 5: + raise TypeError(f"env.{method} must return a gym 5-tuple") + if not isinstance(result[4], dict): + raise TypeError(f"env.{method} info must be a mapping") + return tuple(result) + + def current_observation(self) -> tuple[dict[str, Any], dict[str, Any]]: + result = self._call_backend("current_observation") + if not isinstance(result, (tuple, list)) or len(result) != 2: + raise TypeError("current_observation must return (observation, info)") + observation, info = result + if not isinstance(observation, dict) or not isinstance(info, dict): + raise TypeError("current_observation returned invalid payload") + return observation, info + + def get_task_language(self) -> str: + value = self._call_backend("get_task_language") + if not isinstance(value, str): + raise TypeError("BEHAVIOR task language must be a string") + return value + + def get_camera_meta(self, camera_name: str = "head", **kwargs: Any) -> dict: + value = self._call_backend("get_camera_meta", camera_name, **kwargs) + if not isinstance(value, dict): + raise TypeError("BEHAVIOR camera metadata must be a mapping") + return value + + def render_camera(self, camera_name: str = "head", **kwargs: Any) -> np.ndarray: + image = np.asarray(self._call_backend("render_camera", camera_name, **kwargs)) + if image.dtype != np.uint8 or image.ndim != 3 or image.shape[-1] != 3: + raise ValueError( + f"rendered RGB must be uint8[H,W,3], got {image.dtype}{image.shape}" + ) + return np.ascontiguousarray(image) + + def observe(self, **kwargs: Any) -> dict[str, Any]: + result = self._call_backend("observe", **kwargs) + if not isinstance(result, dict): + raise TypeError("env.observe must return a mapping") + return _encode_observe_images(result) + + def pixel_to_world(self, **kwargs: Any) -> dict[str, Any]: + return self._call_backend("pixel_to_world", **kwargs) + + def navigate_to(self, **kwargs: Any) -> dict[str, Any]: + return self._call_backend("navigate_to", **kwargs) + + def move_to(self, **kwargs: Any) -> dict[str, Any]: + return self._call_backend("move_to", **kwargs) + + def rotate_wrist(self, **kwargs: Any) -> dict[str, Any]: + return self._call_backend("rotate_wrist", **kwargs) + + def close_gripper(self, **kwargs: Any) -> dict[str, Any]: + if not kwargs: + raise ValueError( + "env.close_gripper requires primitive arguments; " + "lifecycle shutdown uses BehaviorEnvFacade.close()" + ) + return self._call_backend("close", **kwargs) + + def open_gripper(self, **kwargs: Any) -> dict[str, Any]: + return self._call_backend("open", **kwargs) + + def press(self, **kwargs: Any) -> dict[str, Any]: + return self._call_backend("press", **kwargs) + + def close(self) -> None: + if self._closed: + return + closer = getattr(self._backend, "close", None) + if callable(closer): + closer() + self._closed = True + + +def _build_meta(args: argparse.Namespace) -> dict[str, Any]: + task_spec = get_task_spec(args.task_name) + return { + "runtime": "behavior_env", + "task_name": task_spec.task_name, + "task": int(args.task_index), + "task_language": task_spec.task_language, + "activity_definition_id": int(args.activity_definition_id), + "activity_instance_id": int(args.activity_instance_id), + "public_seed": int(args.public_seed), + "scene_model": str(args.scene_model), + "max_episode_steps": int(args.max_episode_steps), + "action_dim": ACTION_DIM, + "action_horizon": DEFAULT_ACTION_CHUNK, + "official_success_path": ["info", "done", "success"], + "behavior_repo": str(Path(args.behavior_repo).expanduser().resolve()), + "activity_instance_dir": ( + None + if not args.activity_instance_dir + else str(Path(args.activity_instance_dir).expanduser().resolve()) + ), + "rlinf_env_config_path": ( + None + if not args.env_config_path + else str(Path(args.env_config_path).expanduser().resolve()) + ), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--task-name", required=True) + parser.add_argument("--public-seed", type=int, required=True) + parser.add_argument("--task-index", type=int, required=True) + parser.add_argument("--activity-definition-id", type=int, required=True) + parser.add_argument("--activity-instance-id", type=int, required=True) + parser.add_argument("--scene-model", required=True) + parser.add_argument("--max-episode-steps", type=int, required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--behavior-repo", required=True) + parser.add_argument("--activity-instance-dir", default=None) + parser.add_argument("--env-config-path", default=None) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, required=True) + parser.add_argument("--cuda-device", default=None) + parser.add_argument("--parent-watch", action="store_true") + args = parser.parse_args() + cuda_device = _single_cuda_device(args.cuda_device) + if cuda_device is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = cuda_device + os.environ["RPENT_RLINF_ROOT"] = str( + Path(args.behavior_repo).expanduser().resolve() + ) + + output_dir = Path(args.output_dir).expanduser().resolve() + output_dir.mkdir(parents=True, exist_ok=True) + from robots.behavior.rlinf_env import OfficialBehaviorBackend + + meta = _build_meta(args) + backend = OfficialBehaviorBackend(meta=meta, output_dir=output_dir) + facade = BehaviorEnvFacade(backend=backend, meta=meta) + # ProcessDaemon.stop sends SIGTERM. Let the serving loop finish and close + # the RLinf actor pool on its owning thread, not Ray's exit handler. + previous_sigterm = signal.signal( + signal.SIGTERM, lambda signum, frame: facade._shutdown_event.set() + ) + try: + facade.serve( + transport="http", + host=args.host, + port=args.port, + parent_watch=args.parent_watch, + ) + finally: + signal.signal(signal.SIGTERM, previous_sigterm) + + +if __name__ == "__main__": + main() + + +__all__ = ["BehaviorEnvFacade", "main"] diff --git a/robots/behavior/install_behavior_runtime.sh b/robots/behavior/install_behavior_runtime.sh new file mode 100644 index 000000000..3b396ea84 --- /dev/null +++ b/robots/behavior/install_behavior_runtime.sh @@ -0,0 +1,252 @@ +#!/usr/bin/env bash + +set -Eeuo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +RPENT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" +REPRO_ROOT="${RPENT_REPRO_ROOT:-${RPENT_ROOT}/.behavior-runtime}" +RLINF_ROOT="${RLINF_ROOT:-${REPRO_ROOT}/RLinf}" +RPENT_VENV="${RPENT_VENV:-${REPRO_ROOT}/venvs/rpent}" +BEHAVIOR_VENV="${BEHAVIOR_VENV:-${REPRO_ROOT}/venvs/behavior}" +LOG_DIR="${LOG_DIR:-${REPRO_ROOT}/logs/install}" +TOOLS_DIR="${TOOLS_DIR:-${REPRO_ROOT}/tools}" +export DOWNLOAD_DIR="${DOWNLOAD_DIR:-${REPRO_ROOT}/downloads}" +UV_VERSION="${UV_VERSION:-0.12.7}" +PYTHON_VERSION="${PYTHON_VERSION:-3.10}" +RLINF_REPO_URL="${RLINF_REPO_URL:-https://github.com/RLinf/RLinf.git}" +RLINF_COMMIT="${RLINF_COMMIT:-dd92c62857da4c67aa5e7c36f731c0d6a121f6d7}" + +mkdir -p "${LOG_DIR}" "${TOOLS_DIR}" "${DOWNLOAD_DIR}" "$(dirname "${RPENT_VENV}")" +LOG_FILE="${LOG_DIR}/install-$(date -u +%Y%m%dT%H%M%SZ).log" +exec > >(tee -a "${LOG_FILE}") 2>&1 + +trap 'echo "ERROR: command failed at line ${LINENO}. See ${LOG_FILE}" >&2' ERR + +require_command() { + if ! command -v "$1" >/dev/null 2>&1; then + echo "Missing required host command: $1" >&2 + exit 1 + fi +} + +for command_name in bash curl git nvidia-smi sha256sum; do + require_command "${command_name}" +done + +echo "RPent root: ${RPENT_ROOT}" +echo "RLinf root: ${RLINF_ROOT}" +echo "RPent venv: ${RPENT_VENV}" +echo "BEHAVIOR venv: ${BEHAVIOR_VENV}" +echo "Install log: ${LOG_FILE}" + +UV_BIN="${TOOLS_DIR}/uv" +UV_VERSION_OUTPUT="$("${UV_BIN}" --version 2>/dev/null || true)" +if [[ ! -x "${UV_BIN}" ]] || { + [[ "${UV_VERSION_OUTPUT}" != "uv ${UV_VERSION}" ]] && + [[ "${UV_VERSION_OUTPUT}" != "uv ${UV_VERSION} "* ]] +}; then + UV_INSTALLER="${TOOLS_DIR}/uv-install-${UV_VERSION}.sh" + curl -LsSf "https://astral.sh/uv/${UV_VERSION}/install.sh" -o "${UV_INSTALLER}" + env UV_UNMANAGED_INSTALL="${TOOLS_DIR}" sh "${UV_INSTALLER}" +fi +UV_VERSION_OUTPUT="$("${UV_BIN}" --version)" +if [[ "${UV_VERSION_OUTPUT}" != "uv ${UV_VERSION}" ]] && + [[ "${UV_VERSION_OUTPUT}" != "uv ${UV_VERSION} "* ]]; then + echo "Expected uv ${UV_VERSION}, got ${UV_VERSION_OUTPUT}." >&2 + exit 1 +fi +export PATH="${TOOLS_DIR}:${PATH}" + +if [[ ! -d "${RLINF_ROOT}/.git" ]]; then + mkdir -p "$(dirname "${RLINF_ROOT}")" + git clone "${RLINF_REPO_URL}" "${RLINF_ROOT}" + git -C "${RLINF_ROOT}" checkout --detach "${RLINF_COMMIT}" +else + ACTUAL_RLINF_COMMIT="$(git -C "${RLINF_ROOT}" rev-parse HEAD)" + if [[ "${ACTUAL_RLINF_COMMIT}" != "${RLINF_COMMIT}" ]]; then + echo "Existing RLinf checkout is ${ACTUAL_RLINF_COMMIT}, expected ${RLINF_COMMIT}." >&2 + echo "Use a new RLINF_ROOT; this script will not overwrite a checkout." >&2 + exit 1 + fi + if [[ -n "$(git -C "${RLINF_ROOT}" status --porcelain)" ]]; then + echo "Existing RLinf checkout is dirty; refusing to run its installer." >&2 + exit 1 + fi +fi + +if [[ ! -x "${RPENT_VENV}/bin/python" ]]; then + "${UV_BIN}" venv --python "${PYTHON_VERSION}" "${RPENT_VENV}" +fi +"${UV_BIN}" pip install --python "${RPENT_VENV}/bin/python" -e "${RPENT_ROOT}" +# Episode video writer backend; mirrors the .[behavior] extra pin. +# Do NOT switch the rpent venv to .[behavior] (keeps the CLI venv light). +"${UV_BIN}" pip install --python "${RPENT_VENV}/bin/python" "imageio-ffmpeg>=0.5" + +export UV_TORCH_BACKEND=cu124 +( + cd "${RLINF_ROOT}" + bash requirements/install.sh embodied \ + --model openpi \ + --env behavior \ + --venv "${BEHAVIOR_VENV}" \ + --install-rlinf \ + --no-flash-attn \ + --no-root +) + +BEHAVIOR_PYTHON="${BEHAVIOR_VENV}/bin/python" +if [[ ! -x "${BEHAVIOR_PYTHON}" ]]; then + echo "RLinf installer did not create ${BEHAVIOR_PYTHON}." >&2 + exit 1 +fi + +# Install RPent and all HTTP sidecar dependencies before the final compatibility pins. +# Planner packages in this environment are metadata-only for the sidecars; the +# separate RPent venv remains the canonical planner and Dashboard environment. +"${UV_BIN}" pip install --python "${BEHAVIOR_PYTHON}" -e "${RPENT_ROOT}" + +# The official BEHAVIOR runtime is validated against CUDA 12.4 and torch 2.5.1. +"${UV_BIN}" pip install --python "${BEHAVIOR_PYTHON}" --reinstall \ + --index-url https://download.pytorch.org/whl/cu124 \ + 'torch==2.5.1+cu124' \ + 'torchvision==0.20.1+cu124' \ + 'torchaudio==2.5.1+cu124' + +# Shared constraints protect the simulation stack during motion-planner installation. +FINAL_PINS=( + 'numpy==1.26.4' + 'protobuf==6.33.0' + 'ml-dtypes==0.5.3' + 'click==8.2.1' + 'llvmlite==0.48.0' + 'numba==0.66.0' + 'fastapi==0.110.0' + 'starlette==0.36.3' + 'pydantic==2.9.2' + 'pydantic-core==2.23.4' + 'uvicorn==0.52.4' + 'ray==2.55.1' + 'tensorflow-metadata==1.21.0' + 'typeguard==4.5.2' + 'gdown==6.1.0' + 'pymunk==7.3.0' + 'zarr==3.0.0a5' + 'google-api-core==2.30.3' + 'googleapis-common-protos==1.75.0' + 'proto-plus==1.28.0' + 'beautifulsoup4==4.15.0' + 'soupsieve==2.8.4' + 'asciitree==0.3.3' + 'crc32c==2.8' + 'donfig==0.8.1.post1' + 'numcodecs==0.13.1' + 'torchcodec==0.2.0' + 'rlinf-openpi==0.1.1' + 'rlinf-transformer-openpi==4.53.2' + 'lerobot==0.3.3' + 'openpi-client==0.1.2' +) +CUROBO_COMMIT=4ea77366ca48ee453e7df139e39fa6532af49f3b # v0.8.0 +CUROBO_CONSTRAINTS="${LOG_DIR}/curobo-constraints.txt" +printf '%s\n' "${FINAL_PINS[@]}" \ + 'torch==2.5.1+cu124' 'torchvision==0.20.1+cu124' 'torchaudio==2.5.1+cu124' \ + 'isaacsim==4.5.0.0' 'isaacsim-core==4.5.0.0' 'isaacsim-rl==4.5.0.0' \ + 'omnigibson==3.7.2' 'transformers==4.53.2' 'trimesh==5.1.0' \ + > "${CUROBO_CONSTRAINTS}" +"${UV_BIN}" pip install --python "${BEHAVIOR_PYTHON}" \ + --constraint "${CUROBO_CONSTRAINTS}" \ + "nvidia-curobo[cu12-torch] @ git+https://github.com/NVlabs/curobo.git@${CUROBO_COMMIT}" + +# Final compatibility repin runs after every dependency installer. --no-deps +# prevents a late resolver pass from silently changing Isaac/OpenPI versions. +"${UV_BIN}" pip install --python "${BEHAVIOR_PYTHON}" --no-deps "${FINAL_PINS[@]}" + +# OpenPI's replacement files must match the final transformers build. No package +# installation is allowed after this block. +"${UV_BIN}" pip install --python "${BEHAVIOR_PYTHON}" --no-deps \ + 'transformers==4.53.2' 'tokenizers==0.21.4' 'huggingface-hub==0.36.2' +SITE_PACKAGES="$("${BEHAVIOR_PYTHON}" -c 'import site; print(site.getsitepackages()[0])')" +TRANSFORMERS_REPLACE="${SITE_PACKAGES}/openpi/models_pytorch/transformers_replace" +if [[ ! -d "${TRANSFORMERS_REPLACE}" ]]; then + echo "Missing OpenPI transformer replacement directory: ${TRANSFORMERS_REPLACE}" >&2 + exit 1 +fi +cp -a "${TRANSFORMERS_REPLACE}/." "${SITE_PACKAGES}/transformers/" + +"${BEHAVIOR_PYTHON}" - <<'PY' +from importlib.metadata import version +from pathlib import Path +import hashlib +import site + +expected = { + "torch": "2.5.1+cu124", + "torchvision": "0.20.1+cu124", + "torchaudio": "2.5.1+cu124", + "numpy": "1.26.4", + "fastapi": "0.110.0", + "starlette": "0.36.3", + "pydantic": "2.9.2", + "transformers": "4.53.2", + "tokenizers": "0.21.4", + "rlinf-openpi": "0.1.1", + "rlinf-transformer-openpi": "4.53.2", + "lerobot": "0.3.3", +} +for package, wanted in expected.items(): + actual = version(package) + if actual != wanted: + raise RuntimeError(f"{package}: expected {wanted}, got {actual}") + +root = Path(site.getsitepackages()[0]) +replacement = root / "openpi/models_pytorch/transformers_replace" +target = root / "transformers" +files = sorted(path.relative_to(replacement) for path in replacement.rglob("*.py")) +if not files: + raise RuntimeError("OpenPI transformer replacement contains no Python files") +for relative in files: + source_hash = hashlib.sha256((replacement / relative).read_bytes()).digest() + target_hash = hashlib.sha256((target / relative).read_bytes()).digest() + if source_hash != target_hash: + raise RuntimeError(f"transformer replacement mismatch: {relative}") + +import torch +import omnigibson +import ray +import rlinf +import openpi + +if not torch.cuda.is_available(): + raise RuntimeError("torch.cuda.is_available() is false") +tensor = torch.ones(1, device="cuda") +print("CUDA smoke:", tensor, torch.cuda.get_device_name(0)) +print("Critical imports and transformer replacement: OK") +PY + +MANIFEST_DIR="${REPRO_ROOT}/manifests" +mkdir -p "${MANIFEST_DIR}" +"${UV_BIN}" pip freeze --python "${RPENT_VENV}/bin/python" \ + > "${MANIFEST_DIR}/rpent-venv.freeze.txt" +"${UV_BIN}" pip freeze --python "${BEHAVIOR_PYTHON}" \ + > "${MANIFEST_DIR}/behavior-venv.freeze.txt" +{ + echo "generated_at_utc=$(date -u +%Y-%m-%dT%H:%M:%SZ)" + echo "rpent_commit=$(git -C "${RPENT_ROOT}" rev-parse HEAD)" + echo "rpent_dirty=$(test -n "$(git -C "${RPENT_ROOT}" status --porcelain)" && echo true || echo false)" + echo "rlinf_commit=$(git -C "${RLINF_ROOT}" rev-parse HEAD)" + echo "uv_version=$("${UV_BIN}" --version)" + echo "python_version=$("${BEHAVIOR_PYTHON}" --version 2>&1)" +} > "${MANIFEST_DIR}/source-versions.txt" + +echo "Running uv dependency metadata check (report-only for upstream pin conflicts)." +"${UV_BIN}" pip check --python "${BEHAVIOR_PYTHON}" \ + > "${MANIFEST_DIR}/behavior-pip-check.txt" 2>&1 || true +cat "${MANIFEST_DIR}/behavior-pip-check.txt" + +cd "${RPENT_ROOT}" +"${RPENT_VENV}/bin/python" -m robots.behavior.selfcheck + +echo "Installation complete." +echo "Behavior Python: ${BEHAVIOR_PYTHON}" +echo "Version manifests: ${MANIFEST_DIR}" +echo "Next: export the asset variables and run behavior-download-assets --verify" diff --git a/robots/behavior/install_runtime.py b/robots/behavior/install_runtime.py new file mode 100644 index 000000000..321769715 --- /dev/null +++ b/robots/behavior/install_runtime.py @@ -0,0 +1,49 @@ +# 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. + +"""Console entry point for the reviewed BEHAVIOR dual-venv installer. + +This command requires an RPent source checkout with ``robots.behavior`` +available through editable install. +""" + +from __future__ import annotations + +import argparse +import subprocess +from collections.abc import Sequence +from pathlib import Path + + +def _build_parser() -> argparse.ArgumentParser: + return argparse.ArgumentParser( + prog="behavior-install-runtime", + description=( + "Install the reviewed BEHAVIOR dual-venv runtime. Paths and version " + "pins are controlled by the environment variables documented in " + "the BEHAVIOR usage guide." + ), + ) + + +def main(argv: Sequence[str] | None = None) -> int: + _build_parser().parse_args(argv) + script = Path(__file__).with_name("install_behavior_runtime.sh") + if not script.is_file(): + raise RuntimeError(f"packaged BEHAVIOR installer is missing: {script}") + return subprocess.run(["bash", str(script)], check=False).returncode + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/robots/behavior/memory/__init__.py b/robots/behavior/memory/__init__.py new file mode 100644 index 000000000..641bd1c0e --- /dev/null +++ b/robots/behavior/memory/__init__.py @@ -0,0 +1,75 @@ +# 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. + +"""BEHAVIOR episode-memory index and validation schema.""" + +from typing import TYPE_CHECKING, Any + +from robots.behavior.memory.schema import ( + MemoryValidationError, + canonical_json_bytes, + canonical_json_file_bytes, + require_sha256, + sha256_bytes, +) + +if TYPE_CHECKING: + from robots.behavior.memory.index import ( + EpisodeExperience, + EpisodeFrameKey, + EpisodeMemoryHit, + EpisodeMemoryIndex, + empty_episode_memory_index, + load_current_catalog, + load_revision_dir, + write_candidate_revision, + ) + +_INDEX_EXPORTS = frozenset( + { + "EpisodeExperience", + "EpisodeFrameKey", + "EpisodeMemoryHit", + "EpisodeMemoryIndex", + "empty_episode_memory_index", + "load_current_catalog", + "load_revision_dir", + "write_candidate_revision", + } +) + + +def __getattr__(name: str) -> Any: + if name not in _INDEX_EXPORTS: + raise AttributeError(name) + from robots.behavior.memory import index + + return getattr(index, name) + + +__all__ = [ + "EpisodeExperience", + "EpisodeFrameKey", + "EpisodeMemoryHit", + "EpisodeMemoryIndex", + "MemoryValidationError", + "canonical_json_bytes", + "canonical_json_file_bytes", + "empty_episode_memory_index", + "load_current_catalog", + "load_revision_dir", + "require_sha256", + "sha256_bytes", + "write_candidate_revision", +] diff --git a/robots/behavior/memory/index.py b/robots/behavior/memory/index.py new file mode 100644 index 000000000..b606473a8 --- /dev/null +++ b/robots/behavior/memory/index.py @@ -0,0 +1,767 @@ +# 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. + +"""Production episode-level BEHAVIOR memory index. + +Only head DINOv2 CLS384 keyframes are active. Wrist embeddings may be carried +for audit and shadow distances, but they never decide use vs record. +""" + +from __future__ import annotations + +import io +import json +import os +import tempfile +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from types import MappingProxyType +from typing import Any + +import numpy as np + +from robots.behavior.dino_v2.encoder import ( + DINOV2_DIMENSION, + DISTANCE_METRIC, + l2_matrix, + l2_normalize_row, +) +from robots.behavior.memory.schema import ( + MemoryValidationError, + canonical_json_file_bytes, + fail, + require_exact_keys, + require_sha256, + sha256_bytes, +) + +SCHEMA_ID = "rpent_behavior_episode_memory_index_v1" +REVISION_SCHEMA_ID = "rpent_behavior_episode_memory_revision_v1" +CURRENT_POINTER_SCHEMA_ID = "rpent_behavior_episode_memory_current_v1" +MANIFEST_SCHEMA_ID = "rpent_behavior_episode_memory_manifest_v1" +HEAD_ACTIVE_DISTANCE_MAX = 0.05367707759141922 +ACTIVE_CHANNEL = "head" +SHADOW_CHANNELS = ("left_wrist", "right_wrist") + + +def _nonempty_string(value: Any, *, path: str) -> str: + if not isinstance(value, str) or not value.strip() or "\x00" in value: + fail("MEMORY_EPISODE_SCHEMA_INVALID", path, "expected non-empty string") + return value.strip() + + +@dataclass(frozen=True, slots=True) +class EpisodeFrameKey: + frame_id: str + episode_id: str + experience_id: str + task_name: str + frame_index: int + embedding_row: int + keyframe_kind: str + source_record_id: str + frame_identity: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + for field_name in ( + "frame_id", + "episode_id", + "experience_id", + "task_name", + "keyframe_kind", + "source_record_id", + ): + _nonempty_string(getattr(self, field_name), path=f"frame.{field_name}") + if isinstance(self.frame_index, bool) or self.frame_index < 0: + fail( + "MEMORY_EPISODE_SCHEMA_INVALID", + "frame.frame_index", + "expected non-negative int", + ) + if isinstance(self.embedding_row, bool) or self.embedding_row < 0: + fail( + "MEMORY_EPISODE_SCHEMA_INVALID", + "frame.embedding_row", + "expected non-negative int", + ) + + def to_dict(self) -> dict[str, Any]: + return { + "frame_id": self.frame_id, + "episode_id": self.episode_id, + "experience_id": self.experience_id, + "task_name": self.task_name, + "frame_index": self.frame_index, + "embedding_row": self.embedding_row, + "keyframe_kind": self.keyframe_kind, + "source_record_id": self.source_record_id, + "frame_identity": dict(self.frame_identity), + } + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> "EpisodeFrameKey": + require_exact_keys( + value, + { + "frame_id", + "episode_id", + "experience_id", + "task_name", + "frame_index", + "embedding_row", + "keyframe_kind", + "source_record_id", + "frame_identity", + }, + path="frame", + ) + return cls( + frame_id=str(value["frame_id"]), + episode_id=str(value["episode_id"]), + experience_id=str(value["experience_id"]), + task_name=str(value["task_name"]), + frame_index=int(value["frame_index"]), + embedding_row=int(value["embedding_row"]), + keyframe_kind=str(value["keyframe_kind"]), + source_record_id=str(value["source_record_id"]), + frame_identity=dict(value["frame_identity"]), + ) + + +@dataclass(frozen=True, slots=True) +class EpisodeExperience: + episode_id: str + experience_id: str + logical_experience_id: str + task_name: str + usage: Mapping[str, Any] + outcome: Mapping[str, Any] + frame_keys: tuple[EpisodeFrameKey, ...] + canonical_trajectory_ref: Mapping[str, Any] | None = None + trajectory_refs: tuple[Mapping[str, Any], ...] = () + reproduction_evidence: tuple[Mapping[str, Any], ...] = () + source: Mapping[str, Any] = field(default_factory=dict) + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + for field_name in ( + "episode_id", + "experience_id", + "logical_experience_id", + "task_name", + ): + _nonempty_string(getattr(self, field_name), path=f"experience.{field_name}") + if not self.frame_keys: + fail( + "MEMORY_EPISODE_SCHEMA_INVALID", + "experience.frame_keys", + "at least one head keyframe required", + ) + for frame in self.frame_keys: + if ( + frame.episode_id != self.episode_id + or frame.experience_id != self.experience_id + or frame.task_name != self.task_name + ): + fail( + "MEMORY_EPISODE_SCHEMA_INVALID", + "experience.frame_keys", + "frame identity does not match experience", + ) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_id": SCHEMA_ID, + "episode_id": self.episode_id, + "experience_id": self.experience_id, + "logical_experience_id": self.logical_experience_id, + "task_name": self.task_name, + "usage": dict(self.usage), + "outcome": dict(self.outcome), + "canonical_trajectory_ref": None + if self.canonical_trajectory_ref is None + else dict(self.canonical_trajectory_ref), + "trajectory_refs": [dict(item) for item in self.trajectory_refs], + "reproduction_evidence": [ + dict(item) for item in self.reproduction_evidence + ], + "source": dict(self.source), + "metadata": dict(self.metadata), + "frame_keys": [frame.to_dict() for frame in self.frame_keys], + } + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> "EpisodeExperience": + require_exact_keys( + value, + { + "schema_id", + "episode_id", + "experience_id", + "logical_experience_id", + "task_name", + "usage", + "outcome", + "canonical_trajectory_ref", + "trajectory_refs", + "reproduction_evidence", + "source", + "metadata", + "frame_keys", + }, + path="experience", + ) + if value["schema_id"] != SCHEMA_ID: + fail( + "MEMORY_EPISODE_SCHEMA_INVALID", + "experience.schema_id", + "schema mismatch", + ) + frame_values = value["frame_keys"] + if not isinstance(frame_values, list): + fail( + "MEMORY_EPISODE_SCHEMA_INVALID", + "experience.frame_keys", + "expected list", + ) + return cls( + episode_id=str(value["episode_id"]), + experience_id=str(value["experience_id"]), + logical_experience_id=str(value["logical_experience_id"]), + task_name=str(value["task_name"]), + usage=dict(value["usage"]), + outcome=dict(value["outcome"]), + canonical_trajectory_ref=None + if value["canonical_trajectory_ref"] is None + else dict(value["canonical_trajectory_ref"]), + trajectory_refs=tuple(dict(item) for item in value["trajectory_refs"]), + reproduction_evidence=tuple( + dict(item) for item in value["reproduction_evidence"] + ), + source=dict(value["source"]), + metadata=dict(value["metadata"]), + frame_keys=tuple( + EpisodeFrameKey.from_mapping(item) for item in frame_values + ), + ) + + +@dataclass(frozen=True, slots=True) +class EpisodeMemoryHit: + rank: int + distance: float + matched_frame: EpisodeFrameKey + experience: EpisodeExperience + shadow_distances: Mapping[str, float] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_id": "rpent_behavior_episode_memory_hit_v1", + "rank": self.rank, + "distance": self.distance, + "distance_metric": DISTANCE_METRIC, + "threshold": HEAD_ACTIVE_DISTANCE_MAX, + "episode_id": self.experience.episode_id, + "experience_id": self.experience.experience_id, + "logical_experience_id": self.experience.logical_experience_id, + "task_name": self.experience.task_name, + "usage": dict(self.experience.usage), + "outcome": dict(self.experience.outcome), + "matched_frame": self.matched_frame.to_dict(), + "experience": self.experience.to_dict(), + "returned_scope": "whole_experience", + "stage_inference": None, + "wrist_shadow_only": True, + "shadow_distances": dict(self.shadow_distances), + } + + +class EpisodeMemoryIndex: + def __init__( + self, + *, + experiences: Sequence[EpisodeExperience], + head_embeddings: np.ndarray, + wrist_shadow_embeddings: Mapping[str, np.ndarray] | None = None, + revision: Mapping[str, Any] | None = None, + ) -> None: + self._experiences = tuple(experiences) + self._frames = tuple( + frame for exp in self._experiences for frame in exp.frame_keys + ) + self._head = l2_matrix(head_embeddings, path="head_embeddings") + if self._head.shape[0] != len(self._frames): + fail( + "MEMORY_EPISODE_INDEX_INVALID", + "head_embeddings", + "row count must equal head keyframes", + ) + self._experience_by_id = { + item.experience_id: item for item in self._experiences + } + episode_ids = {item.episode_id for item in self._experiences} + if len(self._experience_by_id) != len(self._experiences) or len( + episode_ids + ) != len(self._experiences): + fail( + "MEMORY_EPISODE_INDEX_INVALID", + "experiences", + "experience and episode IDs must be unique", + ) + by_task: dict[str, list[int]] = {} + for index, frame in enumerate(self._frames): + if frame.embedding_row != index: + fail( + "MEMORY_EPISODE_INDEX_INVALID", + "frames", + "embedding rows must be contiguous", + ) + by_task.setdefault(frame.task_name, []).append(index) + self._by_task = {task: tuple(indices) for task, indices in by_task.items()} + shadow: dict[str, np.ndarray] = {} + for channel, values in (wrist_shadow_embeddings or {}).items(): + name = str(channel) + if name not in SHADOW_CHANNELS: + fail( + "MEMORY_EPISODE_INDEX_INVALID", + f"shadow.{name}", + "only wrist shadow channels are accepted", + ) + matrix = l2_matrix(values, path=f"shadow.{name}") + if matrix.shape[0] != len(self._frames): + fail( + "MEMORY_EPISODE_INDEX_INVALID", + f"shadow.{name}", + "row count mismatch", + ) + shadow[name] = matrix + self._shadow = MappingProxyType(shadow) + self._revision = MappingProxyType(dict(revision or {})) + + @property + def episode_count(self) -> int: + return len(self._experiences) + + @property + def frame_count(self) -> int: + return len(self._frames) + + @property + def experiences(self) -> tuple[EpisodeExperience, ...]: + return self._experiences + + @property + def revision(self) -> Mapping[str, Any]: + return self._revision + + def search( + self, + *, + task_name: str, + head_embedding: np.ndarray, + k: int = 1, + wrist_shadow_embeddings: Mapping[str, np.ndarray] | None = None, + ) -> tuple[EpisodeMemoryHit, ...]: + task = _nonempty_string(task_name, path="query.task_name") + if isinstance(k, bool) or k < 1: + fail("MEMORY_EPISODE_QUERY_INVALID", "query.k", "expected positive int") + candidates = self._by_task.get(task, ()) + if not candidates: + return () + query = l2_normalize_row(head_embedding, path="query.head_embedding")[None, :] + distances = np.asarray( + 1.0 - np.clip(query @ self._head[list(candidates)].T, -1.0, 1.0), + dtype=np.float64, + )[0] + best_by_experience: dict[ + str, tuple[float, EpisodeFrameKey, dict[str, float]] + ] = {} + for offset, frame_index in enumerate(candidates): + frame = self._frames[frame_index] + shadow_distances = self._shadow_distances( + frame_index, wrist_shadow_embeddings + ) + candidate = (float(distances[offset]), frame, shadow_distances) + current = best_by_experience.get(frame.experience_id) + if current is None or (candidate[0], frame.frame_id) < ( + current[0], + current[1].frame_id, + ): + best_by_experience[frame.experience_id] = candidate + hits = [ + EpisodeMemoryHit( + rank=0, + distance=distance, + matched_frame=frame, + experience=self._experience_by_id[frame.experience_id], + shadow_distances=MappingProxyType(shadow), + ) + for distance, frame, shadow in best_by_experience.values() + ] + ordered = sorted( + hits, + key=lambda hit: ( + hit.distance, + hit.experience.experience_id, + hit.matched_frame.frame_id, + ), + ) + return tuple( + EpisodeMemoryHit( + rank=index, + distance=hit.distance, + matched_frame=hit.matched_frame, + experience=hit.experience, + shadow_distances=hit.shadow_distances, + ) + for index, hit in enumerate(ordered[:k], start=1) + ) + + def retrieve( + self, + *, + task_name: str, + head_embedding: np.ndarray, + wrist_shadow_embeddings: Mapping[str, np.ndarray] | None = None, + ) -> Mapping[str, Any]: + hits = self.search( + task_name=task_name, + head_embedding=head_embedding, + k=max(1, self.episode_count), + wrist_shadow_embeddings=wrist_shadow_embeddings, + ) + selected = next( + (hit for hit in hits if hit.distance <= HEAD_ACTIVE_DISTANCE_MAX), None + ) + return MappingProxyType( + { + "schema_id": "rpent_behavior_episode_memory_retrieval_v1", + "decision": "use_experience" if selected is not None else "record_new", + "reason": "head_keyframe_under_active_threshold" + if selected is not None + else "no_same_task_head_keyframe_under_active_threshold", + "task_filter_applied_before_vision": True, + "active_channel": ACTIVE_CHANNEL, + "head_active_distance_max": HEAD_ACTIVE_DISTANCE_MAX, + "wrist_shadow_only": True, + "hit": None if selected is None else selected.to_dict(), + "stage_inference": None, + "candidate_count_after_task_filter": len( + self._by_task.get(str(task_name).strip(), ()) + ), + } + ) + + def _shadow_distances( + self, + frame_index: int, + queries: Mapping[str, np.ndarray] | None, + ) -> dict[str, float]: + result: dict[str, float] = {} + for channel, query in (queries or {}).items(): + name = str(channel) + if name not in self._shadow or query is None: + continue + row = l2_normalize_row(query, path=f"query.{name}")[None, :] + result[name] = float( + 1.0 + - np.clip( + row @ self._shadow[name][frame_index : frame_index + 1].T, -1.0, 1.0 + )[0, 0] + ) + return result + + +def empty_episode_memory_index() -> EpisodeMemoryIndex: + return EpisodeMemoryIndex( + experiences=(), + head_embeddings=np.zeros((0, DINOV2_DIMENSION), dtype=np.float32), + revision={ + "schema_id": REVISION_SCHEMA_ID, + "empty_catalog_reason": "memory_dir_omitted", + "activation_allowed": False, + }, + ) + + +def load_current_catalog(memory_dir: Path | None) -> EpisodeMemoryIndex: + """Load the current catalog; omitted memory_dir is the only legal empty catalog.""" + + if memory_dir is None: + return empty_episode_memory_index() + root = Path(memory_dir) + if not root.is_dir(): + fail( + "MEMORY_EPISODE_CATALOG_MISSING", + str(root), + "explicit memory-dir is missing", + ) + pointer_path = root / "current.json" + pointer = _read_json(pointer_path) + require_exact_keys( + pointer, {"schema_id", "revision_document_sha256"}, path="current.json" + ) + if pointer["schema_id"] != CURRENT_POINTER_SCHEMA_ID: + fail("MEMORY_EPISODE_POINTER_INVALID", "current.json", "schema mismatch") + revision_sha = require_sha256( + pointer["revision_document_sha256"], path="current.revision_document_sha256" + ) + revision_dir = root / "revisions" / revision_sha + return load_revision_dir(revision_dir, expected_revision_sha256=revision_sha) + + +def load_revision_dir( + revision_dir: Path, *, expected_revision_sha256: str | None = None +) -> EpisodeMemoryIndex: + if not revision_dir.is_dir(): + fail( + "MEMORY_EPISODE_REVISION_MISSING", + str(revision_dir), + "revision directory missing", + ) + manifest = _read_json(revision_dir / "manifest.json") + require_exact_keys( + manifest, + { + "schema_id", + "revision_document_sha256", + "catalog_sha256", + "embeddings_npz_sha256", + "experience_count", + "frame_count", + }, + path="manifest.json", + ) + if manifest["schema_id"] != MANIFEST_SCHEMA_ID: + fail("MEMORY_EPISODE_MANIFEST_INVALID", "manifest.schema_id", "schema mismatch") + revision_sha = require_sha256( + manifest["revision_document_sha256"], path="manifest.revision_document_sha256" + ) + if ( + expected_revision_sha256 is not None + and revision_sha != expected_revision_sha256 + ): + fail( + "MEMORY_EPISODE_HASH_MISMATCH", + "manifest.revision_document_sha256", + "current pointer mismatch", + ) + revision_bytes = _read_regular(revision_dir / "revision.json") + if sha256_bytes(revision_bytes) != revision_sha: + fail( + "MEMORY_EPISODE_HASH_MISMATCH", "revision.json", "document digest mismatch" + ) + catalog_bytes = _read_regular(revision_dir / "catalog.jsonl") + if sha256_bytes(catalog_bytes) != require_sha256( + manifest["catalog_sha256"], path="manifest.catalog_sha256" + ): + fail("MEMORY_EPISODE_HASH_MISMATCH", "catalog.jsonl", "catalog digest mismatch") + embeddings_bytes = _read_regular(revision_dir / "embeddings.npz") + if sha256_bytes(embeddings_bytes) != require_sha256( + manifest["embeddings_npz_sha256"], path="manifest.embeddings_npz_sha256" + ): + fail( + "MEMORY_EPISODE_HASH_MISMATCH", + "embeddings.npz", + "embedding digest mismatch", + ) + revision = json.loads(revision_bytes.decode("utf-8")) + experiences = tuple( + EpisodeExperience.from_mapping(json.loads(line.decode("utf-8"))) + for line in catalog_bytes.splitlines() + if line + ) + with np.load(io.BytesIO(embeddings_bytes), allow_pickle=False) as data: + head = np.asarray(data["head"], dtype=np.float32) + shadow = { + name: np.asarray(data[name], dtype=np.float32) + for name in SHADOW_CHANNELS + if name in data.files + } + index = EpisodeMemoryIndex( + experiences=experiences, + head_embeddings=head, + wrist_shadow_embeddings=shadow, + revision=revision, + ) + if index.episode_count != int( + manifest["experience_count"] + ) or index.frame_count != int(manifest["frame_count"]): + fail("MEMORY_EPISODE_MANIFEST_INVALID", "manifest.counts", "count mismatch") + return index + + +def write_candidate_revision( + *, + memory_dir: Path, + experiences: Sequence[EpisodeExperience], + head_embeddings: np.ndarray, + wrist_shadow_embeddings: Mapping[str, np.ndarray] | None = None, + encoder_identity: Mapping[str, Any] | None = None, + parent_revision_document_sha256: str | None = None, + activate_current: bool = True, +) -> Mapping[str, Any]: + """Validate, write content-addressed revision, then atomically advance current.""" + + root = Path(memory_dir) + root.mkdir(parents=True, exist_ok=True) + candidate_index = EpisodeMemoryIndex( + experiences=experiences, + head_embeddings=head_embeddings, + wrist_shadow_embeddings=wrist_shadow_embeddings, + ) + catalog_bytes = b"".join( + canonical_json_file_bytes(exp.to_dict(), path=f"experience[{index}]") + for index, exp in enumerate(candidate_index.experiences) + ) + embedding_payload = _npz_bytes( + {"head": candidate_index._head, **dict(candidate_index._shadow)} + ) + catalog_sha = sha256_bytes(catalog_bytes) + embeddings_sha = sha256_bytes(embedding_payload) + revision = { + "schema_id": REVISION_SCHEMA_ID, + "format_version": 1, + "preliminary": True, + "activation_allowed": False, + "active_thresholds": {"head_distance_max": HEAD_ACTIVE_DISTANCE_MAX}, + "distance_metric": DISTANCE_METRIC, + "active_channel": ACTIVE_CHANNEL, + "wrist_policy": "shadow_only", + "encoder_identity": dict(encoder_identity or {}), + "parent_revision_document_sha256": parent_revision_document_sha256, + "catalog_sha256": catalog_sha, + "embeddings_npz_sha256": embeddings_sha, + "experience_count": candidate_index.episode_count, + "frame_count": candidate_index.frame_count, + } + revision_bytes = canonical_json_file_bytes(revision, path="revision") + revision_sha = sha256_bytes(revision_bytes) + revision_dir = root / "revisions" / revision_sha + _write_revision_dir( + revision_dir, + revision_bytes=revision_bytes, + catalog_bytes=catalog_bytes, + embedding_bytes=embedding_payload, + manifest={ + "schema_id": MANIFEST_SCHEMA_ID, + "revision_document_sha256": revision_sha, + "catalog_sha256": catalog_sha, + "embeddings_npz_sha256": embeddings_sha, + "experience_count": candidate_index.episode_count, + "frame_count": candidate_index.frame_count, + }, + ) + load_revision_dir(revision_dir, expected_revision_sha256=revision_sha) + pointer = { + "schema_id": CURRENT_POINTER_SCHEMA_ID, + "revision_document_sha256": revision_sha, + } + if activate_current: + _atomic_write( + root / "current.json", canonical_json_file_bytes(pointer, path="current") + ) + return MappingProxyType( + { + "revision_document_sha256": revision_sha, + "revision_dir": str(revision_dir), + "current": bool(activate_current), + } + ) + + +def _npz_bytes(arrays: Mapping[str, np.ndarray]) -> bytes: + with io.BytesIO() as buffer: + np.savez( + buffer, + **{ + name: np.asarray(value, dtype=np.float32) + for name, value in arrays.items() + }, + ) + return buffer.getvalue() + + +def _read_regular(path: Path) -> bytes: + if path.is_symlink() or not path.is_file(): + fail("MEMORY_EPISODE_SOURCE_INVALID", str(path), "expected regular file") + return path.read_bytes() + + +def _read_json(path: Path) -> Mapping[str, Any]: + try: + value = json.loads(_read_regular(path).decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + fail("MEMORY_EPISODE_SOURCE_INVALID", str(path), str(exc)) + if not isinstance(value, Mapping): + fail("MEMORY_EPISODE_SOURCE_INVALID", str(path), "expected JSON object") + return value + + +def _atomic_write(path: Path, payload: bytes) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + mode="wb", prefix=f".{path.name}.", dir=path.parent, delete=False + ) as handle: + tmp = Path(handle.name) + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + try: + os.replace(tmp, path) + finally: + if tmp.exists(): + tmp.unlink() + + +def _write_new(path: Path, payload: bytes) -> None: + if path.exists(): + if path.is_file() and not path.is_symlink() and path.read_bytes() == payload: + return + fail("MEMORY_EPISODE_OUTPUT_COLLISION", str(path), "existing bytes differ") + _atomic_write(path, payload) + + +def _write_revision_dir( + revision_dir: Path, + *, + revision_bytes: bytes, + catalog_bytes: bytes, + embedding_bytes: bytes, + manifest: Mapping[str, Any], +) -> None: + revision_dir.mkdir(parents=True, exist_ok=True) + _write_new(revision_dir / "revision.json", revision_bytes) + _write_new(revision_dir / "catalog.jsonl", catalog_bytes) + _write_new(revision_dir / "embeddings.npz", embedding_bytes) + _write_new( + revision_dir / "manifest.json", + canonical_json_file_bytes(dict(manifest), path="manifest"), + ) + + +__all__ = [ + "ACTIVE_CHANNEL", + "HEAD_ACTIVE_DISTANCE_MAX", + "EpisodeExperience", + "EpisodeFrameKey", + "EpisodeMemoryHit", + "EpisodeMemoryIndex", + "MemoryValidationError", + "empty_episode_memory_index", + "load_current_catalog", + "load_revision_dir", + "write_candidate_revision", +] diff --git a/robots/behavior/memory/schema.py b/robots/behavior/memory/schema.py new file mode 100644 index 000000000..d4296f33f --- /dev/null +++ b/robots/behavior/memory/schema.py @@ -0,0 +1,82 @@ +# 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. + +"""Small deterministic schema helpers for BEHAVIOR episode memory.""" + +from __future__ import annotations + +import hashlib +import json +import re +from collections.abc import Mapping +from typing import Any + +SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$") + + +class MemoryValidationError(ValueError): + """Fail-closed validation error with a stable code and path.""" + + def __init__(self, code: str, path: str, detail: str) -> None: + super().__init__(f"{code}: {path}: {detail}") + self.code = code + self.path = path + self.detail = detail + + +def fail(code: str, path: str, detail: str) -> None: + raise MemoryValidationError(code, path, detail) + + +def require_sha256(value: Any, *, path: str) -> str: + if not isinstance(value, str) or SHA256_PATTERN.fullmatch(value) is None: + fail("MEMORY_SCHEMA_INVALID", path, "expected one lowercase SHA-256 digest") + return value + + +def canonical_json_bytes(value: Any, *, path: str = "$") -> bytes: + try: + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ).encode("utf-8") + except (TypeError, ValueError) as exc: + fail("MEMORY_JSON_INVALID", path, f"{type(exc).__name__}: {exc}") + + +def canonical_json_file_bytes(value: Any, *, path: str = "$") -> bytes: + return canonical_json_bytes(value, path=path) + b"\n" + + +def sha256_bytes(payload: bytes) -> str: + return hashlib.sha256(payload).hexdigest() + + +def require_exact_keys( + value: Mapping[str, Any], + expected: set[str] | frozenset[str], + *, + path: str, +) -> None: + actual = set(value) + expected_set = set(expected) + if actual != expected_set: + fail( + "MEMORY_SCHEMA_INVALID", + path, + f"expected keys {sorted(expected_set)}, actual {sorted(actual)}", + ) diff --git a/robots/behavior/motion.py b/robots/behavior/motion.py new file mode 100644 index 000000000..6daa948c5 --- /dev/null +++ b/robots/behavior/motion.py @@ -0,0 +1,382 @@ +# 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. + +"""Read kinematics in the existing RLinf simulation actor's owning thread. + +Ray's built-in actor call executes this importable function serially alongside +reset/step. No simulator object crosses the process boundary. Actions continue +through OfficialBehaviorBackend.chunk_step, including official-success checks. +""" + +from __future__ import annotations + +import numpy as np + + +def get_camera_observation(process, env_index: int = 0) -> dict: + """Capture synchronized RGB, optical-axis depth and calibration without stepping.""" + import omnigibson as og + + from robots.behavior.rlinf_env import _torch_to_numpy + + env = process.env.envs[env_index] + robot = env.robots[0] + cameras = {} + for camera, link in ( + ("head", "zed_link"), + ("left_wrist", "left_realsense_link"), + ("right_wrist", "right_realsense_link"), + ): + sensors = [ + sensor for name, sensor in robot.sensors.items() if f":{link}:" in name + ] + if len(sensors) != 1: + raise ValueError(f"expected one {camera} sensor, found {len(sensors)}") + sensor = sensors[0] + sensor.add_modality("depth_linear") + # Initialize the camera-parameter annotator before flushing render latency. + sensor.camera_parameters + cameras[camera] = sensor + env.load_observation_space() + for _ in range(3): + og.sim.render() + result = {} + for camera, sensor in cameras.items(): + obs, _ = sensor.get_obs() + params = sensor.camera_parameters + view = ( + np.asarray(_torch_to_numpy(params["cameraViewTransform"])).reshape(4, 4).T + ) + result[camera] = { + "rgb": np.asarray(_torch_to_numpy(obs["rgb"]))[..., :3], + "depth": np.asarray(_torch_to_numpy(obs["depth_linear"])).squeeze(), + "intrinsic": np.asarray(_torch_to_numpy(sensor.intrinsic_matrix)), + "camera_to_world": np.linalg.inv(view), + } + return result + + +def get_planning_state(process, env_index: int = 0) -> dict: + """Read the actor's current articulation and conservative collision geometry.""" + from pathlib import Path + + from robots.behavior.rlinf_env import _torch_to_numpy + + env = process.env.envs[env_index] + robot = env.robots[0] + result = get_motion_state(process, env_index) + q = np.asarray(_torch_to_numpy(robot.get_joint_positions())) + result["joint_positions"] = dict(zip(robot.joints, q.tolist())) + result["urdf_path"] = robot.urdf_path + model_dir = Path(robot.urdf_path).parent.parent + result["collision_config_path"] = str( + model_dir / "curobo" / "r1pro_description_curobo_arm_no_torso.yaml" + ) + pos, quat = robot.links["base_link"].get_position_orientation() + result["base_position"] = np.asarray(_torch_to_numpy(pos)) + result["base_quaternion_xyzw"] = np.asarray(_torch_to_numpy(quat)) + result["robot_aabb"] = [np.asarray(_torch_to_numpy(x)) for x in robot.aabb] + result["obstacles"] = {} + for obj in env.scene.objects: + if obj is robot: + continue + for name, link in obj.links.items(): + low, high = (np.asarray(_torch_to_numpy(x)) for x in link.aabb) + if np.all(high > low): + result["obstacles"][f"{obj.name}_{name}"] = {"low": low, "high": high} + return result + + +def navigation_collision(state: dict, goal: np.ndarray) -> str | None: + """Conservative swept footprint for a short holonomic base segment.""" + low, high = map(np.asarray, state["robot_aabb"]) + position = np.asarray(state["base_position"]) + radius = float( + np.linalg.norm(np.maximum(high[:2] - position[:2], position[:2] - low[:2])) + ) + for fraction in np.linspace( + 0, 1, max(2, int(np.linalg.norm(goal[:2] - position[:2]) / 0.01) + 1) + ): + point = position[:2] + fraction * (goal[:2] - position[:2]) + for name, obstacle in state["obstacles"].items(): + a, b = np.asarray(obstacle["low"]), np.asarray(obstacle["high"]) + if b[2] <= low[2] or a[2] >= high[2]: + continue + nearest = np.clip(point, a[:2], b[:2]) + if np.linalg.norm(nearest - point) <= radius: + return name + return None + + +def build_robot_config(state: dict, hands=("left", "right")) -> dict: + """Adapt the installed R1Pro collision model to cuRobo 0.8's URDF loader.""" + import xml.etree.ElementTree as ET + + import yaml + from scipy.spatial.transform import Rotation + + with open(state["collision_config_path"]) as source: + original = yaml.safe_load(source)["robot_cfg"]["kinematics"] + xml_joints = ET.parse(state["urdf_path"]).getroot().findall("joint") + joints = {j.attrib["name"] for j in xml_joints if j.attrib["type"] != "fixed"} + active = [f"{hand}_arm_joint{i}" for hand in hands for i in range(1, 8)] + # R1Pro's holonomic-base import fixes the six URDF wheel/steering joints + # at their zero transform; they are not DOFs in the live articulation. + fixed_wheels = { + f"{kind}_motor_joint{i}" for kind in ("wheel", "steer") for i in range(1, 4) + } + locked = { + name: state["joint_positions"][name] + for name in joints - set(active) - fixed_wheels + } + kinematics = { + name: original[name] + for name in ( + "collision_link_names", + "collision_spheres", + "collision_sphere_buffer", + "self_collision_buffer", + "self_collision_ignore", + "extra_links", + "extra_collision_spheres", + ) + } + # cuRobo 0.8's loader computes padding compensation in a local variable + # but passes the uncompensated dict to SelfCollisionKinematicsCfg. Restore + # its intended semantics here: external collision padding remains intact; + # self-collision uses the original spheres plus the explicit self margins. + buffer = original["collision_sphere_buffer"] + kinematics["self_collision_buffer"] = { + name: original["self_collision_buffer"].get(name, 0.0) + - (buffer if isinstance(buffer, (float, int)) else buffer.get(name, 0.0)) + for name in original["collision_link_names"] + } + for joint in xml_joints: + if joint.attrib["name"] not in fixed_wheels: + continue + name = joint.find("child").attrib["link"] + origin = joint.find("origin") + xyz = [float(x) for x in origin.attrib.get("xyz", "0 0 0").split()] + rpy = [float(x) for x in origin.attrib.get("rpy", "0 0 0").split()] + quat = Rotation.from_euler("xyz", rpy).as_quat()[[3, 0, 1, 2]].tolist() + kinematics["extra_links"][name] = { + "link_name": name, + "parent_link_name": joint.find("parent").attrib["link"], + "joint_name": joint.attrib["name"], + "joint_type": "FIXED", + "fixed_transform": xyz + quat, + } + for hand in ("left", "right"): + name = f"{hand}_eef_link" + kinematics["extra_links"][name] = { + "link_name": name, + "parent_link_name": f"{hand}_gripper_link", + "joint_name": name + "_joint", + "joint_type": "FIXED", + # r1pro_source_cfg.yaml: xyzw [0,1,0,0], translated to wxyz here. + "fixed_transform": [0, 0, -0.06, 0, 0, 1, 0], + } + kinematics.update( + base_link="base_link", + urdf_path=state["urdf_path"], + tool_frames=[f"{hand}_eef_link" for hand in hands], + lock_joints=locked, + cspace={ + "joint_names": active, + "default_joint_position": [state["joint_positions"][n] for n in active], + "cspace_distance_weight": [ + original["cspace"]["cspace_distance_weight"][ + original["cspace"]["joint_names"].index(n) + ] + for n in active + ], + "null_space_weight": [ + original["cspace"]["null_space_weight"][ + original["cspace"]["joint_names"].index(n) + ] + for n in active + ], + "max_acceleration": 5.0, + "max_jerk": 100.0, + "velocity_scale": 0.25, + }, + ) + return {"kinematics": kinematics} + + +class BehaviorMotionPlanner: + """cuRobo 0.8 arm planner, constructed lazily on the ENV facade thread. + + The live torso, unselected arm and grippers are locked. The world uses + conservative link AABBs; a blocked path is a failure, never a teleport. + """ + + def __init__(self): + self._planner = None + self._lock_key = None + + def close(self): + if self._planner is not None: + self._planner.destroy() + self._planner = None + self._lock_key = None + + def plan(self, state: dict, targets: dict) -> dict: + import torch + from curobo.motion_planner import MotionPlanner, MotionPlannerCfg + from curobo.scene import Scene + from curobo.types import GoalToolPose, JointState, Pose + from scipy.spatial.transform import Rotation + + config = build_robot_config(state, tuple(targets)) + kinematics = config["kinematics"] + locks = kinematics["lock_joints"] + key = ( + tuple(targets), + tuple((n, round(v, 3)) for n, v in sorted(locks.items())), + ) + base_rotation = Rotation.from_quat(state["base_quaternion_xyzw"]) + inverse = base_rotation.inv() + orientation = inverse.as_quat()[[3, 0, 1, 2]].tolist() + cuboids = {} + for name, obstacle in state["obstacles"].items(): + low, high = np.asarray(obstacle["low"]), np.asarray(obstacle["high"]) + center = inverse.apply((low + high) / 2 - state["base_position"]) + cuboids[name] = { + "dims": (high - low).tolist(), + "pose": center.tolist() + orientation, + } + scene = {"cuboid": cuboids} + if self._planner is None or self._lock_key != key: + self.close() + self._planner = MotionPlanner( + MotionPlannerCfg.create( + robot={"robot_cfg": config}, + scene_model=scene, + collision_cache={"cuboid": max(len(cuboids), 1)}, + self_collision_check=True, + max_batch_size=1, + max_goalset=1, + ) + ) + self._lock_key = key + else: + self._planner.update_world(Scene.create(scene)) + planner = self._planner + q = torch.tensor( + [[state["joint_positions"][n] for n in planner.joint_names]], + device="cuda", + dtype=torch.float32, + ) + start = JointState.from_position(q, joint_names=planner.joint_names) + fk = planner.compute_kinematics(start) + goals = {} + for hand, target in targets.items(): + name = f"{hand}_eef_link" + actual = state["hands"][hand] + predicted = fk.tool_poses.get_link_pose(name) + predicted_world = ( + base_rotation.apply( + predicted.position.detach().cpu().numpy().reshape(3) + ) + + state["base_position"] + ) + if np.linalg.norm(predicted_world - actual["position"]) > 0.005: + raise ValueError(f"{hand}: URDF/live kinematics mismatch") + relative = inverse.apply( + np.asarray(target["position"]) - state["base_position"] + ) + quat = (inverse * Rotation.from_quat(target["quaternion_xyzw"])).as_quat()[ + [3, 0, 1, 2] + ] + goals[name] = Pose( + torch.tensor(relative[None], device="cuda", dtype=torch.float32), + torch.tensor(quat[None], device="cuda", dtype=torch.float32), + ) + goal = GoalToolPose.from_poses(goals, ordered_tool_frames=planner.tool_frames) + result = planner.plan_pose(goal, start, max_attempts=1) + if result is None or not bool(result.success.all()): + return { + "success": False, + "stop_reason": "planning_failed", + "details": str(getattr(result, "status", "no trajectory")), + } + trajectory = result.get_interpolated_plan().reorder(planner.joint_names) + positions = ( + trajectory.position.detach() + .cpu() + .numpy() + .reshape(-1, len(planner.joint_names)) + ) + dt = float(planner.trajopt_solver.config.interpolation_dt) + return { + "success": True, + "positions": positions, + "joint_names": planner.joint_names, + "dt": dt, + } + + +def get_motion_state(process, env_index: int = 0) -> dict: + import omnigibson as og + from omnigibson.utils.transform_utils import quat2mat + + from robots.behavior.rlinf_env import _torch_to_numpy + + env = process.env.envs[env_index] + robot = env.robots[0] + controls = robot.get_control_dict() + joint_positions = robot.get_joint_positions() + _, base_quat = robot.get_position_orientation() + + def array(value): + return np.asarray(_torch_to_numpy(value)).copy() + + rotation = array(quat2mat(base_quat)) + result = {"control_dt": float(og.sim.get_sim_step_dt()), "hands": {}} + for hand in ("left", "right"): + controller = robot.controllers[f"arm_{hand}"] + if controller.motor_type != "position" or controller.use_delta_commands: + raise RuntimeError("motion requires absolute position arm controllers") + idx = robot.arm_control_idx[hand] + pos, quat = robot.get_eef_pose(hand) + jacobian = array(controls[f"eef_{hand}_jacobian_relative"][:, idx]) + world_jacobian = np.concatenate( + (rotation @ jacobian[:3], rotation @ jacobian[3:]), axis=0 + ) + palm_pos, _ = robot.links[f"{hand}_gripper_link"].get_position_orientation() + direction = array(pos - palm_pos) + length = float(np.linalg.norm(direction)) + if length < 1e-6: + raise RuntimeError("EEF and gripper origins do not define an approach axis") + contacts = sorted( + { + str(body) + for link in robot.finger_links[hand] + for contact in link.contact_list() + for body in (contact.body0, contact.body1) + if not str(body).startswith(robot.prim_path + "/") + } + ) + result["hands"][hand] = { + "position": array(pos), + "quaternion_xyzw": array(quat), + "joint_positions": array(joint_positions[idx]), + "joint_lower_limits": array(robot.joint_lower_limits[idx]), + "joint_upper_limits": array(robot.joint_upper_limits[idx]), + "jacobian": world_jacobian, + "approach_direction": direction / length, + "contacts": contacts, + } + return result diff --git a/robots/behavior/policy_checkpoint.py b/robots/behavior/policy_checkpoint.py new file mode 100644 index 000000000..ffe83aeb8 --- /dev/null +++ b/robots/behavior/policy_checkpoint.py @@ -0,0 +1,230 @@ +# 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. + +"""Identity contract for the shared BEHAVIOR Pi0.5 checkpoint.""" + +from __future__ import annotations + +import hashlib +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping + +POLICY_CHECKPOINT_BINDING_SCHEMA_VERSION = 1 +POLICY_CHECKPOINT_ENV = "PI05_CHECKPOINT_PATH" +SHARED_POLICY_PROFILE_ID = "pi05-b1kpt50-cs32" +SHARED_POLICY_CHECKPOINT_PATH = Path( + os.environ.get(POLICY_CHECKPOINT_ENV, SHARED_POLICY_PROFILE_ID) +) + + +class PolicyCheckpointError(ValueError): + """Raised when a checkpoint violates the shared BEHAVIOR contract.""" + + +@dataclass(frozen=True) +class CheckpointFileRequirement: + relative_path: str + size_bytes: int + sha256: str + + +@dataclass(frozen=True) +class PolicyCheckpointProfile: + profile_id: str + files: tuple[CheckpointFileRequirement, ...] + + +@dataclass(frozen=True) +class PolicyCheckpointBinding: + schema_version: int + profile_id: str + resolved_path: str + files: tuple[CheckpointFileRequirement, ...] + binding_sha256: str + + def as_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "profile_id": self.profile_id, + "resolved_path": self.resolved_path, + "files": { + item.relative_path: { + "size_bytes": item.size_bytes, + "sha256": item.sha256, + } + for item in self.files + }, + "binding_sha256": self.binding_sha256, + } + + +SHARED_POLICY_PROFILE = PolicyCheckpointProfile( + profile_id=SHARED_POLICY_PROFILE_ID, + files=( + CheckpointFileRequirement( + relative_path="model.safetensors", + size_bytes=7_233_650_408, + sha256="7e257666d835f6af701de493676a6c86a0421b2efc737a0f911d782b7a09f635", + ), + CheckpointFileRequirement( + relative_path="config.json", + size_bytes=149, + sha256="a4ae208203adfdd64c5fdbd4b0dc257e4ebbc82e464cb146dd0377051b25fc0a", + ), + CheckpointFileRequirement( + relative_path="assets/behavior-1k/2025-challenge-demos/norm_stats.json", + size_bytes=6_368, + sha256="d66ed16830a98f90dde8a315058b4a0df59f5e05734c1686d8b3f66787d0a929", + ), + ), +) + + +def _canonical_sha256(value: Mapping[str, Any]) -> str: + return hashlib.sha256( + json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + ).hexdigest() + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _binding_payload( + profile: PolicyCheckpointProfile, + resolved_path: Path, +) -> dict[str, Any]: + return { + "schema_version": POLICY_CHECKPOINT_BINDING_SCHEMA_VERSION, + "profile_id": profile.profile_id, + "resolved_path": str(resolved_path), + "files": { + item.relative_path: { + "size_bytes": item.size_bytes, + "sha256": item.sha256, + } + for item in profile.files + }, + } + + +def validate_policy_checkpoint( + path: str | Path = SHARED_POLICY_CHECKPOINT_PATH, +) -> PolicyCheckpointBinding: + """Verify and bind the expected Pi05-Behavior checkpoint files.""" + + profile = SHARED_POLICY_PROFILE + requested = Path(path).expanduser() + try: + resolved = requested.resolve(strict=True) + except OSError as error: + raise PolicyCheckpointError( + f"your Pi05-Behavior model checkpoint is unavailable: {error}" + ) from error + if not resolved.is_dir(): + raise PolicyCheckpointError( + f"your Pi05-Behavior model checkpoint is not a directory: {resolved}" + ) + for requirement in profile.files: + candidate = resolved / requirement.relative_path + if candidate.is_symlink() or not candidate.is_file(): + raise PolicyCheckpointError( + "your Pi05-Behavior model checkpoint file is missing or unsafe: " + f"{candidate}" + ) + size = candidate.stat().st_size + if size != requirement.size_bytes: + raise PolicyCheckpointError( + "your Pi05-Behavior model checkpoint size mismatch for " + f"{requirement.relative_path}: expected {requirement.size_bytes}, " + f"got {size}" + ) + actual_sha256 = _file_sha256(candidate) + if actual_sha256 != requirement.sha256: + raise PolicyCheckpointError( + "your Pi05-Behavior model checkpoint SHA256 mismatch for " + f"{requirement.relative_path}: expected {requirement.sha256}, " + f"got {actual_sha256}" + ) + payload = _binding_payload(profile, resolved) + return PolicyCheckpointBinding( + schema_version=POLICY_CHECKPOINT_BINDING_SCHEMA_VERSION, + profile_id=profile.profile_id, + resolved_path=str(resolved), + files=profile.files, + binding_sha256=_canonical_sha256(payload), + ) + + +def write_policy_checkpoint_manifest( + destination: str | Path, + path: str | Path = SHARED_POLICY_CHECKPOINT_PATH, +) -> PolicyCheckpointBinding: + """Validate the checkpoint and write the shared server manifest.""" + + binding = validate_policy_checkpoint(path) + target = Path(destination).expanduser() + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text( + json.dumps(binding.as_dict(), sort_keys=True, indent=2) + "\n", + encoding="utf-8", + ) + return binding + + +def assert_matching_policy_checkpoint_binding( + actual: Mapping[str, Any] | None, + expected: PolicyCheckpointBinding | Mapping[str, Any], +) -> dict[str, Any]: + expected_value = ( + expected.as_dict() + if isinstance(expected, PolicyCheckpointBinding) + else dict(expected) + ) + if not isinstance(actual, Mapping): + raise PolicyCheckpointError("VLA health metadata lacks checkpoint_binding") + actual_value = dict(actual) + if actual_value != expected_value: + raise PolicyCheckpointError( + "VLA checkpoint binding does not match your Pi05-Behavior model" + ) + return actual_value + + +__all__ = [ + "POLICY_CHECKPOINT_BINDING_SCHEMA_VERSION", + "POLICY_CHECKPOINT_ENV", + "SHARED_POLICY_CHECKPOINT_PATH", + "SHARED_POLICY_PROFILE", + "SHARED_POLICY_PROFILE_ID", + "CheckpointFileRequirement", + "PolicyCheckpointBinding", + "PolicyCheckpointError", + "PolicyCheckpointProfile", + "assert_matching_policy_checkpoint_binding", + "validate_policy_checkpoint", + "write_policy_checkpoint_manifest", +] diff --git a/robots/behavior/prompt_bundle.py b/robots/behavior/prompt_bundle.py new file mode 100644 index 000000000..37ad232bb --- /dev/null +++ b/robots/behavior/prompt_bundle.py @@ -0,0 +1,46 @@ +# 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. + +"""BEHAVIOR prompt bundle assembly.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from robots.behavior.prompts import eval as eval_parts +from robots.behavior.prompts import explore as explore_parts +from robots.behavior.prompts import user as user_parts +from rpent.prompt.utils import PromptNode + + +def system_prompt(variables: Mapping[str, object] | None = None) -> PromptNode: + """Assemble the BEHAVIOR system prompt for the selected run mode.""" + mode = (variables or {}).get("behavior_mode", "eval") + if mode == "eval": + return eval_parts.system_prompt() + if mode == "explore": + return explore_parts.system_prompt() + raise ValueError(f"unsupported BEHAVIOR prompt mode: {mode!r}") + + +def user_prompt(variables: Mapping[str, object] | None = None) -> PromptNode: + """Assemble the BEHAVIOR user prompt tree.""" + return { + "CELL": user_parts.CELL, + "MODE": user_parts.MODE, + "BEGIN": user_parts.BEGIN, + } + + +__all__ = ["system_prompt", "user_prompt"] diff --git a/robots/behavior/prompts/__init__.py b/robots/behavior/prompts/__init__.py new file mode 100644 index 000000000..16fd5e775 --- /dev/null +++ b/robots/behavior/prompts/__init__.py @@ -0,0 +1,15 @@ +# 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. + +"""Prompt sections for the BEHAVIOR robot plugin.""" diff --git a/robots/behavior/prompts/eval.py b/robots/behavior/prompts/eval.py new file mode 100644 index 000000000..3d188b720 --- /dev/null +++ b/robots/behavior/prompts/eval.py @@ -0,0 +1,46 @@ +# 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. + +"""BEHAVIOR single-attempt evaluation prompt.""" + +from __future__ import annotations + +from robots.behavior.prompts import system as base +from rpent.prompt.utils import PromptNode + +ROLE_AND_EVALUATION = """You are the planner for one BEHAVIOR evaluation +episode. The invocation controls only that episode. Pursue the exact task while +remaining honest about official success.""" + +MEMORY = """Memory access is read-only in Eval. Read only relevant material +through the official MemoryManager tools; do not write to the corpus.""" + + +def system_prompt() -> PromptNode: + """Assemble the complete BEHAVIOR Eval prompt.""" + return { + "ROLE AND EVALUATION": ROLE_AND_EVALUATION, + "CURRENT INVOCATION": base.CURRENT_INVOCATION, + "INVOCATION MODEL": base.INVOCATION_MODEL, + "RUNTIME": base.RUNTIME, + "YOUR GOAL": base.GOAL, + "MEMORY": f"{base.MEMORY_CONTEXT}\n\n{MEMORY}", + "EVIDENCE": base.EVIDENCE, + "PLANNER TOOLS": base.PLANNER_TOOLS, + "TERMINATION": base.TERMINATION, + "OUTPUT DISCIPLINE": base.OUTPUT_DISCIPLINE, + } + + +__all__ = ["MEMORY", "ROLE_AND_EVALUATION", "system_prompt"] diff --git a/robots/behavior/prompts/explore.py b/robots/behavior/prompts/explore.py new file mode 100644 index 000000000..7b6cab4be --- /dev/null +++ b/robots/behavior/prompts/explore.py @@ -0,0 +1,56 @@ +# 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. + +"""BEHAVIOR single-attempt exploration prompt.""" + +from __future__ import annotations + +from robots.behavior.prompts import system as base +from rpent.prompt.utils import PromptNode + +ROLE_AND_MODE = """You are the planner for one BEHAVIOR Explore attempt. This +invocation owns exactly one episode. The standard RPent Explore session loop, +not the planner, starts any later attempt.""" + +MEMORY = """Write provisional memory only under `{{memory_inbox}}/wip/`. +After current official success, write `{{output_dir}}/{{recipe_tag}}.json` as +the task audit: task, seed, the verified success receipt, and the winning +session's actual command sequence and strategy notes. This audit is a run +output, not a memory-corpus write, and is separate from the terminal receipt. +Do not claim commands or success absent from the +current receipts. The runner exports `recipe_{{recipe_tag}}.jsonl` after finish. +Only after official success may reusable notes be promoted to root-level +Markdown drafts in `{{memory_inbox}}`; follow the existing MemoryManager +frontmatter schema. Unsolved notes stay in wip and are not published. +The main CLI merges the audit/recipe pair and valid memory drafts after Explore +finishes when `--auto-merge-memory` is enabled. Do not invoke merge yourself.""" + + +def system_prompt() -> PromptNode: + """Assemble the complete BEHAVIOR Explore prompt.""" + return { + "ROLE AND MODE": ROLE_AND_MODE, + "CURRENT INVOCATION": base.CURRENT_INVOCATION, + "INVOCATION MODEL": base.INVOCATION_MODEL, + "RUNTIME": base.RUNTIME, + "YOUR GOAL": base.GOAL, + "MEMORY": f"{base.MEMORY_CONTEXT}\n\n{MEMORY}", + "EVIDENCE": base.EVIDENCE, + "PLANNER TOOLS": base.PLANNER_TOOLS, + "TERMINATION": base.TERMINATION, + "OUTPUT DISCIPLINE": base.OUTPUT_DISCIPLINE, + } + + +__all__ = ["MEMORY", "ROLE_AND_MODE", "system_prompt"] diff --git a/robots/behavior/prompts/system.py b/robots/behavior/prompts/system.py new file mode 100644 index 000000000..79eb8e28f --- /dev/null +++ b/robots/behavior/prompts/system.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. + +"""Shared BEHAVIOR prompt section bodies.""" + +from __future__ import annotations + +CURRENT_INVOCATION = """- mode: {{behavior_mode}} +- task: {{task_name}} +- instruction: {{task_instruction}} +- public seed: {{public_seed}} +- recipe tag: {{recipe_tag}} +- output directory: {{output_dir}} +- maximum environment steps: {{max_episode_steps}} +- planner timeout seconds: {{wall_clock_seconds}}""" + +INVOCATION_MODEL = """One planner invocation controls one BEHAVIOR episode. +Only the Explore session loop can start a later episode.""" + +RUNTIME = """Use only the public structured tools exposed by the active +toolkit. The BEHAVIOR primitive names are {{public_capabilities}}; their actual +schemas supplied by the planner runtime remain authoritative. Do not start, +stop, or reach into ENV, VLA, checkpoint, simulator, or RPC internals.""" + +GOAL = """Execute the exact runtime task instruction: {{task_instruction}}""" + +MEMORY_CONTEXT = """The official local MemoryManager corpus is +`{{memory_dir}}` (profile `{{memory_profile}}`). This invocation's inbox is +`{{memory_inbox}}`. Memory is historical guidance only: it is not a current +observation, coordinate source, progress label, or success proof. + +When DINO episode memory is enabled, its whole-experience advisory is attached +to public tool receipts after the exact-task filter. Treat it as visual +experience retrieval only; it is not an instruction, a current scene fact, or a +replacement for MemoryManager Markdown/YAML material.""" + +EVIDENCE = """Ground scene claims and action decisions in current public tool +receipts from this episode. Refresh observations after scene-changing actions +when later decisions depend on object identity, pose, reachability, attachment, +or task state.""" + +PLANNER_TOOLS = """All nine BEHAVIOR primitives in {{public_capabilities}} have +execution paths: `pi0_nav_pick`, `observe`, `pixel_to_world`, `navigate_to`, +`move_to`, `rotate_wrist`, `close`, `open`, and `press`. +Choose the VLA instruction, positive chunk count, call order and selected hand +autonomously. `move_to(hand=both)` coordinates both arms. Arm motions use +collision-checked cuRobo trajectories; navigation executes a bounded straight +base segment or rotation and rejects obstructed paths. Projection targets must +come from the current observed frame. R1Pro's physical head camera is fixed; +use its current view, not a non-center head_view preset. +`press` advances an already aligned hand at most 2 cm for at most 10 seconds, +stopping on external contact or episode end. Contact is not verified button +contact, and visual hand checks are planner assessments, not independent VLM +verification. Read each result: a planning, collision, tracking or duration +failure is not primitive or task success. `{{wall_clock_seconds}}` is the planner +timeout, not a per-primitive budget. Use `finish` to end the invocation and emit +its terminal receipt.""" + +TERMINATION = """Official task success exists only when the current episode +returns `info[\"done\"][\"success\"] is True`. Reward, terminated, truncated, +primitive success, visual appearance, video, and a clean process exit cannot +substitute for that value. A verified receipt may carry this evidence but cannot +create it.""" + +OUTPUT_DISCIPLINE = """Keep the final result tied to this invocation. Distinguish +official task success from primitive progress, interruption, termination, +truncation, and workflow completion. Never present historical memory as current +evidence.""" + +__all__ = [ + "CURRENT_INVOCATION", + "EVIDENCE", + "GOAL", + "INVOCATION_MODEL", + "MEMORY_CONTEXT", + "OUTPUT_DISCIPLINE", + "PLANNER_TOOLS", + "RUNTIME", + "TERMINATION", +] diff --git a/robots/behavior/prompts/user.py b/robots/behavior/prompts/user.py new file mode 100644 index 000000000..1acde3c0a --- /dev/null +++ b/robots/behavior/prompts/user.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. + +"""User prompt sections for one BEHAVIOR invocation.""" + +from __future__ import annotations + +CELL = """- task: {{task_name}} +- seed: {{public_seed}} +- instruction: {{task_instruction}} +- output_dir: {{output_dir}} +- audit: {{output_dir}}/{{recipe_tag}}.json +- recipe: {{output_dir}}/recipe_{{recipe_tag}}.jsonl +- memory: {{memory_dir}}""" + +MODE = """BEHAVIOR {{behavior_mode}}; one fresh episode per planner invocation; +official success requires current `info[\"done\"][\"success\"] is True`. +DINO episode memory, when enabled, appears only as advisory content in public +tool receipts.""" + +BEGIN = """Execute the selected task using the active public tools. Base each +action on current public evidence and report the outcome honestly.""" + +__all__ = ["BEGIN", "CELL", "MODE"] diff --git a/robots/behavior/rlinf_env.py b/robots/behavior/rlinf_env.py new file mode 100644 index 000000000..8c7c83d29 --- /dev/null +++ b/robots/behavior/rlinf_env.py @@ -0,0 +1,2024 @@ +# 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. + +"""Bundled official BEHAVIOR backend for the RPent env RPC server. + +This module is intentionally independent from the historical RPent BEHAVIOR +runtime helpers. It builds an RLinf ``BehaviorEnv`` config, owns the single +live env instance, and exposes the narrow duck-typed surface consumed by +``robots.behavior.env_server.BehaviorEnvFacade``. +""" + +from __future__ import annotations + +import base64 +import hashlib +import io +import json +import os +import sys +import time +from collections.abc import Mapping +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np + +from robots.behavior.schemas import ENV_ACTION_SEGMENTS, RAW_PROPRIO_SEGMENTS +from robots.behavior.terminal_success import official_success_receipt_sha256 +from rpent.utils.config import get_repo_root, get_rlinf_repo_path + +ACTION_DIM = 23 +ACTION_HORIZON = 32 +GRIPPER_COMMAND_CONTROL_CYCLES = 15 +GRIPPER_OPEN_COMMAND = 1.0 +GRIPPER_CLOSE_COMMAND = -1.0 +PHYSICAL_CAMERAS = ("head", "left_wrist", "right_wrist") +EXACT_OFFICIAL_CONFIG_MODE = "exact_official_v1" +EXACT_OFFICIAL_RUNTIME_SUPPORT_SCHEMA = ( + "rlinf.behavior.exact_official_runtime_support.v1" +) +EXACT_OFFICIAL_OVERLAY_SCHEMA = "rlinf.behavior.exact_official_overlay.v1" +EXACT_OFFICIAL_WRAPPER_SELECTOR = "official_rgb_v1" +RLINF_ROOT_ENV = "RPENT_RLINF_ROOT" +RLINF_ENV_CONFIG_ENV = "RPENT_BEHAVIOR_RLINF_ENV_CONFIG" +ACTIVITY_INSTANCE_DIR_ENV = "RPENT_BEHAVIOR_ACTIVITY_INSTANCE_DIR" +ACTIVITY_INSTANCE_FORMAT_ENV = "RPENT_BEHAVIOR_ACTIVITY_INSTANCE_FORMAT" +EXACT_CONFIG_ENV = "RPENT_BEHAVIOR_EXACT_OFFICIAL_CONFIG" +RESET_TRACE_ENV = "RLINF_BEHAVIOR_RESET_TRACE" +_COMPLETE_EXACT_FIELDS = { + "omni_config_mode", + "omni_config", + "omni_config_semantic_sha256", + "omni_config_runtime_support", + "omni_config_runtime_support_sha256", + "omni_config_effective_overlay", + "omni_config_effective_overlay_sha256", + "omni_config_effective_sha256", +} + + +def discover_rlinf_root() -> Path: + """Return the RLinf checkout that contains the official BehaviorEnv.""" + + root = (get_rlinf_repo_path() or (get_repo_root().parent / "RLinf")).resolve() + if (root / "rlinf" / "envs" / "behavior" / "behavior_env.py").is_file(): + return root + raise FileNotFoundError( + "could not locate RLinf behavior_env.py; set " + f"{RLINF_ROOT_ENV} to the RLinf checkout. searched: {root}" + ) + + +def ensure_rlinf_import_path() -> Path: + """Put the selected RLinf checkout on sys.path and return it.""" + + root = discover_rlinf_root() + root_text = str(root) + if root_text not in sys.path: + sys.path.insert(0, root_text) + return root + + +def _canonical_json_sha256(value: Any) -> str: + payload = json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _read_structured_file(path: Path) -> Any: + suffix = path.suffix.lower() + if suffix == ".json": + return json.loads(path.read_text(encoding="utf-8")) + from omegaconf import OmegaConf + + cfg = OmegaConf.load(path) + return OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True) + + +def _coerce_positive_int(value: Any, *, field: str) -> int: + if isinstance(value, (bool, np.bool_)): + raise ValueError(f"{field} must be a positive integer") + try: + result = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{field} must be a positive integer") from exc + if result <= 0: + raise ValueError(f"{field} must be a positive integer") + return result + + +def _coerce_nonnegative_int(value: Any, *, field: str) -> int: + if isinstance(value, (bool, np.bool_)): + raise ValueError(f"{field} must be a non-negative integer") + try: + result = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{field} must be a non-negative integer") from exc + if result < 0: + raise ValueError(f"{field} must be a non-negative integer") + return result + + +def _require_text(meta: Mapping[str, Any], field: str) -> str: + value = meta.get(field) + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"meta[{field!r}] must be a non-empty string") + return value.strip() + + +def _task_identity(meta: Mapping[str, Any]) -> dict[str, Any]: + return { + "task_name": _require_text(meta, "task_name"), + "task_language": _require_text(meta, "task_language"), + "activity_definition_id": _coerce_nonnegative_int( + meta.get("activity_definition_id"), + field="activity_definition_id", + ), + "activity_instance_id": _coerce_nonnegative_int( + meta.get("activity_instance_id"), + field="activity_instance_id", + ), + "public_seed": _coerce_nonnegative_int( + meta.get("public_seed", 0), + field="public_seed", + ), + "scene_model": _require_text(meta, "scene_model"), + "max_episode_steps": _coerce_positive_int( + meta.get("max_episode_steps"), + field="max_episode_steps", + ), + } + + +def _resolution(value: Any, default: tuple[int, int]) -> list[int]: + if value is None: + return [int(default[0]), int(default[1])] + if not isinstance(value, (list, tuple)) or len(value) != 2: + raise ValueError("camera resolution must contain two positive integers") + return [ + _coerce_positive_int(value[0], field="camera resolution"), + _coerce_positive_int(value[1], field="camera resolution"), + ] + + +def _exact_runtime_support( + *, + official: Mapping[str, Any], + meta: Mapping[str, Any], +) -> dict[str, Any]: + camera_cfg = ( + official.get("camera") if isinstance(official.get("camera"), Mapping) else {} + ) + return { + "schema_version": EXACT_OFFICIAL_RUNTIME_SUPPORT_SCHEMA, + "source_profile_sha256": _canonical_json_sha256(official), + "wrapper_selector": EXACT_OFFICIAL_WRAPPER_SELECTOR, + "macro": { + "use_gpu_dynamics": bool(meta.get("use_gpu_dynamics", False)), + "headless": bool(meta.get("headless", True)), + "enable_flatcache": bool(meta.get("enable_flatcache", True)), + "enable_object_states": bool(meta.get("enable_object_states", True)), + "enable_transition_rules": bool(meta.get("enable_transition_rules", True)), + "render_viewer_camera": bool(meta.get("render_viewer_camera", False)), + "use_numpy_controller_backend": bool( + meta.get("use_numpy_controller_backend", True) + ), + }, + "camera": { + "head_resolution": _resolution( + meta.get("head_resolution", camera_cfg.get("head_resolution")), + (720, 720), + ), + "wrist_resolution": _resolution( + meta.get("wrist_resolution", camera_cfg.get("wrist_resolution")), + (480, 480), + ), + }, + } + + +def _exact_overlay( + official: Mapping[str, Any], +) -> tuple[dict[str, Any], dict[str, Any]]: + official_copy = json.loads( + json.dumps(official, ensure_ascii=False, allow_nan=False) + ) + env_cfg = official_copy.setdefault("env", {}) + task_cfg = official_copy.setdefault("task", {}) + termination = task_cfg.setdefault("termination_config", {}) + source_max_steps = _coerce_positive_int( + termination.get("max_steps"), + field="task.termination_config.max_steps", + ) + flatten_source = env_cfg.get("flatten_obs_space") + if flatten_source is not True: + raise ValueError( + "exact official omni_config.env.flatten_obs_space must be True " + "so RLinf can apply its reviewed runtime overlay" + ) + overlay = { + "schema_version": EXACT_OFFICIAL_OVERLAY_SCHEMA, + "changes": { + "env.flatten_obs_space": {"source": True, "effective": False}, + "task.termination_config.max_steps": { + "source": source_max_steps, + "effective": source_max_steps - 1, + }, + }, + } + effective = json.loads( + json.dumps(official_copy, ensure_ascii=False, allow_nan=False) + ) + effective["env"]["flatten_obs_space"] = False + effective["task"]["termination_config"]["max_steps"] = source_max_steps - 1 + return overlay, effective + + +def _assert_official_identity( + official: Mapping[str, Any], + meta: Mapping[str, Any], +) -> None: + identity = _task_identity(meta) + task_cfg = official.get("task") + scene_cfg = official.get("scene") + if not isinstance(task_cfg, Mapping) or not isinstance(scene_cfg, Mapping): + raise ValueError("exact official config must contain task and scene mappings") + mismatches = { + "task.activity_name": ( + identity["task_name"], + task_cfg.get("activity_name"), + ), + "task.activity_definition_id": ( + identity["activity_definition_id"], + task_cfg.get("activity_definition_id"), + ), + "task.activity_instance_id": ( + identity["activity_instance_id"], + task_cfg.get("activity_instance_id"), + ), + "scene.scene_model": ( + identity["scene_model"], + scene_cfg.get("scene_model"), + ), + } + bad = { + key: {"expected": expected, "actual": actual} + for key, (expected, actual) in mismatches.items() + if actual != expected + } + if bad: + raise ValueError(f"exact official config identity mismatch: {bad}") + + +def _exact_config_from_official( + official: Mapping[str, Any], + meta: Mapping[str, Any], + output_dir: Path, +) -> dict[str, Any]: + official_dict = json.loads( + json.dumps(official, ensure_ascii=False, allow_nan=False) + ) + identity = _task_identity(meta) + _assert_official_identity(official_dict, meta) + support = dict( + meta.get("omni_config_runtime_support") or {} + ) or _exact_runtime_support( + official=official_dict, + meta=meta, + ) + overlay = meta.get("omni_config_effective_overlay") + effective: dict[str, Any] | None = None + if isinstance(overlay, Mapping): + effective = None + overlay = json.loads(json.dumps(overlay, ensure_ascii=False, allow_nan=False)) + else: + overlay, effective = _exact_overlay(official_dict) + if effective is None: + effective = json.loads( + json.dumps(official_dict, ensure_ascii=False, allow_nan=False) + ) + changes = dict(overlay["changes"]) + effective["env"]["flatten_obs_space"] = changes["env.flatten_obs_space"][ + "effective" + ] + effective["task"]["termination_config"]["max_steps"] = changes[ + "task.termination_config.max_steps" + ]["effective"] + return { + "env_type": "behavior", + "total_num_envs": 1, + "auto_reset": False, + "ignore_terminations": False, + "use_rel_reward": True, + "seed": identity["public_seed"], + "group_size": 1, + "use_fixed_reset_state_ids": False, + "max_steps_per_rollout_epoch": identity["max_episode_steps"], + "max_episode_steps": identity["max_episode_steps"], + "skip_intermediate_obs_in_chunk": True, + "num_env_subprocess": 1, + "direct_omnigibson_env": True, + "video_cfg": { + "save_video": False, + "info_on_video": True, + "video_base_dir": str(output_dir / "video"), + }, + "base_config_name": "r1pro_behavior", + "use_eval_utils_cfg": False, + "policy_wrapper": None, + "omni_config_mode": EXACT_OFFICIAL_CONFIG_MODE, + "omni_config": official_dict, + "omni_config_semantic_sha256": str( + meta.get("omni_config_semantic_sha256") + or _canonical_json_sha256(official_dict) + ), + "omni_config_runtime_support": support, + "omni_config_runtime_support_sha256": str( + meta.get("omni_config_runtime_support_sha256") + or _canonical_json_sha256(support) + ), + "omni_config_effective_overlay": overlay, + "omni_config_effective_overlay_sha256": str( + meta.get("omni_config_effective_overlay_sha256") + or _canonical_json_sha256(overlay) + ), + "omni_config_effective_sha256": str( + meta.get("omni_config_effective_sha256") + or _canonical_json_sha256(effective) + ), + "action_trace_path": str(output_dir / "behavior_action_trace.jsonl"), + "action_trace_interval": 1, + } + + +def _load_exact_official_config(meta: Mapping[str, Any]) -> Mapping[str, Any] | None: + if meta.get("omni_config_mode") == EXACT_OFFICIAL_CONFIG_MODE and isinstance( + meta.get("omni_config"), + Mapping, + ): + return meta + + path_value = ( + meta.get("exact_official_config_path") + or meta.get("official_omni_config_path") + or os.environ.get(EXACT_CONFIG_ENV) + ) + if not path_value: + return None + loaded = _read_structured_file(Path(str(path_value)).expanduser().resolve()) + if not isinstance(loaded, Mapping): + raise ValueError("exact official config file must contain a mapping") + return loaded + + +def _default_env_config_path(rlinf_root: Path, meta: Mapping[str, Any]) -> Path: + path_value = meta.get("rlinf_env_config_path") or os.environ.get( + RLINF_ENV_CONFIG_ENV + ) + if path_value: + return Path(str(path_value)).expanduser().resolve() + return ( + rlinf_root + / "examples" + / "embodiment" + / "config" + / "env" + / "behavior_r1pro.yaml" + ) + + +def _bootstrap_template_path( + instance_dir: Path, + *, + scene_model: str, + task_name: str, + activity_definition_id: int, +) -> Path: + """Resolve the full instance-0 scene used before applying a TRO delta. + + RLinf's ``ActivityInstanceLoader`` applies ``*_template-tro_state.json`` + only immediately before reset. OmniGibson therefore needs a complete + same-task template to construct the object scope first. The official + challenge dataset stores that bootstrap template beside the task-specific + ``*_instances`` directory. + """ + + template_name = ( + f"{scene_model}_task_{task_name}_{activity_definition_id}_0_template.json" + ) + candidates = (instance_dir / template_name, instance_dir.parent / template_name) + template_path = next((path for path in candidates if path.is_file()), None) + if template_path is None: + raise FileNotFoundError( + "BEHAVIOR bootstrap scene template not found: " + + " or ".join(str(path) for path in candidates) + ) + return template_path + + +def _resolve_activity_instance_dir( + activity_dir: Path, + *, + scene_model: str, + task_name: str, + activity_definition_id: int, + activity_instance_id: int, +) -> Path: + """Resolve a task instance directory from either supported CLI shape. + + The public launcher accepts the downloaded challenge dataset root, while + RLinf's loader requires the task-specific directory that directly contains + the cached instance JSON files. An already task-specific directory remains + valid for callers that supply one explicitly. + """ + + instance_name = ( + f"{scene_model}_task_{task_name}_{activity_definition_id}_" + f"{activity_instance_id}_template-tro_state.json" + ) + task_dir_name = f"{scene_model}_task_{task_name}_instances" + candidates = ( + activity_dir, + activity_dir / task_dir_name, + activity_dir / "scenes" / scene_model / "json" / task_dir_name, + ) + instance_dir = next( + (path for path in candidates if (path / instance_name).is_file()), + None, + ) + if instance_dir is None: + raise FileNotFoundError( + "BEHAVIOR activity instance not found: " + + " or ".join(str(path / instance_name) for path in candidates) + ) + return instance_dir + + +def _apply_default_config_identity( + cfg: Any, + *, + identity: Mapping[str, Any], + output_dir: Path, + meta: Mapping[str, Any], +) -> Any: + from omegaconf import OmegaConf + + cfg.env_type = "behavior" + cfg.total_num_envs = 1 + cfg.auto_reset = False + cfg.ignore_terminations = False + cfg.use_fixed_reset_state_ids = False + cfg.seed = int(identity["public_seed"]) + cfg.direct_omnigibson_env = True + cfg.num_env_subprocess = 1 + cfg.max_episode_steps = int(identity["max_episode_steps"]) + cfg.max_steps_per_rollout_epoch = int(identity["max_episode_steps"]) + cfg.skip_intermediate_obs_in_chunk = True + cfg.video_cfg.save_video = False + cfg.video_cfg.video_base_dir = str(output_dir / "video") + cfg.omni_config.env.env_wrapper = str(meta.get("env_wrapper") or "rgb") + cfg.omni_config.env.flatten_obs_space = False + cfg.omni_config.env.flatten_action_space = False + cfg.omni_config.env.automatic_reset = False + cfg.omni_config.task.activity_name = str(identity["task_name"]) + cfg.omni_config.task.activity_definition_id = int( + identity["activity_definition_id"] + ) + cfg.omni_config.task.activity_instance_id = int(identity["activity_instance_id"]) + cfg.omni_config.task.online_object_sampling = False + cfg.omni_config.task.termination_config.max_steps = int( + identity["max_episode_steps"] + ) + cfg.omni_config.scene.scene_model = str(identity["scene_model"]) + + activity_dir = meta.get("activity_instance_dir") or os.environ.get( + ACTIVITY_INSTANCE_DIR_ENV + ) + if activity_dir: + instance_dir = _resolve_activity_instance_dir( + Path(str(activity_dir)).expanduser().resolve(), + scene_model=str(identity["scene_model"]), + task_name=str(identity["task_name"]), + activity_definition_id=int(identity["activity_definition_id"]), + activity_instance_id=int(identity["activity_instance_id"]), + ) + cfg.omni_config.task.activity_instance_dir = str(instance_dir) + cfg.omni_config.task.instance_resample_mode = "disabled" + instance_file_format = str( + meta.get("activity_instance_file_format") + or os.environ.get(ACTIVITY_INSTANCE_FORMAT_ENV) + or "tro_state" + ) + cfg.omni_config.task.instance_file_format = instance_file_format + cfg.omni_config.task.use_presampled_robot_pose = bool( + meta.get("use_presampled_robot_pose", True) + ) + if instance_file_format == "tro_state": + # A TRO-state file is a delta, not an OmniGibson scene template. + # Bootstrap the same task's object scope from the official full + # instance-0 template; ActivityInstanceLoader applies the selected + # native instance immediately before the first reset. + cfg.omni_config.scene.scene_file = str( + _bootstrap_template_path( + instance_dir, + scene_model=str(identity["scene_model"]), + task_name=str(identity["task_name"]), + activity_definition_id=int(identity["activity_definition_id"]), + ) + ) + cfg.omni_config.scene.scene_instance = None + + for key, default in ( + ("head_resolution", (720, 720)), + ("wrist_resolution", (480, 480)), + ): + if meta.get(key) is not None: + OmegaConf.update( + cfg, + f"omni_config.camera.{key}", + _resolution(meta.get(key), default), + merge=False, + ) + + cfg.action_trace_path = str(output_dir / "behavior_action_trace.jsonl") + cfg.action_trace_interval = 1 + return cfg + + +def build_behavior_env_config(meta: Mapping[str, Any], output_dir: str | Path) -> Any: + """Build the RLinf BehaviorEnv config without launching simulation. + + If an exact official config is supplied through ``meta`` or + ``RPENT_BEHAVIOR_EXACT_OFFICIAL_CONFIG``, this returns an + ``exact_official_v1`` RLinf config. Otherwise it loads RLinf's canonical + ``behavior_r1pro.yaml`` and applies the task/instance identity from RPent. + """ + + from omegaconf import OmegaConf + + output_path = Path(output_dir).expanduser().resolve() + identity = _task_identity(meta) + exact_loaded = _load_exact_official_config(meta) + if exact_loaded is not None: + if exact_loaded.get( + "omni_config_mode" + ) == EXACT_OFFICIAL_CONFIG_MODE and _COMPLETE_EXACT_FIELDS.issubset( + exact_loaded + ): + cfg_dict = dict(exact_loaded) + official = cfg_dict.get("omni_config") + if not isinstance(official, Mapping): + raise ValueError("exact official omni_config must be a mapping") + _assert_official_identity(official, meta) + cfg_dict.setdefault("seed", identity["public_seed"]) + cfg_dict.setdefault("max_episode_steps", identity["max_episode_steps"]) + cfg_dict.setdefault( + "max_steps_per_rollout_epoch", + identity["max_episode_steps"], + ) + cfg_dict.setdefault("auto_reset", False) + cfg_dict.setdefault("ignore_terminations", False) + cfg_dict.setdefault("use_fixed_reset_state_ids", False) + cfg_dict.setdefault("skip_intermediate_obs_in_chunk", True) + cfg_dict.setdefault("num_env_subprocess", 1) + cfg_dict.setdefault("direct_omnigibson_env", True) + cfg_dict.setdefault( + "video_cfg", + { + "save_video": False, + "info_on_video": True, + "video_base_dir": str(output_path / "video"), + }, + ) + cfg_dict.setdefault( + "action_trace_path", + str(output_path / "behavior_action_trace.jsonl"), + ) + cfg_dict.setdefault("action_trace_interval", 1) + return OmegaConf.create(cfg_dict) + + official = exact_loaded.get("omni_config", exact_loaded) + if not isinstance(official, Mapping): + raise ValueError("exact official omni_config must be a mapping") + return OmegaConf.create( + _exact_config_from_official(official, meta, output_path) + ) + + rlinf_root = ensure_rlinf_import_path() + config_path = _default_env_config_path(rlinf_root, meta) + if not config_path.is_file(): + raise FileNotFoundError(f"RLinf BEHAVIOR env config not found: {config_path}") + cfg = OmegaConf.load(config_path) + return _apply_default_config_identity( + cfg, + identity=identity, + output_dir=output_path, + meta=meta, + ) + + +def _torch_to_numpy(value: Any) -> Any: + if hasattr(value, "detach") and hasattr(value, "cpu") and hasattr(value, "numpy"): + return value.detach().cpu().numpy() + return value + + +def _jsonable(value: Any) -> Any: + value = _torch_to_numpy(value) + if isinstance(value, np.ndarray): + return value + if isinstance(value, np.generic): + return value.item() + if isinstance(value, Mapping): + return {str(key): _jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(item) for item in value] + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return repr(value) + + +def _strict_public_json(value: Any) -> Any: + value = _torch_to_numpy(value) + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, Mapping): + return {str(key): _strict_public_json(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_strict_public_json(item) for item in value] + if isinstance(value, bytes): + return { + "format": "png", + "data": base64.b64encode(value).decode("ascii"), + } + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return repr(value) + + +def _reset_trace_enabled() -> bool: + return os.environ.get(RESET_TRACE_ENV) == "1" + + +def _emit_reset_trace_marker( + event: str, + *, + elapsed_s: float | None = None, + **fields: Any, +) -> None: + if not _reset_trace_enabled(): + return + payload = { + "schema_version": 1, + "component": "OfficialBehaviorBackend", + "event": event, + **fields, + } + if elapsed_s is not None: + payload["elapsed_s"] = float(elapsed_s) + try: + print( + json.dumps( + _strict_public_json(payload), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ), + flush=True, + ) + except Exception: + # Tracing must never change reset / timeout / exception semantics. + pass + + +def _image_uint8(value: Any) -> np.ndarray: + arr = np.asarray(_torch_to_numpy(value)) + if arr.ndim != 3 or arr.shape[-1] not in {3, 4}: + raise ValueError(f"image must be [H,W,3 or 4], got {arr.shape}") + arr = arr[..., :3] + if arr.dtype == np.uint8: + return np.ascontiguousarray(arr) + if np.issubdtype(arr.dtype, np.floating): + max_value = float(np.nanmax(arr)) if arr.size else 1.0 + if max_value <= 1.0 + 1e-6: + arr = arr * 255.0 + arr = np.rint(arr).clip(0, 255).astype(np.uint8) + return np.ascontiguousarray(arr) + + +def _first_batch(value: Any) -> Any: + arr = np.asarray(_torch_to_numpy(value)) + if arr.ndim >= 1 and arr.shape[0] == 1: + return arr[0] + return arr + + +def _task_description(value: Any, default: str) -> str: + value = _jsonable(value) + if isinstance(value, (list, tuple)): + for item in value: + if isinstance(item, str) and item.strip(): + return item.strip() + return default + if isinstance(value, str) and value.strip(): + return value.strip() + return default + + +def _extract_raw_observation(raw_obs: Mapping[str, Any]) -> dict[str, Any]: + main_image = None + left_image = None + right_image = None + proprio = None + for sensor_data in raw_obs.values(): + if not isinstance(sensor_data, Mapping): + continue + for key, value in sensor_data.items(): + if not isinstance(key, str): + continue + if "proprio" in key: + # RLinf exposes proprio as a tensor / ndarray, whereas camera + # observations are mappings containing an ``rgb`` value. + proprio = value + elif not isinstance(value, Mapping): + continue + elif "left_realsense_link:Camera:0" in key and "rgb" in value: + left_image = value["rgb"] + elif "right_realsense_link:Camera:0" in key and "rgb" in value: + right_image = value["rgb"] + elif "zed_link:Camera:0" in key and "rgb" in value: + main_image = value["rgb"] + if ( + main_image is None + or left_image is None + or right_image is None + or proprio is None + ): + raise ValueError("raw BEHAVIOR observation lacks main/wrist RGB or proprio") + return { + "main_images": main_image, + "wrist_images": np.stack( + [_image_uint8(left_image), _image_uint8(right_image)], + axis=0, + ), + "states": proprio, + } + + +def _normalize_single_observation( + obs: Mapping[str, Any], *, task_language: str +) -> dict[str, Any]: + if "main_images" not in obs or "wrist_images" not in obs or "states" not in obs: + obs = _extract_raw_observation(obs) + + main = _image_uint8(_first_batch(obs["main_images"])) + wrists_value = _first_batch(obs["wrist_images"]) + wrists = np.asarray(_torch_to_numpy(wrists_value)) + if wrists.ndim == 5 and wrists.shape[0] == 1: + wrists = wrists[0] + if wrists.ndim != 4 or wrists.shape[0] != 2: + raise ValueError(f"wrist_images must be [2,H,W,3], got {wrists.shape}") + left = _image_uint8(wrists[0]) + right = _image_uint8(wrists[1]) + states = np.asarray(_first_batch(obs["states"]), dtype=np.float32) + if states.ndim != 1: + raise ValueError(f"states must be [raw_proprio_dim], got {states.shape}") + if not np.isfinite(states).all(): + raise ValueError("states contains NaN or infinity") + return { + "main_images": main, + "wrist_images": np.ascontiguousarray(np.stack([left, right], axis=0)), + "states": np.ascontiguousarray(states.astype(np.float32, copy=False)), + "task_descriptions": _task_description( + obs.get("task_descriptions"), + task_language, + ), + "extra_view_images": None, + } + + +def _validate_action_chunk(actions: Any) -> np.ndarray: + arr = np.asarray(actions, dtype=np.float32) + if arr.ndim != 2 or arr.shape[1] != ACTION_DIM or arr.shape[0] < 1: + raise ValueError(f"BEHAVIOR actions must be [T,{ACTION_DIM}], got {arr.shape}") + if not np.isfinite(arr).all(): + raise ValueError("BEHAVIOR actions contain NaN or infinity") + return np.ascontiguousarray(arr) + + +def _raw_success(info: Any) -> bool: + done = info.get("done") if isinstance(info, Mapping) else None + value = done.get("success") if isinstance(done, Mapping) else None + return isinstance(value, (bool, np.bool_)) and bool(value) + + +def _receipt_from_info( + info: Mapping[str, Any], *, env_step: int +) -> dict[str, Any] | None: + if not _raw_success(info) or type(env_step) is not int or env_step < 0: + return None + material = { + "schema_version": 1, + "source": 'info["done"]["success"]', + "env_step": env_step, + "raw_done": {"success": True}, + } + return { + **material, + "receipt_sha256": official_success_receipt_sha256(material), + } + + +def _png_bytes(image: np.ndarray) -> bytes: + import imageio.v2 as imageio + + buf = io.BytesIO() + imageio.imwrite(buf, _image_uint8(image), format="png") + return buf.getvalue() + + +def _write_frame_files( + frames: Mapping[str, bytes], + *, + output_dir: Path, + group_id: str, +) -> dict[str, str]: + capture_dir = output_dir / "captures" + capture_dir.mkdir(parents=True, exist_ok=True) + paths: dict[str, str] = {} + for camera, payload in frames.items(): + filename = f"{group_id}_{camera}.png" + path = capture_dir / filename + path.write_bytes(payload) + try: + paths[str(camera)] = str(path.relative_to(output_dir)) + except ValueError: + paths[str(camera)] = str(path) + return paths + + +class OfficialBehaviorBackend: + """Duck-typed backend around one official RLinf BehaviorEnv.""" + + def __init__( + self, + *, + meta: Mapping[str, Any], + output_dir: str | Path, + behavior_env_cls: Any | None = None, + cfg: Any | None = None, + ) -> None: + self.meta = dict(meta) + self.identity = _task_identity(self.meta) + self.output_dir = Path(output_dir).expanduser().resolve() + self.output_dir.mkdir(parents=True, exist_ok=True) + self._last_obs: dict[str, Any] | None = None + self._last_info: dict[str, Any] = {} + self._last_raw_obs: Any = None + self._closed = False + self._episode_ended = False + self._total_env_steps = 0 + self._official_success_latched = False + self._official_success_receipt: dict[str, Any] | None = None + self._camera_frames: dict[str, Any] = {} + self._camera_frame_step = -1 + self._projections: dict[str, Any] = {} + self._motion_planner = None + self._gripper_latch = { + "left": GRIPPER_OPEN_COMMAND, + "right": GRIPPER_OPEN_COMMAND, + } + self.cfg = ( + cfg + if cfg is not None + else build_behavior_env_config(self.meta, self.output_dir) + ) + if behavior_env_cls is None: + ensure_rlinf_import_path() + from rlinf.envs.behavior.behavior_env import BehaviorEnv + + behavior_env_cls = BehaviorEnv + self._env = behavior_env_cls( + self.cfg, + num_envs=1, + seed_offset=0, + total_num_processes=1, + worker_info=SimpleNamespace(group_world_size=1), + record_metrics=False, + ) + + @property + def total_env_steps(self) -> int: + return self._total_env_steps + + @property + def official_success_latched(self) -> bool: + return self._official_success_latched + + @property + def official_success_receipt(self) -> dict[str, Any] | None: + if self._official_success_receipt is None: + return None + return dict(self._official_success_receipt) + + def _wrap_raw_obs(self, raw_obs: Any) -> dict[str, Any]: + if isinstance(raw_obs, Mapping) and { + "main_images", + "wrist_images", + "states", + }.issubset(raw_obs): + return _normalize_single_observation( + raw_obs, + task_language=self.identity["task_language"], + ) + wrapper = getattr(self._env, "_wrap_obs", None) + if callable(wrapper): + try: + wrapped = wrapper([raw_obs]) + return _normalize_single_observation( + wrapped, + task_language=self.identity["task_language"], + ) + except Exception: + pass + if isinstance(raw_obs, Mapping): + return _normalize_single_observation( + raw_obs, + task_language=self.identity["task_language"], + ) + raise TypeError("BEHAVIOR raw observation is not a mapping") + + def _note_info( + self, + info: Any, + *, + telemetry: Mapping[str, Any] | None = None, + ) -> dict[str, Any]: + info_dict = dict(_jsonable(info)) if isinstance(info, Mapping) else {} + runtime = info_dict.get("_rpent") + if not isinstance(runtime, dict): + runtime = {} + runtime["total_env_steps"] = int(self._total_env_steps) + runtime["global_env_steps"] = int(self._total_env_steps) + if telemetry is not None: + for field in ( + "executed_steps", + "stop_reason", + "success_step_in_chunk", + ): + value = telemetry.get(field) + if value is not None: + info_dict[field] = _strict_public_json(value) + if _raw_success(info_dict): + self._official_success_latched = True + receipt = _receipt_from_info(info_dict, env_step=self._total_env_steps) + if receipt is not None: + self._official_success_receipt = receipt + runtime["official_success_receipt"] = dict(receipt) + info_dict["_rpent"] = runtime + self._last_info = info_dict + return info_dict + + def _reset_raw(self) -> tuple[Any, dict[str, Any]]: + reset_raw = getattr(self._env, "reset_raw", None) + env_reset = getattr(self._env, "env_reset", None) + if callable(reset_raw): + branch = "reset_raw" + elif callable(env_reset): + branch = "env_reset" + else: + branch = "reset_fallback" + started_at = time.monotonic() + _emit_reset_trace_marker( + "official_behavior_backend._reset_raw.enter", + branch=branch, + ) + try: + if callable(reset_raw): + obs, info = reset_raw(env_idx=0) + elif callable(env_reset): + observations, infos = env_reset() + if not isinstance(observations, (list, tuple)) or not observations: + raise TypeError( + "RLinf BehaviorEnv.env_reset returned no observations" + ) + obs = observations[0] + info = infos[0] if isinstance(infos, (list, tuple)) and infos else {} + else: + ret = self._env.reset() + if isinstance(ret, (tuple, list)) and len(ret) == 2: + obs, info = ret + else: + obs, info = ret, {} + info_out = dict(_jsonable(info)) if isinstance(info, Mapping) else {} + except Exception as exc: + _emit_reset_trace_marker( + "official_behavior_backend._reset_raw.exit", + branch=branch, + status="error", + error_type=type(exc).__name__, + error=str(exc), + elapsed_s=time.monotonic() - started_at, + ) + raise + _emit_reset_trace_marker( + "official_behavior_backend._reset_raw.exit", + branch=branch, + status="ok", + info_is_mapping=isinstance(info, Mapping), + observation_type=type(obs).__name__, + elapsed_s=time.monotonic() - started_at, + ) + return obs, info_out + + def _step_one_raw( + self, action: np.ndarray + ) -> tuple[Any, float, bool, bool, dict[str, Any]]: + step_raw = getattr(self._env, "step_raw", None) + if callable(step_raw): + obs, reward, terminated, truncated, info = step_raw(action, env_idx=0) + return ( + obs, + float(np.asarray(_torch_to_numpy(reward)).reshape(-1)[0]), + bool(np.asarray(_torch_to_numpy(terminated)).reshape(-1)[0]), + bool(np.asarray(_torch_to_numpy(truncated)).reshape(-1)[0]), + dict(_jsonable(info)) if isinstance(info, Mapping) else {}, + ) + + env_chunk_step = getattr(self._env, "env_chunk_step", None) + if callable(env_chunk_step): + import torch + + chunk_action = torch.as_tensor( + action.reshape(1, 1, ACTION_DIM), dtype=torch.float32 + ) + raw_obs_list, rewards, terms, truncs, infos = env_chunk_step(chunk_action) + obs = raw_obs_list[-1][0] if raw_obs_list[-1] is not None else None + info = infos[-1][0] if infos[-1] else {} + return ( + obs, + float(np.asarray(_torch_to_numpy(rewards[-1])).reshape(-1)[0]), + bool(np.asarray(_torch_to_numpy(terms[-1])).reshape(-1)[0]), + bool(np.asarray(_torch_to_numpy(truncs[-1])).reshape(-1)[0]), + dict(_jsonable(info)) if isinstance(info, Mapping) else {}, + ) + + chunk_step = getattr(self._env, "chunk_step", None) + if callable(chunk_step): + obs_list, rewards, terms, truncs, infos = chunk_step( + action.reshape(1, 1, ACTION_DIM) + ) + obs = obs_list[-1] if isinstance(obs_list, (list, tuple)) else obs_list + info = infos[-1] if isinstance(infos, (list, tuple)) and infos else {} + if isinstance(info, list) and info: + info = info[0] + return ( + obs, + float(np.asarray(_torch_to_numpy(rewards)).reshape(-1)[-1]), + bool(np.asarray(_torch_to_numpy(terms)).reshape(-1)[-1]), + bool(np.asarray(_torch_to_numpy(truncs)).reshape(-1)[-1]), + dict(_jsonable(info)) if isinstance(info, Mapping) else {}, + ) + + raise RuntimeError("RLinf BehaviorEnv exposes no raw step interface") + + def reset(self) -> tuple[dict[str, Any], dict[str, Any]]: + started_at = time.monotonic() + _emit_reset_trace_marker( + "official_behavior_backend.reset.enter", + total_env_steps_before=int(self._total_env_steps), + ) + try: + self._total_env_steps = 0 + self._episode_ended = False + self._camera_frames = {} + self._camera_frame_step = -1 + self._projections = {} + self._gripper_latch = { + "left": GRIPPER_OPEN_COMMAND, + "right": GRIPPER_OPEN_COMMAND, + } + raw_obs, info = self._reset_raw() + self._last_raw_obs = raw_obs + self._last_obs = self._wrap_raw_obs(raw_obs) + info_out = self._note_info(info) + except Exception as exc: + _emit_reset_trace_marker( + "official_behavior_backend.reset.exit", + status="error", + error_type=type(exc).__name__, + error=str(exc), + elapsed_s=time.monotonic() - started_at, + ) + raise + _emit_reset_trace_marker( + "official_behavior_backend.reset.exit", + status="ok", + total_env_steps=int(self._total_env_steps), + observation_keys=sorted(self._last_obs), + elapsed_s=time.monotonic() - started_at, + ) + return self._last_obs, info_out + + def current_observation(self) -> tuple[dict[str, Any], dict[str, Any]]: + if self._last_obs is None: + raise RuntimeError("no BEHAVIOR observation is available before reset") + return self._last_obs, self._last_info + + def _remember_gripper_commands(self, action: np.ndarray) -> None: + for hand in ("left", "right"): + segment = ENV_ACTION_SEGMENTS[f"{hand}_gripper"] + value = float(np.asarray(action[segment], dtype=np.float32).reshape(-1)[0]) + if np.isfinite(value): + self._gripper_latch[hand] = value + + def _latest_raw_proprio(self) -> np.ndarray: + obs, _info = self.current_observation() + raw = np.asarray(obs.get("states"), dtype=np.float32) + required = max(segment.stop or 0 for segment in RAW_PROPRIO_SEGMENTS.values()) + if raw.ndim != 1 or raw.shape[0] < required: + raise ValueError( + "raw R1Pro proprio must be a vector with at least " + f"{required} values, got {raw.shape}" + ) + if not np.isfinite(raw).all(): + raise ValueError("raw R1Pro proprio contains NaN or infinity") + return raw + + def _hold_action_from_current_proprio(self) -> np.ndarray: + raw = self._latest_raw_proprio() + action = np.zeros(ACTION_DIM, dtype=np.float32) + action[ENV_ACTION_SEGMENTS["base"]] = 0.0 + for segment_name in ("trunk", "left_arm", "right_arm"): + action[ENV_ACTION_SEGMENTS[segment_name]] = raw[ + RAW_PROPRIO_SEGMENTS[segment_name] + ] + action[ENV_ACTION_SEGMENTS["left_gripper"]] = self._gripper_latch["left"] + action[ENV_ACTION_SEGMENTS["right_gripper"]] = self._gripper_latch["right"] + return _validate_action_chunk(action[None, :])[0] + + def _motion_error( + self, + name: str, + kwargs: Mapping[str, Any], + *, + stop_reason: str, + error: str, + ) -> dict[str, Any]: + return { + "status": "failed", + "name": name, + "primitive_success": False, + "task_success": self.official_success_latched, + "stop_reason": stop_reason, + "error": error, + "request": _strict_public_json(dict(kwargs)), + "info": self._last_info, + } + + def _gripper_command( + self, + name: str, + kwargs: Mapping[str, Any], + *, + command: float, + ) -> dict[str, Any]: + request = dict(kwargs) + hand = request.get("hand") + if hand not in {"left", "right"}: + raise ValueError("hand must be 'left' or 'right'") + if "visual_hand_check" not in request: + raise ValueError("visual_hand_check is required") + if self.official_success_latched: + return { + "status": "skipped", + "name": name, + "primitive_success": False, + "task_success": True, + "stop_reason": "already_officially_successful", + "request": _strict_public_json(request), + "info": self._last_info, + } + if self._episode_ended: + return self._motion_error( + name, + request, + stop_reason="episode_ended", + error="BEHAVIOR episode already terminated or truncated", + ) + + try: + action = self._hold_action_from_current_proprio() + action[ENV_ACTION_SEGMENTS[f"{hand}_gripper"]] = float(command) + chunk = np.repeat(action[None, :], GRIPPER_COMMAND_CONTROL_CYCLES, axis=0) + _obs, reward, terminated, truncated, info = self.chunk_step( + chunk, + return_all_frames=False, + ) + except Exception as exc: + return self._motion_error( + name, + request, + stop_reason="error", + error=str(exc), + ) + executed_steps = int(info.get("executed_steps") or 0) + stop_reason = str(info.get("stop_reason") or "requested_actions_completed") + result: dict[str, Any] = { + "status": "ok" if executed_steps > 0 else "failed", + "name": name, + "primitive_success": executed_steps > 0, + "task_success": self.official_success_latched, + "stop_reason": stop_reason, + "hand": hand, + "gripper_command": float(command), + "requested_steps": GRIPPER_COMMAND_CONTROL_CYCLES, + "executed_steps": executed_steps, + "total_env_steps": int(self.total_env_steps), + "reward": float(reward), + "terminated": bool(terminated), + "truncated": bool(truncated), + "action_shape": [1, ACTION_DIM], + "action_chunk_shape": [int(chunk.shape[0]), int(chunk.shape[1])], + "hold_action_source": "raw_proprio_reordered_with_gripper_command_latches", + "visual_hand_check": _strict_public_json(request["visual_hand_check"]), + "visual_hand_check_verification": "not_verified", + "request": _strict_public_json(request), + "info": info, + "_observation": self._last_obs, + } + if "release_visual_check" in request: + result["release_visual_check"] = _strict_public_json( + request["release_visual_check"] + ) + result["release_visual_check_verification"] = "not_verified" + if self.official_success_latched: + result["official_success_receipt"] = self.official_success_receipt + return result + + def step( + self, + action: Any, + ) -> tuple[dict[str, Any] | None, float, bool, bool, dict[str, Any]]: + array = np.asarray(action, dtype=np.float32) + if array.shape != (ACTION_DIM,) or not np.isfinite(array).all(): + raise ValueError(f"BEHAVIOR action must be finite [{ACTION_DIM}]") + return self.chunk_step(array[None, :], return_all_frames=False) + + def chunk_step( + self, + actions: Any, + *, + return_all_frames: bool = False, + ) -> tuple[Any, float, bool, bool, dict[str, Any]]: + action_array = _validate_action_chunk(actions) + last_obs: Any = None + last_reward = 0.0 + terminated = False + truncated = False + last_info: dict[str, Any] = {} + success_step: int | None = None + executed_steps = 0 + stop_reason = "requested_actions_completed" + frames: list[dict[str, Any]] = [] + + for step_offset, action in enumerate(action_array): + raw_obs, reward, step_terminated, step_truncated, info = self._step_one_raw( + action + ) + self._remember_gripper_commands(action) + executed_steps = step_offset + 1 + self._total_env_steps += 1 + last_obs = raw_obs + last_reward = float(reward) + last_info = info + terminated = bool(step_terminated) + truncated = bool(step_truncated) + if return_all_frames and raw_obs is not None: + frames.append(self._wrap_raw_obs(raw_obs)) + if _raw_success(info): + success_step = step_offset + stop_reason = "official_task_success" + break + if terminated: + stop_reason = "terminated" + break + if truncated: + stop_reason = "truncated" + break + + if last_obs is not None: + self._last_raw_obs = last_obs + self._last_obs = self._wrap_raw_obs(last_obs) + if terminated or truncated or success_step is not None: + self._episode_ended = True + telemetry = { + "executed_steps": int(executed_steps), + "stop_reason": stop_reason, + "success_step_in_chunk": success_step, + } + info_out = self._note_info(last_info, telemetry=telemetry) + observation: Any = frames if return_all_frames else self._last_obs + return observation, last_reward, terminated, truncated, info_out + + def get_task_language(self) -> str: + return str(self.identity["task_language"]) + + def healthz(self) -> dict[str, Any]: + return { + "status": "ok", + "runtime": "behavior_rlinf_env", + "pid": os.getpid(), + "total_env_steps": self.total_env_steps, + "official_success_latched": self.official_success_latched, + } + + def get_env_meta(self) -> dict[str, Any]: + return dict(self.meta) + + def render_camera(self, camera_name: str = "head", **_kwargs: Any) -> np.ndarray: + obs, _info = self.current_observation() + camera = _physical_camera(camera_name) + if camera == "head": + return np.asarray(obs["main_images"], dtype=np.uint8) + index = 0 if camera == "left_wrist" else 1 + return np.asarray(obs["wrist_images"][index], dtype=np.uint8) + + def get_camera_meta( + self, + camera_name: str = "head", + **_kwargs: Any, + ) -> dict[str, Any]: + camera = _physical_camera(camera_name) + frame = self._get_camera_frames()[camera] + return { + "camera_name": camera, + "available": True, + "rgb_shape": list(frame["rgb"].shape), + "rgb_dtype": str(frame["rgb"].dtype), + "calibration_available": True, + "depth_available": True, + "intrinsic": frame["intrinsic"], + "camera_to_world": frame["camera_to_world"], + } + + def observe(self, camera: str = "head", **kwargs: Any) -> dict[str, Any]: + from robots.behavior.schemas import ( + FRAME_REVIEW_ASSESSMENTS, + validate_observe_request, + ) + + request = validate_observe_request(camera=camera, **kwargs) + camera = _physical_camera(camera) + if request.get("head_view", "center") != "center": + return self._motion_error( + "observe", + request, + stop_reason="head_view_unavailable", + error="R1Pro has no movable head camera; use the current physical view", + ) + review = request.get("frame_review") + probe = request.get("depth_probe") + if review is not None or probe is not None: + value = review if review is not None else probe + frame_id = f"behavior-{self.total_env_steps}-{camera}" + if ( + not isinstance(value, Mapping) + or value.get("frame_id") != frame_id + or self._camera_frame_step != self.total_env_steps + ): + return self._motion_error( + "observe", + request, + stop_reason="stale_frame", + error="review/probe requires the current observed frame", + ) + if review is not None: + if review.get("assessment") not in FRAME_REVIEW_ASSESSMENTS: + raise ValueError("invalid frame_review assessment") + return { + "status": "ok", + "primitive_success": True, + "frame_review": dict(review), + "verification": "planner_assessment_not_independently_verified", + "info": self._last_info, + } + if probe.get("assessment") != "target_point_visually_confirmed": + raise ValueError("invalid depth_probe assessment") + result = self.pixel_to_world( + camera=camera, **{k: v for k, v in probe.items() if k != "assessment"} + ) + result.pop("projection_id", None) + return { + **result, + "depth_probe": dict(probe), + "verification": "depth_measured_visual_assessment_not_verified", + "info": self._last_info, + } + frames = self._get_camera_frames() + payloads = {name: _png_bytes(frame["rgb"]) for name, frame in frames.items()} + depths = {} + for name, frame in frames.items(): + gray = np.nan_to_num(frame["depth"], nan=0, posinf=0, neginf=0) + gray = np.rint(np.clip(gray / 5.0, 0, 1) * 255).astype(np.uint8) + depths[name] = _png_bytes(np.repeat(gray[..., None], 3, axis=-1)) + frame_id = f"behavior-{self.total_env_steps}-{camera}" + return { + "status": "ok", + "camera": camera, + "paired_hand": request.get("paired_hand"), + "frame_id": frame_id, + "step": self.total_env_steps, + "_image_bytes": payloads["head"], + "_depth_image_bytes": depths["head"], + "_image_left_wrist_bytes": payloads["left_wrist"], + "_depth_left_wrist_bytes": depths["left_wrist"], + "_image_right_wrist_bytes": payloads["right_wrist"], + "_depth_right_wrist_bytes": depths["right_wrist"], + "depth_display_range_m": [0, 5], + "frames": _write_frame_files( + payloads, + output_dir=self.output_dir, + group_id=frame_id, + ), + "info": self._last_info, + } + + def move_to(self, **kwargs: Any) -> dict[str, Any]: + if kwargs.get("hand") == "both": + return self._move_both_hands_to(kwargs) + return self._move_single_hand_to(kwargs) + + def _move_single_hand_to(self, kwargs: Mapping[str, Any]) -> dict[str, Any]: + return self._plan_motion("move_to", kwargs, {kwargs["hand"]: kwargs["target"]}) + + def _move_both_hands_to(self, kwargs: Mapping[str, Any]) -> dict[str, Any]: + from robots.behavior.schemas import ( + validate_move_both_targets, + validate_move_both_visual_hand_checks, + ) + + validate_move_both_visual_hand_checks(kwargs.get("visual_hand_checks")) + return self._plan_motion( + "move_to", kwargs, validate_move_both_targets(kwargs.get("targets")) + ) + + def _plan_motion( + self, name: str, request: Mapping[str, Any], targets: dict + ) -> dict: + from scipy.spatial.transform import Rotation + + from robots.behavior.motion import BehaviorMotionPlanner, get_planning_state + + started = self.total_env_steps + try: + if self._episode_ended: + return self._motion_error( + name, + request, + stop_reason="episode_ended", + error="episode has ended", + ) + state = self._call_actor(get_planning_state) + poses = {} + for hand, target in targets.items(): + live = state["hands"][hand] + position = np.array(live["position"], dtype=np.float64, copy=True) + rotation = Rotation.from_quat(live["quaternion_xyzw"]) + orientation = live["quaternion_xyzw"] + if name == "rotate_wrist": + angle = float(request["angle_deg"]) + if not np.isfinite(angle) or abs(angle) > 180: + raise ValueError( + "angle_deg must be finite and within [-180,180]" + ) + direction = request.get("direction", "counterclockwise") + if direction not in {"clockwise", "counterclockwise"}: + raise ValueError("invalid wrist rotation direction") + angle *= -1 if direction == "clockwise" else 1 + orientation = ( + rotation * Rotation.from_rotvec([0, 0, np.deg2rad(angle)]) + ).as_quat() + elif "projection_id" in target: + point = self._projection(target["projection_id"]) + goal = np.asarray(point["world_xyz"]) + delta = goal - position + standoff = float(target.get("standoff_m", 0)) + if not np.isfinite(standoff) or standoff < 0: + raise ValueError("standoff_m must be finite and nonnegative") + position = ( + goal - delta / max(np.linalg.norm(delta), 1e-12) * standoff + ) + else: + delta = np.asarray(target["delta_xyz"], dtype=np.float64) + if delta.shape != (3,) or not np.isfinite(delta).all(): + raise ValueError("delta_xyz must be finite [3]") + if target["frame"] == "eef": + delta = rotation.apply(delta) + elif target["frame"] != "world": + raise ValueError("frame must be world or eef") + position += delta + poses[hand] = {"position": position, "quaternion_xyzw": orientation} + if self._motion_planner is None: + self._motion_planner = BehaviorMotionPlanner() + planned = self._motion_planner.plan(state, poses) + if not planned["success"]: + return self._motion_error( + name, + request, + stop_reason=planned["stop_reason"], + error=planned["details"], + ) + positions = planned["positions"] + times = np.arange(len(positions)) * planned["dt"] + dt = state["control_dt"] + if times[-1] > 30 or not np.isfinite(positions).all(): + raise ValueError( + "planned trajectory exceeds 30 seconds or is nonfinite" + ) + sample_times = np.minimum( + np.arange(int(np.ceil(times[-1] / dt)) + 1) * dt, times[-1] + ) + positions = np.stack( + [ + np.interp(sample_times, times, positions[:, i]) + for i in range(positions.shape[1]) + ], + axis=1, + ) + hold = self._hold_action_from_current_proprio() + actions = np.repeat(hold[None, :], len(positions), axis=0) + for i, joint_name in enumerate(planned["joint_names"]): + hand, _, number = joint_name.partition("_arm_joint") + actions[ + :, ENV_ACTION_SEGMENTS[f"{hand}_arm"].start + int(number) - 1 + ] = positions[:, i] + reason = "trajectory_completed" + for offset in range(0, len(actions), 32): + _, _, terminated, truncated, info = self.chunk_step( + actions[offset : offset + 32] + ) + if self.official_success_latched or terminated or truncated: + reason = str(info["stop_reason"]) + break + # Terminal observations are already captured by chunk_step. Never + # issue another actor query after official success or truncation. + final = self._get_motion_state() if not self._episode_ended else None + errors = { + hand: float( + np.linalg.norm( + np.asarray(final["hands"][hand]["position"]) + - target["position"] + ) + ) + for hand, target in poses.items() + if final is not None + } + angles = { + hand: float( + ( + Rotation.from_quat( + final["hands"][hand]["quaternion_xyzw"] + ).inv() + * Rotation.from_quat(target["quaternion_xyzw"]) + ).magnitude() + ) + for hand, target in poses.items() + if final is not None + } + reached = ( + final is not None + and all(x <= 0.01 for x in errors.values()) + and all(x <= 0.1 for x in angles.values()) + ) + succeeded = self.official_success_latched or ( + reason == "trajectory_completed" and reached + ) + if reason == "trajectory_completed" and not reached: + reason = "tracking_error" + return { + "status": "ok" if succeeded else "failed", + "name": name, + "primitive_success": succeeded, + "task_success": self.official_success_latched, + "stop_reason": reason, + "executed_steps": self.total_env_steps - started, + "total_env_steps": self.total_env_steps, + "position_error_m": errors, + "orientation_error_rad": angles, + "_observation": self._last_obs, + "visual_hand_check_verification": "not_verified", + "request": _strict_public_json(request), + "info": self._last_info, + } + except Exception as exc: + result = self._motion_error( + name, request, stop_reason="error", error=str(exc) + ) + result["executed_steps"] = self.total_env_steps - started + return result + + def _projection(self, projection_id: str) -> dict: + if ( + self._camera_frame_step != self.total_env_steps + or projection_id not in self._projections + ): + raise ValueError("projection is not from the current observed frame") + return self._projections[projection_id] + + def navigate_to(self, **kwargs: Any) -> dict[str, Any]: + from scipy.spatial.transform import Rotation + + from robots.behavior.motion import get_planning_state, navigation_collision + from robots.behavior.schemas import validate_relative_navigation_motion + + started = self.total_env_steps + try: + if self._episode_ended: + return self._motion_error( + "navigate_to", + kwargs, + stop_reason="episode_ended", + error="episode has ended", + ) + state = self._call_actor(get_planning_state) + position = np.asarray(state["base_position"]) + yaw = Rotation.from_quat(state["base_quaternion_xyzw"]).as_euler("xyz")[2] + target, target_yaw = position.copy(), yaw + if "relative_motion" in kwargs: + motion = validate_relative_navigation_motion(kwargs["relative_motion"]) + if motion["kind"] == "translation": + distance = motion["distance_m"] * ( + 1 if motion["direction"] == "forward" else -1 + ) + target[:2] += distance * np.array([np.cos(yaw), np.sin(yaw)]) + else: + target_yaw += np.deg2rad(motion["angle_deg"]) * ( + 1 if motion["direction"] == "left" else -1 + ) + else: + check = kwargs.get("navigation_visual_check") + if ( + not isinstance(check, Mapping) + or check.get("camera") != "head" + or check.get("frame_id") != f"behavior-{self.total_env_steps}-head" + or check.get("assessment") != "navigation_target_visually_confirmed" + ): + raise ValueError( + "navigation_visual_check must confirm the head target" + ) + goal = np.asarray( + self._projection(kwargs["projection_id"])["world_xyz"] + ) + standoff = float(kwargs.get("standoff_m", 0.85)) + if not 0.45 <= standoff <= 1.5: + raise ValueError("standoff_m must be within [0.45,1.5]") + delta = goal[:2] - position[:2] + length = np.linalg.norm(delta) + target[:2] += delta / max(length, 1e-12) * max(0, length - standoff) + target_yaw = np.arctan2(delta[1], delta[0]) + obstacle = navigation_collision(state, target) + if obstacle: + return self._motion_error( + "navigate_to", + kwargs, + stop_reason="collision", + error=f"swept footprint intersects {obstacle}", + ) + dt = state["control_dt"] + reason = "duration_limit" + for _ in range(int(np.ceil(30 / dt))): + live = state["base_position"] + yaw = Rotation.from_quat(state["base_quaternion_xyzw"]).as_euler("xyz")[ + 2 + ] + delta = target[:2] - np.asarray(live)[:2] + angle = (target_yaw - yaw + np.pi) % (2 * np.pi) - np.pi + if np.linalg.norm(delta) < 0.01 and abs(angle) < np.deg2rad(1): + reason = "target_reached" + break + world_velocity = delta * min( + 2.0, 0.2 / max(np.linalg.norm(delta), 1e-12) + ) + local_velocity = ( + np.array([[np.cos(yaw), np.sin(yaw)], [-np.sin(yaw), np.cos(yaw)]]) + @ world_velocity + ) + action = self._hold_action_from_current_proprio() + # RLinf base controller scales normalized x/y commands by 0.75 m/s. + action[:2] = local_velocity / 0.75 + action[2] = np.clip(2 * angle, -0.4, 0.4) + _, _, terminated, truncated, info = self.chunk_step(action[None, :]) + if self.official_success_latched or terminated or truncated: + reason = str(info["stop_reason"]) + break + state = self._call_actor(get_planning_state) + obstacle = navigation_collision(state, target) + if obstacle: + reason = "collision" + break + if not self._episode_ended: + _, _, terminated, truncated, info = self.chunk_step( + self._hold_action_from_current_proprio()[None, :] + ) + if self.official_success_latched or terminated or truncated: + reason = str(info["stop_reason"]) + succeeded = self.official_success_latched or reason == "target_reached" + return { + "status": "ok" if succeeded else "failed", + "name": "navigate_to", + "primitive_success": succeeded, + "_observation": self._last_obs, + "task_success": self.official_success_latched, + "stop_reason": reason, + "executed_steps": self.total_env_steps - started, + "total_env_steps": self.total_env_steps, + "request": _strict_public_json(kwargs), + "info": self._last_info, + } + except Exception as exc: + # If a read/planning error followed a base command, release that + # command through the same monitored action channel before returning. + brake_error = None + if self.total_env_steps > started and not self._episode_ended: + try: + self.chunk_step(self._hold_action_from_current_proprio()[None, :]) + except Exception as brake_exc: + brake_error = str(brake_exc) + result = self._motion_error( + "navigate_to", kwargs, stop_reason="error", error=str(exc) + ) + result["executed_steps"] = self.total_env_steps - started + if brake_error is not None: + result["brake_error"] = brake_error + return result + + def rotate_wrist(self, **kwargs: Any) -> dict[str, Any]: + check = kwargs.get("visual_hand_check") + if ( + not isinstance(check, Mapping) + or set(check) != {"camera", "frame_id", "selected_hand", "assessment"} + or check.get("camera") not in PHYSICAL_CAMERAS + or not isinstance(check.get("frame_id"), str) + or not check.get("frame_id") + or check.get("assessment") != "selected_hand_visually_confirmed" + or check.get("selected_hand") != kwargs.get("hand") + ): + raise ValueError("visual_hand_check must identify the selected hand") + return self._plan_motion("rotate_wrist", kwargs, {kwargs["hand"]: {}}) + + def open(self, **kwargs: Any) -> dict[str, Any]: + return self._gripper_command("open", kwargs, command=GRIPPER_OPEN_COMMAND) + + def close(self, **kwargs: Any) -> dict[str, Any]: + if kwargs: + return self._gripper_command("close", kwargs, command=GRIPPER_CLOSE_COMMAND) + if self._closed: + return {"status": "ok", "closed": True, "already_closed": True} + if self._motion_planner is not None: + self._motion_planner.close() + closer = getattr(self._env, "close", None) + if callable(closer): + closer() + self._closed = True + return {"status": "ok", "closed": True} + + def press(self, **kwargs: Any) -> dict[str, Any]: + """Press along the currently aligned hand axis, at most 2 cm. + + Differential IK keeps orientation fixed. This local contact motion does + not choose a button or claim official success from contact alone. + """ + request = dict(kwargs) + hand = request.get("hand") + if hand not in {"left", "right"}: + raise ValueError("hand must be 'left' or 'right'") + check = request.get("visual_hand_check") + if not isinstance(check, Mapping) or check.get("selected_hand") != hand: + raise ValueError("visual_hand_check must identify the selected hand") + duration = float(request.get("duration_s", 1.0)) + if not np.isfinite(duration) or not 0.0 < duration <= 10.0: + raise ValueError("duration_s must be finite and in (0, 10]") + if self._episode_ended or self.official_success_latched: + return self._motion_error( + "press", request, stop_reason="episode_ended", error="episode has ended" + ) + + started = self.total_env_steps + try: + state = self._get_motion_state() + initial = state["hands"][hand] + origin = np.asarray(initial["position"], dtype=np.float64) + direction = np.asarray(initial["approach_direction"], dtype=np.float64) + dt = float(state["control_dt"]) + if not np.isfinite(dt) or dt <= 0: + raise ValueError("invalid environment control timestep") + target = origin + direction * 0.02 + stop_reason = "duration_limit" + contacts = [] + travel = 0.0 + for _ in range(int(np.ceil(duration / dt))): + live = state["hands"][hand] + position = np.asarray(live["position"], dtype=np.float64) + travel = float(np.linalg.norm(position - origin)) + contacts = live["contacts"] + if contacts: + stop_reason = "contact" + break + if travel >= 0.02: + stop_reason = "travel_limit" + break + error = target - position + if np.linalg.norm(error) <= 0.001: + stop_reason = "target_reached" + break + delta = error * min(1.0, 0.02 * dt / np.linalg.norm(error)) + jacobian = np.asarray(live["jacobian"], dtype=np.float64) + if jacobian.shape != (6, 7) or not np.isfinite(jacobian).all(): + raise ValueError("expected a finite [6,7] arm Jacobian") + twist = np.concatenate((delta, np.zeros(3))) + dq = jacobian.T @ np.linalg.solve( + jacobian @ jacobian.T + 1e-4 * np.eye(6), twist + ) + # Cap joint speed at 0.5 rad/s; preserve unselected joints/latches. + dq *= min(1.0, 0.5 * dt / max(float(np.max(np.abs(dq))), 1e-12)) + q = np.asarray(live["joint_positions"]) + dq + if np.any(q < live["joint_lower_limits"]) or np.any( + q > live["joint_upper_limits"] + ): + stop_reason = "joint_limit" + break + action = self._hold_action_from_current_proprio() + action[ENV_ACTION_SEGMENTS[f"{hand}_arm"]] = q + _, _, terminated, truncated, info = self.chunk_step(action[None, :]) + if self.official_success_latched or terminated or truncated: + stop_reason = str(info["stop_reason"]) + break + state = self._get_motion_state() + live = state["hands"][hand] + travel = float(np.linalg.norm(np.asarray(live["position"]) - origin)) + succeeded = stop_reason in { + "contact", + "target_reached", + "official_task_success", + } + return { + "status": "ok" if succeeded else "failed", + "name": "press", + "hand": hand, + "_observation": self._last_obs, + "primitive_success": succeeded, + "task_success": self.official_success_latched, + "stop_reason": stop_reason, + "executed_steps": self.total_env_steps - started, + "total_env_steps": self.total_env_steps, + "travel_m": travel, + "travel_limit_m": 0.02, + "contacts": contacts, + "contact_verification": "any_non_robot_contact_not_button_verified", + "visual_hand_check": _strict_public_json(check), + "visual_hand_check_verification": "not_verified", + "request": _strict_public_json(request), + "info": self._last_info, + } + except Exception as exc: + result = self._motion_error( + "press", request, stop_reason="error", error=str(exc) + ) + result["executed_steps"] = self.total_env_steps - started + return result + + def _get_motion_state(self) -> dict[str, Any]: + from robots.behavior.motion import get_motion_state + + # RLinf owns the OG actor; query it on its existing serial execution lane. + return self._call_actor(get_motion_state) + + def pixel_to_world(self, **kwargs: Any) -> dict[str, Any]: + camera = _physical_camera(kwargs.get("camera")) + expected = f"behavior-{self.total_env_steps}-{camera}" + if ( + kwargs.get("frame_id") != expected + or self._camera_frame_step != self.total_env_steps + ): + return self._motion_error( + "pixel_to_world", + kwargs, + stop_reason="stale_frame", + error="observe the current frame before projection", + ) + frame = self._camera_frames[camera] + u, v = kwargs["u"], kwargs["v"] + depth = frame["depth"] + if ( + type(u) is not int + or type(v) is not int + or not 0 <= u < depth.shape[1] + or not 0 <= v < depth.shape[0] + ): + raise ValueError("pixel must be an integer inside the observed image") + window = kwargs.get("depth_window_px", 7) + if type(window) is not int or not 1 <= window <= 31: + raise ValueError("depth_window_px must be in [1,31]") + radius = window // 2 + samples = depth[ + max(0, v - radius) : v + radius + 1, max(0, u - radius) : u + radius + 1 + ] + valid = samples[np.isfinite(samples) & (samples > 0)] + if not valid.size: + return self._motion_error( + "pixel_to_world", + kwargs, + stop_reason="invalid_depth", + error="no finite positive depth at this pixel", + ) + z = float(np.median(valid)) + k = frame["intrinsic"] + # USD cameras face -Z, +Y up; image rows increase downwards. + optical = np.array( + [(u - k[0, 2]) * z / k[0, 0], -(v - k[1, 2]) * z / k[1, 1], -z, 1.0] + ) + xyz = (frame["camera_to_world"] @ optical)[:3] + value = { + "camera": camera, + "frame_id": expected, + "u": u, + "v": v, + "world_xyz": xyz.tolist(), + "depth_m": z, + } + projection_id = _canonical_json_sha256(value) + self._projections[projection_id] = value + return { + "status": "ok", + "primitive_success": True, + "task_success": self.official_success_latched, + "projection_id": projection_id, + **value, + } + + def _get_camera_frames(self) -> dict[str, Any]: + from robots.behavior.motion import get_camera_observation + + if self._camera_frame_step != self.total_env_steps: + self._camera_frames = self._call_actor(get_camera_observation) + self._camera_frame_step = self.total_env_steps + self._projections = {} + return self._camera_frames + + def _call_actor(self, function): + import ray + + pool = self._env.pool + index = self._env.pool_offset + return ray.get( + pool.env_processes[index % pool.num_env_subprocess].__ray_call__.remote( + function, index // pool.num_env_subprocess + ) + ) + + +def _physical_camera(value: Any) -> str: + camera = str(value or "head") + aliases = { + "main": "head", + "zed": "head", + "left": "left_wrist", + "right": "right_wrist", + } + camera = aliases.get(camera, camera) + if camera not in PHYSICAL_CAMERAS: + raise ValueError("camera must be head, left_wrist, or right_wrist") + return camera + + +__all__ = [ + "ACTION_DIM", + "ACTION_HORIZON", + "PHYSICAL_CAMERAS", + "OfficialBehaviorBackend", + "build_behavior_env_config", + "discover_rlinf_root", + "ensure_rlinf_import_path", +] diff --git a/robots/behavior/robot_spec.py b/robots/behavior/robot_spec.py new file mode 100644 index 000000000..f5ee36ad0 --- /dev/null +++ b/robots/behavior/robot_spec.py @@ -0,0 +1,128 @@ +# 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. + +"""BEHAVIOR robot extension: RobotSpec factory and toolkit bridge.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from robots.behavior.prompt_bundle import system_prompt, user_prompt +from rpent.dashboard.events import DashboardEventSink +from rpent.memory import MemoryManager +from rpent.robots.prompt_bundle import PromptBundle +from rpent.robots.robot_spec import RobotSpec, RunConfig + +BEHAVIOR_DASHBOARD_SPEC = { + "task": { + "command": "/rpent-task", + "usage": "/rpent-task ", + "fields": ( + { + "name": "task_name", + "suggestions": ("turning_on_radio", "picking_up_trash"), + }, + {"name": "public_seed", "kind": "integer", "minimum": 0}, + ), + "display": "{task_name} / s{public_seed}", + "output_slug": "{task_name}_s{public_seed}", + }, + "runtime_components": ( + {"name": "env", "label": "ENV", "scope": "unique"}, + {"name": "vla", "label": "VLA", "scope": "shared"}, + {"name": "dino", "label": "DINO", "scope": "shared"}, + ), + "frame_channels": ( + { + "name": "head", + "label": "head camera", + "result_key": "_image_bytes", + "legacy_path_key": "image_cam_path", + }, + { + "name": "left_wrist", + "label": "left wrist", + "result_key": "_image_left_wrist_bytes", + }, + { + "name": "right_wrist", + "label": "right wrist", + "result_key": "_image_right_wrist_bytes", + }, + ), +} + + +def get_robot_spec() -> RobotSpec: + from robots.behavior import runtime + + return RobotSpec( + name="behavior", + prompts=PromptBundle(system=system_prompt, user=user_prompt), + add_cli_args=runtime.add_cli_args, + parse_config=runtime.parse_config, + init_runtime=runtime.init_runtime, + on_explore_session=runtime.on_explore_session, + dashboard=BEHAVIOR_DASHBOARD_SPEC, + ) + + +def get_toolkit( + *, + primitives_kwargs: dict[str, Any], + dashboard_events: DashboardEventSink, + config: RunConfig, + mode: str | None = None, + attempts_per_session: int = 0, + state_output_dir: Path | str | None = None, +): + """Return the BEHAVIOR toolkit through the standard main contract.""" + + from robots.behavior.toolkit import BehaviorToolkit + + toolkit_kwargs = dict(primitives_kwargs) + if attempts_per_session > 0: + raise ValueError( + "BEHAVIOR explore runs one attempt per session; use --explore-sessions" + ) + if mode is None: + behavior_mode = str(config.prompt_vars.get("behavior_mode", "eval")) + elif mode == "exploration": + behavior_mode = "explore" + elif mode == "evaluation": + behavior_mode = "eval" + else: + raise ValueError(f"unsupported BEHAVIOR toolkit mode: {mode!r}") + if behavior_mode not in {"eval", "explore"}: + raise ValueError(f"unsupported BEHAVIOR toolkit mode: {behavior_mode!r}") + toolkit_kwargs["behavior_phase"] = behavior_mode + memory_dir = config.prompt_vars.get("memory_dir") + if not memory_dir: + raise ValueError("BEHAVIOR RunConfig is missing memory_dir") + memory = MemoryManager( + root=Path(memory_dir), + memory_access="inbox_write" if behavior_mode == "explore" else "read_only", + inbox_cell_tag=config.recipe_tag if behavior_mode == "explore" else None, + ) + return BehaviorToolkit( + primitives_kwargs=toolkit_kwargs, + dashboard_events=dashboard_events, + memory=memory, + config=config, + state_output_dir=state_output_dir, + ) + + +__all__ = ["BEHAVIOR_DASHBOARD_SPEC", "get_robot_spec", "get_toolkit"] diff --git a/robots/behavior/runtime.py b/robots/behavior/runtime.py new file mode 100644 index 000000000..836544480 --- /dev/null +++ b/robots/behavior/runtime.py @@ -0,0 +1,860 @@ +# 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. + +"""Standard RPent runtime hooks for BEHAVIOR.""" + +from __future__ import annotations + +import argparse +import os +import re +import time +from datetime import datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from robots.behavior.policy_checkpoint import ( + POLICY_CHECKPOINT_ENV, + SHARED_POLICY_CHECKPOINT_PATH, + SHARED_POLICY_PROFILE_ID, + write_policy_checkpoint_manifest, +) +from robots.behavior.schemas import ( + ACTION_DIM, + DEFAULT_ACTION_CHUNK, + behavior_tool_specs_for_task, +) +from robots.behavior.task_specs import ( + BehaviorTaskSpec, + get_task_spec, + get_task_spec_by_index, +) +from rpent.dashboard.events import DashboardEventSink +from rpent.robots.robot_spec import RunConfig +from rpent.robots.runtime import stop_owned_daemons, try_spawn_server, try_wait_server +from rpent.utils.config import get_memory_dir, get_repo_root +from rpent.utils.daemon import ProcessDaemon, pick_free_port +from rpent.utils.rpc import make_rpc_client +from rpent.utils.rpc.http_rpc import HttpRpcClient + +if TYPE_CHECKING: + from rpent.utils.rpc import RpcClient + +BEHAVIOR_MODES = ("eval", "explore") +BEHAVIOR_COMPONENTS = {"env", "vla", "dino"} +DEFAULT_EVAL_COMPONENTS = {"env", "vla", "dino"} +DEFAULT_MAX_EPISODE_STEPS = 43_200 +DEFAULT_PLANNER_TIMEOUT_S = 7_200 +RLINF_ROOT_ENV = "RPENT_RLINF_ROOT" +BEHAVIOR_PYTHON_ENV = "RPENT_BEHAVIOR_PYTHON" + + +def _default_behavior_repo() -> Path: + configured = os.environ.get(RLINF_ROOT_ENV) + if configured: + return Path(configured).expanduser() + return get_repo_root().parent / "RLinf" + + +def _default_behavior_python(behavior_repo: Path) -> str: + configured = os.environ.get(BEHAVIOR_PYTHON_ENV) + if configured: + return str(Path(configured).expanduser()) + return str(behavior_repo / ".venv-behavior" / "bin" / "python") + + +def _single_cuda_device(value: Any) -> str | None: + if value in (None, ""): + return None + device = str(value) + if re.fullmatch(r"[0-9]+", device) is None: + raise ValueError("CUDA device must be a single physical GPU ordinal") + return device + + +def _component_cuda_device( + args: argparse.Namespace, + component: str, +) -> str | None: + if component == "env": + specific = getattr(args, "behavior_env_cuda_device", None) + elif component in {"vla", "dino"}: + specific = getattr(args, "behavior_model_cuda_device", None) + else: + raise ValueError(f"unsupported CUDA component: {component}") + return _single_cuda_device( + specific if specific not in (None, "") else getattr(args, "cuda_device", None) + ) + + +def _task_from_args(args: argparse.Namespace) -> BehaviorTaskSpec: + task_name = getattr(args, "task_name", None) + task_value = getattr(args, "task", None) + if task_name: + spec = get_task_spec(str(task_name)) + if task_value is not None and str(task_value).strip().isdigit(): + by_index = get_task_spec_by_index(int(task_value)) + if by_index is not spec: + raise ValueError( + f"BEHAVIOR task identity mismatch: {task_name!r} != {task_value!r}" + ) + return spec + if task_value is None: + raise ValueError("--task-name is required") + if isinstance(task_value, str) and not task_value.strip().isdigit(): + return get_task_spec(task_value.strip()) + return get_task_spec_by_index(int(task_value)) + + +def _public_seed_from_args(args: argparse.Namespace) -> int: + public_seed = getattr(args, "public_seed", None) + seed = getattr(args, "seed", None) + if public_seed is not None and seed is not None and int(public_seed) != int(seed): + raise ValueError("--public-seed and --seed disagree") + value = public_seed if public_seed is not None else seed + if value is None: + raise ValueError("--public-seed is required") + if isinstance(value, bool): + raise ValueError("public seed must be an integer") + return int(value) + + +def add_cli_args(parser: argparse.ArgumentParser, use_dashboard: bool) -> None: + required = not use_dashboard + parser.set_defaults(planner="codex", planner_timeout_s=DEFAULT_PLANNER_TIMEOUT_S) + parser.set_defaults(memory_profile="local") + parser.add_argument( + "--task-name", + required=required, + choices=("turning_on_radio", "picking_up_trash"), + ) + parser.add_argument( + "--task", + default=None, + help="Task name or task index alias; task-name is preferred.", + ) + parser.add_argument("--public-seed", type=int, required=required) + parser.add_argument( + "--seed", + type=int, + default=None, + help="Alias for --public-seed for dashboard and legacy launchers.", + ) + parser.add_argument( + "--behavior-mode", + choices=BEHAVIOR_MODES, + default="eval", + help="BEHAVIOR prompt/runtime mode: eval or explore.", + ) + parser.add_argument( + "--auto-merge-memory", + action=argparse.BooleanOptionalAction, + default=True, + help="Merge exploration output into layered memory (default: enabled).", + ) + parser.add_argument( + "--explore-attempts-per-session", + type=int, + default=0, + help="Unsupported for BEHAVIOR; use --explore-sessions instead.", + ) + parser.add_argument( + "--explore-sessions", + type=int, + default=1, + help="Fresh one-attempt BEHAVIOR sessions per exploration run.", + ) + parser.add_argument( + "--max-episode-steps", + type=int, + default=DEFAULT_MAX_EPISODE_STEPS, + ) + parser.add_argument("--env-endpoint", default=None) + parser.add_argument("--vla-endpoint", default=None) + parser.add_argument("--dino-endpoint", default=None) + default_behavior_repo = _default_behavior_repo() + parser.add_argument( + "--behavior-repo", + default=str(default_behavior_repo), + help=( + "Source checkout containing the RLinf BEHAVIOR integration. " + f"Can also be set with {RLINF_ROOT_ENV}." + ), + ) + parser.add_argument( + "--behavior-python", + default=_default_behavior_python(default_behavior_repo), + help=( + "Python executable for the official BEHAVIOR/OmniGibson env process. " + f"Can also be set with {BEHAVIOR_PYTHON_ENV}." + ), + ) + parser.add_argument( + "--activity-instance-dir", + default=None, + help="Optional explicit official BEHAVIOR task-instance JSON directory.", + ) + parser.add_argument( + "--env-config-path", + default=None, + help="Optional explicit RLinf BEHAVIOR env YAML; the bundled adapter otherwise resolves the pinned template.", + ) + parser.add_argument( + "--policy-checkpoint", + default=str(SHARED_POLICY_CHECKPOINT_PATH), + help=( + "Path to your Pi05-Behavior model checkpoint. " + f"Can also be set with {POLICY_CHECKPOINT_ENV}." + ), + ) + parser.add_argument( + "--cuda-device", + default=None, + help=( + "Shared fallback physical GPU ordinal. Component-specific BEHAVIOR " + "GPU flags take precedence." + ), + ) + parser.add_argument( + "--behavior-env-cuda-device", + default=None, + help="Physical GPU ordinal exposed only to the BEHAVIOR env process.", + ) + parser.add_argument( + "--behavior-model-cuda-device", + default=None, + help="Physical GPU ordinal shared by the BEHAVIOR VLA and DINO processes.", + ) + parser.add_argument( + "--behavior-memory-dir", + default=None, + help="Explicit DINO episode-memory catalog root. Omission selects a legal empty catalog.", + ) + parser.add_argument("--dino-source-archive", default=None) + parser.add_argument("--dino-weights", default=None) + parser.add_argument("--dino-cache-dir", default=None) + parser.add_argument("--vla-ready-timeout-s", type=float, default=900.0) + + +def parse_config(args: argparse.Namespace) -> RunConfig: + mode = str(getattr(args, "behavior_mode", None) or "eval") + if mode not in BEHAVIOR_MODES: + raise ValueError(f"unsupported --behavior-mode {mode!r}") + if getattr(args, "explore", False): + if mode != "explore": + raise ValueError("BEHAVIOR --explore requires --behavior-mode explore") + if getattr(args, "dashboard", False): + raise ValueError( + "BEHAVIOR --explore is CLI-only; use --behavior-mode explore " + "without --explore for Dashboard TaskRuns" + ) + if getattr(args, "env_endpoint", None) is not None: + raise ValueError( + "BEHAVIOR explore requires an owned env sidecar; omit --env-endpoint" + ) + if getattr(args, "explore_attempts_per_session", 0) > 0: + raise ValueError( + "BEHAVIOR explore runs one attempt per session; use --explore-sessions" + ) + spec = _task_from_args(args) + public_seed = _public_seed_from_args(args) + activity_instance_id = spec.instance_for_public_seed(public_seed, phase=mode) + if getattr(args, "activity_instance_id", None) is not None: + requested = int(args.activity_instance_id) + if requested != activity_instance_id: + raise ValueError( + "activity_instance_id must match the task public-seed mapping: " + f"expected {activity_instance_id}, got {requested}" + ) + cuda_device = _single_cuda_device(getattr(args, "cuda_device", None)) + env_cuda_device = _component_cuda_device(args, "env") + model_cuda_device = _component_cuda_device(args, "vla") + if int(getattr(args, "max_episode_steps", 0) or 0) <= 0: + raise ValueError("--max-episode-steps must be positive") + + args.task_name = spec.task_name + args.task = spec.task_index + args.public_seed = public_seed + args.seed = public_seed + args.behavior_phase = mode + args.activity_definition_id = spec.activity_definition_id + args.activity_instance_id = activity_instance_id + args.scene_model = spec.scene_model + args.cuda_device = cuda_device + args.behavior_env_cuda_device = env_cuda_device + args.behavior_model_cuda_device = model_cuda_device + + recipe_tag = spec.tag(public_seed) + output_dir = getattr(args, "output_dir", None) + if output_dir is None: + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S-%f") + output_dir = get_repo_root() / "logs" / f"{timestamp}_behavior_{recipe_tag}" + output_dir = Path(output_dir).expanduser().resolve() + requested_memory_profile = getattr(args, "memory_profile", None) + memory_profile = str(requested_memory_profile or "local") + if memory_profile != "local": + raise ValueError("BEHAVIOR requires --memory-profile local") + configured_memory_dir = getattr(args, "memory_dir", None) + memory_dir = ( + Path(configured_memory_dir).expanduser().resolve() + if configured_memory_dir + else get_memory_dir("behavior").resolve() + ) + args.memory_profile = memory_profile + args.memory_dir = str(memory_dir) + configured_episode_memory_dir = getattr(args, "behavior_memory_dir", None) + episode_memory_dir = ( + Path(configured_episode_memory_dir).expanduser().resolve() + if configured_episode_memory_dir + else None + ) + args.behavior_memory_dir = ( + str(episode_memory_dir) if episode_memory_dir is not None else None + ) + args.behavior_memory_dir_explicit = episode_memory_dir is not None + episode_memory_profile = ( + "explicit" if episode_memory_dir is not None else "empty_episode_catalog" + ) + return RunConfig( + recipe_tag=recipe_tag, + output_dir=output_dir, + prompt_vars={ + "task_name": spec.task_name, + "task": spec.task_index, + "task_language": spec.task_language, + "instruction": spec.task_language, + "task_instruction": spec.task_language, + "public_seed": public_seed, + "activity_definition_id": spec.activity_definition_id, + "activity_instance_id": activity_instance_id, + "scene_model": spec.scene_model, + "behavior_mode": mode, + "behavior_phase": mode, + "max_episode_steps": int(args.max_episode_steps), + "wall_clock_seconds": int(getattr(args, "planner_timeout_s", 7200) or 7200), + "public_capabilities": [ + item["name"] for item in behavior_tool_specs_for_task(spec) + ], + "memory_dir": str(memory_dir), + "memory_profile": memory_profile, + "memory_inbox": str(memory_dir / "_internal" / "inbox" / recipe_tag), + "behavior_episode_memory": episode_memory_profile, + "behavior_memory_dir": str(episode_memory_dir) + if episode_memory_dir is not None + else "", + "recipe_tag": recipe_tag, + "output_dir": str(output_dir), + }, + task_desc={ + "env": "behavior", + "task_name": spec.task_name, + "task": spec.task_index, + "public_seed": public_seed, + "activity_definition_id": spec.activity_definition_id, + "activity_instance_id": activity_instance_id, + "scene_model": spec.scene_model, + "mapping_version": spec.mapping_version, + "behavior_mode": mode, + "policy_profile_id": SHARED_POLICY_PROFILE_ID, + "action_dim": ACTION_DIM, + "action_horizon": DEFAULT_ACTION_CHUNK, + "cuda_device": cuda_device, + "behavior_env_cuda_device": env_cuda_device, + "behavior_model_cuda_device": model_cuda_device, + "memory_profile": memory_profile, + "memory_dir": str(memory_dir), + "behavior_episode_memory": episode_memory_profile, + "behavior_memory_dir_explicit": episode_memory_dir is not None, + "behavior_memory_dir": str(episode_memory_dir) + if episode_memory_dir is not None + else None, + }, + ) + + +def env_runtime_contract(args: argparse.Namespace) -> dict[str, Any]: + return { + "runtime": "behavior_env", + "task_name": args.task_name, + "task": int(args.task), + "task_language": get_task_spec(args.task_name).task_language, + "activity_definition_id": int(args.activity_definition_id), + "activity_instance_id": int(args.activity_instance_id), + "public_seed": int(args.public_seed), + "scene_model": str(args.scene_model), + "max_episode_steps": int(args.max_episode_steps), + "action_dim": ACTION_DIM, + "action_horizon": DEFAULT_ACTION_CHUNK, + "official_success_path": ["info", "done", "success"], + "behavior_repo": str(Path(args.behavior_repo).expanduser().resolve()), + "activity_instance_dir": ( + None + if not getattr(args, "activity_instance_dir", None) + else str(Path(args.activity_instance_dir).expanduser().resolve()) + ), + "rlinf_env_config_path": ( + None + if not getattr(args, "env_config_path", None) + else str(Path(args.env_config_path).expanduser().resolve()) + ), + } + + +def _behavior_python_path(value: str | Path) -> Path: + """Return an absolute executable path without dereferencing venv symlinks.""" + + return Path(value).expanduser().absolute() + + +def _behavior_subprocess_env( + *, + cuda_device: str | None = None, + **overrides: str, +) -> dict[str, str]: + repo_root = str(get_repo_root()) + pythonpath = os.environ.get("PYTHONPATH") + return { + **overrides, + "PYTHONPATH": ( + repo_root if not pythonpath else os.pathsep.join([repo_root, pythonpath]) + ), + **({"CUDA_VISIBLE_DEVICES": cuda_device} if cuda_device is not None else {}), + } + + +def vla_runtime_contract(args: argparse.Namespace) -> dict[str, Any]: + return { + "runtime": "pi05_vla", + "embodiment": "behavior", + "config_name": "pi05_behavior", + "action_dim": ACTION_DIM, + "action_horizon": DEFAULT_ACTION_CHUNK, + "policy_profile_id": SHARED_POLICY_PROFILE_ID, + "checkpoint": str(Path(args.policy_checkpoint).expanduser()), + } + + +def _wait_for_server_endpoint( + daemon: ProcessDaemon, + *, + log_offset: int = 0, + timeout_s: float = 1800.0, +) -> str: + """Read this launch's bound RPC address, never a preselected free port. + + ProcessDaemon appends logs, so callers supply the pre-launch byte offset. + Application readiness is still checked by the shared RPC health probe. + """ + if daemon.log_path is None: + raise ValueError("endpoint discovery requires a daemon log_path") + deadline = time.monotonic() + timeout_s + with Path(daemon.log_path).open("rb") as log: + log.seek(log_offset) + while time.monotonic() < deadline: + exit_code = daemon.poll() + if exit_code is not None: + raise RuntimeError( + f"{daemon.name} exited with code {exit_code}; see {daemon.log_path}" + ) + position = log.tell() + line = log.readline() + if not line.endswith(b"\n"): + log.seek(position) + time.sleep(0.1) + continue + match = re.fullmatch( + rb"RPC server listening on (http://[^\s:]+:[1-9][0-9]*)\r?\n", + line, + ) + if match is not None: + return match.group(1).decode("ascii") + raise TimeoutError( + f"{daemon.name} did not announce an RPC endpoint within {timeout_s}s; " + f"see {daemon.log_path}" + ) + + +def _spawn_env_server( + args: argparse.Namespace, + output_dir: Path, +) -> tuple[ProcessDaemon | None, "RpcClient"]: + output_dir.mkdir(parents=True, exist_ok=True) + if args.env_endpoint is not None: + return None, make_rpc_client(args.env_endpoint) + # Bind only after Ray/OmniGibson initialization; the OS chooses the port. + host, port = "127.0.0.1", 0 + cuda_device = _component_cuda_device(args, "env") + # Keep the virtualenv launcher path intact. Resolving ``bin/python`` + # follows its symlink to the system interpreter and silently drops the + # virtualenv's site-packages (OmniGibson, OpenPI, OmegaConf, ...). + behavior_python = _behavior_python_path(args.behavior_python) + if not behavior_python.is_file(): + raise RuntimeError(f"BEHAVIOR Python executable is missing: {behavior_python}") + cmd = [ + str(behavior_python), + str(get_repo_root() / "robots" / "behavior" / "env_server.py"), + "--task-name", + str(args.task_name), + "--public-seed", + str(args.public_seed), + "--task-index", + str(args.task), + "--activity-definition-id", + str(args.activity_definition_id), + "--activity-instance-id", + str(args.activity_instance_id), + "--scene-model", + str(args.scene_model), + "--max-episode-steps", + str(args.max_episode_steps), + "--output-dir", + str(output_dir), + "--behavior-repo", + str(Path(args.behavior_repo).expanduser().resolve()), + "--host", + host, + "--port", + str(port), + "--parent-watch", + ] + if getattr(args, "activity_instance_dir", None): + cmd.extend( + [ + "--activity-instance-dir", + str(Path(args.activity_instance_dir).expanduser().resolve()), + ] + ) + if getattr(args, "env_config_path", None): + cmd.extend( + [ + "--env-config-path", + str(Path(args.env_config_path).expanduser().resolve()), + ] + ) + if cuda_device is not None: + cmd.extend(["--cuda-device", cuda_device]) + daemon = ProcessDaemon( + name="behavior_env_server", + cmd=cmd, + env_overrides=_behavior_subprocess_env( + cuda_device=cuda_device, + **{ + "ROBOT_PLATFORM": "BEHAVIOR", + "OMNIGIBSON_HEADLESS": "1", + # Ray otherwise clears CUDA_VISIBLE_DEVICES for the zero-GPU actor + # that owns the single OmniGibson subprocess. + "RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO": "0", + }, + ), + log_path=str(output_dir / "behavior_env_server.log"), + ) + log_path = Path(daemon.log_path) + log_offset = log_path.stat().st_size if log_path.exists() else 0 + daemon.start() + try: + endpoint = _wait_for_server_endpoint(daemon, log_offset=log_offset) + return daemon, HttpRpcClient(endpoint) + except BaseException: + # try_spawn_server cannot own this daemon until this function returns. + daemon.stop() + raise + + +def _spawn_vla_server( + args: argparse.Namespace, + output_dir: Path, +) -> tuple[ProcessDaemon | None, "RpcClient"]: + output_dir.mkdir(parents=True, exist_ok=True) + if args.vla_endpoint is not None: + return None, make_rpc_client(args.vla_endpoint) + host, port = "127.0.0.1", pick_free_port() + cuda_device = _component_cuda_device(args, "vla") + # See _spawn_env_server: do not dereference a virtualenv's python symlink. + behavior_python = _behavior_python_path(args.behavior_python) + if not behavior_python.is_file(): + raise RuntimeError(f"BEHAVIOR Python executable is missing: {behavior_python}") + checkpoint_manifest = output_dir / "policy_checkpoint_manifest.json" + write_policy_checkpoint_manifest( + checkpoint_manifest, + Path(args.policy_checkpoint).expanduser(), + ) + cmd = [ + str(behavior_python), + str(get_repo_root() / "rpent" / "robots" / "components" / "pi05_vla_server.py"), + "--embodiment", + "behavior", + "--transport", + "http", + "--host", + host, + "--port", + str(port), + "--model-path", + str(Path(args.policy_checkpoint).expanduser()), + "--checkpoint-manifest", + str(checkpoint_manifest), + "--parent-watch", + ] + if cuda_device is not None: + cmd.extend(["--cuda-device", cuda_device]) + daemon = ProcessDaemon( + name="behavior_vla_server", + cmd=cmd, + env_overrides=_behavior_subprocess_env(cuda_device=cuda_device), + log_path=str(output_dir / "behavior_vla_server.log"), + ) + daemon.start() + return daemon, make_rpc_client(f"http://{host}:{port}") + + +def _spawn_dino_server( + args: argparse.Namespace, + output_dir: Path, +) -> tuple[ProcessDaemon | None, "RpcClient"]: + output_dir.mkdir(parents=True, exist_ok=True) + if args.dino_endpoint is not None: + return None, make_rpc_client(args.dino_endpoint) + host, port = "127.0.0.1", pick_free_port() + cuda_device = _component_cuda_device(args, "dino") + behavior_python = _behavior_python_path(args.behavior_python) + if not behavior_python.is_file(): + raise RuntimeError(f"BEHAVIOR Python executable is missing: {behavior_python}") + cmd = [ + str(behavior_python), + str(get_repo_root() / "robots" / "behavior" / "dino_v2" / "server.py"), + "--host", + host, + "--port", + str(port), + "--parent-watch", + ] + if getattr(args, "dino_source_archive", None): + cmd.extend( + [ + "--source-archive", + str(Path(args.dino_source_archive).expanduser().resolve()), + ] + ) + if getattr(args, "dino_weights", None): + cmd.extend(["--weights", str(Path(args.dino_weights).expanduser().resolve())]) + if getattr(args, "dino_cache_dir", None): + cmd.extend( + ["--cache-dir", str(Path(args.dino_cache_dir).expanduser().resolve())] + ) + if cuda_device is not None: + cmd.extend(["--cuda-device", cuda_device]) + daemon = ProcessDaemon( + name="behavior_dino_server", + cmd=cmd, + env_overrides=_behavior_subprocess_env(cuda_device=cuda_device), + log_path=str(output_dir / "behavior_dino_server.log"), + ) + daemon.start() + return daemon, HttpRpcClient(f"http://{host}:{port}") + + +def _connect_env( + args: argparse.Namespace, + rpc: "RpcClient", + output_dir: Path, +) -> dict[str, Any]: + from robots.behavior.env_client import BehaviorEnvClient + + env = BehaviorEnvClient(rpc, expected_meta=env_runtime_contract(args)) + initial_observation, initial_info = env.reset() + task_language = initial_observation.get("task_descriptions") + if isinstance(task_language, (list, tuple)): + task_language = next( + (item for item in task_language if isinstance(item, str)), None + ) + if task_language is not None and str(task_language).strip(): + expected = get_task_spec(args.task_name).task_language + if str(task_language).strip() != expected: + raise RuntimeError("environment task language does not match TaskSpec") + return { + "env": env, + "task_name": args.task_name, + "behavior_phase": args.behavior_phase, + "public_seed": int(args.public_seed), + "max_episode_steps": int(args.max_episode_steps), + "action_horizon": DEFAULT_ACTION_CHUNK, + "initial_observation": initial_observation, + "initial_info": initial_info, + "output_dir": Path(output_dir), + } + + +def _connect_vla(args: argparse.Namespace, rpc: "RpcClient") -> dict[str, Any]: + from robots.behavior.policy_checkpoint import ( + assert_matching_policy_checkpoint_binding, + validate_policy_checkpoint, + ) + from rpent.robots.components.pi05_vla_client import Pi05VLAClient + + expected_binding = validate_policy_checkpoint(args.policy_checkpoint) + server_meta = rpc.call( + "healthz", + timeout_s=min(float(getattr(args, "vla_ready_timeout_s", 900.0)), 30.0), + ) + if not isinstance(server_meta, dict): + raise TypeError("VLA healthz must return a mapping") + assert_matching_policy_checkpoint_binding( + server_meta.get("checkpoint_binding"), + expected_binding, + ) + return { + "model": Pi05VLAClient(rpc, embodiment="behavior"), + } + + +def _connect_dino(args: argparse.Namespace, rpc: "RpcClient") -> dict[str, Any]: + from robots.behavior.dino_v2.client import BehaviorDinoClient + from robots.behavior.memory.index import load_current_catalog + + client = BehaviorDinoClient(rpc, expected_meta={"runtime": "behavior_dino"}) + configured_memory_dir = getattr(args, "behavior_memory_dir", None) + explicit_marker = getattr(args, "behavior_memory_dir_explicit", None) + explicit = ( + bool(configured_memory_dir) + if explicit_marker is None + else bool(explicit_marker) + ) + if explicit and not configured_memory_dir: + raise ValueError("explicit BEHAVIOR episode-memory catalog path is missing") + memory_dir = ( + Path(configured_memory_dir).expanduser().resolve() if explicit else None + ) + return { + "dino_component": client, + "episode_memory_index": load_current_catalog(memory_dir), + } + + +def init_runtime( + args: argparse.Namespace, + output_dir: Path, + dashboard_events: DashboardEventSink, + components: set[str] | None, +) -> tuple[list[ProcessDaemon], dict[str, Any]]: + """Initialize requested BEHAVIOR components under the RobotSpec contract.""" + + selected = set(DEFAULT_EVAL_COMPONENTS if components is None else components) + if components is None and getattr(args, "explore", False): + # The CLI session hook owns a fresh ENV; VLA/DINO survive handoffs. + selected.discard("env") + unknown = selected.difference(BEHAVIOR_COMPONENTS) + if unknown: + raise ValueError(f"unknown BEHAVIOR runtime components: {sorted(unknown)}") + + owned_daemons: dict[str, ProcessDaemon] = {} + primitives_kwargs: dict[str, Any] = {} + pending_env: tuple[ProcessDaemon | None, RpcClient] | None = None + pending_vla: tuple[ProcessDaemon | None, RpcClient] | None = None + pending_dino: tuple[ProcessDaemon | None, RpcClient] | None = None + if "env" in selected: + pending_env = try_spawn_server( + owned_daemons, + dashboard_events, + "env", + lambda: _spawn_env_server(args, output_dir), + ) + if "vla" in selected: + pending_vla = try_spawn_server( + owned_daemons, + dashboard_events, + "vla", + lambda: _spawn_vla_server(args, output_dir), + ) + if "dino" in selected: + pending_dino = try_spawn_server( + owned_daemons, + dashboard_events, + "dino", + lambda: _spawn_dino_server(args, output_dir), + ) + if pending_env is not None: + daemon, rpc = pending_env + primitives_kwargs.update( + try_wait_server( + owned_daemons, + dashboard_events, + "env", + rpc, + daemon, + 1800.0 if daemon is not None else 120.0, + post_fn=lambda: _connect_env(args, rpc, output_dir), + ) + ) + if pending_vla is not None: + daemon, rpc = pending_vla + primitives_kwargs.update( + try_wait_server( + owned_daemons, + dashboard_events, + "vla", + rpc, + daemon, + float(getattr(args, "vla_ready_timeout_s", 900.0)), + post_fn=lambda: _connect_vla(args, rpc), + ) + ) + if pending_dino is not None: + daemon, rpc = pending_dino + primitives_kwargs.update( + try_wait_server( + owned_daemons, + dashboard_events, + "dino", + rpc, + daemon, + 600.0 if daemon is not None else 120.0, + post_fn=lambda: _connect_dino(args, rpc), + ) + ) + return list(owned_daemons.values()), primitives_kwargs + + +def on_explore_session( + args: argparse.Namespace, + state_output_dir: Path, + dashboard_events: DashboardEventSink, + daemons: list[ProcessDaemon], +) -> dict[str, Any]: + """Replace only this run's ENV sidecar for a fresh Explore episode.""" + for daemon in tuple(daemons): + if daemon.name == "behavior_env_server": + stop_owned_daemons({"env": daemon}, dashboard_events) + daemons.remove(daemon) + env_daemons, env_kwargs = init_runtime( + args, state_output_dir, dashboard_events, {"env"} + ) + daemons.extend(env_daemons) + if len(env_daemons) != 1: + raise RuntimeError("BEHAVIOR explore requires one owned env daemon per session") + return env_kwargs + + +__all__ = [ + "BEHAVIOR_COMPONENTS", + "BEHAVIOR_MODES", + "DEFAULT_EVAL_COMPONENTS", + "DEFAULT_MAX_EPISODE_STEPS", + "DEFAULT_PLANNER_TIMEOUT_S", + "add_cli_args", + "env_runtime_contract", + "init_runtime", + "on_explore_session", + "parse_config", + "vla_runtime_contract", +] diff --git a/robots/behavior/schemas.py b/robots/behavior/schemas.py new file mode 100644 index 000000000..73dcfaaae --- /dev/null +++ b/robots/behavior/schemas.py @@ -0,0 +1,848 @@ +# 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. + +"""Validated BEHAVIOR/R1Pro observation, action, and public tool contracts.""" + +from __future__ import annotations + +import copy +import math +from collections.abc import Mapping +from typing import Any + +import numpy as np + +from robots.behavior.task_specs import BehaviorTaskSpec, get_task_spec + +ACTION_DIM = 23 +DEFAULT_ACTION_CHUNK = 32 +CAMERA_KEYS = ("main", "left_wrist", "right_wrist") +PHYSICAL_CAMERAS = ("head", "left_wrist", "right_wrist") +HEAD_VIEW_PRESETS = ( + "center", + "up", + "down", + "left", + "right", + "down_left", + "down_right", +) +FRAME_REVIEW_ASSESSMENTS = ( + "target_bearing_surface_confirmed", + "opposite_surface_confirmed", + "side_or_indeterminate", +) + +PUBLIC_TOOL_CONTRACTS: dict[int, tuple[str, ...]] = { + 1: ( + "pi0_nav_pick", + "observe", + "pixel_to_world", + "move_to", + "rotate_wrist", + "close", + "open", + "press", + "save_robot_state_checkpoint", + ), + 2: ( + "pi0_nav_pick", + "observe", + "pixel_to_world", + "move_to", + "rotate_wrist", + "close", + "open", + "press", + "save_robot_state_checkpoint", + "navigate_to", + ), + 3: ( + "pi0_nav_pick", + "observe", + "pixel_to_world", + "move_to", + "rotate_wrist", + "close", + "open", + "press", + "save_robot_state_checkpoint", + "navigate_to", + "move_both_to", + ), + 4: ( + "pi0_nav_pick", + "observe", + "pixel_to_world", + "move_to", + "rotate_wrist", + "close", + "open", + "press", + "save_robot_state_checkpoint", + "navigate_to", + "move_both_to", + "get_prepared_motion_status", + ), + 5: ( + "pi0_nav_pick", + "observe", + "pixel_to_world", + "navigate_to", + "move_to", + "rotate_wrist", + "close", + "open", + "press", + ), +} +CURRENT_PUBLIC_TOOL_CONTRACT_VERSION = 5 +BEHAVIOR_TOOL_NAMES = PUBLIC_TOOL_CONTRACTS[CURRENT_PUBLIC_TOOL_CONTRACT_VERSION] +PUBLIC_PRIMITIVE_ENTRYPOINTS = { + "pi0_nav_pick": "BehaviorPrimitives.pi0_nav_pick", + "observe": "BehaviorPrimitives.observe", + "pixel_to_world": "BehaviorPrimitives.pixel_to_world", + "navigate_to": "BehaviorPrimitives.navigate_to", + "move_to": "BehaviorPrimitives.move_to", + "rotate_wrist": "BehaviorPrimitives.rotate_wrist", + "close": "BehaviorPrimitives.close", + "open": "BehaviorPrimitives.open", + "press": "BehaviorPrimitives.press", +} +if tuple(PUBLIC_TOOL_CONTRACTS) != (1, 2, 3, 4, 5): + raise ValueError("BEHAVIOR public tool contract versions must be contiguous") +if tuple(PUBLIC_PRIMITIVE_ENTRYPOINTS) != BEHAVIOR_TOOL_NAMES: + raise ValueError("BEHAVIOR primitive entrypoints must match the public contract") +if len(BEHAVIOR_TOOL_NAMES) != 9 or len(set(BEHAVIOR_TOOL_NAMES)) != 9: + raise ValueError("BEHAVIOR toolkit must expose 9 unique public primitives") + +POLICY_STATE_SEGMENTS: dict[str, slice] = { + "base": slice(0, 3), + "trunk": slice(3, 7), + "left_arm": slice(7, 14), + "right_arm": slice(14, 21), + "left_gripper": slice(21, 22), + "right_gripper": slice(22, 23), +} +ENV_ACTION_SEGMENTS: dict[str, slice] = { + "base": slice(0, 3), + "trunk": slice(3, 7), + "left_arm": slice(7, 14), + "left_gripper": slice(14, 15), + "right_arm": slice(15, 22), + "right_gripper": slice(22, 23), +} +RAW_PROPRIO_SEGMENTS: dict[str, slice] = { + "left_arm": slice(158, 165), + "left_gripper": slice(193, 195), + "right_arm": slice(197, 204), + "right_gripper": slice(232, 234), + "trunk": slice(236, 240), + "base": slice(253, 256), +} + + +def _validate_segments(name: str, segments: Mapping[str, slice]) -> None: + covered: list[int] = [] + for segment, indices in segments.items(): + if ( + indices.start is None + or indices.stop is None + or indices.step not in (None, 1) + ): + raise ValueError(f"{name}.{segment} must be a contiguous slice") + covered.extend(range(indices.start, indices.stop)) + if covered != list(range(ACTION_DIM)): + raise ValueError( + f"{name} must cover 0..{ACTION_DIM - 1} exactly, got {covered}" + ) + + +_validate_segments("POLICY_STATE_SEGMENTS", POLICY_STATE_SEGMENTS) +_validate_segments("ENV_ACTION_SEGMENTS", ENV_ACTION_SEGMENTS) +if POLICY_STATE_SEGMENTS == ENV_ACTION_SEGMENTS: + raise ValueError("policy state and env action layouts must remain distinct") + + +def segment_ranges(segments: Mapping[str, slice]) -> dict[str, list[int]]: + return {name: [part.start, part.stop] for name, part in segments.items()} + + +def validate_policy_state(state: Any) -> np.ndarray: + array = np.asarray(state, dtype=np.float32) + if array.shape != (ACTION_DIM,): + raise ValueError( + f"compact policy state must be [{ACTION_DIM}], got {array.shape}" + ) + if not np.isfinite(array).all(): + raise ValueError("compact policy state contains NaN or infinity") + return array + + +def extract_policy_state(raw_proprio: Any) -> np.ndarray: + raw = np.asarray(raw_proprio, dtype=np.float32) + if raw.ndim != 1 or raw.shape[0] < RAW_PROPRIO_SEGMENTS["base"].stop: + raise ValueError( + "raw R1Pro proprio must be a vector with at least " + f"{RAW_PROPRIO_SEGMENTS['base'].stop} values, got {raw.shape}" + ) + compact = np.concatenate( + [ + raw[RAW_PROPRIO_SEGMENTS["base"]], + raw[RAW_PROPRIO_SEGMENTS["trunk"]], + raw[RAW_PROPRIO_SEGMENTS["left_arm"]], + raw[RAW_PROPRIO_SEGMENTS["right_arm"]], + np.asarray([raw[RAW_PROPRIO_SEGMENTS["left_gripper"]].sum()]), + np.asarray([raw[RAW_PROPRIO_SEGMENTS["right_gripper"]].sum()]), + ] + ) + return validate_policy_state(compact) + + +def validate_action_chunk( + actions: Any, *, max_horizon: int | None = None +) -> np.ndarray: + array = np.asarray(actions, dtype=np.float32) + if array.ndim != 2 or array.shape[1] != ACTION_DIM or array.shape[0] < 1: + raise ValueError( + f"BEHAVIOR actions must be [T,{ACTION_DIM}], got {array.shape}" + ) + if not np.isfinite(array).all(): + raise ValueError("BEHAVIOR actions contain NaN or infinity") + if max_horizon is not None and array.shape[0] > int(max_horizon): + raise ValueError( + f"BEHAVIOR action horizon {array.shape[0]} exceeds {int(max_horizon)}" + ) + return array + + +ENV_WIRE_SCHEMA: dict[str, Any] = { + "name": "behavior_env_rpc", + "version": 1, + "observation": { + "main_images": "uint8[H,W,3]", + "wrist_images": "uint8[2,H,W,3]", + "states": "float[raw_proprio_dim]", + "task_descriptions": "str", + }, + "action": { + "shape": f"float[T,{ACTION_DIM}]", + "segments": segment_ranges(ENV_ACTION_SEGMENTS), + }, + "official_success_path": ["info", "done", "success"], +} + +VLA_WIRE_SCHEMA: dict[str, Any] = { + "name": "pi05_vla_rpc_behavior", + "version": 2, + "request": { + "method": "vla.predict", + "main_images": "uint8[1,H,W,3]", + "wrist_images": "uint8[1,2,H,W,3]", + "states": "float[1,raw_proprio_dim]", + "task_descriptions": "list[str]", + "extra_view_images": "None", + "compact_state_segments": segment_ranges(POLICY_STATE_SEGMENTS), + "mode": "eval", + }, + "response": { + "actions": f"float[1,T,{ACTION_DIM}]", + "env_action_segments": segment_ranges(ENV_ACTION_SEGMENTS), + }, +} + + +def _planner_spec( + name: str, + description: str, + properties: dict[str, Any], + *, + required: list[str] | None = None, + one_of: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + schema = { + "type": "object", + "properties": properties, + "required": required or [], + "additionalProperties": False, + } + if one_of is not None: + schema["oneOf"] = one_of + return {"name": name, "description": description, "input_schema": schema} + + +_CAMERA_ROLE_SCHEMA = { + "type": "string", + "enum": ["head", "left_wrist", "right_wrist"], +} +_HAND_SCHEMA = {"type": "string", "enum": ["left", "right"]} +_VISUAL_HAND_CHECK_SCHEMA = { + "type": "object", + "properties": { + "camera": {"type": "string", "enum": ["head", "left_wrist", "right_wrist"]}, + "frame_id": {"type": "string", "minLength": 1}, + "selected_hand": {"type": "string", "enum": ["left", "right"]}, + "assessment": {"type": "string", "const": "selected_hand_visually_confirmed"}, + }, + "required": ["camera", "frame_id", "selected_hand", "assessment"], + "additionalProperties": False, +} +_RELEASE_VISUAL_CHECK_SCHEMA = { + "type": "object", + "properties": { + "camera": {"type": "string", "enum": ["head", "left_wrist", "right_wrist"]}, + "frame_id": {"type": "string", "minLength": 1}, + "selected_hand": {"type": "string", "enum": ["left", "right"]}, + "assessment": { + "type": "string", + "const": "attached_object_fully_inside_receptacle_opening", + }, + }, + "required": ["camera", "frame_id", "selected_hand", "assessment"], + "additionalProperties": False, +} +_DELTA_XYZ_SCHEMA = { + "type": "array", + "items": {"type": "number"}, + "minItems": 3, + "maxItems": 3, +} + +PI0_NAV_PICK_SPEC = _planner_spec( + "pi0_nav_pick", + ( + "Invoke a Pi0.5 planner tool supporting navigation, grasping, and " + "pressing. chunks is a positive requested work bound with no fixed " + "maximum; execution is bounded only by the remaining episode steps, raw " + "official success, termination, truncation, or infrastructure failure. " + "task_success reports only raw info.done.success." + ), + { + "instruction": { + "type": "string", + "minLength": 1, + "description": "Exact VLA task language for this invocation.", + }, + "chunks": { + "type": "integer", + "minimum": 1, + "description": "Positive [32,23] action chunk request count.", + }, + }, + required=["instruction", "chunks"], +) + +OBSERVE_SPEC = _planner_spec( + "observe", + "Capture or review a public BEHAVIOR RGB-D observation without advancing physics.", + { + "camera": _CAMERA_ROLE_SCHEMA, + "head_view": {"type": "string", "enum": list(HEAD_VIEW_PRESETS)}, + "paired_hand": _HAND_SCHEMA, + "frame_review": { + "type": "object", + "properties": { + "frame_id": {"type": "string", "minLength": 1}, + "assessment": { + "type": "string", + "enum": list(FRAME_REVIEW_ASSESSMENTS), + }, + }, + "required": ["frame_id", "assessment"], + "additionalProperties": False, + }, + "depth_probe": { + "type": "object", + "properties": { + "frame_id": {"type": "string", "minLength": 1}, + "u": {"type": "integer"}, + "v": {"type": "integer"}, + "depth_window_px": {"type": "integer", "minimum": 1, "maximum": 31}, + "assessment": { + "type": "string", + "const": "target_point_visually_confirmed", + }, + }, + "required": ["frame_id", "u", "v", "depth_window_px", "assessment"], + "additionalProperties": False, + }, + }, + required=["camera"], +) +OBSERVE_SPEC["input_schema"]["allOf"] = [ + {"not": {"required": ["frame_review", "depth_probe"]}}, + { + "if": {"required": ["head_view"]}, + "then": { + "properties": {"camera": {"const": "head"}}, + "not": { + "anyOf": [ + {"required": ["frame_review"]}, + {"required": ["depth_probe"]}, + ] + }, + }, + }, + { + "if": {"required": ["paired_hand"]}, + "then": { + "properties": {"camera": {"const": "head"}}, + "not": { + "anyOf": [ + {"required": ["frame_review"]}, + {"required": ["depth_probe"]}, + ] + }, + }, + }, +] + +PIXEL_TO_WORLD_SPEC = _planner_spec( + "pixel_to_world", + "Back-project one pixel from a fresh public RGB-D frame.", + { + "camera": _CAMERA_ROLE_SCHEMA, + "frame_id": {"type": "string", "minLength": 1}, + "u": {"type": "integer"}, + "v": {"type": "integer"}, + "depth_window_px": { + "type": "integer", + "default": 7, + "minimum": 1, + "maximum": 31, + }, + "target_fact": {"type": "string", "const": "soda_can_floor_outside_receptacle"}, + }, + required=["camera", "frame_id", "u", "v"], +) + +_MOVE_TARGET_SCHEMA = { + "type": "object", + "properties": { + "projection_id": {"type": "string", "minLength": 1}, + "standoff_m": {"type": "number", "minimum": 0.0}, + "delta_xyz": _DELTA_XYZ_SCHEMA, + "frame": {"type": "string", "enum": ["world", "eef"]}, + }, + "oneOf": [ + { + "required": ["projection_id"], + "not": { + "anyOf": [ + {"required": ["delta_xyz"]}, + {"required": ["frame"]}, + ] + }, + }, + { + "required": ["delta_xyz", "frame"], + "not": { + "anyOf": [ + {"required": ["projection_id"]}, + {"required": ["standoff_m"]}, + ] + }, + }, + ], + "additionalProperties": False, +} + +_MOVE_TO_HAND_SCHEMA = {"type": "string", "enum": ["left", "right", "both"]} +_MOVE_TO_BOTH_TARGETS_SCHEMA = { + "type": "object", + "properties": { + "left": { + "type": "object", + "properties": { + "delta_xyz": _DELTA_XYZ_SCHEMA, + "frame": {"type": "string", "enum": ["world", "eef"]}, + }, + "required": ["delta_xyz", "frame"], + "additionalProperties": False, + }, + "right": { + "type": "object", + "properties": { + "delta_xyz": _DELTA_XYZ_SCHEMA, + "frame": {"type": "string", "enum": ["world", "eef"]}, + }, + "required": ["delta_xyz", "frame"], + "additionalProperties": False, + }, + }, + "required": ["left", "right"], + "additionalProperties": False, +} +_MOVE_TO_BOTH_VISUAL_HAND_CHECKS_SCHEMA = { + "type": "object", + "properties": { + "left": _VISUAL_HAND_CHECK_SCHEMA, + "right": _VISUAL_HAND_CHECK_SCHEMA, + }, + "required": ["left", "right"], + "additionalProperties": False, +} + +MOVE_TO_SPEC = _planner_spec( + "move_to", + "Move the selected BEHAVIOR hand, or coordinate both hands when hand is both.", + { + "hand": _MOVE_TO_HAND_SCHEMA, + "visual_hand_check": _VISUAL_HAND_CHECK_SCHEMA, + "target": _MOVE_TARGET_SCHEMA, + "targets": _MOVE_TO_BOTH_TARGETS_SCHEMA, + "visual_hand_checks": _MOVE_TO_BOTH_VISUAL_HAND_CHECKS_SCHEMA, + "support_motion_phase": { + "type": "string", + "enum": ["carry_can", "transit_next_can"], + }, + }, + one_of=[ + { + "properties": {"hand": {"enum": ["left", "right"]}}, + "required": ["hand", "target"], + "not": { + "anyOf": [ + {"required": ["targets"]}, + {"required": ["visual_hand_checks"]}, + ] + }, + }, + { + "properties": {"hand": {"const": "both"}}, + "required": ["hand", "targets", "visual_hand_checks"], + "not": { + "anyOf": [ + {"required": ["target"]}, + {"required": ["visual_hand_check"]}, + {"required": ["support_motion_phase"]}, + ] + }, + }, + ], +) + +ROTATE_WRIST_SPEC = _planner_spec( + "rotate_wrist", + "Rotate one selected BEHAVIOR wrist.", + { + "hand": _HAND_SCHEMA, + "visual_hand_check": _VISUAL_HAND_CHECK_SCHEMA, + "angle_deg": {"type": "number"}, + "direction": {"type": "string", "enum": ["clockwise", "counterclockwise"]}, + }, + required=["hand", "visual_hand_check", "angle_deg"], +) + +CLOSE_SPEC = _planner_spec( + "close", + "Close one selected BEHAVIOR gripper.", + { + "hand": _HAND_SCHEMA, + "visual_hand_check": _VISUAL_HAND_CHECK_SCHEMA, + }, + required=["hand", "visual_hand_check"], +) + +OPEN_SPEC = _planner_spec( + "open", + "Open one selected BEHAVIOR gripper.", + { + "hand": _HAND_SCHEMA, + "visual_hand_check": _VISUAL_HAND_CHECK_SCHEMA, + "release_visual_check": _RELEASE_VISUAL_CHECK_SCHEMA, + }, + required=["hand", "visual_hand_check"], +) + +PRESS_SPEC = _planner_spec( + "press", + "Press with one selected BEHAVIOR hand.", + { + "hand": _HAND_SCHEMA, + "visual_hand_check": _VISUAL_HAND_CHECK_SCHEMA, + "duration_s": {"type": "number", "exclusiveMinimum": 0.0, "maximum": 10.0}, + }, + required=["hand", "visual_hand_check"], +) + +_NAVIGATION_VISUAL_CHECK_SCHEMA = { + "type": "object", + "properties": { + "camera": {"type": "string", "const": "head"}, + "frame_id": {"type": "string", "minLength": 1}, + "assessment": { + "type": "string", + "const": "navigation_target_visually_confirmed", + }, + }, + "required": ["camera", "frame_id", "assessment"], + "additionalProperties": False, +} +_RELATIVE_NAVIGATION_MOTION_SCHEMA = { + "type": "object", + "properties": { + "kind": {"type": "string", "enum": ["translation", "rotation"]}, + "direction": { + "type": "string", + "enum": ["forward", "backward", "left", "right"], + }, + "distance_m": {"type": "number", "exclusiveMinimum": 0.0, "maximum": 1.5}, + "angle_deg": {"type": "number", "exclusiveMinimum": 0.0, "maximum": 180.0}, + }, + "additionalProperties": False, +} +NAVIGATE_TO_SPEC = _planner_spec( + "navigate_to", + "Request a base navigation goal from a projection or explicit relative motion.", + { + "projection_id": {"type": "string", "minLength": 1}, + "navigation_visual_check": _NAVIGATION_VISUAL_CHECK_SCHEMA, + "relative_motion": _RELATIVE_NAVIGATION_MOTION_SCHEMA, + "standoff_m": { + "type": "number", + "default": 0.85, + "minimum": 0.45, + "maximum": 1.5, + }, + }, + one_of=[ + { + "required": ["projection_id", "navigation_visual_check"], + "not": {"required": ["relative_motion"]}, + }, + { + "required": ["relative_motion"], + "not": { + "anyOf": [ + {"required": ["projection_id"]}, + {"required": ["navigation_visual_check"]}, + {"required": ["standoff_m"]}, + ] + }, + }, + ], +) + + +def _non_bool_number(value: Any, *, field: str) -> float: + if isinstance(value, (bool, np.bool_)) or not isinstance( + value, + (int, float, np.integer, np.floating), + ): + raise ValueError(f"{field} must be a finite number") + value = float(value) + if not math.isfinite(value): + raise ValueError(f"{field} must be a finite number") + return value + + +def validate_observe_request( + *, + camera: Any, + head_view: Any = None, + paired_hand: Any = None, + frame_review: Any = None, + depth_probe: Any = None, +) -> dict[str, Any]: + if not isinstance(camera, str) or camera not in PHYSICAL_CAMERAS: + raise ValueError("camera must be head, left_wrist, or right_wrist") + if frame_review is not None and depth_probe is not None: + raise ValueError("frame_review and depth_probe are mutually exclusive") + if head_view is not None: + if not isinstance(head_view, str) or head_view not in HEAD_VIEW_PRESETS: + raise ValueError("head_view must be a supported preset") + if camera != "head": + raise ValueError("head_view is available only for camera='head'") + if frame_review is not None or depth_probe is not None: + raise ValueError("head_view cannot be combined with review/probe") + if paired_hand is not None: + if paired_hand not in {"left", "right"}: + raise ValueError("paired_hand must be 'left' or 'right'") + if camera != "head": + raise ValueError("paired_hand is available only for camera='head'") + if frame_review is not None or depth_probe is not None: + raise ValueError("paired_hand cannot be combined with review/probe") + request: dict[str, Any] = {"camera": camera} + if head_view is not None: + request["head_view"] = head_view + if paired_hand is not None: + request["paired_hand"] = paired_hand + if frame_review is not None: + request["frame_review"] = frame_review + if depth_probe is not None: + request["depth_probe"] = depth_probe + return request + + +def _identifier(value: Any, *, name: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{name} must be a non-empty string") + return value.strip() + + +def validate_relative_navigation_motion(value: Any) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise ValueError("relative_motion must be an object") + motion = dict(value) + kind = motion.get("kind") + direction = motion.get("direction") + if kind == "translation": + expected = {"kind", "direction", "distance_m"} + if direction not in {"forward", "backward"}: + raise ValueError("translation direction must be forward or backward") + amount_name = "distance_m" + maximum = 1.5 + elif kind == "rotation": + expected = {"kind", "direction", "angle_deg"} + if direction not in {"left", "right"}: + raise ValueError("rotation direction must be left or right") + amount_name = "angle_deg" + maximum = 180.0 + else: + raise ValueError("relative_motion.kind must be translation or rotation") + if set(motion) != expected: + raise ValueError(f"relative_motion.{kind} requires exactly {sorted(expected)}") + amount = _non_bool_number( + motion[amount_name], field=f"relative_motion.{amount_name}" + ) + if amount <= 0.0 or amount > maximum: + raise ValueError(f"relative_motion.{amount_name} must be within (0,{maximum}]") + return {"kind": str(kind), "direction": str(direction), amount_name: amount} + + +def validate_visibility_recovery_check(value: Any, *, hand: Any) -> dict[str, str]: + if hand not in {"left", "right"}: + raise ValueError("hand must be 'left' or 'right'") + required = {"view_pair_id", "selected_hand", "assessment"} + if not isinstance(value, Mapping) or set(value) != required: + raise ValueError( + "visibility_recovery_check requires exactly view_pair_id, selected_hand, " + "and assessment" + ) + if value["selected_hand"] != hand: + raise ValueError("hand must equal visibility_recovery_check.selected_hand") + if value["assessment"] != "global_target_and_matching_wrist_observed": + raise ValueError( + "visibility_recovery_check.assessment must be " + "global_target_and_matching_wrist_observed" + ) + return { + "view_pair_id": _identifier(value["view_pair_id"], name="view_pair_id"), + "selected_hand": str(hand), + "assessment": "global_target_and_matching_wrist_observed", + } + + +def validate_move_both_targets(value: Any) -> dict[str, dict[str, Any]]: + if not isinstance(value, Mapping) or set(value) != {"left", "right"}: + raise ValueError("targets requires exactly left and right") + normalized: dict[str, dict[str, Any]] = {} + for hand in ("left", "right"): + target = value[hand] + if not isinstance(target, Mapping) or set(target) != {"delta_xyz", "frame"}: + raise ValueError(f"targets.{hand} requires exactly delta_xyz and frame") + frame = target["frame"] + if frame not in {"world", "eef"}: + raise ValueError(f"targets.{hand}.frame must be world or eef") + delta = target["delta_xyz"] + delta_items = delta.tolist() if isinstance(delta, np.ndarray) else delta + if not isinstance(delta_items, (list, tuple)) or len(delta_items) != 3: + raise ValueError(f"targets.{hand}.delta_xyz must contain three numbers") + values = np.asarray(delta_items, dtype=np.float64) + if values.shape != (3,) or not np.isfinite(values).all(): + raise ValueError(f"targets.{hand}.delta_xyz must contain finite numbers") + normalized[hand] = {"delta_xyz": values.tolist(), "frame": str(frame)} + return normalized + + +def validate_move_both_visual_hand_checks(value: Any) -> dict[str, dict[str, str]]: + if not isinstance(value, Mapping) or set(value) != {"left", "right"}: + raise ValueError("visual_hand_checks requires exactly left and right") + normalized: dict[str, dict[str, str]] = {} + required = {"camera", "frame_id", "selected_hand", "assessment"} + for hand in ("left", "right"): + check = value[hand] + if not isinstance(check, Mapping) or set(check) != required: + raise ValueError(f"visual_hand_checks.{hand} has invalid keys") + if check["camera"] not in {"head", f"{hand}_wrist"}: + raise ValueError(f"visual_hand_checks.{hand}.camera is invalid") + if check["selected_hand"] != hand: + raise ValueError(f"visual_hand_checks.{hand}.selected_hand must be {hand}") + if check["assessment"] != "selected_hand_visually_confirmed": + raise ValueError(f"visual_hand_checks.{hand}.assessment is invalid") + normalized[hand] = { + "camera": str(check["camera"]), + "frame_id": _identifier(check["frame_id"], name="frame_id"), + "selected_hand": hand, + "assessment": "selected_hand_visually_confirmed", + } + return normalized + + +def behavior_tool_specs_for_task( + task: str | BehaviorTaskSpec, +) -> tuple[dict[str, Any], ...]: + task_spec = get_task_spec(task) if isinstance(task, str) else task + specs = { + "pi0_nav_pick": copy.deepcopy(PI0_NAV_PICK_SPEC), + "observe": copy.deepcopy(OBSERVE_SPEC), + "pixel_to_world": copy.deepcopy(PIXEL_TO_WORLD_SPEC), + "navigate_to": copy.deepcopy(NAVIGATE_TO_SPEC), + "move_to": copy.deepcopy(MOVE_TO_SPEC), + "rotate_wrist": copy.deepcopy(ROTATE_WRIST_SPEC), + "close": copy.deepcopy(CLOSE_SPEC), + "open": copy.deepcopy(OPEN_SPEC), + "press": copy.deepcopy(PRESS_SPEC), + } + if task_spec.release_visual_policy is None: + specs["open"]["input_schema"]["properties"].pop("release_visual_check", None) + return tuple(specs[name] for name in BEHAVIOR_TOOL_NAMES) + + +__all__ = [ + "ACTION_DIM", + "BEHAVIOR_TOOL_NAMES", + "CAMERA_KEYS", + "CLOSE_SPEC", + "CURRENT_PUBLIC_TOOL_CONTRACT_VERSION", + "DEFAULT_ACTION_CHUNK", + "ENV_ACTION_SEGMENTS", + "ENV_WIRE_SCHEMA", + "FRAME_REVIEW_ASSESSMENTS", + "HEAD_VIEW_PRESETS", + "MOVE_TO_SPEC", + "NAVIGATE_TO_SPEC", + "OBSERVE_SPEC", + "OPEN_SPEC", + "PI0_NAV_PICK_SPEC", + "PHYSICAL_CAMERAS", + "PIXEL_TO_WORLD_SPEC", + "POLICY_STATE_SEGMENTS", + "PRESS_SPEC", + "PUBLIC_PRIMITIVE_ENTRYPOINTS", + "PUBLIC_TOOL_CONTRACTS", + "ROTATE_WRIST_SPEC", + "VLA_WIRE_SCHEMA", + "behavior_tool_specs_for_task", + "extract_policy_state", + "segment_ranges", + "validate_action_chunk", + "validate_move_both_targets", + "validate_move_both_visual_hand_checks", + "validate_observe_request", + "validate_policy_state", + "validate_relative_navigation_motion", + "validate_visibility_recovery_check", +] diff --git a/robots/behavior/selfcheck.py b/robots/behavior/selfcheck.py new file mode 100644 index 000000000..de23fc75c --- /dev/null +++ b/robots/behavior/selfcheck.py @@ -0,0 +1,71 @@ +# 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. + +"""Minimal import/runtime-contract selfcheck for the BEHAVIOR robot plugin.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + + +def run_import_selfcheck() -> dict[str, Any]: + from robots.behavior.robot_spec import get_robot_spec + from robots.behavior.schemas import BEHAVIOR_TOOL_NAMES + from robots.behavior.task_specs import get_task_spec + + if len(BEHAVIOR_TOOL_NAMES) != 9: + raise RuntimeError("BEHAVIOR must expose exactly 9 public primitives") + + spec = get_robot_spec() + parser = argparse.ArgumentParser(prog="behavior-selfcheck") + spec.add_cli_args(parser, use_dashboard=False) + args = parser.parse_args( + [ + "--task-name", + "turning_on_radio", + "--public-seed", + "1", + ] + ) + args.output_dir = str(Path("/tmp/rpent_behavior_selfcheck")) + config = spec.parse_config(args) + return { + "robot": spec.name, + "robot_spec_fields": sorted(spec.__dataclass_fields__), + "task_name": config.prompt_vars["task_name"], + "activity_instance_id": config.prompt_vars["activity_instance_id"], + "behavior_mode": config.prompt_vars["behavior_mode"], + "memory_profile": config.prompt_vars["memory_profile"], + "memory_dir": config.prompt_vars["memory_dir"], + "behavior_episode_memory": config.prompt_vars["behavior_episode_memory"], + "runtime_components": [ + item["name"] for item in spec.dashboard["runtime_components"] + ], + "tool_count_without_finish": len(BEHAVIOR_TOOL_NAMES), + "radio_task_language": get_task_spec("turning_on_radio").task_language, + } + + +def main() -> None: + print(json.dumps(run_import_selfcheck(), indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() + + +__all__ = ["run_import_selfcheck"] diff --git a/robots/behavior/sft_offline_converter.py b/robots/behavior/sft_offline_converter.py new file mode 100644 index 000000000..f9f7dfddd --- /dev/null +++ b/robots/behavior/sft_offline_converter.py @@ -0,0 +1,663 @@ +# 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. + +"""BEHAVIOR episode-memory library for SFT rollups and runtime catalogs. + +``behavior-build-memory`` consumes ``compile_runtime_catalog``; callers can use +``write_content_addressed_rollup`` to prepare immutable expert episode rollups. +This module has no command-line entry point and does not activate memory. +""" + +from __future__ import annotations + +import hashlib +import io +import json +import os +import re +import tempfile +from collections.abc import Mapping, Sequence +from pathlib import Path +from types import MappingProxyType +from typing import Any + +from robots.behavior.memory.schema import ( + canonical_json_file_bytes, + fail, + require_exact_keys, + require_sha256, + sha256_bytes, +) + +SELECTION_SCHEMA_ID = "rpent_behavior_sft_expert_selection_v1" +ROLLED_ARTIFACT_SCHEMA_ID = "rpent_behavior_sft_offline_rollup_v1" +EXPECTED_TASK_IDS = ("task-0000", "task-0001", "task-0010", "task-0034", "task-0040") +ACTIVE_VIEW_TASKS = {"turning_on_radio", "picking_up_trash"} +EXPECTED_EPISODES = 10 +EXPECTED_SEGMENTS = 91 + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def load_selection(path: Path) -> Mapping[str, Any]: + if not path.is_file() or path.is_symlink(): + fail( + "MEMORY_SFT_SELECTION_MISSING", + str(path), + "selection manifest must be a regular file", + ) + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: + fail("MEMORY_SFT_SELECTION_INVALID", str(path), f"{type(exc).__name__}: {exc}") + if not isinstance(value, Mapping): + fail("MEMORY_SFT_SELECTION_INVALID", str(path), "expected JSON object") + validate_selection(value) + return MappingProxyType(dict(value)) + + +def validate_selection(document: Mapping[str, Any]) -> None: + require_exact_keys( + document, + { + "schema_id", + "created_at", + "preliminary", + "activation_allowed", + "active", + "formal_compiler_admission", + "contract_status", + "source_release", + "coverage", + "evidence_boundaries", + "episodes", + }, + path="$", + ) + if ( + document["schema_id"] != SELECTION_SCHEMA_ID + or document["preliminary"] is not True + or document["activation_allowed"] is not False + or document["active"] is not False + or document["formal_compiler_admission"] is not False + ): + fail("MEMORY_SFT_SELECTION_INVALID", "$", "non-activation identity mismatch") + coverage = document["coverage"] + if not isinstance(coverage, Mapping): + fail("MEMORY_SFT_SELECTION_INVALID", "coverage", "expected object") + expected_coverage = { + "selected_episode_count": EXPECTED_EPISODES, + "selected_segment_count": EXPECTED_SEGMENTS, + "catalog_episode_count": 5, + "query_episode_count": 5, + } + for key, expected in expected_coverage.items(): + if coverage.get(key) != expected: + fail( + "MEMORY_SFT_SELECTION_INVALID", + f"coverage.{key}", + f"expected {expected}", + ) + episodes = document["episodes"] + if not isinstance(episodes, list) or len(episodes) != EXPECTED_EPISODES: + fail("MEMORY_SFT_SELECTION_INVALID", "episodes", "expected 10 episodes") + task_ids = {str(row.get("task_id")) for row in episodes if isinstance(row, Mapping)} + if task_ids != set(EXPECTED_TASK_IDS): + fail( + "MEMORY_SFT_SELECTION_INVALID", + "episodes.task_id", + "expected exact five-task coverage", + ) + segment_count = 0 + for index, episode in enumerate(episodes): + if not isinstance(episode, Mapping): + fail( + "MEMORY_SFT_SELECTION_INVALID", f"episodes[{index}]", "expected object" + ) + for file_key in ("annotation", "metadata", "parquet"): + entry = episode.get(file_key) + if not isinstance(entry, Mapping): + fail( + "MEMORY_SFT_SELECTION_INVALID", + f"episodes[{index}].{file_key}", + "expected object", + ) + require_sha256( + entry.get("sha256"), path=f"episodes[{index}].{file_key}.sha256" + ) + videos = episode.get("videos") + if not isinstance(videos, Mapping) or set(videos) != { + "head", + "left_wrist", + "right_wrist", + }: + fail( + "MEMORY_SFT_SELECTION_INVALID", + f"episodes[{index}].videos", + "expected three camera pins", + ) + for camera, entry in videos.items(): + if not isinstance(entry, Mapping): + fail( + "MEMORY_SFT_SELECTION_INVALID", + f"episodes[{index}].videos.{camera}", + "expected object", + ) + require_sha256( + entry.get("sha256"), path=f"episodes[{index}].videos.{camera}.sha256" + ) + segments = episode.get("segments") + if not isinstance(segments, list) or not segments: + fail( + "MEMORY_SFT_SELECTION_INVALID", + f"episodes[{index}].segments", + "expected non-empty list", + ) + segment_count += len(segments) + if segment_count != EXPECTED_SEGMENTS: + fail( + "MEMORY_SFT_SELECTION_INVALID", "segments", "expected 91 selected segments" + ) + + +def keyframes_for_episode(episode: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]: + frames: dict[int, dict[str, Any]] = {} + segments = episode["segments"] + for segment in segments: + start = int(segment["start_frame"]) + end_exclusive = int(segment["end_frame_exclusive"]) + end = max(start, end_exclusive - 1) + _add_frame(frames, start, "segment_start", segment) + _add_frame(frames, end, "segment_end", segment) + if end_exclusive - start >= 96: + _add_frame( + frames, + start + (end_exclusive - start) // 2, + "long_segment_midpoint", + segment, + ) + first_start = min(int(segment["start_frame"]) for segment in segments) + last_end = max(int(segment["end_frame_exclusive"]) - 1 for segment in segments) + _add_frame(frames, first_start, "episode_first", segments[0]) + _add_frame(frames, last_end, "episode_last", segments[-1]) + return tuple(frames[index] for index in sorted(frames)) + + +def build_rollup( + selection: Mapping[str, Any], *, selection_sha256: str +) -> Mapping[str, Any]: + active: list[Mapping[str, Any]] = [] + sealed: list[Mapping[str, Any]] = [] + for episode in selection["episodes"]: + row = { + "episode_id": episode["episode_id"], + "task_id": episode["task_id"], + "task_name": episode["task_name"], + "split": episode["split"], + "segments": episode["segments"], + "keyframes": list(keyframes_for_episode(episode)), + "source_refs": { + "annotation": episode["annotation"], + "metadata": episode["metadata"], + "parquet": episode["parquet"], + "videos": episode["videos"], + }, + "usage": { + "source": "official_sft_offline", + "active_view": episode["task_name"] in ACTIVE_VIEW_TASKS, + "wrist_policy": "shadow_only", + }, + "outcome": { + "success": None, + "authority": "official_sft_demonstration_without_runtime_success_receipt", + }, + } + if episode["task_name"] in ACTIVE_VIEW_TASKS: + active.append(row) + else: + sealed.append(row) + if len(active) != 4 or len(sealed) != 6: + fail( + "MEMORY_SFT_ROLLUP_INVALID", + "active_view", + "expected Radio/Trash 4 active-view episodes and 6 sealed episodes", + ) + return { + "schema_id": ROLLED_ARTIFACT_SCHEMA_ID, + "preliminary": True, + "activation_allowed": False, + "selection_manifest_sha256": selection_sha256, + "task_count": 5, + "episode_count": 10, + "segment_count": 91, + "keyframe_policy": "episode first/last, segment start/end, long midpoint, dedupe by frame index", + "active_view_policy": "turning_on_radio and picking_up_trash only", + "active_view": active, + "sealed_archive": sealed, + } + + +def write_content_addressed_rollup( + *, selection_manifest: Path, output_dir: Path +) -> Mapping[str, Any]: + raw = selection_manifest.read_bytes() + selection_sha = sha256_bytes(raw) + selection = load_selection(selection_manifest) + artifact = build_rollup(selection, selection_sha256=selection_sha) + payload = canonical_json_file_bytes(artifact, path="rollup") + digest = sha256_bytes(payload) + object_dir = output_dir / "objects" + object_path = object_dir / f"{digest}.json" + _write_once(object_path, payload) + pointer = { + "schema_id": "rpent_behavior_sft_offline_rollup_pointer_v1", + "artifact_sha256": digest, + "artifact_path": str(object_path), + "preliminary": True, + "activation_allowed": False, + } + _atomic_write( + output_dir / "latest.json", canonical_json_file_bytes(pointer, path="pointer") + ) + return MappingProxyType(pointer) + + +def _resolve_source_file( + relative_path: str, roots: Sequence[Path], *, expected_sha256: str +) -> Path: + matches = [ + root / relative_path for root in roots if (root / relative_path).is_file() + ] + if len(matches) != 1: + fail( + "MEMORY_SFT_SOURCE_RESOLUTION_INVALID", + relative_path, + f"expected one source under configured roots, found {len(matches)}", + ) + path = matches[0].resolve() + actual = _sha256_file(path) + if actual != expected_sha256: + fail( + "MEMORY_SFT_SOURCE_HASH_MISMATCH", + relative_path, + f"expected {expected_sha256}, actual {actual}", + ) + return path + + +def _decode_video_frames(path: Path, frame_indices: Sequence[int]) -> list[Any]: + # imageio-ffmpeg is already part of the Behavior optional extra and avoids + # adding OpenCV to the source-plugin contract. + import imageio.v2 as imageio + + try: + reader = imageio.get_reader(str(path), format="ffmpeg") + except Exception as exc: + fail("MEMORY_SFT_VIDEO_INVALID", str(path), f"reader open failed: {exc}") + decoded: list[Any] = [] + try: + for frame_index in frame_indices: + try: + frame = reader.get_data(int(frame_index)) + except Exception as exc: + fail( + "MEMORY_SFT_VIDEO_INVALID", + str(path), + f"cannot decode frame {frame_index}: {type(exc).__name__}: {exc}", + ) + decoded.append(frame) + finally: + reader.close() + return decoded + + +def _encode_in_batches(encoder: Any, images: Sequence[Any], *, batch_size: int) -> Any: + import numpy as np + + rows: list[Any] = [] + for offset in range(0, len(images), batch_size): + batch = encoder.encode_batch(list(images[offset : offset + batch_size])) + rows.extend(item for item in batch if item is not None) + if len(rows) != len(images): + fail("MEMORY_SFT_EMBEDDING_INVALID", "encoder", "missing embedding row") + return np.stack(rows, axis=0).astype(np.float32, copy=False) + + +def _load_episode_rollups(rollups_dir: Path) -> Mapping[str, tuple[Path, str]]: + result: dict[str, tuple[Path, str]] = {} + pattern = re.compile(r"^Episode id: `([^`]+)`\.$", re.MULTILINE) + for path in sorted(rollups_dir.glob("*.memory.md")): + text = path.read_text(encoding="utf-8") + match = pattern.search(text) + if match: + result[match.group(1)] = (path, text) + if len(result) != EXPECTED_EPISODES: + fail( + "MEMORY_SFT_ROLLUP_INVALID", + str(rollups_dir), + "expected 10 episode memory.md rollups", + ) + return MappingProxyType(result) + + +def compile_runtime_catalog( + *, + selection_manifest: Path, + output_dir: Path, + video_roots: Sequence[Path], + rollups_dir: Path, + source_archive: Path, + weights: Path, + cache_dir: Path | None, + batch_size: int, +) -> Mapping[str, Any]: + """Compile all ten official SFT episodes and a four-episode runtime view.""" + + if output_dir.exists(): + fail( + "MEMORY_SFT_OUTPUT_COLLISION", + str(output_dir), + "output directory already exists", + ) + if batch_size < 1 or batch_size > 32: + fail("MEMORY_SFT_BATCH_INVALID", "batch_size", "expected 1..32") + selection_raw = selection_manifest.read_bytes() + selection = load_selection(selection_manifest) + rollups = _load_episode_rollups(rollups_dir) + + # CUDA visibility is set by main() before these imports. + import numpy as np + import torch + import torchvision + + from robots.behavior.dino_v2.encoder import ( + EXPECTED_SOURCE_COMMIT, + MODEL_ID, + MODEL_REVISION, + Dinov2DeploymentPaths, + Dinov2Engine, + Dinov2RevisionIdentity, + ) + from robots.behavior.memory.index import ( + EpisodeExperience, + EpisodeFrameKey, + write_candidate_revision, + ) + + if not torch.cuda.is_available(): + fail( + "MEMORY_SFT_CUDA_UNAVAILABLE", + "cuda", + "compiler requires one visible CUDA device", + ) + identity = Dinov2RevisionIdentity( + model_id=MODEL_ID, + model_revision=MODEL_REVISION, + source_commit=EXPECTED_SOURCE_COMMIT, + source_archive_sha256=_sha256_file(source_archive), + weights_sha256=_sha256_file(weights), + torch_version=str(torch.__version__), + torchvision_version=str(torchvision.__version__), + device="cuda", + ) + encoder = Dinov2Engine( + identity, + Dinov2DeploymentPaths( + source_archive_path=source_archive.resolve(), + weights_path=weights.resolve(), + cache_dir=None if cache_dir is None else cache_dir.resolve(), + ), + ) + active_experiences: list[Any] = [] + active_head: list[Any] = [] + active_left: list[Any] = [] + active_right: list[Any] = [] + all_inventory: list[dict[str, Any]] = [] + all_head: list[Any] = [] + all_left: list[Any] = [] + all_right: list[Any] = [] + try: + for episode in selection["episodes"]: + episode_id = str(episode["episode_id"]) + keyframes = keyframes_for_episode(episode) + frame_indices = [int(item["frame_index"]) for item in keyframes] + encoded_channels: dict[str, Any] = {} + source_videos: dict[str, dict[str, Any]] = {} + for channel in ("head", "left_wrist", "right_wrist"): + video = episode["videos"][channel] + path = _resolve_source_file( + str(video["relative_path"]), + video_roots, + expected_sha256=str(video["sha256"]), + ) + images = _decode_video_frames(path, frame_indices) + encoded_channels[channel] = _encode_in_batches( + encoder, images, batch_size=batch_size + ) + source_videos[channel] = { + "relative_path": str(video["relative_path"]), + "sha256": str(video["sha256"]), + } + all_offset = sum(array.shape[0] for array in all_head) + all_head.append(encoded_channels["head"]) + all_left.append(encoded_channels["left_wrist"]) + all_right.append(encoded_channels["right_wrist"]) + rollup_source, memory_markdown = rollups[episode_id] + active = str(episode["task_name"]) in ACTIVE_VIEW_TASKS + all_inventory.append( + { + "episode_id": episode_id, + "task_name": episode["task_name"], + "active_view": active, + "sealed": not active, + "frame_count": len(keyframes), + "all_embedding_rows": [all_offset, all_offset + len(keyframes)], + "memory_markdown": f"episode_rollups/{rollup_source.name}", + "source_videos": source_videos, + } + ) + if not active: + continue + active_offset = sum(array.shape[0] for array in active_head) + frames = tuple( + EpisodeFrameKey( + frame_id=f"{episode_id}:head:{item['frame_index']}", + episode_id=episode_id, + experience_id=f"episode:{episode_id}", + task_name=str(episode["task_name"]), + frame_index=int(item["frame_index"]), + embedding_row=active_offset + index, + keyframe_kind="+".join(item["keyframe_kinds"]), + source_record_id="+".join(item["source_segment_ids"]), + frame_identity={ + "camera": "head", + "keyframe_kinds": list(item["keyframe_kinds"]), + "source_segment_ids": list(item["source_segment_ids"]), + }, + ) + for index, item in enumerate(keyframes) + ) + active_experiences.append( + EpisodeExperience( + episode_id=episode_id, + experience_id=f"episode:{episode_id}", + logical_experience_id=f"official-sft:{episode_id}", + task_name=str(episode["task_name"]), + usage={ + "returned_scope": "whole_experience", + "episode_memory_markdown": memory_markdown, + "stage_inference": None, + "summary_status": "builder_generated_pending_phase6_review", + }, + outcome={ + "success": None, + "authority": "user_authorized_official_sft_expert_demonstration", + "raw_done_success": None, + }, + frame_keys=frames, + canonical_trajectory_ref={ + "kind": "official_sft_parquet", + **dict(episode["parquet"]), + }, + trajectory_refs=tuple( + {"kind": f"official_sft_{channel}_video", **video} + for channel, video in source_videos.items() + ), + source={ + "selection_manifest_sha256": sha256_bytes(selection_raw), + "episode_split": episode["split"], + "layout_fingerprint_sha256": episode[ + "layout_fingerprint_sha256" + ], + }, + metadata={ + "preliminary": True, + "activation_allowed": False, + "segments": episode["segments"], + "wrist_policy": "shadow_only", + }, + ) + ) + active_head.append(encoded_channels["head"]) + active_left.append(encoded_channels["left_wrist"]) + active_right.append(encoded_channels["right_wrist"]) + finally: + encoder.close() + + output_dir.mkdir(parents=True, exist_ok=False) + episode_rollup_output = output_dir / "episode_rollups" + episode_rollup_output.mkdir() + for _, (source_path, text) in sorted(rollups.items()): + _write_once(episode_rollup_output / source_path.name, text.encode("utf-8")) + all_embedding_bytes = io.BytesIO() + np.savez( + all_embedding_bytes, + head=np.concatenate(all_head, axis=0), + left_wrist=np.concatenate(all_left, axis=0), + right_wrist=np.concatenate(all_right, axis=0), + ) + all_embedding_payload = all_embedding_bytes.getvalue() + _write_once(output_dir / "all_episode_embeddings.npz", all_embedding_payload) + _write_once( + output_dir / "all_episode_inventory.json", + canonical_json_file_bytes( + {"episodes": all_inventory}, path="all_episode_inventory" + ), + ) + candidate = write_candidate_revision( + memory_dir=output_dir / "active_catalog", + experiences=active_experiences, + head_embeddings=np.concatenate(active_head, axis=0), + wrist_shadow_embeddings={ + "left_wrist": np.concatenate(active_left, axis=0), + "right_wrist": np.concatenate(active_right, axis=0), + }, + encoder_identity=identity.as_dict(), + activate_current=True, + ) + manifest = { + "schema_id": "rpent_behavior_sft_episode_catalog_artifact_v1", + "preliminary": True, + "activation_allowed": False, + "connected_to_active_runtime": False, + "source_kind": "user_authorized_official_behavior_sft_training_data", + "selection_manifest_sha256": sha256_bytes(selection_raw), + "task_count": 5, + "episode_count": 10, + "segment_count": 91, + "active_view_episode_count": 4, + "sealed_episode_count": 6, + "keyframe_policy": "episode first/last, segment start/end, long-segment midpoint, deduplicated", + "active_channel": "head", + "wrist_policy": "shadow_only_pending_fresh_policy_query_review", + "stage_evidence_boundary": "SFT expert annotations are preserved as content; runtime retrieval makes no stage inference", + "held_out_observed": False, + "batch_size": batch_size, + "cuda_visible_device_count": int(torch.cuda.device_count()), + "encoder_identity": identity.as_dict(), + "all_episode_embeddings_sha256": sha256_bytes(all_embedding_payload), + "active_catalog_revision_document_sha256": candidate[ + "revision_document_sha256" + ], + "active_catalog_path": "active_catalog", + "sealed_tasks": sorted( + {row["task_name"] for row in all_inventory if row["sealed"]} + ), + } + manifest_payload = canonical_json_file_bytes(manifest, path="artifact_manifest") + _write_once(output_dir / "manifest.json", manifest_payload) + return MappingProxyType( + { + "artifact_dir": str(output_dir), + "manifest_sha256": sha256_bytes(manifest_payload), + "active_catalog_revision_document_sha256": candidate[ + "revision_document_sha256" + ], + "preliminary": True, + "activation_allowed": False, + } + ) + + +def _add_frame( + frames: dict[int, dict[str, Any]], + frame_index: int, + kind: str, + segment: Mapping[str, Any], +) -> None: + frames.setdefault( + frame_index, + { + "frame_index": frame_index, + "keyframe_kinds": [], + "source_segment_ids": [], + }, + ) + row = frames[frame_index] + if kind not in row["keyframe_kinds"]: + row["keyframe_kinds"].append(kind) + segment_id = str(segment["segment_id"]) + if segment_id not in row["source_segment_ids"]: + row["source_segment_ids"].append(segment_id) + + +def _atomic_write(path: Path, payload: bytes) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + mode="wb", prefix=f".{path.name}.", dir=path.parent, delete=False + ) as handle: + tmp = Path(handle.name) + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + try: + os.replace(tmp, path) + finally: + if tmp.exists(): + tmp.unlink() + + +def _write_once(path: Path, payload: bytes) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists(): + if path.is_file() and not path.is_symlink() and path.read_bytes() == payload: + return + fail("MEMORY_SFT_OUTPUT_COLLISION", str(path), "existing bytes differ") + _atomic_write(path, payload) diff --git a/robots/behavior/task_specs.py b/robots/behavior/task_specs.py new file mode 100644 index 000000000..1f839fa7c --- /dev/null +++ b/robots/behavior/task_specs.py @@ -0,0 +1,325 @@ +# 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. + +"""Immutable task-scoped facts for the supported BEHAVIOR tasks.""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal, Mapping + +BehaviorPhase = Literal["explore", "eval"] +InstanceKind = Literal["explore", "eval", "candidate"] + + +@dataclass(frozen=True) +class ReleaseVisualPolicy: + """One task-specific visual authorization contract for object release.""" + + camera: str + assessment: str + + +@dataclass(frozen=True) +class BehaviorInstanceClassification: + """Task-scoped classification of one native activity instance.""" + + task_name: str + instance_id: int + kind: InstanceKind + public_seed: int | None + + +@dataclass(frozen=True) +class BehaviorTaskSpec: + """Immutable identity, mapping, and task policy for one BEHAVIOR task.""" + + task_index: int + task_name: str + task_language: str + prompt_profile_id: str + activity_definition_id: int + scene_model: str + public_seed_to_instance: Mapping[int, int] + mapping_version: str + candidate_mapping_version: str + explore_public_seeds: tuple[int, ...] + eval_public_seeds: tuple[int, ...] + release_visual_policy: ReleaseVisualPolicy | None = None + + def __post_init__(self) -> None: + if ( + isinstance(self.task_index, bool) + or not isinstance(self.task_index, int) + or self.task_index < 0 + ): + raise ValueError("task_index must be a non-negative integer") + if not self.task_name or not self.task_language: + raise ValueError("task_name and task_language must be non-empty") + if self.prompt_profile_id != self.task_name: + raise ValueError("prompt_profile_id must exactly match task_name") + if ( + isinstance(self.activity_definition_id, bool) + or not isinstance(self.activity_definition_id, int) + or self.activity_definition_id < 0 + ): + raise ValueError("activity_definition_id must be non-negative") + if not self.scene_model: + raise ValueError("scene_model must be non-empty") + + mapping = dict(self.public_seed_to_instance) + if not mapping or any( + isinstance(seed, bool) + or not isinstance(seed, int) + or seed < 0 + or isinstance(instance_id, bool) + or not isinstance(instance_id, int) + or instance_id <= 0 + for seed, instance_id in mapping.items() + ): + raise ValueError("public seed mapping must contain positive instance IDs") + if tuple(sorted(mapping)) != tuple(range(len(mapping))): + raise ValueError("public seed mapping must be contiguous from seed 0") + if len(set(mapping.values())) != len(mapping): + raise ValueError("public seed mapping must not reuse an instance") + object.__setattr__(self, "public_seed_to_instance", MappingProxyType(mapping)) + + explore = tuple(self.explore_public_seeds) + evaluate = tuple(self.eval_public_seeds) + if ( + not explore + or len(set(explore)) != len(explore) + or len(set(evaluate)) != len(evaluate) + or set(explore).intersection(evaluate) + or set(explore).union(evaluate) != set(mapping) + ): + raise ValueError( + "Explore and Eval public seeds must partition the public mapping" + ) + object.__setattr__(self, "explore_public_seeds", explore) + object.__setattr__(self, "eval_public_seeds", evaluate) + + if not self.mapping_version or not self.candidate_mapping_version: + raise ValueError("mapping versions must be non-empty") + if self.release_visual_policy is not None: + if self.release_visual_policy.camera != "head": + raise ValueError("release visual policy camera must be head") + if not self.release_visual_policy.assessment: + raise ValueError("release visual assessment must be non-empty") + + @property + def state_dir_name(self) -> str: + return f"{self.scene_model}_task_{self.task_name}_instances" + + def tag(self, public_seed: int) -> str: + self.instance_for_public_seed(public_seed) + return f"{self.task_name}_s{public_seed}" + + def instance_for_public_seed( + self, + public_seed: int, + *, + phase: BehaviorPhase | None = None, + ) -> int: + if isinstance(public_seed, bool) or not isinstance(public_seed, int): + raise ValueError("public_seed must be an integer") + try: + instance_id = self.public_seed_to_instance[public_seed] + except KeyError as error: + raise ValueError( + f"{self.task_name} has no public seed s{public_seed}" + ) from error + if phase is not None: + if phase == "explore": + allowed = self.explore_public_seeds + elif phase == "eval": + allowed = self.eval_public_seeds + else: + raise ValueError(f"unsupported BEHAVIOR phase: {phase!r}") + if public_seed not in allowed: + raise ValueError( + f"{self.task_name} does not allow s{public_seed} in {phase}" + ) + return instance_id + + def public_seed_for_instance(self, instance_id: int) -> int | None: + if ( + isinstance(instance_id, bool) + or not isinstance(instance_id, int) + or instance_id <= 0 + ): + raise ValueError("instance_id must be a positive integer") + return next( + ( + seed + for seed, mapped_instance in self.public_seed_to_instance.items() + if mapped_instance == instance_id + ), + None, + ) + + def classify_instance(self, instance_id: int) -> BehaviorInstanceClassification: + public_seed = self.public_seed_for_instance(instance_id) + if public_seed in self.explore_public_seeds: + kind: InstanceKind = "explore" + elif public_seed in self.eval_public_seeds: + kind = "eval" + else: + kind = "candidate" + return BehaviorInstanceClassification( + task_name=self.task_name, + instance_id=instance_id, + kind=kind, + public_seed=public_seed, + ) + + +_TRASH_RELEASE_VISUAL_POLICY: Final = ReleaseVisualPolicy( + camera="head", + assessment="attached_object_fully_inside_receptacle_opening", +) + +TURNING_ON_RADIO_TASK_SPEC: Final = BehaviorTaskSpec( + task_index=0, + task_name="turning_on_radio", + task_language="Turn on the radio receiver that's on the table in the living room.", + prompt_profile_id="turning_on_radio", + activity_definition_id=0, + scene_model="house_double_floor_lower", + public_seed_to_instance={ + 0: 242, + 1: 109, + 2: 181, + 3: 187, + 4: 197, + 5: 203, + 6: 211, + 7: 212, + 8: 295, + 9: 298, + }, + mapping_version="turning_on_radio_public_seed_v1", + candidate_mapping_version="turning_on_radio_candidate_instance_v1", + explore_public_seeds=(0,), + eval_public_seeds=tuple(range(1, 10)), +) + +PICKING_UP_TRASH_TASK_SPEC: Final = BehaviorTaskSpec( + task_index=1, + task_name="picking_up_trash", + # Match the pinned RLinf behavior_task.jsonl verbatim, including its spelling. + task_language=( + "Put the three can of soda from the living room inside the tash can " + "in the kitchen." + ), + prompt_profile_id="picking_up_trash", + activity_definition_id=0, + scene_model="house_double_floor_lower", + public_seed_to_instance={ + 0: 196, + 1: 67, + 2: 155, + 3: 106, + 4: 161, + 5: 245, + 6: 171, + 7: 156, + 8: 162, + 9: 246, + 10: 108, + 11: 152, + 12: 84, + 13: 198, + 14: 199, + 15: 100, + 16: 111, + 17: 151, + 18: 130, + 19: 168, + }, + mapping_version="picking_up_trash_public_seed_v1", + candidate_mapping_version="picking_up_trash_candidate_instance_v1", + explore_public_seeds=tuple(range(10)), + eval_public_seeds=tuple(range(10, 20)), + release_visual_policy=_TRASH_RELEASE_VISUAL_POLICY, +) + +_TASK_SPECS_BY_NAME: Final[Mapping[str, BehaviorTaskSpec]] = MappingProxyType( + { + spec.task_name: spec + for spec in (TURNING_ON_RADIO_TASK_SPEC, PICKING_UP_TRASH_TASK_SPEC) + } +) +_TASK_SPECS_BY_INDEX: Final[Mapping[int, BehaviorTaskSpec]] = MappingProxyType( + {spec.task_index: spec for spec in _TASK_SPECS_BY_NAME.values()} +) + + +def get_task_spec(task_name: str) -> BehaviorTaskSpec: + try: + return _TASK_SPECS_BY_NAME[task_name] + except (KeyError, TypeError) as error: + raise ValueError(f"unsupported BEHAVIOR task name: {task_name!r}") from error + + +def get_task_spec_by_index(task_index: int) -> BehaviorTaskSpec: + if isinstance(task_index, bool) or not isinstance(task_index, int): + raise ValueError("task_index must be an integer") + try: + return _TASK_SPECS_BY_INDEX[task_index] + except KeyError as error: + raise ValueError(f"unsupported BEHAVIOR task index: {task_index!r}") from error + + +def resolve_task_spec(*, task_name: str, task_index: int) -> BehaviorTaskSpec: + by_name = get_task_spec(task_name) + by_index = get_task_spec_by_index(task_index) + if by_name is not by_index: + raise ValueError( + f"BEHAVIOR task identity mismatch: {task_name!r} != index {task_index}" + ) + return by_name + + +def instance_for_public_seed( + task_name: str, + public_seed: int, + *, + phase: BehaviorPhase | None = None, +) -> int: + return get_task_spec(task_name).instance_for_public_seed(public_seed, phase=phase) + + +def classify_instance( + task_name: str, + instance_id: int, +) -> BehaviorInstanceClassification: + return get_task_spec(task_name).classify_instance(instance_id) + + +__all__ = [ + "BehaviorInstanceClassification", + "BehaviorPhase", + "BehaviorTaskSpec", + "InstanceKind", + "PICKING_UP_TRASH_TASK_SPEC", + "ReleaseVisualPolicy", + "TURNING_ON_RADIO_TASK_SPEC", + "classify_instance", + "get_task_spec", + "get_task_spec_by_index", + "instance_for_public_seed", + "resolve_task_spec", +] diff --git a/robots/behavior/terminal_success.py b/robots/behavior/terminal_success.py new file mode 100644 index 000000000..e96f3d6dd --- /dev/null +++ b/robots/behavior/terminal_success.py @@ -0,0 +1,136 @@ +# 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. + +"""Raw BEHAVIOR success helpers. + +Official task success is only the exact boolean at ``info["done"]["success"]``. +Later snapshots, visual state, videos, and planner ``finish`` status do not +create or revoke that bit. +""" + +from __future__ import annotations + +import copy +import hashlib +import hmac +import json +from collections.abc import Mapping +from typing import Any + +import numpy as np + + +def _canonical_json_bytes(value: Any) -> bytes: + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ).encode("utf-8") + + +def official_success_receipt_sha256(receipt: Mapping[str, Any]) -> str: + """Return the canonical digest after excluding ``receipt_sha256``.""" + + material = {key: value for key, value in receipt.items() if key != "receipt_sha256"} + return hashlib.sha256(_canonical_json_bytes(material)).hexdigest() + + +def validate_official_success_receipt(value: Any) -> dict[str, Any] | None: + """Return a validated official-success receipt copy, or ``None``.""" + + if not isinstance(value, Mapping): + return None + receipt = dict(value) + schema_version = receipt.get("schema_version") + raw_done = receipt.get("raw_done") + env_step = receipt.get("env_step") + digest = receipt.get("receipt_sha256") + if ( + type(schema_version) is not int + or schema_version != 1 + or receipt.get("source") != 'info["done"]["success"]' + or not isinstance(raw_done, Mapping) + or raw_done.get("success") is not True + or isinstance(env_step, bool) + or not isinstance(env_step, int) + or env_step < 0 + or not isinstance(digest, str) + ): + return None + try: + expected = official_success_receipt_sha256(receipt) + except (TypeError, ValueError): + return None + if not hmac.compare_digest(digest, expected): + return None + return copy.deepcopy(receipt) + + +def official_task_success(info: Any) -> bool: + """Return only the raw official BEHAVIOR success bit.""" + + done = info.get("done") if isinstance(info, dict) else None + value = done.get("success") if isinstance(done, dict) else None + return isinstance(value, (bool, np.bool_)) and bool(value) + + +def official_success_receipt_from_info(info: Any) -> dict[str, Any] | None: + """Extract a trusted runtime receipt if one is present and source-matched.""" + + runtime = info.get("_rpent") if isinstance(info, dict) else None + if not isinstance(runtime, dict): + return None + return validate_official_success_receipt(runtime.get("official_success_receipt")) + + +def make_raw_success_receipt( + info: Any, *, env_step: int | None = None +) -> dict[str, Any] | None: + """Return a deterministic receipt for raw success when the env did not provide one.""" + + if not official_task_success(info): + return None + runtime = info.get("_rpent") if isinstance(info, dict) else {} + if isinstance(runtime, dict): + step_value = runtime.get("total_env_steps", runtime.get("global_env_steps")) + else: + step_value = None + if type(step_value) is not int or step_value < 0: + if env_step is None: + step_value = 0 + elif type(env_step) is int and env_step >= 0: + step_value = env_step + else: + return None + material = { + "schema_version": 1, + "source": 'info["done"]["success"]', + "env_step": step_value, + "raw_done": {"success": True}, + } + return { + **material, + "receipt_sha256": official_success_receipt_sha256(material), + } + + +__all__ = [ + "make_raw_success_receipt", + "official_success_receipt_sha256", + "official_success_receipt_from_info", + "official_task_success", + "validate_official_success_receipt", +] diff --git a/robots/behavior/toolkit.py b/robots/behavior/toolkit.py new file mode 100644 index 000000000..093627fa9 --- /dev/null +++ b/robots/behavior/toolkit.py @@ -0,0 +1,219 @@ +# 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. + +"""Standard RPent Toolkit implementation for BEHAVIOR.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np + +from robots.behavior.schemas import behavior_tool_specs_for_task +from robots.behavior.task_specs import get_task_spec +from robots.behavior.tools import BehaviorPrimitives +from rpent.dashboard.events import ( + DashboardEventSink, + NullDashboardEventSink, + ToolResultEvent, +) +from rpent.memory import MemoryManager +from rpent.session import EnvState, write_command_recipe_from_states +from rpent.tools import common +from rpent.tools.toolkit import Toolkit, ToolResult +from rpent.utils.templates import substitute + + +class BehaviorToolkit(Toolkit): + """Expose BEHAVIOR primitives through the latest standard-main contract.""" + + _FRAME_ARTIFACTS = { + "head": "head_rgb.png", + "left_wrist": "left_wrist_rgb.png", + "right_wrist": "right_wrist_rgb.png", + } + + def __init__( + self, + *, + primitives_kwargs: dict[str, Any], + dashboard_events: DashboardEventSink | None = None, + memory: MemoryManager, + config: Any = None, + state_output_dir: str | Path | None = None, + ) -> None: + values = dict(primitives_kwargs) + if config is not None: + prompt_vars = dict(getattr(config, "prompt_vars", {}) or {}) + values.setdefault("task_name", prompt_vars.get("task_name")) + values.setdefault("public_seed", prompt_vars.get("public_seed")) + behavior_phase = prompt_vars.get("behavior_phase") or prompt_vars.get( + "behavior_mode" + ) + if behavior_phase is not None: + values.setdefault("behavior_phase", behavior_phase) + values.setdefault("max_episode_steps", prompt_vars.get("max_episode_steps")) + values.setdefault("output_dir", getattr(config, "output_dir", None)) + output_dir = Path( + values.get("output_dir") or getattr(config, "output_dir", Path.cwd()) + ) + self._run_output_dir = Path(getattr(config, "output_dir", output_dir)) + self._state_output_dir = Path(state_output_dir or output_dir) + values["output_dir"] = self._state_output_dir + + super().__init__( + dashboard_events=dashboard_events or NullDashboardEventSink(), + state=EnvState(self._state_output_dir), + memory=memory, + ) + self._run_state = EnvState(self._run_output_dir) + self._task_spec = get_task_spec( + str(values.get("task_name") or "turning_on_radio") + ) + values["episode_video_writer"] = self._state.open_video_writer( + "episode.mp4", + step=None, + fps=20, + max_frames=2000, + ) + self._primitives = BehaviorPrimitives(**values) + for spec in behavior_tool_specs_for_task(self._task_spec): + if values.get("env") is None: + continue + if spec["name"] == "pi0_nav_pick" and values.get("model") is None: + continue + self.add_tool(spec["name"], spec, getattr(self._primitives, spec["name"])) + finish_spec = next( + spec for spec in common.TOOLS_SPEC if spec["name"] == "finish" + ) + self.add_tool("finish", finish_spec, self._primitives.finish) + + @property + def primitives(self) -> BehaviorPrimitives: + return self._primitives + + def get_tools_spec(self) -> list[dict[str, Any]]: + return substitute( + [spec for spec, _ in self._tools.values()], + variables={"output_dir": str(self._primitives.output_dir)}, + ) + + def execute_tool(self, name: str, input_dict: dict[str, Any]) -> ToolResult: + result = super().execute_tool(name, input_dict) + if self._dashboard_result_has_frames(result.result): + try: + self._dashboard_events.emit( + ToolResultEvent(name=name, result=result.result) + ) + except Exception: + pass + if ( + name == "finish" + and isinstance(result.result, dict) + and result.result.get("_finish") is True + ): + saved = self._state.save("terminal_receipt.json", result.result, step=None) + if saved is None: + raise RuntimeError("failed to write terminal_receipt.json") + return result + + @staticmethod + def _dashboard_result_has_frames(result: Any) -> bool: + if not isinstance(result, dict): + return False + for key in ( + "_image_bytes", + "_depth_image_bytes", + "_image_left_wrist_bytes", + "_depth_left_wrist_bytes", + "_image_right_wrist_bytes", + "_depth_right_wrist_bytes", + ): + if result.get(key): + return True + return False + + def _save_observation_images( + self, observation: dict[str, Any], *, step: int + ) -> None: + head = observation.get("main_images") + wrists = observation.get("wrist_images") + if head is not None: + image = np.asarray(head) + if image.ndim == 3: + self._state.save("head_rgb.png", image[..., :3], step=step) + if wrists is not None: + wrist_array = np.asarray(wrists) + if wrist_array.ndim == 4 and wrist_array.shape[0] >= 2: + self._state.save( + "left_wrist_rgb.png", wrist_array[0, ..., :3], step=step + ) + self._state.save( + "right_wrist_rgb.png", wrist_array[1, ..., :3], step=step + ) + + def get_env_state( + self, + *, + command: dict[str, Any], + result: dict[str, Any], + elapsed_s: float, + ) -> dict[str, Any]: + snapshot = self._primitives.snapshot() + terminated = bool(snapshot.get("task_success")) + with self._state.record_step( + state=snapshot, + terminated=terminated, + truncated=False, + command=command, + result=result, + elapsed_s=elapsed_s, + ) as step: + observation = self._primitives.current_observation + if isinstance(observation, dict): + self._save_observation_images(observation, step=step) + step_idx = step + record = self._state.get(step_idx) + return { + **snapshot, + "step_idx": step_idx, + "artifacts": sorted(record.artifacts), + "command": command, + "result": result, + } + + def close(self) -> None: + """Release clients/transports only; never synthesize task success.""" + + self._primitives.shutdown() + + def solved(self) -> bool: + return self._primitives.solved() + + def write_recipe(self, recipe_tag: str) -> str | None: + """Write an idempotent public recipe JSONL for a solved session.""" + + if not self.solved(): + return None + if not isinstance(recipe_tag, str) or not recipe_tag.strip(): + recipe_tag = self._task_spec.tag(self._primitives.public_seed) + return write_command_recipe_from_states( + self._state, + recipe_tag.strip(), + output_state=self._run_state, + ) + + +__all__ = ["BehaviorToolkit"] diff --git a/robots/behavior/tools.py b/robots/behavior/tools.py new file mode 100644 index 000000000..ef2c3b091 --- /dev/null +++ b/robots/behavior/tools.py @@ -0,0 +1,711 @@ +# 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. + +"""Primitive handlers for the standard-main BEHAVIOR toolkit.""" + +from __future__ import annotations + +import json +import time +from pathlib import Path +from typing import Any + +import numpy as np + +from robots.behavior.schemas import ( + DEFAULT_ACTION_CHUNK, + validate_action_chunk, + validate_move_both_targets, + validate_move_both_visual_hand_checks, + validate_observe_request, + validate_relative_navigation_motion, +) +from robots.behavior.task_specs import get_task_spec +from robots.behavior.terminal_success import ( + make_raw_success_receipt, + official_success_receipt_from_info, + official_task_success, +) +from rpent.tools.toolkit import readonly +from rpent.utils.logging import get_logger + +logger = get_logger("behavior_tools") + +_PRIVATE_RESULT_KEYS = { + "_memory_source", + "activity_instance_id", + "ground_truth", + "gt", + "hidden_state", + "native_instance", + "private_environment_metadata", + "simulator_state", + "suggested_next_action", + "suggested_next_tool", +} +_PUBLIC_IMAGE_BYTE_FIELDS = { + "_image_bytes", + "_depth_image_bytes", + "_image_left_wrist_bytes", + "_depth_left_wrist_bytes", + "_image_right_wrist_bytes", + "_depth_right_wrist_bytes", +} + + +def _jsonable(value: Any) -> Any: + try: + import torch + + if torch.is_tensor(value): + value = value.detach().cpu().numpy() + except Exception: + pass + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, dict): + return {str(key): _jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(item) for item in value] + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return repr(value) + + +def _sanitize_public_result(value: Any) -> Any: + """Remove privileged diagnostics while preserving public image byte fields.""" + + if isinstance(value, dict): + public: dict[str, Any] = {} + for key, item in value.items(): + lowered = str(key).lower() + if lowered == "info": + public[str(key)] = _public_info_summary(item) + continue + if ( + lowered in _PRIVATE_RESULT_KEYS + or lowered.startswith(("ground_truth_", "gt_", "private_")) + or lowered.endswith("_ground_truth") + ): + continue + if lowered in _PUBLIC_IMAGE_BYTE_FIELDS: + public[str(key)] = None if item is None else bytes(item) + continue + public[str(key)] = _sanitize_public_result(item) + return public + if isinstance(value, (list, tuple)): + return [_sanitize_public_result(item) for item in value] + return _jsonable(value) + + +def _public_info_summary(info: Any) -> dict[str, Any]: + if not isinstance(info, dict): + return {} + public: dict[str, Any] = {} + if isinstance(info.get("done"), dict): + public["done"] = _jsonable(info["done"]) + for field in ("executed_steps", "stop_reason", "success_step_in_chunk"): + if field in info: + public[field] = _jsonable(info[field]) + runtime = info.get("_rpent") + if isinstance(runtime, dict): + allowed = { + "total_env_steps", + "global_env_steps", + "attempt_index", + "attempt_nonce", + "run_nonce", + "official_success_receipt", + } + public["_rpent"] = { + key: _jsonable(runtime[key]) for key in allowed if key in runtime + } + return public + + +def _info_from_result(value: Any) -> dict[str, Any] | None: + if ( + isinstance(value, (tuple, list)) + and len(value) == 5 + and isinstance(value[4], dict) + ): + return value[4] + if isinstance(value, dict): + info = value.get("info") + if isinstance(info, dict): + return info + if isinstance(value.get("done"), dict): + return value + return None + + +def _terminal_capture_pointer_from_info(info: Any) -> dict[str, Any] | None: + """Reduce a terminal capture to public linkage fields only.""" + + runtime = info.get("_rpent") if isinstance(info, dict) else None + if not isinstance(runtime, dict): + return None + capture = runtime.get("terminal_capture") + if not isinstance(capture, dict): + capture = runtime.get("vla_after_capture") + if not isinstance(capture, dict): + return None + group_id = capture.get("capture_group_id") + step = capture.get( + "capture_env_step", capture.get("env_step", capture.get("simulator_step")) + ) + frame_ids = capture.get("frame_ids") + if ( + not isinstance(group_id, str) + or not group_id + or isinstance(step, bool) + or not isinstance(step, (int, np.integer)) + or not isinstance(frame_ids, dict) + ): + return None + cameras = ("head", "left_wrist", "right_wrist") + if any( + not isinstance(frame_ids.get(camera), str) or not frame_ids[camera] + for camera in cameras + ): + return None + return { + "capture_group_id": group_id, + "capture_env_step": int(step), + "simulator_step": int(step), + "frame_ids": {camera: frame_ids[camera] for camera in cameras}, + } + + +def _observation_summary(observation: Any) -> dict[str, Any]: + if not isinstance(observation, dict): + return {} + summary: dict[str, Any] = {} + for key, value in observation.items(): + if key in {"main_images", "wrist_images", "states"}: + array = np.asarray(value) + summary[key] = { + "shape": list(array.shape), + "dtype": str(array.dtype), + } + elif key == "task_descriptions": + summary[key] = _jsonable(value) + elif isinstance(value, (str, int, float, bool)) or value is None: + summary[str(key)] = value + return summary + + +class BehaviorPrimitives: + """Handlers registered by :class:`robots.behavior.toolkit.BehaviorToolkit`.""" + + def __init__( + self, + *, + env: Any = None, + model: Any = None, + max_episode_steps: int | None = None, + output_dir: str | Path | None = None, + episode_video_writer: Any = None, + action_horizon: int = DEFAULT_ACTION_CHUNK, + initial_observation: dict[str, Any] | None = None, + initial_info: Any = None, + behavior_phase: str = "eval", + task_name: str = "turning_on_radio", + public_seed: int = 0, + episode_memory_index: Any = None, + dino_component: Any = None, + ) -> None: + self.env = env + self.model = model + self.max_episode_steps = ( + None if max_episode_steps is None else int(max_episode_steps) + ) + self.output_dir = Path(output_dir) if output_dir else Path.cwd() + self._episode_video_writer = episode_video_writer + self._recording = episode_video_writer is not None + self.action_horizon = int(action_horizon) + self._current_observation = initial_observation + self._current_info = initial_info if isinstance(initial_info, dict) else {} + self.behavior_phase = str(behavior_phase) + if self.behavior_phase not in {"eval", "explore"}: + raise ValueError("behavior_phase must be 'eval' or 'explore'") + self.task_spec = get_task_spec(str(task_name)) + self.task_name = self.task_spec.task_name + self.public_seed = int(public_seed) + self.task_spec.instance_for_public_seed(self.public_seed, phase=None) + self.episode_memory_index = episode_memory_index + self.dino_component = dino_component + self._episode_memory_decision = self._retrieve_episode_memory( + self._current_observation + ) + self.started_monotonic = time.monotonic() + self.last_result: dict[str, Any] | None = None + self._local_env_steps = 0 + self._official_success_latched = official_task_success(self._current_info) + self._official_success_receipt = official_success_receipt_from_info( + self._current_info + ) or make_raw_success_receipt(self._current_info, env_step=self.total_env_steps) + if self._recording: + self.record_frame(self._current_observation) + + @property + def elapsed_wall_clock_s(self) -> float: + return max(0.0, time.monotonic() - self.started_monotonic) + + @property + def total_env_steps(self) -> int: + reported = getattr(self.env, "total_env_steps", None) + if isinstance(reported, (int, np.integer)) and not isinstance( + reported, (bool, np.bool_) + ): + return max(self._local_env_steps, int(reported)) + return self._local_env_steps + + @property + def current_observation(self) -> dict[str, Any] | None: + return self._current_observation + + def solved(self) -> bool: + env_solved = bool(getattr(self.env, "official_success_latched", False)) + return bool(self._official_success_latched or env_solved) + + def official_success_receipt(self) -> dict[str, Any] | None: + env_receipt = getattr(self.env, "official_success_receipt", None) + if isinstance(env_receipt, dict): + return _jsonable(env_receipt) + return ( + _jsonable(self._official_success_receipt) + if self._official_success_receipt + else None + ) + + def _remaining_steps(self) -> int | None: + if self.max_episode_steps is None: + return None + return max(0, int(self.max_episode_steps) - self.total_env_steps) + + def _require_env(self) -> Any: + if self.env is None: + raise RuntimeError("BEHAVIOR env component is unavailable") + return self.env + + def _require_model(self) -> Any: + if self.model is None: + raise RuntimeError("BEHAVIOR VLA component is unavailable") + return self.model + + def _note_info(self, info: Any) -> None: + if not isinstance(info, dict): + return + self._current_info = info + runtime = info.get("_rpent") + if isinstance(runtime, dict): + steps = runtime.get("total_env_steps", runtime.get("global_env_steps")) + if isinstance(steps, (int, np.integer)) and not isinstance( + steps, (bool, np.bool_) + ): + self._local_env_steps = max(self._local_env_steps, int(steps)) + if official_task_success(info): + self._official_success_latched = True + self._official_success_receipt = official_success_receipt_from_info( + info + ) or make_raw_success_receipt(info, env_step=self.total_env_steps) + + def recorded_frame_count(self) -> int: + writer = self._episode_video_writer + return int(getattr(writer, "frames_written", 0) or 0) + + def record_frame(self, observation: Any) -> None: + if not self._recording or self._episode_video_writer is None: + return + if not isinstance(observation, dict): + return + image = self._rgb8(observation.get("main_images")) + if image is None: + return + try: + appended = bool(self._episode_video_writer.append(image)) + except Exception as exc: + self._recording = False + logger.warning("failed to append BEHAVIOR episode frame: %s", exc) + return + if not appended: + self._recording = False + logger.warning("stopped BEHAVIOR episode recording after frame limit") + + def stop_recording(self) -> str | None: + writer = self._episode_video_writer + self._recording = False + if writer is None: + return None + try: + return writer.close() + except Exception as exc: + logger.warning("failed to save BEHAVIOR episode video: %s", exc) + return None + + @staticmethod + def _rgb8(value: Any, *, first: int | None = None) -> np.ndarray | None: + if value is None: + return None + image = np.asarray(value) + if first is not None: + if image.ndim != 4 or image.shape[0] <= first: + return None + image = image[first] + elif image.ndim == 4 and image.shape[0] == 1: + image = image[0] + if image.ndim != 3 or image.shape[2] < 3: + return None + image = image[..., :3] + if image.dtype != np.uint8: + if ( + np.issubdtype(image.dtype, np.floating) + and image.size + and float(np.nanmax(image)) <= 1.0 + ): + image = np.rint(np.clip(image, 0.0, 1.0) * 255.0) + image = np.clip(image, 0, 255).astype(np.uint8) + return np.ascontiguousarray(image) + + def _retrieve_episode_memory(self, observation: Any) -> dict[str, Any] | None: + if ( + self.episode_memory_index is None + or self.dino_component is None + or not isinstance(observation, dict) + ): + return None + head = self._rgb8(observation.get("main_images")) + if head is None: + return None + wrists = observation.get("wrist_images") + left = self._rgb8(wrists, first=0) + right = self._rgb8(wrists, first=1) + encoded = self.dino_component.encode_batch([head, left, right]) + head_embedding = encoded[0] + if head_embedding is None: + raise RuntimeError( + "DINO returned no head embedding for episode-memory retrieval" + ) + shadow = { + channel: vector + for channel, vector in zip(("left_wrist", "right_wrist"), encoded[1:]) + if vector is not None + } + decision = self.episode_memory_index.retrieve( + task_name=self.task_name, + head_embedding=head_embedding, + wrist_shadow_embeddings=shadow, + ) + return _jsonable(dict(decision)) + + def _envelope( + self, + name: str, + payload: Any, + *, + primitive_success: bool | None = None, + stop_reason: str | None = None, + ) -> dict[str, Any]: + info = _info_from_result(payload) + self._note_info(info) + if isinstance(payload, dict) and "_observation" in payload: + payload = dict(payload) + observation = payload.pop("_observation") + if isinstance(observation, dict): + self._current_observation = observation + # Backend primitives return their final physical observation; + # do not issue a new RPC after official success just for video. + self.record_frame(observation) + if self.solved() or payload.get("stop_reason") in { + "terminated", + "truncated", + }: + self.stop_recording() + public_payload = _sanitize_public_result(payload) + result: dict[str, Any] = { + "name": name, + "primitive_success": ( + bool(primitive_success) + if primitive_success is not None + else not ( + isinstance(public_payload, dict) and public_payload.get("error") + ) + ), + "task_success": self.solved(), + "official_success_source": 'info["done"]["success"]', + "total_env_steps": self.total_env_steps, + "max_episode_steps": self.max_episode_steps, + } + if stop_reason is not None: + result["stop_reason"] = stop_reason + if isinstance(public_payload, dict): + result.update(public_payload) + else: + result["value"] = public_payload + if self._episode_memory_decision is not None: + result["episode_memory"] = self._episode_memory_decision + if self.solved(): + result["official_success_receipt"] = self.official_success_receipt() + terminal_capture = _terminal_capture_pointer_from_info(self._current_info) + if terminal_capture is not None: + result["terminal_capture"] = terminal_capture + self.last_result = result + return result + + def snapshot(self) -> dict[str, Any]: + return { + "task_name": self.task_name, + "public_seed": self.public_seed, + "behavior_phase": self.behavior_phase, + "task_success": self.solved(), + "official_success_source": 'info["done"]["success"]', + "official_success_receipt": self.official_success_receipt(), + "total_env_steps": self.total_env_steps, + "max_episode_steps": self.max_episode_steps, + "elapsed_wall_clock_s": round(self.elapsed_wall_clock_s, 3), + "observation": _observation_summary(self._current_observation), + "episode_memory": self._episode_memory_decision, + } + + def pi0_nav_pick(self, *, instruction: str, chunks: int) -> dict[str, Any]: + if not isinstance(instruction, str) or not instruction.strip(): + raise ValueError("instruction must be a non-empty string") + if isinstance(chunks, bool) or not isinstance(chunks, int) or chunks <= 0: + raise ValueError("chunks must be a positive integer") + env = self._require_env() + model = self._require_model() + if self.solved(): + return self._envelope( + "pi0_nav_pick", + {}, + primitive_success=True, + stop_reason="already_officially_successful", + ) + + started_steps = self.total_env_steps + chunks_used = 0 + full_chunks = 0 + stop_reason = "exact_requested_chunks" + last_info: dict[str, Any] | None = self._current_info + started = time.monotonic() + + for _ in range(chunks): + remaining = self._remaining_steps() + if remaining is not None and remaining <= 0: + stop_reason = "episode_step_budget_exhausted" + break + if self._current_observation is None: + self._current_observation, last_info = env.current_observation() + self._note_info(last_info) + if self.solved(): + stop_reason = "official_task_success" + break + env_obs = dict(self._current_observation) + env_obs["task_descriptions"] = instruction.strip() + actions = model.predict(env_obs, options={"mode": "eval"}) + action_array = validate_action_chunk(actions) + if remaining is not None: + action_array = action_array[:remaining] + if action_array.shape[0] <= 0: + stop_reason = "episode_step_budget_exhausted" + break + ret = env.chunk_step(action_array, return_all_frames=self._recording) + chunks_used += 1 + obs, _reward, terminated, truncated, info = ret + if isinstance(obs, list): + for frame_obs in obs: + self.record_frame(frame_obs) + if obs and isinstance(obs[-1], dict): + self._current_observation = obs[-1] + elif isinstance(obs, dict): + self._current_observation = obs + self.record_frame(obs) + last_info = info if isinstance(info, dict) else {} + self._note_info(last_info) + executed_steps = None + if isinstance(last_info, dict): + value = last_info.get("executed_steps") + if isinstance(value, (int, np.integer)) and not isinstance( + value, (bool, np.bool_) + ): + executed_steps = int(value) + if executed_steps is None: + executed_steps = int(action_array.shape[0]) + self._local_env_steps = max( + self._local_env_steps, + started_steps + executed_steps, + ) + if executed_steps >= int(action_array.shape[0]): + full_chunks += 1 + if self.solved(): + stop_reason = "official_task_success" + self.stop_recording() + break + if bool(terminated): + stop_reason = "terminated" + self.stop_recording() + break + if bool(truncated): + stop_reason = "truncated" + self.stop_recording() + break + + env_steps_used = max(0, self.total_env_steps - started_steps) + result = { + "terminated": stop_reason == "terminated", + "truncated": stop_reason == "truncated", + "stop_reason": stop_reason, + "requested_chunks": int(chunks), + "chunks_used": chunks_used, + "full_chunks_executed": full_chunks, + "exact_requested_chunks_completed": chunks_used == int(chunks) + and stop_reason == "exact_requested_chunks", + "env_steps_used": env_steps_used, + "total_env_steps": self.total_env_steps, + "max_episode_steps": self.max_episode_steps, + "action_horizon": self.action_horizon, + "required_action_shape": [None, 23], + "elapsed_s": round(time.monotonic() - started, 3), + "info": last_info or {}, + } + return self._envelope( + "pi0_nav_pick", + result, + primitive_success=bool(chunks_used > 0 or self.solved()), + stop_reason=stop_reason, + ) + + @readonly + def observe(self, **kwargs: Any) -> dict[str, Any]: + env = self._require_env() + request = validate_observe_request(**kwargs) + result = env.observe(**request) + info = _info_from_result(result) + if info is not None: + self._note_info(info) + return self._envelope("observe", result, primitive_success=True) + + @readonly + def pixel_to_world(self, **kwargs: Any) -> dict[str, Any]: + env = self._require_env() + result = env.pixel_to_world(**kwargs) + return self._envelope("pixel_to_world", result, primitive_success=True) + + def navigate_to(self, **kwargs: Any) -> dict[str, Any]: + env = self._require_env() + if "relative_motion" in kwargs and kwargs["relative_motion"] is not None: + kwargs = { + **kwargs, + "relative_motion": validate_relative_navigation_motion( + kwargs["relative_motion"] + ), + } + return self._envelope("navigate_to", env.navigate_to(**kwargs)) + + def move_to(self, **kwargs: Any) -> dict[str, Any]: + env = self._require_env() + hand = kwargs.get("hand") + if hand == "both": + kwargs = { + **kwargs, + "targets": validate_move_both_targets(kwargs.get("targets")), + "visual_hand_checks": validate_move_both_visual_hand_checks( + kwargs.get("visual_hand_checks") + ), + } + elif hand not in {"left", "right"}: + raise ValueError("hand must be 'left', 'right', or 'both'") + return self._envelope("move_to", env.move_to(**kwargs)) + + def rotate_wrist(self, **kwargs: Any) -> dict[str, Any]: + env = self._require_env() + return self._envelope("rotate_wrist", env.rotate_wrist(**kwargs)) + + def close(self, **kwargs: Any) -> dict[str, Any]: + env = self._require_env() + return self._envelope("close", env.close_gripper(**kwargs)) + + def open(self, **kwargs: Any) -> dict[str, Any]: + env = self._require_env() + return self._envelope("open", env.open_gripper(**kwargs)) + + def press(self, **kwargs: Any) -> dict[str, Any]: + env = self._require_env() + return self._envelope("press", env.press(**kwargs)) + + @readonly + def finish(self, *, status: str, summary: str) -> dict[str, Any]: + if not isinstance(status, str) or not status.strip(): + raise ValueError("status must be a non-empty string") + if not isinstance(summary, str) or not summary.strip(): + raise ValueError("summary must be a non-empty string") + episode_video_artifact = self.stop_recording() + receipt = { + "schema_version": 1, + "kind": "behavior_finish_terminal_receipt", + "planner_status": status.strip(), + "summary": summary.strip(), + "task_success": self.solved(), + "official_success_source": 'info["done"]["success"]', + "official_success_receipt": self.official_success_receipt(), + "total_env_steps": self.total_env_steps, + "max_episode_steps": self.max_episode_steps, + } + if episode_video_artifact: + receipt["episode_video_artifact"] = episode_video_artifact + receipt["episode_video_frames"] = self.recorded_frame_count() + result = {"_finish": True, **receipt} + self.last_result = result + return result + + def shutdown(self) -> None: + # VLA and DINO belong to the Dashboard Session shared runtime. A + # TaskRun only releases its task-scoped ENV transport; shared daemons + # and clients are released by the runtime owner after the session. + self.stop_recording() + for candidate in (self.env,): + if candidate is None: + continue + closer = getattr(candidate, "close_transport", None) + if not callable(closer): + closer = getattr(candidate, "close", None) + if callable(closer): + try: + closer() + except Exception: + pass + + def recipe_records(self) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + if self.last_result is None: + return records + records.append( + { + "task_name": self.task_name, + "public_seed": self.public_seed, + "last_result": _sanitize_public_result(self.last_result), + } + ) + return json.loads(json.dumps(records, default=str)) + + +__all__ = [ + "BehaviorPrimitives", + "_sanitize_public_result", + "_terminal_capture_pointer_from_info", + "official_task_success", +] diff --git a/rpent/cli/dashboard.py b/rpent/cli/dashboard.py index a90c44fcc..aaf6863a4 100644 --- a/rpent/cli/dashboard.py +++ b/rpent/cli/dashboard.py @@ -179,6 +179,17 @@ def _run_dashboard_task( started = time.time() solved = False memory_manager = None + # BEHAVIOR TaskRuns use behavior_mode; LIBERO uses the shared explore flag. + # A BEHAVIOR TaskRun already owns its fresh ENV (one episode per TaskRun). + toolkit_mode = ( + "exploration" + if ( + getattr(task_args, "behavior_mode", "eval") == "explore" + if args.robot_name == "behavior" + else getattr(task_args, "explore", False) + ) + else "evaluation" + ) try: task_daemons, task_primitives_kwargs = robot_spec.init_runtime( task_args, @@ -231,13 +242,13 @@ def _run_dashboard_task( state.begin_planner_session( video_path=state_output_dir / "episode.mp4", ) - if args.robot_name == "libero": + if args.robot_name in ("libero", "behavior"): toolkit = get_toolkit( args.robot_name, primitives_kwargs=primitives_kwargs, dashboard_events=state, config=run_config, - mode="exploration" if task_args.explore else "evaluation", + mode=toolkit_mode, attempts_per_session=getattr( task_args, "explore_attempts_per_session", 0 ), @@ -277,7 +288,7 @@ def _run_dashboard_task( messages += result.messages stats = result.stats agent_error = result.error - if args.robot_name == "libero": + if args.robot_name in ("libero", "behavior"): solved = toolkit.solved() if solved: recipe_path = toolkit.write_recipe(recipe_tag) @@ -334,7 +345,7 @@ def _run_dashboard_task( init_output_dir(session_root, verbose=args.verbose) if ( - getattr(task_args, "explore", False) + toolkit_mode == "exploration" and getattr(task_args, "auto_merge_memory", False) and not agent_error and not state.task_replacement_requested @@ -351,6 +362,6 @@ def _run_dashboard_task( except Exception as exc: warning = f"memory finalization failed: {type(exc).__name__}: {exc}" logger.warning("%s", warning) - state.report_task_warning(f"Task succeeded, but {warning}") + state.report_task_warning(f"Memory finalization warning: {warning}") return agent_error diff --git a/rpent/cli/main.py b/rpent/cli/main.py index e2faf5cf7..5c7f8070f 100644 --- a/rpent/cli/main.py +++ b/rpent/cli/main.py @@ -320,10 +320,11 @@ def main() -> int: ) args = parser.parse_args() args.robot_name = early.robot_name + on_explore_session = getattr(robot_spec, "on_explore_session", None) if args.dashboard and args.interactive: parser.error("--dashboard and --interactive cannot be used together") - if args.explore and args.robot_name != "libero": - parser.error("--explore is currently supported only for LIBERO") + if args.explore and args.robot_name != "libero" and on_explore_session is None: + parser.error(f"--explore is not supported for {args.robot_name}") if args.explore and args.memory_profile == "hf": parser.error("--explore cannot be used with --memory-profile hf") if args.explore and getattr(args, "explore_sessions", 1) <= 0: @@ -331,12 +332,21 @@ def main() -> int: args.memory_profile = args.memory_profile or ("local" if args.explore else "hf") if args.memory_profile == "hf" and args.memory_dir is not None: parser.error("--memory-dir requires --memory-profile local or --explore") + run_config = None + if args.explore and on_explore_session is not None: + # Let session-owning plugins validate their constraints before startup, + # including requests routed to the Dashboard launcher. + try: + run_config = robot_spec.parse_config(args) + except ValueError as exc: + parser.error(str(exc)) if args.dashboard: from rpent.cli.dashboard import run_dashboard_session return run_dashboard_session(args, robot_spec, parser=parser) - run_config = robot_spec.parse_config(args) + if run_config is None: + run_config = robot_spec.parse_config(args) recipe_tag = run_config.recipe_tag output_dir = run_config.output_dir prompt_vars = run_config.prompt_vars @@ -449,7 +459,13 @@ def main() -> int: state_output_dir = ( output_dir / "sessions" / f"session_{session_number:03d}" ) - if robot_name == "libero": + if args.explore and on_explore_session is not None: + primitives_kwargs.update( + on_explore_session( + args, state_output_dir, dashboard_events, daemons + ) + ) + if robot_name == "libero" or on_explore_session is not None: toolkit = get_toolkit( robot_name, primitives_kwargs=primitives_kwargs, @@ -481,7 +497,7 @@ def main() -> int: messages += result.messages stats = result.stats agent_error = result.error - if robot_name == "libero": + if robot_name == "libero" or on_explore_session is not None: solved = toolkit.solved() if solved: recipe_path = toolkit.write_recipe(recipe_tag) diff --git a/rpent/dashboard/state.py b/rpent/dashboard/state.py index 3904ccd44..90baf243e 100644 --- a/rpent/dashboard/state.py +++ b/rpent/dashboard/state.py @@ -54,6 +54,7 @@ TaskRequest = dict[str, Any] _INTEGER = re.compile(r"-?[0-9]+") _UNSAFE_SLUG = re.compile(r"[^A-Za-z0-9_.-]+") +_OFFICIAL_SUCCESS_SOURCE = 'info["done"]["success"]' def _to_json_safe(value: Any) -> Any: @@ -71,6 +72,46 @@ def _to_json_safe(value: Any) -> Any: return value +def _channel_result_keys(channel: dict[str, Any]) -> tuple[str, ...]: + explicit = channel.get("result_key", channel.get("result_bytes_key")) + if isinstance(explicit, str) and explicit: + return (explicit,) + if isinstance(explicit, (list, tuple)): + return tuple(str(key) for key in explicit if key) + name = str(channel.get("name", "")) + if name == "camera": + return ("_image_cam_bytes", "_image_bytes") + if name == "wrist": + return ("_image_wrist_bytes",) + if name == "head": + return ("_image_bytes",) + if name == "left_wrist": + return ("_image_left_wrist_bytes", "_image_cam_bytes") + if name == "right_wrist": + return ("_image_right_wrist_bytes", "_image_wrist_bytes") + return () + + +def _result_official_success(result: dict[str, Any], *, terminated: bool) -> bool: + """Return whether a tool result contains robot-official success evidence.""" + + if ( + result.get("official_success_source") == _OFFICIAL_SUCCESS_SOURCE + or "official_success_receipt" in result + ): + receipt = result.get("official_success_receipt") + if not isinstance(receipt, dict): + return False + raw_done = receipt.get("raw_done") + return ( + result.get("task_success") is True + and receipt.get("source") == _OFFICIAL_SUCCESS_SOURCE + and isinstance(raw_done, dict) + and raw_done.get("success") is True + ) + return bool(terminated) + + def _parse_task(task_spec: dict[str, Any], text: str) -> TaskRequest | None: tokens = text.split() command = task_spec["command"] @@ -147,6 +188,7 @@ def __init__( self._task_state: str | None = None self._terminated = False self._truncated = False + self._official_success = False self._error: str | None = None self._usage = {"in": 0, "out": 0, "tool_calls": 0} self._planner_usage_base = {"in": 0, "out": 0, "tool_calls": 0} @@ -157,6 +199,10 @@ def __init__( self._events: list[dict[str, Any]] = [] self._timeline: list[dict[str, Any]] = [] self._frames: dict[str, bytes] = {} + self._frame_result_keys = { + channel["name"]: _channel_result_keys(channel) + for channel in self._frame_channels + } self._frame_idx = -1 self.env_state: EnvState | None = None self.frame_artifacts: dict[str, str] = {} @@ -323,6 +369,9 @@ def complete_task( self._task_state = state self._terminated = any(item.get("terminated") for item in self._timeline) self._truncated = any(item.get("truncated") for item in self._timeline) + self._official_success = any( + item.get("official_success") for item in self._timeline + ) self._error = None if error is None else str(error) self._task_replacement_requested = False self._seal_interaction_locked() @@ -349,6 +398,7 @@ def _begin_task_locked( self._task_replacement_requested = False self._terminated = False self._truncated = False + self._official_success = False self._error = None self._usage = {"in": 0, "out": 0, "tool_calls": 0} self._planner_usage_base = {"in": 0, "out": 0, "tool_calls": 0} @@ -655,10 +705,12 @@ def _apply_tool_result(self, event: ToolResultEvent) -> None: result = event.result if not isinstance(result, dict): return - frames = { - "camera": result.get("_image_cam_bytes") or result.get("_image_bytes"), - "wrist": result.get("_image_wrist_bytes"), - } + frames: dict[str, Any] = {} + for kind, keys in self._frame_result_keys.items(): + for key in keys: + if result.get(key): + frames[kind] = result[key] + break self._update_frames( step=result.get("step"), frames={kind: data for kind, data in frames.items() if data}, @@ -676,6 +728,7 @@ def _apply_tool_result(self, event: ToolResultEvent) -> None: action = str(command.get("action", name)) terminated = bool(result.get("terminated")) truncated = bool(result.get("truncated")) + official_success = _result_official_success(result, terminated=terminated) action_video_path = self._action_video_from_result( result, step=step, @@ -691,6 +744,7 @@ def _apply_tool_result(self, event: ToolResultEvent) -> None: "elapsed_s": log.get("elapsed_s"), "terminated": terminated, "truncated": truncated, + "official_success": official_success, "action_video_path": action_video, "action_video_artifact": action_video_artifact, "has_action_video": bool(action_video_artifact or action_video_path), @@ -699,6 +753,7 @@ def _apply_tool_result(self, event: ToolResultEvent) -> None: self._timeline.append(item) self._terminated = self._terminated or terminated self._truncated = self._truncated or truncated + self._official_success = self._official_success or official_success def _action_video_from_result( self, @@ -735,6 +790,10 @@ def on_step(self, record: StepRecord, *, step_offset: int = 0) -> None: "elapsed_s": record.elapsed_s, "terminated": record.terminated, "truncated": record.truncated, + "official_success": _result_official_success( + record.result if isinstance(record.result, dict) else {}, + terminated=record.terminated, + ), "action_video_artifact": action_video, "has_action_video": action_video is not None, } @@ -748,6 +807,9 @@ def on_step(self, record: StepRecord, *, step_offset: int = 0) -> None: ) self._terminated = self._terminated or record.terminated self._truncated = self._truncated or record.truncated + self._official_success = self._official_success or bool( + item["official_success"] + ) def _update_step_frames(self, record: StepRecord, *, display_step: int) -> None: """Load dashboard frame bytes from the step's canonical artifacts.""" @@ -962,6 +1024,7 @@ def snapshot(self) -> dict[str, Any]: "state": self._visible_state_locked(), "terminated": self._terminated, "truncated": self._truncated, + "official_success": self._official_success, "error": self._error, "usage": dict(self._usage), "runtime": self._runtime_snapshot(), @@ -985,6 +1048,7 @@ def run_detail(self) -> dict[str, Any]: "state": self._visible_state_locked(), "terminated": self._terminated, "truncated": self._truncated, + "official_success": self._official_success, "error": self._error, "usage": dict(self._usage), "runtime": self._runtime_snapshot(), diff --git a/rpent/dashboard/static/dashboard.js b/rpent/dashboard/static/dashboard.js index dbe4b7985..e1ad577f1 100644 --- a/rpent/dashboard/static/dashboard.js +++ b/rpent/dashboard/static/dashboard.js @@ -745,12 +745,12 @@ function renderRuntimeStatus(runtime) { container.replaceChildren(...items); } -function setResult(terminated, state) { +function setResult(terminated, state, officialSuccess = false) { const b = $("#resultBadge"); if (state === "succeeded" || terminated) { b.style.display = ""; - b.className = "badge " + (terminated ? "b-ok" : "b-fail"); - b.textContent = terminated ? copy.solved : copy.notSolved; + b.className = "badge " + (officialSuccess ? "b-ok" : "b-fail"); + b.textContent = officialSuccess ? copy.solved : copy.notSolved; } else { b.style.display = "none"; } @@ -1128,7 +1128,7 @@ async function refreshMeta(opts = {}) { const generationState = syncTaskGeneration(r); if (generationState === "stale") return; setBadge(r.state, r.control_error || r.error); - setResult(r.terminated, r.state); + setResult(r.terminated, r.state, r.official_success); renderRuntimeStatus(r.runtime); interactionController.applySnapshot(r); const currentTask = r.current_task; @@ -1164,7 +1164,7 @@ function connectSSE() { const generationState = syncTaskGeneration(sig); if (generationState === "stale") return; setBadge(sig.state, sig.control_error || sig.error); - setResult(sig.terminated, sig.state); + setResult(sig.terminated, sig.state, sig.official_success); renderRuntimeStatus(sig.runtime); interactionController.applySnapshot(sig); mediaState.frameAvailable = sig.frame_available || null; diff --git a/rpent/memory/manager.py b/rpent/memory/manager.py index 53c08b0be..6b9743416 100644 --- a/rpent/memory/manager.py +++ b/rpent/memory/manager.py @@ -275,9 +275,12 @@ def merge_memory( audit = run_dir / f"{cell_tag}.json" recipe = run_dir / f"{cell_tag}_recipe.jsonl" + if not recipe.exists(): + # The generic EnvState command writer uses this filename. + recipe = run_dir / f"recipe_{cell_tag}.jsonl" if solved and audit.exists() and recipe.exists(): audit_target = tiers["task"] / audit.name - recipe_target = tiers["task"] / recipe.name + recipe_target = tiers["task"] / f"{cell_tag}_recipe.jsonl" if not audit_target.exists() and not recipe_target.exists(): shutil.copy2(audit, audit_target) shutil.copy2(recipe, recipe_target) diff --git a/rpent/planner/codex.py b/rpent/planner/codex.py index cf97f3b2c..232457344 100644 --- a/rpent/planner/codex.py +++ b/rpent/planner/codex.py @@ -219,7 +219,7 @@ def solve( elapsed = time.time() - started text = state.get("text", "") or output_path.read_text(errors="replace") - error = error or recorder.error + error = _planner_result_error(error=error, recorder=recorder) logger.info("Codex SDK finished in %.1fs", elapsed) logger.info("output: %s", output_path) @@ -236,6 +236,7 @@ def solve( "raw_stream_path": str(raw_stream_path), "last_message_path": str(last_message_path), "last_message_chars": len(recorder.final_response or ""), + "nonfatal_recorder_errors": _nonfatal_recorder_error_count(recorder), **recorder.stats(), }, error=error, @@ -445,9 +446,14 @@ def emit_user(text: str, *, initial: bool = False) -> None: "raw_stream_path": str(raw_stream_path), "last_message_path": str(last_message_path), "last_message_chars": len(recorder.final_response or ""), + "nonfatal_recorder_errors": _nonfatal_recorder_error_count(recorder), **recorder.stats(), }, - error=error or session.error or recorder.error, + error=_planner_result_error( + error=error, + session_error=session.error, + recorder=recorder, + ), ) # -- config builder ---------------------------------------------------- @@ -640,6 +646,7 @@ class _Recorder: final_response: str | None = None finish_result: dict[str, Any] | None = None error: str | None = None + turn_failed: bool = False def stats(self) -> dict[str, int]: return {"turns_used": self.turns, "tool_calls": self.tool_calls, **self.usage} @@ -726,6 +733,8 @@ def _render_turn_completed(self, turn: Any) -> str: duration_ms = _get(turn, "duration_ms") if error := _get(turn, "error"): self.error = str(_get(error, "message", str(error))) + if status == "failed": + self.turn_failed = True parts = ["[codex-result]", status] if duration_ms is not None: @@ -782,6 +791,35 @@ def _maybe_capture_finish(self, name: str, item: Any) -> None: self.finish_result = {"_finish": True, **args} +def _nonfatal_recorder_error_count(recorder: _Recorder) -> int: + return int( + recorder.finish_result is not None + and recorder.error is not None + and not recorder.turn_failed + ) + + +def _planner_result_error( + *, + error: str | None, + recorder: _Recorder, + session_error: str | None = None, +) -> str | None: + """Return fatal planner error without treating post-finish stream noise as fatal.""" + + if error is not None: + return error + if session_error is not None: + return session_error + if _nonfatal_recorder_error_count(recorder): + logger.warning( + "Codex stream reported a non-fatal event after finish: %s", + recorder.error, + ) + return None + return recorder.error + + # --------------------------------------------------------------------------- # Codex config overrides # --------------------------------------------------------------------------- diff --git a/rpent/robots/components/pi05_vla_client.py b/rpent/robots/components/pi05_vla_client.py index d0fee3519..7d641d8e0 100644 --- a/rpent/robots/components/pi05_vla_client.py +++ b/rpent/robots/components/pi05_vla_client.py @@ -30,6 +30,16 @@ from rpent.robots.components.vla_client_base import BaseVLAClient +_BEHAVIOR_ACTION_DIM = 23 +_BEHAVIOR_RAW_PROPRIO_SEGMENTS: dict[str, slice] = { + "left_arm": slice(158, 165), + "left_gripper": slice(193, 195), + "right_arm": slice(197, 204), + "right_gripper": slice(232, 234), + "trunk": slice(236, 240), + "base": slice(253, 256), +} + # --------------------------------------------------------------------------- # Obs encoder registry # --------------------------------------------------------------------------- @@ -65,10 +75,101 @@ def _batch_view(v): } +def _encode_obs_behavior(env_obs: dict) -> dict: + """BEHAVIOR/R1Pro single-env obs -> openpi batched wire obs.""" + + if not isinstance(env_obs, dict): + raise TypeError("BEHAVIOR observation must be a mapping") + + main = np.asarray(env_obs.get("main_images")) + if main.ndim != 3 or main.shape[-1] != 3: + raise ValueError(f"main_images must be [H,W,3], got {main.shape}") + if main.dtype != np.uint8: + raise TypeError(f"main_images must have dtype uint8, got {main.dtype}") + + wrists = np.asarray(env_obs.get("wrist_images")) + if wrists.ndim != 4 or wrists.shape[0] != 2 or wrists.shape[-1] != 3: + raise ValueError(f"wrist_images must be [2,H,W,3], got {wrists.shape}") + if wrists.dtype != np.uint8: + raise TypeError(f"wrist_images must have dtype uint8, got {wrists.dtype}") + + states = np.asarray(env_obs.get("states"), dtype=np.float32) + if states.ndim != 1: + raise ValueError(f"states must be [raw_proprio_dim], got {states.shape}") + if not np.isfinite(states).all(): + raise ValueError("states contains NaN or infinity") + _extract_behavior_policy_state(states) + + task_description = env_obs.get("task_descriptions") + if isinstance(task_description, (list, tuple)): + instruction = next( + ( + item.strip() + for item in task_description + if isinstance(item, str) and item.strip() + ), + "", + ) + else: + instruction = str(task_description or "") + + return { + "main_images": np.ascontiguousarray(main)[None], + "wrist_images": np.ascontiguousarray(wrists)[None], + "extra_view_images": None, + "states": np.ascontiguousarray(states)[None], + "task_descriptions": [instruction], + } + + +def _validate_behavior_action_chunk( + actions: Any, *, max_horizon: int | None = None +) -> np.ndarray: + array = np.asarray(actions, dtype=np.float32) + if array.ndim != 2 or array.shape[1] != _BEHAVIOR_ACTION_DIM or array.shape[0] < 1: + raise ValueError( + f"BEHAVIOR actions must be [T,{_BEHAVIOR_ACTION_DIM}], got {array.shape}" + ) + if not np.isfinite(array).all(): + raise ValueError("BEHAVIOR actions contain NaN or infinity") + if max_horizon is not None and array.shape[0] > int(max_horizon): + raise ValueError( + f"BEHAVIOR action horizon {array.shape[0]} exceeds {int(max_horizon)}" + ) + return array + + +def _extract_behavior_policy_state(raw_proprio: Any) -> np.ndarray: + raw = np.asarray(raw_proprio, dtype=np.float32) + if raw.ndim != 1 or raw.shape[0] < _BEHAVIOR_RAW_PROPRIO_SEGMENTS["base"].stop: + raise ValueError( + "raw R1Pro proprio must be a vector with at least " + f"{_BEHAVIOR_RAW_PROPRIO_SEGMENTS['base'].stop} values, got {raw.shape}" + ) + compact = np.concatenate( + [ + raw[_BEHAVIOR_RAW_PROPRIO_SEGMENTS["base"]], + raw[_BEHAVIOR_RAW_PROPRIO_SEGMENTS["trunk"]], + raw[_BEHAVIOR_RAW_PROPRIO_SEGMENTS["left_arm"]], + raw[_BEHAVIOR_RAW_PROPRIO_SEGMENTS["right_arm"]], + np.asarray([raw[_BEHAVIOR_RAW_PROPRIO_SEGMENTS["left_gripper"]].sum()]), + np.asarray([raw[_BEHAVIOR_RAW_PROPRIO_SEGMENTS["right_gripper"]].sum()]), + ] + ) + if compact.shape != (_BEHAVIOR_ACTION_DIM,): + raise ValueError( + f"compact policy state must be [{_BEHAVIOR_ACTION_DIM}], got {compact.shape}" + ) + if not np.isfinite(compact).all(): + raise ValueError("compact policy state contains NaN or infinity") + return compact + + # NOTE: an embodiment registered here must also exist in the server's # ``PI05_EMBODIMENTS`` (and ``PI05_ROBOT_PLATFORMS`` if it sets ROBOT_PLATFORM); # the two registries are kept in sync manually. _ENCODE_OBS: dict[str, Any] = { + "behavior": _encode_obs_behavior, "libero": _encode_obs_libero, } @@ -93,6 +194,8 @@ def __init__(self, client, *, embodiment: str): f"registered={list(_ENCODE_OBS)}" ) self._embodiment = embodiment + if embodiment == "behavior": + self._TIMEOUT_S = {**self._TIMEOUT_S, "predict": 600.0} # ---- obs encode (symmetric with server decode_obs_) ---- @@ -105,5 +208,11 @@ def encode_obs(self, env_obs: dict) -> dict: def predict(self, env_obs: dict, options: dict | None = None) -> np.ndarray: """Encode obs, request ``vla.predict``, strip batch dim, return ``[chunk, action_dim]``.""" openpi_obs = self.encode_obs(env_obs) - actions = super().predict(openpi_obs, options) - return np.asarray(actions)[0] + actions = np.asarray(super().predict(openpi_obs, options)) + if self._embodiment == "behavior": + if actions.ndim != 3 or actions.shape[0] != 1: + raise ValueError( + f"BEHAVIOR Pi0.5 actions must be [1,T,23], got {actions.shape}" + ) + return _validate_behavior_action_chunk(actions[0]) + return actions[0] diff --git a/rpent/robots/components/pi05_vla_server.py b/rpent/robots/components/pi05_vla_server.py index d0df87e81..78233b309 100644 --- a/rpent/robots/components/pi05_vla_server.py +++ b/rpent/robots/components/pi05_vla_server.py @@ -22,14 +22,17 @@ from __future__ import annotations import argparse +import hashlib +import json import os import sys +import threading import time -from typing import Any +from contextlib import nullcontext +from pathlib import Path +from typing import Any, Mapping import numpy as np -import torch -from omegaconf import OmegaConf from rpent.robots.components.vla_facade_base import BaseVLAFacade from rpent.utils.config import ( @@ -50,9 +53,128 @@ # Embodiment registry # --------------------------------------------------------------------------- +_BEHAVIOR_ACTION_DIM = 23 + + +def _checkpoint_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _validate_checkpoint_manifest( + path: str | Path, + manifest_path: str | Path, +) -> tuple[str, dict[str, Any]]: + requested = Path(path).expanduser() + try: + resolved = requested.resolve(strict=True) + except OSError as error: + raise ValueError( + f"your Pi0.5 model checkpoint is unavailable: {error}" + ) from error + if not resolved.is_dir(): + raise ValueError(f"your Pi0.5 model checkpoint is not a directory: {resolved}") + + manifest_file = Path(manifest_path).expanduser() + try: + manifest = json.loads(manifest_file.read_text(encoding="utf-8")) + except OSError as error: + raise ValueError(f"checkpoint manifest is unavailable: {error}") from error + except json.JSONDecodeError as error: + raise ValueError(f"checkpoint manifest is invalid JSON: {error}") from error + if not isinstance(manifest, dict): + raise ValueError("checkpoint manifest must be a JSON object") + + manifest_resolved_path = manifest.get("resolved_path") + if manifest_resolved_path is not None: + manifest_resolved = Path(str(manifest_resolved_path)).expanduser() + try: + manifest_resolved = manifest_resolved.resolve(strict=True) + except OSError as error: + raise ValueError( + f"checkpoint manifest resolved_path is unavailable: {error}" + ) from error + if manifest_resolved != resolved: + raise ValueError( + "checkpoint manifest resolved_path does not match --model-path: " + f"{manifest_resolved} != {resolved}" + ) + + files = manifest.get("files") + if not isinstance(files, dict) or not files: + raise ValueError("checkpoint manifest files must be a non-empty object") + for relative_path, requirement in files.items(): + if not isinstance(relative_path, str) or not relative_path: + raise ValueError("checkpoint manifest file paths must be non-empty strings") + if Path(relative_path).is_absolute() or ".." in Path(relative_path).parts: + raise ValueError( + f"checkpoint manifest file path is not a safe relative path: {relative_path}" + ) + if not isinstance(requirement, Mapping): + raise ValueError( + f"checkpoint manifest entry must be an object: {relative_path}" + ) + try: + expected_size = int(requirement["size_bytes"]) + expected_sha256 = str(requirement["sha256"]) + except (KeyError, TypeError, ValueError) as error: + raise ValueError( + "checkpoint manifest entries require size_bytes and sha256: " + f"{relative_path}" + ) from error + candidate = resolved / relative_path + if candidate.is_symlink() or not candidate.is_file(): + raise ValueError(f"checkpoint file is missing or unsafe: {candidate}") + size = candidate.stat().st_size + if size != expected_size: + raise ValueError( + "checkpoint file size mismatch for " + f"{relative_path}: expected {expected_size}, got {size}" + ) + actual_sha256 = _checkpoint_sha256(candidate) + if actual_sha256 != expected_sha256: + raise ValueError( + "checkpoint file SHA256 mismatch for " + f"{relative_path}: expected {expected_sha256}, " + f"got {actual_sha256}" + ) + return str(resolved), manifest + + # NOTE: an embodiment added here must also be registered in the client's # ``_ENCODE_OBS`` (obs encoding); the two registries are kept in sync manually. PI05_EMBODIMENTS: dict[str, dict] = { + "behavior": { + "seed": 0, + "num_action_chunks": 32, + "action_dim": 32, + "use_proprio": True, + "num_steps": 4, + "add_value_head": False, + "openpi_data": { + "norm_stats_path": "assets/behavior-1k/2025-challenge-demos/norm_stats.json", + "extra_delta_transform": False, + "extract_state_from_proprio": True, + "use_all_wrist_images": True, + "use_quantile_norm": True, + }, + "openpi": { + "config_name": "pi05_behavior", + "num_images_in_input": 3, + "action_dim": 32, + "action_horizon": 32, + "action_chunk": 32, + "action_env_dim": 23, + "num_steps": 4, + "add_value_head": False, + "noise_level": 0.0, + "noise_method": "flow_sde", + "joint_logprob": False, + }, + }, "libero": { "num_action_chunks": 5, "action_dim": 7, @@ -71,6 +193,7 @@ } PI05_ROBOT_PLATFORMS: dict[str, str] = { + "behavior": "BEHAVIOR", "libero": "LIBERO", } @@ -109,8 +232,18 @@ def build_model_cfg(model_path: str, emb_cfg: dict) -> Any: for k, v in emb_cfg.items(): if k == "openpi": cfg["openpi"].update(v) + elif k == "openpi_data": + cfg["openpi_data"] = dict(v) else: cfg[k] = v + openpi_data = cfg.get("openpi_data") + if isinstance(openpi_data, dict): + norm_stats_path = openpi_data.get("norm_stats_path") + if norm_stats_path is not None and not Path(norm_stats_path).is_absolute(): + openpi_data["norm_stats_path"] = os.fspath( + Path(model_path) / norm_stats_path + ) + from omegaconf import OmegaConf return OmegaConf.create(cfg) @@ -130,7 +263,13 @@ class Pi05VLAFacade(BaseVLAFacade): Session-isolation is not supported (``reset_session`` is not registered). """ - def __init__(self, *, model_path: str, embodiment: str): + def __init__( + self, + *, + model_path: str, + embodiment: str, + checkpoint_manifest: str | None = None, + ): if embodiment not in PI05_EMBODIMENTS: raise ValueError( f"unknown pi05 server embodiment: {embodiment!r}; " @@ -138,15 +277,30 @@ def __init__(self, *, model_path: str, embodiment: str): ) emb_cfg = PI05_EMBODIMENTS[embodiment] self._embodiment = embodiment + self._model_path = os.fspath(Path(model_path).expanduser()) + self._checkpoint_binding: dict[str, Any] | None = None + self._predict_lock = threading.Lock() super().__init__() + import torch from rlinf.models.embodiment.openpi import get_model as get_openpi_model platform = PI05_ROBOT_PLATFORMS.get(embodiment) if platform is not None: os.environ.setdefault("ROBOT_PLATFORM", platform) - cfg = build_model_cfg(model_path=model_path, emb_cfg=emb_cfg) + if checkpoint_manifest is not None: + self._model_path, self._checkpoint_binding = _validate_checkpoint_manifest( + model_path, + checkpoint_manifest, + ) + seed = emb_cfg.get("seed") + if seed is not None: + torch.manual_seed(int(seed)) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(int(seed)) + + cfg = build_model_cfg(model_path=self._model_path, emb_cfg=emb_cfg) t0 = time.time() logger.info( "loading Pi0.5 (embodiment=%s, model_path=%s) ...", @@ -154,7 +308,29 @@ def __init__(self, *, model_path: str, embodiment: str): cfg["model_path"], ) self._model = get_openpi_model(cfg, torch_dtype=None).cuda().eval() - logger.info("model ready in %.1fs", time.time() - t0) + self._model_meta = { + "status": "ok", + "runtime": "pi05_vla", + "embodiment": embodiment, + "config_name": emb_cfg.get("openpi", {}).get("config_name"), + "action_horizon": emb_cfg.get("openpi", {}).get("action_chunk") + or emb_cfg.get("num_action_chunks"), + "action_dim": emb_cfg.get("openpi", {}).get("action_env_dim") + or emb_cfg.get("action_dim"), + "model_path": self._model_path, + "device": "cuda", + "load_elapsed_s": round(time.time() - t0, 2), + } + if self._checkpoint_binding is not None: + self._model_meta["checkpoint_binding"] = self._checkpoint_binding + logger.info("model ready in %.1fs", self._model_meta["load_elapsed_s"]) + + def _builtin_dispatch( + self, method: str, args: tuple[Any, ...], kwargs: dict[str, Any] + ) -> Any: + if method == "healthz": + return {**self._model_meta, "pid": os.getpid()} + return super()._builtin_dispatch(method, args, kwargs) # ---- inference ---- @@ -164,10 +340,27 @@ def predict(self, obs: dict, options: dict | None = None) -> np.ndarray: The caller (client) is responsible for encoding env-native obs into the openpi wire format (see ``Pi05VLAClient.encode_obs``). """ + if self._embodiment == "behavior": + if options is None: + options = {} + if not isinstance(options, dict): + raise TypeError("VLA options must be a mapping") + unexpected = set(options) - {"mode"} + if unexpected: + raise ValueError(f"unsupported VLA options: {sorted(unexpected)!r}") mode = (options or {}).get("mode", "eval") - with torch.no_grad(): - actions, _ = self._model.predict_action_batch(obs, mode=mode) - return ( + if self._embodiment == "behavior": + if mode != "eval": + raise ValueError("BEHAVIOR VLA inference mode must be 'eval'") + predict_kwargs = {"mode": mode} + if self._embodiment == "behavior": + predict_kwargs["compute_values"] = False + lock = self._predict_lock if self._embodiment == "behavior" else nullcontext() + import torch + + with lock, torch.no_grad(): + actions, _ = self._model.predict_action_batch(obs, **predict_kwargs) + result = ( actions.detach().cpu().numpy() if ( hasattr(actions, "detach") @@ -176,6 +369,19 @@ def predict(self, obs: dict, options: dict | None = None) -> np.ndarray: ) else np.asarray(actions) ).astype(np.float32) + if self._embodiment == "behavior": + if ( + result.ndim != 3 + or result.shape[0] != 1 + or result.shape[1] < 1 + or result.shape[2] != _BEHAVIOR_ACTION_DIM + or not np.isfinite(result).all() + ): + raise ValueError( + "Pi0.5 returned invalid " + f"[1,T,{_BEHAVIOR_ACTION_DIM}] shape {result.shape}" + ) + return result # --------------------------------------------------------------------------- @@ -209,6 +415,14 @@ def main() -> None: default=None, help="Pi0.5 checkpoint (defaults to PI05_CHECKPOINT_PATH env)", ) + p.add_argument( + "--checkpoint-manifest", + default=None, + help=( + "Optional generic JSON manifest describing required checkpoint files " + "and SHA-256 values." + ), + ) args = p.parse_args() if args.cuda_device is not None: @@ -229,7 +443,11 @@ def main() -> None: "path via --model-path or the environment." ) - facade = Pi05VLAFacade(model_path=model_path, embodiment=args.embodiment) + facade = Pi05VLAFacade( + model_path=model_path, + embodiment=args.embodiment, + checkpoint_manifest=args.checkpoint_manifest, + ) facade.serve( transport=args.transport, host=args.host, diff --git a/rpent/robots/robot_spec.py b/rpent/robots/robot_spec.py index 21ec56b13..11b8a25aa 100644 --- a/rpent/robots/robot_spec.py +++ b/rpent/robots/robot_spec.py @@ -65,3 +65,14 @@ class RobotSpec: dashboard: dict[str, Any] | None = None memory_repo_id: str = "RLinf/RPent-memory" finalize_run: RunFinalizer | None = None + # Optional session-scoped runtime refresh before Explore toolkit creation. + # Mutate the runner's owned-daemon list and return primitive kwargs updates. + # init_runtime(..., None) remains responsible for shared runtime startup. + # Opts into the CLI Explore toolkit mode/state-dir and solved/recipe contract. + on_explore_session: ( + Callable[ + [argparse.Namespace, Path, DashboardEventSink, list["ProcessDaemon"]], + dict[str, Any], + ] + | None + ) = None diff --git a/rpent/session/__init__.py b/rpent/session/__init__.py index 56e6d989f..4e51548e0 100644 --- a/rpent/session/__init__.py +++ b/rpent/session/__init__.py @@ -14,6 +14,18 @@ """Single-session, mutable state (EnvState).""" -from rpent.session.base import EnvState, StepRecord +from rpent.session.base import ( + EnvState, + StepRecord, + VideoArtifactWriter, + recipe_commands_from_states, + write_command_recipe_from_states, +) -__all__ = ["EnvState", "StepRecord"] +__all__ = [ + "EnvState", + "StepRecord", + "VideoArtifactWriter", + "write_command_recipe_from_states", + "recipe_commands_from_states", +] diff --git a/rpent/session/base.py b/rpent/session/base.py index 1e8b3146b..57e9d0ee7 100644 --- a/rpent/session/base.py +++ b/rpent/session/base.py @@ -292,6 +292,36 @@ def exists(self, name: str, *, step: int | None = -1) -> bool: return False return self._artifact_file(name, resolved_step).exists() + def _register_artifact(self, name: str, step: int | None) -> None: + if step is not None: + self._record_for(step).artifacts.add(name) + else: + self._run_artifacts.add(name) + if not self._step_open: + self._write_manifest() + + def open_video_writer( + self, + name: str, + *, + step: int | None = None, + fps: int = 20, + max_frames: int | None = None, + ) -> "VideoArtifactWriter": + """Open a streaming MP4 writer for an EnvState artifact.""" + + resolved_step = self._resolve_read_step(step) + path = self._artifact_file(name, resolved_step) + if path.suffix.lower() != ".mp4": + raise ValueError("video artifacts must use .mp4") + return VideoArtifactWriter( + self, + name=name, + step=resolved_step, + fps=fps, + max_frames=max_frames, + ) + # -- step records ---------------------------------------------------- @contextmanager @@ -359,3 +389,137 @@ def get(self, step: int = -1) -> StepRecord: def records(self) -> list[StepRecord]: return copy.deepcopy(self._steps) + + +def recipe_commands_from_states(env_state: EnvState) -> list[dict[str, Any]]: + """Return replayable top-level tool commands from stateful records.""" + + commands: list[dict[str, Any]] = [] + for record in env_state.records(): + command = record.command + if not isinstance(command, dict) or command.get("action") is None: + continue + if isinstance(record.result, dict) and record.result.get("error"): + continue + commands.append(copy.deepcopy(command)) + return commands + + +def write_command_recipe_from_states( + env_state: EnvState, + recipe_tag: str, + *, + output_state: EnvState | None = None, +) -> str | None: + """Write a generic ``recipe_.jsonl`` command-sequence artifact. + + This shared helper exports top-level replay commands from generic + :class:`EnvState` records. It is intentionally distinct from robot-specific + recipe writers such as LIBERO's ``write_recipe_from_states``, which also + applies reset-window filtering, segment artifacts, and solved gating. + """ + + tag = str(recipe_tag).strip() + if not tag: + raise ValueError("recipe_tag must be a non-empty string") + commands = recipe_commands_from_states(env_state) + if not commands: + return None + target_state = output_state or env_state + name = f"recipe_{tag}.jsonl" + saved = target_state.save(name, commands, step=None) + if saved is None: + raise RuntimeError(f"failed to write {name}") + return str(target_state.artifact_path(saved, step=None)) + + +class VideoArtifactWriter: + """Streaming MP4 writer that registers the artifact only on successful close.""" + + def __init__( + self, + env_state: EnvState, + *, + name: str, + step: int | None, + fps: int, + max_frames: int | None, + ) -> None: + if not isinstance(fps, int) or fps <= 0: + raise ValueError("fps must be a positive integer") + if max_frames is not None and int(max_frames) <= 0: + raise ValueError("max_frames must be positive when provided") + self._env_state = env_state + self._name = env_state._validate_name(name) + self._step = step + self._fps = fps + self._max_frames = None if max_frames is None else int(max_frames) + self._destination = env_state._artifact_file(self._name, step) + self._temporary = self._destination.with_name( + f".{self._destination.stem}.{os.getpid()}.{id(self)}.tmp" + f"{self._destination.suffix}" + ) + self._writer: Any | None = None + self._closed = False + self._aborted = False + self.frames_written = 0 + self.frames_dropped = 0 + + def append(self, frame: Any) -> bool: + """Append one RGB frame; return False when the configured cap is reached.""" + + if self._closed: + raise RuntimeError("video writer is closed") + if self._aborted: + return False + if self._max_frames is not None and self.frames_written >= self._max_frames: + self.frames_dropped += 1 + return False + array = np.asarray(frame) + if array.ndim != 3 or array.shape[2] < 3: + raise ValueError("video frame must have shape [H, W, C>=3]") + array = array[..., :3] + if array.dtype != np.uint8: + array = np.clip(array, 0, 255).astype(np.uint8) + if self._writer is None: + self._destination.parent.mkdir(parents=True, exist_ok=True) + self._writer = imageio.get_writer(self._temporary, fps=self._fps) + self._writer.append_data(np.ascontiguousarray(array)) + self.frames_written += 1 + return True + + def close(self) -> str | None: + """Finalize the video and register it in the EnvState manifest.""" + + if self._closed: + return self._name if self._destination.exists() else None + self._closed = True + try: + if self._writer is not None: + self._writer.close() + self._writer = None + if self._aborted or self.frames_written <= 0: + self._temporary.unlink(missing_ok=True) + return None + os.replace(self._temporary, self._destination) + self._env_state._register_artifact(self._name, self._step) + return self._name + except Exception: + self.abort() + raise + finally: + self._temporary.unlink(missing_ok=True) + + def abort(self) -> None: + """Close and remove the temporary file without publishing the artifact.""" + + if self._aborted: + return + self._aborted = True + try: + if self._writer is not None: + self._writer.close() + self._writer = None + finally: + self._closed = True + self._temporary.unlink(missing_ok=True) diff --git a/rpent/tools/toolkit.py b/rpent/tools/toolkit.py index 871673470..21ca334c1 100644 --- a/rpent/tools/toolkit.py +++ b/rpent/tools/toolkit.py @@ -134,9 +134,14 @@ def _build_content_blocks(self) -> list[dict[str, Any]]: result_for_text = dict(result) image = result_for_text.pop("_image_bytes", None) + depth_image = result_for_text.pop("_depth_image_bytes", None) image_cam = result_for_text.pop("_image_cam_bytes", None) image_nav = result_for_text.pop("_image_nav_bytes", None) image_wrist = result_for_text.pop("_image_wrist_bytes", None) + image_left_wrist = result_for_text.pop("_image_left_wrist_bytes", None) + depth_left_wrist = result_for_text.pop("_depth_left_wrist_bytes", None) + image_right_wrist = result_for_text.pop("_image_right_wrist_bytes", None) + depth_right_wrist = result_for_text.pop("_depth_right_wrist_bytes", None) text = json.dumps(result_for_text, indent=2, default=str) text = _truncate_utf8( text, @@ -161,12 +166,22 @@ def _add_image_bytes(data_bytes: bytes) -> None: if image: _add_image_bytes(image) + if depth_image: + _add_image_bytes(depth_image) if image_cam: _add_image_bytes(image_cam) if image_nav: _add_image_bytes(image_nav) if image_wrist: _add_image_bytes(image_wrist) + if image_left_wrist: + _add_image_bytes(image_left_wrist) + if depth_left_wrist: + _add_image_bytes(depth_left_wrist) + if image_right_wrist: + _add_image_bytes(image_right_wrist) + if depth_right_wrist: + _add_image_bytes(depth_right_wrist) return blocks diff --git a/rpent/utils/daemon.py b/rpent/utils/daemon.py index da7e53a70..30d6eaac9 100644 --- a/rpent/utils/daemon.py +++ b/rpent/utils/daemon.py @@ -45,7 +45,11 @@ def watch_parent_death(on_death: Callable[[], None]) -> None: def _watch() -> None: try: - sys.stdin.buffer.read() + # A daemon thread must not hold BufferedReader's lock while Python + # finalizes stdin (that aborts the interpreter on RPC shutdown). + fd = sys.stdin.fileno() + while os.read(fd, 65536): + pass except Exception: pass on_death() diff --git a/tests/unit_tests/robots/behavior/test_behavior_contracts.py b/tests/unit_tests/robots/behavior/test_behavior_contracts.py new file mode 100644 index 000000000..03e07989f --- /dev/null +++ b/tests/unit_tests/robots/behavior/test_behavior_contracts.py @@ -0,0 +1,1387 @@ +# 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. + +"""Offline contracts for the BEHAVIOR toolkit.""" + +from __future__ import annotations + +import json +import os +import signal +import socket +import sys +import threading +import time +from pathlib import Path +from typing import Any + +import numpy as np +import pytest + +from robots.behavior import build_memory_cli +from robots.behavior.dino_v2.client import BehaviorDinoClient +from robots.behavior.dino_v2.encoder import DINOV2_DIMENSION +from robots.behavior.dino_v2.server import BehaviorDinoFacade +from robots.behavior.env_client import BehaviorEnvClient +from robots.behavior.env_server import BehaviorEnvFacade +from robots.behavior.rlinf_env import ( + GRIPPER_CLOSE_COMMAND, + GRIPPER_COMMAND_CONTROL_CYCLES, + GRIPPER_OPEN_COMMAND, + OfficialBehaviorBackend, +) +from robots.behavior.robot_spec import get_toolkit +from robots.behavior.schemas import ( + BEHAVIOR_TOOL_NAMES, + ENV_ACTION_SEGMENTS, + MOVE_TO_SPEC, + RAW_PROPRIO_SEGMENTS, + validate_action_chunk, +) +from robots.behavior.toolkit import BehaviorToolkit +from robots.behavior.tools import BehaviorPrimitives +from rpent.dashboard.events import NullDashboardEventSink +from rpent.memory import MemoryManager +from rpent.robots import RunConfig +from rpent.utils.daemon import ProcessDaemon, pick_free_port +from rpent.utils.rpc.http_rpc import HttpRpcClient + +EXPECTED_TOOLS = ( + "pi0_nav_pick", + "observe", + "pixel_to_world", + "navigate_to", + "move_to", + "rotate_wrist", + "close", + "open", + "press", +) + + +def test_episode_memory_is_structured_in_public_results(): + from robots.behavior.memory.index import empty_episode_memory_index + + class Encoder: + def encode_batch(self, images): + vector = np.zeros(DINOV2_DIMENSION, dtype=np.float32) + vector[0] = 1 + return [vector if image is not None else None for image in images] + + class Env: + def observe(self, **kwargs): + return {"status": "ok", "info": {"done": {"success": False}}} + + primitives = BehaviorPrimitives( + env=Env(), + initial_observation={"main_images": np.zeros((8, 8, 3), dtype=np.uint8)}, + episode_memory_index=empty_episode_memory_index(), + dino_component=Encoder(), + ) + decision = primitives.snapshot()["episode_memory"] + assert isinstance(decision, dict) + assert decision["decision"] == "record_new" + assert decision["candidate_count_after_task_filter"] == 0 + assert decision["stage_inference"] is None + assert primitives.observe(camera="head")["episode_memory"] == decision + assert json.loads(json.dumps(decision)) == decision + + +@pytest.mark.parametrize( + ("task_name", "mode", "seed", "instruction"), + [ + ( + "turning_on_radio", + "explore", + 0, + "Turn on the radio receiver that's on the table in the living room.", + ), + ( + "picking_up_trash", + "explore", + 0, + "Put the three can of soda from the living room inside the tash can " + "in the kitchen.", + ), + ( + "picking_up_trash", + "eval", + 10, + "Put the three can of soda from the living room inside the tash can " + "in the kitchen.", + ), + ], +) +@pytest.mark.parametrize("batched", [False, True]) +def test_runtime_preserves_rlinf_task_language( + tmp_path: Path, + task_name: str, + mode: str, + seed: int, + instruction: str, + batched: bool, +) -> None: + import argparse + + from robots.behavior import runtime + + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir") + runtime.add_cli_args(parser, use_dashboard=False) + args = parser.parse_args( + [ + "--task-name", + task_name, + "--public-seed", + str(seed), + "--behavior-mode", + mode, + "--behavior-repo", + str(tmp_path / "RLinf"), + "--output-dir", + str(tmp_path / "output"), + ] + ) + config = runtime.parse_config(args) + meta = runtime.env_runtime_contract(args) + assert meta["task_language"] == instruction + for key in ("task_language", "task_instruction", "instruction"): + assert config.prompt_vars[key] == instruction + + class Rpc: + language = instruction + + def call(self, method, **kwargs): + if method == "env.get_env_meta": + return meta + assert method == "env.reset" + text = [self.language] if batched else self.language + return {"task_descriptions": text}, {"done": {"success": False}} + + rpc = Rpc() + connected = runtime._connect_env(args, rpc, config.output_dir) + expected = [instruction] if batched else instruction + assert connected["initial_observation"]["task_descriptions"] == expected + assert connected["env"].official_success_latched is False + + rpc.language = "Put the three soda cans from the living room inside the trash can in the kitchen." + with pytest.raises(RuntimeError, match="task language does not match TaskSpec"): + runtime._connect_env(args, rpc, config.output_dir) + + +def test_env_endpoint_discovery_uses_actual_bind_and_ignores_old_log(tmp_path: Path): + from robots.behavior.runtime import _wait_for_server_endpoint + + log = tmp_path / "env.log" + log.write_text("RPC server listening on http://127.0.0.1:1\n") + offset = log.stat().st_size + daemon = ProcessDaemon( + "test_env", + [ + sys.executable, + "-c", + "from types import SimpleNamespace; " + "from robots.behavior.env_server import BehaviorEnvFacade; " + "print('Ray started; no application endpoint yet', flush=True); " + "BehaviorEnvFacade(backend=SimpleNamespace(close=lambda: None), " + "meta={'task_language': 'test'}).serve(" + "transport='http', host='127.0.0.1', port=0, parent_watch=True)", + ], + log_path=str(log), + ) + daemon.start() + rpc = None + try: + endpoint = _wait_for_server_endpoint(daemon, log_offset=offset) + assert endpoint != "http://127.0.0.1:1" + rpc = HttpRpcClient(endpoint) + assert rpc.call("healthz") == {"status": "ok"} + assert rpc.call("env.get_env_meta") == {"task_language": "test"} + assert rpc.call("shutdown") == {"ok": True} + finally: + if rpc is not None: + rpc.close() + daemon.stop() + + +def test_env_endpoint_discovery_reports_early_exit(tmp_path: Path): + from types import SimpleNamespace + + from robots.behavior.runtime import _wait_for_server_endpoint + + log = tmp_path / "env.log" + log.write_text("initialization failed\n") + daemon = SimpleNamespace(log_path=str(log), name="test_env", poll=lambda: 2) + with pytest.raises(RuntimeError, match="exited with code 2"): + _wait_for_server_endpoint(daemon) + + +class _FakeEnv: + total_env_steps = 0 + official_success_latched = False + official_success_receipt = None + + def __init__(self) -> None: + self.last_move: dict[str, Any] | None = None + self.gripper_calls: list[tuple[str, dict[str, Any]]] = [] + + def move_to(self, **kwargs: Any) -> dict[str, Any]: + self.last_move = kwargs + return {"status": "failed", "stop_reason": "motion_unavailable"} + + def close_gripper(self, **kwargs: Any) -> dict[str, Any]: + self.gripper_calls.append(("close_gripper", kwargs)) + return {"status": "success"} + + def open_gripper(self, **kwargs: Any) -> dict[str, Any]: + self.gripper_calls.append(("open_gripper", kwargs)) + return {"status": "success"} + + +class _FakeRpcClient: + def __init__(self) -> None: + self.calls: list[tuple[str, dict[str, Any]]] = [] + + def call( + self, + method: str, + *, + args: tuple = (), + kwargs: dict[str, Any] | None = None, + timeout_s: float | None = None, + ) -> dict[str, Any]: + self.calls.append((method, kwargs or {})) + if method == "env.get_env_meta": + return {} + if method == "dino.get_meta": + return { + "runtime": "behavior_dino", + "dimension": DINOV2_DIMENSION, + } + return {"status": "success"} + + +class _FakeModel: + def predict(self, observation: dict[str, Any], *, options: dict[str, Any]) -> Any: + assert observation["task_descriptions"] == "turn on the radio" + assert options == {"mode": "eval"} + return np.zeros((32, 23), dtype=np.float32) + + +class _FakeChunkEnv: + total_env_steps = 0 + official_success_latched = False + official_success_receipt = None + + def __init__(self) -> None: + self.return_all_frames: list[bool] = [] + + def chunk_step( + self, + actions: Any, + *, + return_all_frames: bool = False, + ) -> tuple[Any, float, bool, bool, dict[str, Any]]: + array = np.asarray(actions) + self.return_all_frames.append(bool(return_all_frames)) + self.total_env_steps += int(array.shape[0]) + frames = [ + { + "main_images": np.full((16, 16, 3), idx, dtype=np.uint8), + "task_descriptions": "turn on the radio", + } + for idx in range(int(array.shape[0])) + ] + obs: Any = frames if return_all_frames else frames[-1] + return ( + obs, + 0.0, + False, + False, + { + "executed_steps": int(array.shape[0]), + "_rpent": {"total_env_steps": self.total_env_steps}, + }, + ) + + +class _FakeOfficialBehaviorEnv: + def __init__( + self, + *_args: Any, + terminated_on_step: int | None = None, + success_on_step: int | None = None, + **_kwargs: Any, + ) -> None: + self.actions: list[np.ndarray] = [] + self.closed = False + self.terminated_on_step = terminated_on_step + self.success_on_step = success_on_step + self.raw = np.zeros(256, dtype=np.float32) + self.raw[RAW_PROPRIO_SEGMENTS["trunk"]] = np.asarray( + [0.11, 0.12, 0.13, 0.14], + dtype=np.float32, + ) + self.raw[RAW_PROPRIO_SEGMENTS["left_arm"]] = np.linspace( + 0.21, + 0.27, + 7, + dtype=np.float32, + ) + self.raw[RAW_PROPRIO_SEGMENTS["right_arm"]] = np.linspace( + -0.31, + -0.37, + 7, + dtype=np.float32, + ) + self.raw[RAW_PROPRIO_SEGMENTS["left_gripper"]] = 0.05 + self.raw[RAW_PROPRIO_SEGMENTS["right_gripper"]] = 0.06 + + def _obs(self) -> dict[str, Any]: + return { + "main_images": np.zeros((8, 8, 3), dtype=np.uint8), + "wrist_images": np.zeros((2, 8, 8, 3), dtype=np.uint8), + "states": self.raw.copy(), + "task_descriptions": "turn on the radio", + } + + def reset_raw(self, *, env_idx: int = 0) -> tuple[dict[str, Any], dict[str, Any]]: + assert env_idx == 0 + return self._obs(), {"done": {"success": False}} + + def step_raw( + self, + action: Any, + *, + env_idx: int = 0, + ) -> tuple[dict[str, Any], float, bool, bool, dict[str, Any]]: + assert env_idx == 0 + action_array = np.asarray(action, dtype=np.float32) + self.actions.append(action_array.copy()) + step_index = len(self.actions) + terminated = self.terminated_on_step == step_index + success = self.success_on_step == step_index + return self._obs(), 0.0, terminated, False, {"done": {"success": success}} + + def close(self) -> None: + self.closed = True + + +class _ThreadRecordingBehaviorEnvFacade(BehaviorEnvFacade): + def __init__(self) -> None: + super().__init__(backend=object(), meta={"task_language": "test"}) + self.serve_thread_id: int | None = None + self.business_thread_id: int | None = None + + def serve(self, **kwargs: Any) -> None: + self.serve_thread_id = threading.get_ident() + super().serve(**kwargs) + + def get_env_meta(self) -> dict[str, Any]: + self.business_thread_id = threading.get_ident() + return super().get_env_meta() + + +def _both_hand_request() -> dict[str, Any]: + return { + "hand": "both", + "targets": { + "left": {"delta_xyz": [0.01, 0.0, 0.0], "frame": "world"}, + "right": {"delta_xyz": [-0.01, 0.0, 0.0], "frame": "eef"}, + }, + "visual_hand_checks": { + "left": { + "camera": "left_wrist", + "frame_id": "left-frame", + "selected_hand": "left", + "assessment": "selected_hand_visually_confirmed", + }, + "right": { + "camera": "right_wrist", + "frame_id": "right-frame", + "selected_hand": "right", + "assessment": "selected_hand_visually_confirmed", + }, + }, + } + + +def _visual_check(hand: str) -> dict[str, str]: + return { + "camera": f"{hand}_wrist", + "frame_id": f"{hand}-frame", + "selected_hand": hand, + "assessment": "selected_hand_visually_confirmed", + } + + +def _official_backend( + tmp_path: Path, + fake_env: _FakeOfficialBehaviorEnv | None = None, +) -> tuple[OfficialBehaviorBackend, _FakeOfficialBehaviorEnv]: + env = fake_env or _FakeOfficialBehaviorEnv() + backend = OfficialBehaviorBackend( + meta={ + "task_name": "turning_on_radio", + "task_language": "turn on the radio", + "activity_definition_id": 0, + "activity_instance_id": 242, + "public_seed": 0, + "scene_model": "house_double_floor_lower", + "max_episode_steps": 64, + }, + output_dir=tmp_path, + behavior_env_cls=lambda *_args, **_kwargs: env, + cfg=object(), + ) + backend.reset() + return backend, env + + +def _port_accepts_connections(port: int) -> bool: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.settimeout(0.1) + return probe.connect_ex(("127.0.0.1", port)) == 0 + + +def test_public_behavior_surface_is_exactly_nine_tools() -> None: + assert BEHAVIOR_TOOL_NAMES == EXPECTED_TOOLS + + +def test_pixel_projection_uses_observed_depth_and_rejects_old_frame(tmp_path): + backend, _ = _official_backend(tmp_path) + frames = { + camera: { + "rgb": np.zeros((3, 3, 3), dtype=np.uint8), + "depth": np.ones((3, 3), dtype=np.float32) * 2, + "intrinsic": np.array([[2.0, 0, 1], [0, 2, 1], [0, 0, 1]]), + "camera_to_world": np.eye(4), + } + for camera in ("head", "left_wrist", "right_wrist") + } + backend._call_actor = lambda _function: frames + observation = backend.observe(camera="head") + result = backend.pixel_to_world( + camera="head", frame_id=observation["frame_id"], u=1, v=1 + ) + assert result["world_xyz"] == [0.0, 0.0, -2.0] + assert result["task_success"] is False + assert ( + backend._projection(result["projection_id"])["world_xyz"] == result["world_xyz"] + ) + backend.chunk_step(backend._hold_action_from_current_proprio()[None, :]) + assert ( + backend.pixel_to_world( + camera="head", frame_id=observation["frame_id"], u=1, v=1 + )["stop_reason"] + == "stale_frame" + ) + + +@pytest.mark.parametrize("hand", ["left", "both"]) +def test_planned_motion_preserves_unselected_joint_commands(tmp_path, hand): + from types import SimpleNamespace + + backend, env = _official_backend(tmp_path) + hold = backend._hold_action_from_current_proprio().copy() + hands = ("left", "right") if hand == "both" else (hand,) + poses = { + h: {"position": np.zeros(3), "quaternion_xyzw": [0, 0, 0, 1]} + for h in ("left", "right") + } + state = {"hands": poses, "control_dt": 1 / 60} + backend._call_actor = lambda _: state + backend._get_motion_state = lambda: state + names = [f"{h}_arm_joint{i}" for h in hands for i in range(1, 8)] + trajectory = np.ones((2, len(names)), dtype=np.float32) * 0.1 + backend._motion_planner = SimpleNamespace( + plan=lambda *_: { + "success": True, + "joint_names": names, + "positions": trajectory, + "dt": 1 / 60, + } + ) + request = ( + _both_hand_request() + if hand == "both" + else {"hand": hand, "target": {"delta_xyz": [0, 0, 0], "frame": "world"}} + ) + if hand == "both": + for target in request["targets"].values(): + target["delta_xyz"] = [0, 0, 0] + result = backend.move_to(**request) + assert result["primitive_success"] is True + assert result["task_success"] is False + for action in env.actions: + expected = hold.copy() + for selected in hands: + expected[ENV_ACTION_SEGMENTS[f"{selected}_arm"]] = 0.1 + np.testing.assert_allclose(action, expected) + + +def test_failed_plan_does_not_execute_actions(tmp_path): + from types import SimpleNamespace + + backend, env = _official_backend(tmp_path) + backend._call_actor = lambda _: { + "hands": {"left": {"position": np.zeros(3), "quaternion_xyzw": [0, 0, 0, 1]}} + } + backend._motion_planner = SimpleNamespace( + plan=lambda *_: { + "success": False, + "stop_reason": "planning_failed", + "details": "collision", + } + ) + result = backend.move_to( + hand="left", target={"delta_xyz": [0, 0, 0.03], "frame": "world"} + ) + assert result["stop_reason"] == "planning_failed" + assert not env.actions + + +def test_navigation_swept_footprint_rejects_obstacles(): + from robots.behavior.motion import navigation_collision + + state = { + "base_position": np.zeros(3), + "robot_aabb": [[-0.1, -0.1, 0], [0.1, 0.1, 0.5]], + "obstacles": {"box": {"low": [0.4, -0.1, 0.1], "high": [0.6, 0.1, 0.3]}}, + } + assert navigation_collision(state, np.array([1.0, 0, 0])) == "box" + assert navigation_collision(state, np.array([-0.1, 0, 0])) is None + + +def test_rotate_requires_selected_hand_confirmation_before_motion(tmp_path): + backend, env = _official_backend(tmp_path) + with pytest.raises(ValueError, match="visual_hand_check"): + backend.rotate_wrist(hand="left", angle_deg=5) + with pytest.raises(ValueError, match="visual_hand_check"): + backend.rotate_wrist( + hand="left", angle_deg=5, visual_hand_check=_visual_check("right") + ) + assert not env.actions + + +def test_navigation_brake_propagates_raw_success(tmp_path): + backend, env = _official_backend( + tmp_path, _FakeOfficialBehaviorEnv(success_on_step=1) + ) + backend._call_actor = lambda _: { + "base_position": np.zeros(3), + "base_quaternion_xyzw": [0, 0, 0, 1], + "robot_aabb": [[-0.1, -0.1, 0], [0.1, 0.1, 0.5]], + "obstacles": {}, + "control_dt": 1 / 60, + } + result = backend.navigate_to( + relative_motion={ + "kind": "translation", + "direction": "forward", + "distance_m": 0.001, + } + ) + assert len(env.actions) == 1 + assert result["task_success"] is True + assert result["stop_reason"] == "official_task_success" + + +def test_motion_result_refreshes_policy_observation_without_extra_rpc(tmp_path): + observation = { + "main_images": np.zeros((16, 16, 3), dtype=np.uint8), + "states": np.arange(256, dtype=np.float32), + } + + class Env: + def close_gripper(self, **kwargs): + return { + "primitive_success": True, + "task_success": False, + "_observation": observation, + "info": {"done": {"success": False}}, + } + + def current_observation(self): + raise AssertionError("post-action frame must come from the executed action") + + primitives = BehaviorPrimitives( + env=Env(), task_name="turning_on_radio", output_dir=tmp_path + ) + result = primitives.close(hand="left", visual_hand_check=_visual_check("left")) + assert primitives.current_observation is observation + assert "_observation" not in result + assert result["task_success"] is False + + +def test_robot_config_separates_world_and_self_collision_padding(tmp_path): + import yaml + + from robots.behavior.motion import build_robot_config + + names = [f"{h}_arm_joint{i}" for h in ("left", "right") for i in range(1, 8)] + urdf = tmp_path / "robot.urdf" + urdf.write_text( + '' + + "".join(f'' for n in names) + + "" + ) + source = { + "collision_link_names": ["base_link", "left_arm_link1"], + "collision_spheres": {}, + "collision_sphere_buffer": 0.002, + "self_collision_buffer": {"base_link": 0.02}, + "self_collision_ignore": {}, + "extra_links": {}, + "extra_collision_spheres": {}, + "cspace": { + "joint_names": names, + "cspace_distance_weight": [1] * 14, + "null_space_weight": [1] * 14, + }, + } + config_file = tmp_path / "collision.yaml" + config_file.write_text(yaml.safe_dump({"robot_cfg": {"kinematics": source}})) + result = build_robot_config( + { + "collision_config_path": str(config_file), + "urdf_path": str(urdf), + "joint_positions": dict.fromkeys(names, 0), + } + )["kinematics"] + assert result["collision_sphere_buffer"] == 0.002 + assert result["self_collision_buffer"]["base_link"] == pytest.approx(0.018) + assert result["self_collision_buffer"]["left_arm_link1"] == -0.002 + assert result["self_collision_ignore"] == source["self_collision_ignore"] + + +def test_gripper_close_builds_hold_action_and_target_command(tmp_path: Path) -> None: + backend, env = _official_backend(tmp_path) + + result = backend.close(hand="left", visual_hand_check=_visual_check("left")) + chunk = validate_action_chunk(np.stack(env.actions, axis=0)) + first = chunk[0] + + assert result["status"] == "ok" + assert result["primitive_success"] is True + assert result["task_success"] is False + assert result["stop_reason"] == "requested_actions_completed" + assert result["visual_hand_check"] == _visual_check("left") + assert result["visual_hand_check_verification"] == "not_verified" + assert result["action_shape"] == [1, 23] + assert result["action_chunk_shape"] == [GRIPPER_COMMAND_CONTROL_CYCLES, 23] + assert chunk.shape == (GRIPPER_COMMAND_CONTROL_CYCLES, 23) + assert np.allclose(chunk, first) + assert np.allclose(first[ENV_ACTION_SEGMENTS["base"]], 0.0) + assert np.allclose( + first[ENV_ACTION_SEGMENTS["trunk"]], + env.raw[RAW_PROPRIO_SEGMENTS["trunk"]], + ) + assert np.allclose( + first[ENV_ACTION_SEGMENTS["left_arm"]], + env.raw[RAW_PROPRIO_SEGMENTS["left_arm"]], + ) + assert np.allclose( + first[ENV_ACTION_SEGMENTS["right_arm"]], + env.raw[RAW_PROPRIO_SEGMENTS["right_arm"]], + ) + assert first[ENV_ACTION_SEGMENTS["left_gripper"]][0] == GRIPPER_CLOSE_COMMAND + assert first[ENV_ACTION_SEGMENTS["right_gripper"]][0] == GRIPPER_OPEN_COMMAND + + +def test_gripper_latch_holds_non_target_hand_and_open_echoes_release_check( + tmp_path: Path, +) -> None: + backend, env = _official_backend(tmp_path) + + backend.close(hand="left", visual_hand_check=_visual_check("left")) + env.actions.clear() + release_check = { + "camera": "head", + "frame_id": "release-frame", + "assessment": "target_visibly_released", + } + result = backend.open( + hand="right", + visual_hand_check=_visual_check("right"), + release_visual_check=release_check, + ) + first = validate_action_chunk(np.stack(env.actions, axis=0))[0] + + assert result["status"] == "ok" + assert result["release_visual_check"] == release_check + assert result["release_visual_check_verification"] == "not_verified" + assert first[ENV_ACTION_SEGMENTS["left_gripper"]][0] == GRIPPER_CLOSE_COMMAND + assert first[ENV_ACTION_SEGMENTS["right_gripper"]][0] == GRIPPER_OPEN_COMMAND + + +def test_gripper_task_success_uses_only_official_done_success(tmp_path: Path) -> None: + terminated_env = _FakeOfficialBehaviorEnv(terminated_on_step=1) + backend, env = _official_backend(tmp_path / "terminated", terminated_env) + + result = backend.close(hand="left", visual_hand_check=_visual_check("left")) + + assert result["stop_reason"] == "terminated" + assert result["primitive_success"] is True + assert result["task_success"] is False + assert backend.official_success_latched is False + assert len(env.actions) == 1 + rejected = backend.open(hand="left", visual_hand_check=_visual_check("left")) + assert rejected["stop_reason"] == "episode_ended" + assert len(env.actions) == 1 + + success_env = _FakeOfficialBehaviorEnv(success_on_step=1) + success_backend, _success_env = _official_backend(tmp_path / "success", success_env) + success = success_backend.open( + hand="right", + visual_hand_check=_visual_check("right"), + ) + + assert success["stop_reason"] == "official_task_success" + assert success["task_success"] is True + assert success["official_success_receipt"]["source"] == 'info["done"]["success"]' + + +def test_gripper_execution_error_uses_motion_error_envelope(tmp_path: Path) -> None: + backend, env = _official_backend(tmp_path) + assert backend._last_obs is not None + backend._last_obs["states"][RAW_PROPRIO_SEGMENTS["trunk"]] = np.nan + + result = backend.open(hand="left", visual_hand_check=_visual_check("left")) + + assert result["status"] == "failed" + assert result["primitive_success"] is False + assert result["task_success"] is False + assert result["stop_reason"] == "error" + assert "NaN or infinity" in result["error"] + assert env.actions == [] + + +def test_backend_lifecycle_close_still_closes_env_without_primitive_args( + tmp_path: Path, +) -> None: + backend, env = _official_backend(tmp_path) + + assert backend.close() == {"status": "ok", "closed": True} + assert env.closed is True + assert backend.close() == { + "status": "ok", + "closed": True, + "already_closed": True, + } + + +@pytest.mark.parametrize( + "ending", ["contact", "terminated", "official_task_success", "duration_limit"] +) +def test_press_executes_bounded_hold_actions_and_reports_raw_success( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ending: str +) -> None: + backend, env = _official_backend( + tmp_path, + _FakeOfficialBehaviorEnv( + terminated_on_step=1 if ending == "terminated" else None, + success_on_step=1 if ending == "official_task_success" else None, + ), + ) + + def state(): + return { + "control_dt": 0.05, + "hands": { + "left": { + "position": np.array([0.0, 0.0, len(env.actions) * 0.0005]), + "approach_direction": np.array([0.0, 0.0, 1.0]), + "joint_positions": env.raw[RAW_PROPRIO_SEGMENTS["left_arm"]], + "joint_lower_limits": np.full(7, -3.0), + "joint_upper_limits": np.full(7, 3.0), + "jacobian": np.eye(6, 7), + "contacts": ["button"] + if ending == "contact" and env.actions + else [], + } + }, + } + + monkeypatch.setattr(backend, "_get_motion_state", state) + hold = backend._hold_action_from_current_proprio() + result = backend.press( + hand="left", visual_hand_check=_visual_check("left"), duration_s=0.1 + ) + assert result["stop_reason"] == ending + assert result["task_success"] is (ending == "official_task_success") + assert result["primitive_success"] is ( + ending in {"contact", "official_task_success"} + ) + assert result["executed_steps"] == (2 if ending == "duration_limit" else 1) + untouched = np.ones(23, dtype=bool) + untouched[ENV_ACTION_SEGMENTS["left_arm"]] = False + for action in env.actions: + assert validate_action_chunk(action[None, :]).shape == (1, 23) + np.testing.assert_array_equal(action[untouched], hold[untouched]) + assert result["visual_hand_check_verification"] == "not_verified" + + +@pytest.mark.parametrize("duration", [0, -1, 11, float("nan"), float("inf")]) +def test_press_rejects_invalid_duration_before_motion(tmp_path: Path, duration: float): + backend, env = _official_backend(tmp_path) + with pytest.raises(ValueError, match="duration_s"): + backend.press( + hand="left", visual_hand_check=_visual_check("left"), duration_s=duration + ) + assert not env.actions + + +@pytest.mark.parametrize( + ("extra_args", "message"), + [ + ( + [ + "--behavior-mode", + "explore", + "--explore-attempts-per-session", + "1", + ], + "BEHAVIOR explore runs one attempt per session; use --explore-sessions", + ), + ([], "BEHAVIOR --explore requires --behavior-mode explore"), + ( + ["--behavior-mode", "explore", "--dashboard"], + "BEHAVIOR --explore is CLI-only", + ), + ( + ["--behavior-mode", "explore", "--env-endpoint", "127.0.0.1:1"], + "BEHAVIOR explore requires an owned env sidecar; omit --env-endpoint", + ), + ], +) +def test_behavior_explore_rejects_incompatible_options( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + extra_args: list[str], + message: str, +) -> None: + from rpent.cli import main as cli + + monkeypatch.setattr( + sys, + "argv", + [ + "rpent", + "--robot", + "behavior", + "--task-name", + "turning_on_radio", + "--public-seed", + "0", + "--explore", + *extra_args, + ], + ) + + with pytest.raises(SystemExit) as exc_info: + cli.main() + + assert exc_info.value.code == 2 + assert message in capsys.readouterr().err + + +def test_behavior_toolkit_factory_maps_shared_modes(tmp_path: Path) -> None: + config = RunConfig( + recipe_tag="turning_on_radio_s0", + output_dir=tmp_path / "run", + prompt_vars={ + "behavior_mode": "eval", + "memory_dir": str(tmp_path / "memory"), + "task_name": "turning_on_radio", + "public_seed": 0, + "max_episode_steps": 64, + }, + task_desc={}, + ) + dashboard_events = NullDashboardEventSink() + + evaluation = get_toolkit( + primitives_kwargs={}, + dashboard_events=dashboard_events, + config=config, + mode="evaluation", + ) + exploration = get_toolkit( + primitives_kwargs={"output_dir": tmp_path / "session"}, + dashboard_events=dashboard_events, + config=config, + mode="exploration", + ) + config.prompt_vars["behavior_mode"] = "explore" + inherited = get_toolkit( + primitives_kwargs={}, + dashboard_events=dashboard_events, + config=config, + ) + + assert evaluation.primitives.behavior_phase == "eval" + assert exploration.primitives.behavior_phase == "explore" + assert inherited.primitives.behavior_phase == "explore" + assert evaluation.memory._memory_access == "read_only" + assert exploration.memory._memory_access == "inbox_write" + assert exploration.primitives.output_dir == tmp_path / "session" + with pytest.raises(ValueError, match="one attempt per session"): + get_toolkit( + primitives_kwargs={}, + dashboard_events=dashboard_events, + config=config, + mode="exploration", + attempts_per_session=1, + ) + with pytest.raises(ValueError, match="unsupported BEHAVIOR toolkit mode"): + get_toolkit( + primitives_kwargs={}, + dashboard_events=dashboard_events, + config=config, + mode="unsupported", + ) + + +def test_behavior_external_sidecar_python_gets_source_checkout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from robots.behavior import runtime + + monkeypatch.setenv("PYTHONPATH", "/tmp/existing-pythonpath") + + env = runtime._behavior_subprocess_env( + cuda_device="2", + ROBOT_PLATFORM="BEHAVIOR", + ) + + assert env["ROBOT_PLATFORM"] == "BEHAVIOR" + assert env["CUDA_VISIBLE_DEVICES"] == "2" + assert env["PYTHONPATH"].split(os.pathsep)[:2] == [ + str(runtime.get_repo_root()), + "/tmp/existing-pythonpath", + ] + + +def test_behavior_facades_use_default_healthz_and_registered_metadata() -> None: + facade = BehaviorEnvFacade(backend=object(), meta={"task_language": "test"}) + dino = BehaviorDinoFacade( + encoder=object(), + meta={"runtime": "behavior_dino", "dimension": DINOV2_DIMENSION}, + ) + + assert facade._dispatch("healthz", (), {}) == {"status": "ok"} + assert facade._dispatch("env.get_env_meta", (), {}) == {"task_language": "test"} + assert "env.close_gripper" in facade._rpc + assert "env.open_gripper" in facade._rpc + assert "env.close" not in facade._rpc + assert "env.open" not in facade._rpc + with pytest.raises(ValueError, match="requires primitive arguments"): + facade.close_gripper() + assert dino._dispatch("healthz", (), {}) == {"status": "ok"} + dino_meta = dino._dispatch("dino.get_meta", (), {}) + assert dino_meta["runtime"] == "behavior_dino" + assert dino_meta["dimension"] == DINOV2_DIMENSION + assert isinstance(dino_meta["pid"], int) + + +def test_behavior_env_facade_serve_dispatches_business_calls_on_serving_thread() -> ( + None +): + facade = _ThreadRecordingBehaviorEnvFacade() + port = pick_free_port() + thread = threading.Thread( + target=facade.serve, + kwargs={ + "transport": "http", + "host": "127.0.0.1", + "port": port, + }, + daemon=True, + ) + thread.start() + client = HttpRpcClient(f"http://127.0.0.1:{port}") + + try: + deadline = time.monotonic() + 3.0 + while True: + try: + assert client.call("healthz", timeout_s=0.5) == {"status": "ok"} + break + except Exception: + if time.monotonic() >= deadline: + raise + time.sleep(0.01) + + assert client.call("env.get_env_meta", timeout_s=1.0) == { + "task_language": "test" + } + assert facade.business_thread_id == facade.serve_thread_id + assert facade.business_thread_id != threading.get_ident() + assert client.call("shutdown", timeout_s=1.0) == {"ok": True} + finally: + client.close() + facade._shutdown_event.set() + thread.join(timeout=3.0) + + assert not thread.is_alive() + assert facade._closed + assert not _port_accepts_connections(port) + + +def test_env_main_sigterm_uses_owning_thread_cleanup(monkeypatch, tmp_path): + from robots.behavior import env_server, rlinf_env + + closed_on = [] + + class Backend: + def __init__(self, **kwargs): + pass + + def close(self): + closed_on.append(threading.get_ident()) + + def serve(facade, **kwargs): + try: + signal.raise_signal(signal.SIGTERM) + assert facade._shutdown_event.is_set() + finally: + facade.close() + + previous = signal.getsignal(signal.SIGTERM) + monkeypatch.setattr(rlinf_env, "OfficialBehaviorBackend", Backend) + monkeypatch.setattr(env_server.BehaviorEnvFacade, "serve", serve) + monkeypatch.setattr(env_server, "_build_meta", lambda args: {}) + monkeypatch.setattr( + sys, + "argv", + [ + "env_server", + "--task-name", + "turning_on_radio", + "--public-seed", + "0", + "--task-index", + "0", + "--activity-definition-id", + "0", + "--activity-instance-id", + "242", + "--scene-model", + "test", + "--max-episode-steps", + "32", + "--output-dir", + str(tmp_path), + "--behavior-repo", + str(tmp_path), + "--port", + "0", + ], + ) + env_server.main() + assert closed_on == [threading.get_ident()] + assert signal.getsignal(signal.SIGTERM) == previous + + +def test_behavior_clients_and_tools_use_explicit_component_rpc_names() -> None: + rpc = _FakeRpcClient() + client = BehaviorEnvClient(rpc, expected_meta={}) + dino_client = BehaviorDinoClient(rpc, expected_meta={"runtime": "behavior_dino"}) + env = _FakeEnv() + primitives = BehaviorPrimitives(env=env, task_name="turning_on_radio") + + client.close_gripper(hand="right") + client.open_gripper(hand="right") + primitives.close(hand="right") + primitives.open(hand="right") + + assert rpc.calls == [ + ("env.get_env_meta", {}), + ("dino.get_meta", {}), + ("env.close_gripper", {"hand": "right"}), + ("env.open_gripper", {"hand": "right"}), + ] + assert dino_client.server_meta["dimension"] == DINOV2_DIMENSION + assert env.gripper_calls == [ + ("close_gripper", {"hand": "right"}), + ("open_gripper", {"hand": "right"}), + ] + assert "env.close_gripper" in BehaviorEnvClient._TIMEOUT_S + assert "env.open_gripper" in BehaviorEnvClient._TIMEOUT_S + assert "env.close" not in BehaviorEnvClient._TIMEOUT_S + assert "env.open" not in BehaviorEnvClient._TIMEOUT_S + + +def test_move_to_contract_separates_single_and_dual_hand_branches() -> None: + schema = MOVE_TO_SPEC["input_schema"] + single, dual = schema["oneOf"] + assert schema["properties"]["hand"]["enum"] == ["left", "right", "both"] + assert single["properties"]["hand"] == {"enum": ["left", "right"]} + assert single["required"] == ["hand", "target"] + assert dual["properties"]["hand"] == {"const": "both"} + assert dual["required"] == ["hand", "targets", "visual_hand_checks"] + + env = _FakeEnv() + primitives = BehaviorPrimitives(env=env, task_name="turning_on_radio") + for hand in ("left", "right"): + request = {"hand": hand, "target": {"delta_xyz": [0.0, 0.0, 0.01]}} + primitives.move_to(**request) + assert env.last_move == request + + both = _both_hand_request() + primitives.move_to(**both) + assert env.last_move == both + invalid = _both_hand_request() + invalid["targets"] = {"left": invalid["targets"]["left"]} + with pytest.raises(ValueError, match="exactly left and right"): + primitives.move_to(**invalid) + + +def test_finish_writes_terminal_receipt(tmp_path: Path) -> None: + output_dir = tmp_path / "run" + toolkit = BehaviorToolkit( + primitives_kwargs={ + "task_name": "turning_on_radio", + "output_dir": output_dir, + }, + dashboard_events=NullDashboardEventSink(), + memory=MemoryManager(tmp_path / "memory"), + ) + + result = toolkit.execute_tool( + "finish", {"status": "incomplete", "summary": "bounded test"} + ) + receipt = json.loads((output_dir / "terminal_receipt.json").read_text()) + + assert result.is_finish is True + assert receipt["_finish"] is True + assert receipt["kind"] == "behavior_finish_terminal_receipt" + assert receipt["planner_status"] == "incomplete" + assert receipt["summary"] == "bounded test" + + +def test_receipt_is_session_artifact_and_recipe_is_run_artifact( + tmp_path: Path, +) -> None: + run_dir = tmp_path / "run" + session_dir = run_dir / "sessions" / "session_001" + toolkit = BehaviorToolkit( + primitives_kwargs={ + "task_name": "turning_on_radio", + "output_dir": run_dir, + }, + dashboard_events=NullDashboardEventSink(), + memory=MemoryManager(tmp_path / "memory"), + config=RunConfig( + recipe_tag="turning_on_radio_s0", + output_dir=run_dir, + prompt_vars={ + "task_name": "turning_on_radio", + "public_seed": 0, + "memory_dir": str(tmp_path / "memory"), + }, + task_desc={}, + ), + state_output_dir=session_dir, + ) + run_audit = run_dir / "turning_on_radio_s0.json" + run_audit.parent.mkdir(parents=True, exist_ok=True) + run_audit.write_text('{"audit": true}\n') + + toolkit.execute_tool("finish", {"status": "incomplete", "summary": "done"}) + assert (session_dir / "terminal_receipt.json").is_file() + assert json.loads((session_dir / "states.json").read_text())["run_artifacts"] == [ + "terminal_receipt.json" + ] + assert run_audit.read_text() == '{"audit": true}\n' + + toolkit.primitives._official_success_latched = True + with toolkit.state.record_step( + state={"task_success": True}, + terminated=True, + command={"action": "future_stateful_command", "arg": 1}, + result={"ok": True}, + elapsed_s=0.1, + ): + pass + with toolkit.state.record_step( + state={"task_success": True}, + terminated=True, + command={"action": "bad_command"}, + result={"error": "failed"}, + elapsed_s=0.1, + ): + pass + + recipe_path = Path(toolkit.write_recipe("turning_on_radio_s0") or "") + assert recipe_path == run_dir / "recipe_turning_on_radio_s0.jsonl" + assert not (session_dir / "recipe_turning_on_radio_s0.jsonl").exists() + commands = [ + json.loads(line) + for line in recipe_path.read_text().splitlines() + if line.strip() + ] + assert commands == [{"action": "future_stateful_command", "arg": 1}] + assert "command" not in commands[0] + + +def test_unsolved_behavior_session_does_not_write_recipe(tmp_path: Path) -> None: + toolkit = BehaviorToolkit( + primitives_kwargs={ + "task_name": "turning_on_radio", + "output_dir": tmp_path / "run", + }, + dashboard_events=NullDashboardEventSink(), + memory=MemoryManager(tmp_path / "memory"), + ) + with toolkit.state.record_step( + state={"task_success": False}, + command={"action": "pi0_nav_pick", "instruction": "turn on the radio"}, + result={"ok": True}, + elapsed_s=0.1, + ): + pass + + assert toolkit.write_recipe("turning_on_radio_s0") is None + assert not (tmp_path / "run" / "recipe_turning_on_radio_s0.jsonl").exists() + + +def test_behavior_pi0_chunk_records_streaming_episode_video(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + session_dir = run_dir / "sessions" / "session_001" + toolkit = BehaviorToolkit( + primitives_kwargs={ + "env": _FakeChunkEnv(), + "model": _FakeModel(), + "task_name": "turning_on_radio", + "public_seed": 0, + "max_episode_steps": 64, + "initial_observation": { + "main_images": np.zeros((16, 16, 3), dtype=np.uint8), + "task_descriptions": "turn on the radio", + }, + }, + dashboard_events=NullDashboardEventSink(), + memory=MemoryManager(tmp_path / "memory"), + config=RunConfig( + recipe_tag="turning_on_radio_s0", + output_dir=run_dir, + prompt_vars={ + "behavior_mode": "explore", + "task_name": "turning_on_radio", + "public_seed": 0, + "memory_dir": str(tmp_path / "memory"), + }, + task_desc={}, + ), + state_output_dir=session_dir, + ) + + result = toolkit.primitives.pi0_nav_pick( + instruction="turn on the radio", + chunks=1, + ) + toolkit.execute_tool("finish", {"status": "incomplete", "summary": "done"}) + + assert result["chunks_used"] == 1 + assert result["env_steps_used"] == 32 + assert toolkit.primitives.env.return_all_frames == [True] + video_path = session_dir / "episode.mp4" + assert video_path.is_file() + assert video_path.stat().st_size > 0 + assert ( + "episode.mp4" + in json.loads((session_dir / "states.json").read_text())["run_artifacts"] + ) + + +def test_behavior_build_memory_cli_wraps_existing_catalog_compiler( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + captured: dict[str, Any] = {} + + def fake_compile_runtime_catalog(**kwargs: Any) -> dict[str, Any]: + captured.update(kwargs) + return {"artifact_dir": str(kwargs["output_dir"]), "preliminary": True} + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "7") + monkeypatch.setattr( + build_memory_cli, + "compile_runtime_catalog", + fake_compile_runtime_catalog, + ) + + rc = build_memory_cli.main( + [ + "--selection-manifest", + str(tmp_path / "selection.json"), + "--output-dir", + str(tmp_path / "catalog"), + "--video-root", + str(tmp_path / "videos"), + "--rollups-dir", + str(tmp_path / "rollups"), + "--source-archive", + str(tmp_path / "dinov2.tar.gz"), + "--weights", + str(tmp_path / "dinov2.pth"), + "--cache-dir", + str(tmp_path / "cache"), + "--cuda-device", + "2", + "--batch-size", + "8", + ] + ) + + assert rc == 0 + assert json.loads(capsys.readouterr().out) == { + "artifact_dir": str((tmp_path / "catalog").resolve()), + "preliminary": True, + } + assert captured["selection_manifest"] == (tmp_path / "selection.json").resolve() + assert captured["output_dir"] == (tmp_path / "catalog").resolve() + assert captured["video_roots"] == ((tmp_path / "videos").resolve(),) + assert captured["rollups_dir"] == (tmp_path / "rollups").resolve() + assert captured["source_archive"] == (tmp_path / "dinov2.tar.gz").resolve() + assert captured["weights"] == (tmp_path / "dinov2.pth").resolve() + assert captured["cache_dir"] == (tmp_path / "cache").resolve() + assert captured["batch_size"] == 8 + assert os.environ["CUDA_VISIBLE_DEVICES"] == "2" + + +def test_behavior_build_memory_cli_rejects_multi_cuda_device() -> None: + with pytest.raises(SystemExit) as exc_info: + build_memory_cli.main( + [ + "--selection-manifest", + "selection.json", + "--output-dir", + "catalog", + "--video-root", + "videos", + "--rollups-dir", + "rollups", + "--source-archive", + "dinov2.tar.gz", + "--weights", + "dinov2.pth", + "--cuda-device", + "2,7", + ] + ) + assert exc_info.value.code == 2 diff --git a/tests/unit_tests/robots/test_toolkit_contracts.py b/tests/unit_tests/robots/test_toolkit_contracts.py index 1656e1f1b..f96221e73 100644 --- a/tests/unit_tests/robots/test_toolkit_contracts.py +++ b/tests/unit_tests/robots/test_toolkit_contracts.py @@ -20,6 +20,8 @@ import pytest +from robots.behavior import robot_spec as behavior_robot_spec +from robots.behavior import toolkit as behavior_toolkit from robots.libero import robot_spec as libero_robot_spec from robots.libero import toolkit as libero_toolkit from robots.robocasa import robot_spec as robocasa_robot_spec @@ -120,3 +122,56 @@ def test_toolkit_factories_fall_back_to_each_robot_memory_root( ) assert toolkit.memory.root == default_memory.resolve() + + +@pytest.mark.parametrize( + ("mode", "write_allowed"), + [("eval", False), ("explore", True)], +) +def test_behavior_toolkit_uses_one_official_memory_manager( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + mode: str, + write_allowed: bool, +) -> None: + captured: dict[str, Any] = {} + + def fake_toolkit(**kwargs: Any) -> SimpleNamespace: + captured.update(kwargs) + return SimpleNamespace(**kwargs) + + monkeypatch.setattr(behavior_toolkit, "BehaviorToolkit", fake_toolkit) + memory_dir = tmp_path / "memory" + recipe_tag = "turning_on_radio_s1" + config = RunConfig( + recipe_tag=recipe_tag, + output_dir=tmp_path / "run", + prompt_vars={ + "behavior_mode": mode, + "memory_dir": str(memory_dir), + }, + task_desc={}, + ) + + toolkit = behavior_robot_spec.get_toolkit( + primitives_kwargs={}, + dashboard_events=NullDashboardEventSink(), + config=config, + ) + + assert toolkit.memory is captured["memory"] + assert toolkit.memory.root == memory_dir.resolve() + write = toolkit.memory.get_common_tool_bindings()["write_text_file"][1] + destination = memory_dir / "_internal" / "inbox" / recipe_tag / "wip" / "notes.md" + if write_allowed: + write(str(destination), "evidence") + assert destination.read_text() == "evidence" + else: + with pytest.raises(PermissionError, match="writing to memory is denied"): + write(str(destination), "evidence") + + component_names = { + item["name"] + for item in behavior_robot_spec.BEHAVIOR_DASHBOARD_SPEC["runtime_components"] + } + assert component_names == {"env", "vla", "dino"} diff --git a/tests/unit_tests/rpent/cli/test_main_contracts.py b/tests/unit_tests/rpent/cli/test_main_contracts.py index abf1a1a75..3185167d0 100644 --- a/tests/unit_tests/rpent/cli/test_main_contracts.py +++ b/tests/unit_tests/rpent/cli/test_main_contracts.py @@ -184,7 +184,7 @@ def test_robot_and_env_aliases_are_mutually_exclusive( ["--robot", "libero", "--dashboard", "--interactive"], "cannot be used together", ), - (["--robot", "robocasa", "--explore"], "supported only for LIBERO"), + (["--robot", "robocasa", "--explore"], "--explore is not supported"), ( ["--robot", "libero", "--explore", "--memory-profile", "hf"], "cannot be used with --memory-profile hf", @@ -473,6 +473,8 @@ def reject_memory_sync(*args: Any, **kwargs: Any) -> None: assert calls["get_toolkit"][1]["primitives_kwargs"] == {"runtime": "simulated"} assert calls["get_toolkit"][1]["mode"] == "exploration" assert calls["get_toolkit"][1]["attempts_per_session"] == 2 + assert robot_spec.on_explore_session is None + assert calls["init_runtime"][3] is None assert calls["write_recipe"] == "libero_s0" assert calls["merge_memory"] == { "cell_tag": "libero_s0", @@ -622,3 +624,127 @@ def finalize_run(context: RunFinalizationContext) -> Path: } assert robot_toolkit.closed is True assert daemon.stopped is True + + +def test_behavior_explore_hook_restarts_only_env_between_cli_sessions( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + from dataclasses import replace + + from robots.behavior import runtime + from robots.behavior.robot_spec import get_robot_spec + from rpent.planner.base import PlannerResult + from rpent.robots.prompt_bundle import PromptBundle + + cli = _cli_module() + spawned = [] + stopped = [] + toolkit_calls = [] + runtime_calls = [] + merge_calls = [] + solve_calls = [] + + class FakeDaemon: + def __init__(self, component): + self.name = f"behavior_{component}_server" + self.component = component + + def stop(self): + stopped.append(self) + + def spawn(owned, events, component, spawn_fn): + daemon = FakeDaemon(component) + owned[component] = daemon + spawned.append(daemon) + return daemon, daemon + + def wait(owned, events, component, rpc, daemon, timeout_s, **kwargs): + return {component: daemon} + + original_init_runtime = runtime.init_runtime + + def init_runtime(args, output_dir, events, components): + runtime_calls.append((output_dir, components)) + return original_init_runtime(args, output_dir, events, components) + + def get_toolkit(*args, **kwargs): + toolkit_calls.append( + {**kwargs, "primitives_kwargs": dict(kwargs["primitives_kwargs"])} + ) + return SimpleNamespace( + memory=SimpleNamespace(merge_memory=lambda **kw: merge_calls.append(kw)), + solved=lambda: False, + close=lambda: None, + ) + + def solve(**kwargs): + solve_calls.append(kwargs) + assert spawned[-1] not in stopped + if len(solve_calls) == 2: + assert spawned[2] in stopped + return PlannerResult( + finish_result={"status": "incomplete"}, messages=[], stats={} + ) + + monkeypatch.setattr(runtime, "try_spawn_server", spawn) + monkeypatch.setattr(runtime, "try_wait_server", wait) + monkeypatch.setattr(runtime, "init_runtime", init_runtime) + spec = replace( + get_robot_spec(), + prompts=PromptBundle( + system=lambda variables: "system", user=lambda variables: "user" + ), + ) + monkeypatch.setattr(cli, "get_robot_spec", lambda name: spec) + monkeypatch.setattr(cli, "get_toolkit", get_toolkit) + monkeypatch.setattr( + cli, "build_planner", lambda *args, **kwargs: SimpleNamespace(solve=solve) + ) + monkeypatch.setattr( + sys, + "argv", + [ + "rpent", + "--robot", + "behavior", + "--behavior-mode", + "explore", + "--explore", + "--explore-sessions", + "2", + "--task-name", + "turning_on_radio", + "--public-seed", + "0", + "--output-dir", + str(tmp_path), + "--memory-dir", + str(tmp_path / "memory"), + ], + ) + + assert cli.main() == 0 + session_dirs = [tmp_path / "sessions" / f"session_{n:03d}" for n in (1, 2)] + assert runtime_calls == [ + (tmp_path, None), + *[(path, {"env"}) for path in session_dirs], + ] + assert [daemon.component for daemon in spawned] == ["vla", "dino", "env", "env"] + assert len(toolkit_calls) == len(solve_calls) == 2 + for index, call in enumerate(toolkit_calls): + assert call["state_output_dir"] == session_dirs[index] + assert call["mode"] == "exploration" + assert call["attempts_per_session"] == 0 + assert call["primitives_kwargs"] == { + "vla": spawned[0], + "dino": spawned[1], + "env": spawned[index + 2], + } + assert stopped == [spawned[2], spawned[0], spawned[1], spawned[3]] + assert merge_calls == [ + { + "cell_tag": "turning_on_radio_s0", + "run_state_dir": tmp_path, + "solved": False, + } + ] diff --git a/tests/unit_tests/rpent/dashboard/test_session_contracts.py b/tests/unit_tests/rpent/dashboard/test_session_contracts.py index 6f3a78d33..b61b663b5 100644 --- a/tests/unit_tests/rpent/dashboard/test_session_contracts.py +++ b/tests/unit_tests/rpent/dashboard/test_session_contracts.py @@ -20,6 +20,7 @@ import pytest +from rpent.dashboard import session as dashboard_session from rpent.dashboard.session import DashboardSessionController from rpent.dashboard.state import ClaimedTask from rpent.planner.base import PlannerResult @@ -152,8 +153,14 @@ def run_task(task: ClaimedTask, shared_kwargs: dict[str, Any]): def test_dashboard_session_stops_shared_daemons_in_reverse_after_cleanup_error( - caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, ) -> None: + warnings: list[str] = [] + + def fake_warning(message: str, *args: Any) -> None: + warnings.append(message % args if args else message) + + monkeypatch.setattr(dashboard_session.logger, "warning", fake_warning) stopped: list[str] = [] daemons = [ FakeDaemon("first", stopped), @@ -170,18 +177,28 @@ def test_dashboard_session_stops_shared_daemons_in_reverse_after_cleanup_error( controller.run() assert stopped == ["third", "second", "first"] - assert "shared runtime cleanup failed: stop failed" in caplog.text + assert warnings == ["shared runtime cleanup failed: stop failed"] +@pytest.mark.parametrize("robot_name", ["libero", "behavior"]) @pytest.mark.parametrize("merge_fails", [False, True]) +@pytest.mark.parametrize("auto_merge", [False, True]) +@pytest.mark.parametrize("solved", [False, True]) +@pytest.mark.parametrize("mode_explore", [False, True]) def test_dashboard_exploration_finalizes_memory_and_reports_merge_failures( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, + robot_name: str, merge_fails: bool, + auto_merge: bool, + solved: bool, + mode_explore: bool, ) -> None: from rpent.cli import dashboard as dashboard_cli merge_calls: list[dict[str, Any]] = [] + toolkit_calls: list[dict[str, Any]] = [] + recipe_calls: list[str] = [] class FakeMemoryManager: def merge_memory(self, **kwargs: Any) -> dict[str, int]: @@ -194,10 +211,11 @@ class FakeToolkit: memory = FakeMemoryManager() def solved(self) -> bool: - return True + return solved def write_recipe(self, recipe_tag: str) -> str: - return str(tmp_path / f"{recipe_tag}_recipe.jsonl") + recipe_calls.append(recipe_tag) + return str(tmp_path / f"recipe_{recipe_tag}.jsonl") def close(self) -> None: pass @@ -221,17 +239,17 @@ class FakePlanner: def solve(self, **kwargs: Any) -> PlannerResult: del kwargs return PlannerResult( - finish_result={"status": "success"}, + finish_result={"status": "success" if solved else "failure"}, messages=[], stats={}, ) output_dir = tmp_path / "task" run_config = RunConfig( - recipe_tag="libero_s0", + recipe_tag=f"{robot_name}_s0", output_dir=output_dir, prompt_vars={}, - task_desc={"robot": "libero"}, + task_desc={"robot": robot_name}, ) robot_spec = SimpleNamespace( parse_config=lambda args: run_config, @@ -243,11 +261,12 @@ def solve(self, **kwargs: Any) -> PlannerResult: ) args = SimpleNamespace( verbose=False, - robot_name="libero", - explore=True, - auto_merge_memory=True, + robot_name=robot_name, + explore=robot_name == "libero" and mode_explore, + behavior_mode="explore" if robot_name == "libero" or mode_explore else "eval", + auto_merge_memory=auto_merge, explore_sessions=1, - explore_attempts_per_session=2, + explore_attempts_per_session=0, planner="api", base_url=None, model="offline", @@ -260,9 +279,12 @@ def solve(self, **kwargs: Any) -> PlannerResult: ) claimed = ClaimedTask(number=1, request={}, output_dir=output_dir) state = FakeState() - monkeypatch.setattr( - dashboard_cli, "get_toolkit", lambda *args, **kwargs: FakeToolkit() - ) + + def fake_get_toolkit(*args: Any, **kwargs: Any) -> FakeToolkit: + toolkit_calls.append({"args": args, "kwargs": kwargs}) + return FakeToolkit() + + monkeypatch.setattr(dashboard_cli, "get_toolkit", fake_get_toolkit) monkeypatch.setattr( dashboard_cli, "build_planner", lambda *args, **kwargs: FakePlanner() ) @@ -278,18 +300,32 @@ def solve(self, **kwargs: Any) -> PlannerResult: ) assert error is None - assert merge_calls == [ - { - "cell_tag": "libero_s0", - "run_state_dir": output_dir, - "solved": True, - } - ] - if merge_fails: + assert recipe_calls == ([f"{robot_name}_s0"] if solved else []) + assert toolkit_calls[0]["kwargs"]["mode"] == ( + "exploration" if mode_explore else "evaluation" + ) + assert toolkit_calls[0]["kwargs"]["state_output_dir"] == ( + output_dir / "sessions" / "session_001" + if robot_name == "libero" and mode_explore + else output_dir + ) + assert merge_calls == ( + [ + { + "cell_tag": f"{robot_name}_s0", + "run_state_dir": output_dir, + "solved": solved, + } + ] + if auto_merge and mode_explore + else [] + ) + if merge_fails and auto_merge and mode_explore: assert len(state.warnings) == 1 assert ( "memory finalization failed: RuntimeError: merge exploded" in state.warnings[0] ) + assert "succeeded" not in state.warnings[0].lower() else: assert state.warnings == [] diff --git a/tests/unit_tests/rpent/dashboard/test_state_contracts.py b/tests/unit_tests/rpent/dashboard/test_state_contracts.py index d0bfe7bfa..5e6dd0a2a 100644 --- a/tests/unit_tests/rpent/dashboard/test_state_contracts.py +++ b/tests/unit_tests/rpent/dashboard/test_state_contracts.py @@ -269,6 +269,7 @@ def test_dashboard_events_project_runtime_usage_timeline_and_newest_frames( "elapsed_s": 0.5, "terminated": True, "truncated": False, + "official_success": True, "action_video_path": None, "action_video_artifact": None, "has_action_video": False, @@ -332,5 +333,144 @@ def test_dashboard_step_events_offset_new_traces_and_resolve_action_video( assert [item["step"] for item in detail["timeline"]] == [0, 1] assert detail["timeline"][0]["result"] == {"position": [1, 2, 3]} assert detail["timeline"][1]["terminated"] is True + assert detail["timeline"][1]["official_success"] is True assert state.frame("camera") == b"second-frame" assert state.action_video_path(0) == first_env.artifact_path("action.mp4", step=0) + + +def test_dashboard_maps_declared_result_keys_to_physical_frames( + tmp_path: Path, +) -> None: + spec = { + **DASHBOARD_SPEC, + "frame_channels": ( + {"name": "head", "label": "head", "result_key": "_image_bytes"}, + { + "name": "left_wrist", + "label": "left wrist", + "result_key": "_image_left_wrist_bytes", + }, + { + "name": "right_wrist", + "label": "right wrist", + "result_key": "_image_right_wrist_bytes", + }, + ), + } + state = DashboardState( + run_id="three-camera", + output_dir=tmp_path, + dashboard_spec=spec, + ) + + state.emit( + ToolResultEvent( + "observe", + { + "step": 4, + "_image_bytes": b"head", + "_image_left_wrist_bytes": b"left", + "_image_right_wrist_bytes": b"right", + }, + ) + ) + + snapshot = state.snapshot() + assert snapshot["frame_available"] == { + "head": True, + "left_wrist": True, + "right_wrist": True, + } + assert state.frame("head") == b"head" + assert state.frame("left_wrist") == b"left" + assert state.frame("right_wrist") == b"right" + + +def test_dashboard_legacy_frame_key_fallbacks_remain_supported( + tmp_path: Path, +) -> None: + spec = { + **DASHBOARD_SPEC, + "frame_channels": ( + {"name": "head", "label": "head"}, + {"name": "left_wrist", "label": "left wrist"}, + {"name": "right_wrist", "label": "right wrist"}, + {"name": "camera", "label": "camera"}, + {"name": "wrist", "label": "wrist"}, + ), + } + state = DashboardState( + run_id="legacy-fallbacks", + output_dir=tmp_path, + dashboard_spec=spec, + ) + + state.emit( + ToolResultEvent( + "view_env_state", + { + "step": 1, + "_image_bytes": b"head", + "_image_cam_bytes": b"left-and-camera", + "_image_wrist_bytes": b"right-and-wrist", + }, + ) + ) + + assert state.frame("head") == b"head" + assert state.frame("left_wrist") == b"left-and-camera" + assert state.frame("right_wrist") == b"right-and-wrist" + assert state.frame("camera") == b"left-and-camera" + assert state.frame("wrist") == b"right-and-wrist" + + +def test_behavior_official_success_fields_gate_success_projection( + tmp_path: Path, +) -> None: + state = _ready_state(tmp_path) + _claim_started_task(state) + state.emit( + ToolResultEvent( + "finish", + { + "step": 1, + "terminated": True, + "task_success": False, + "official_success_source": 'info["done"]["success"]', + "official_success_receipt": None, + "log": { + "command": {"action": "finish"}, + "result": {"task_success": False}, + }, + }, + ) + ) + detail = state.run_detail() + assert detail["terminated"] is True + assert detail["official_success"] is False + assert detail["timeline"][0]["official_success"] is False + + state = _ready_state(tmp_path / "success") + _claim_started_task(state) + state.emit( + ToolResultEvent( + "finish", + { + "step": 1, + "terminated": True, + "task_success": True, + "official_success_source": 'info["done"]["success"]', + "official_success_receipt": { + "source": 'info["done"]["success"]', + "raw_done": {"success": True}, + }, + "log": { + "command": {"action": "finish"}, + "result": {"task_success": True}, + }, + }, + ) + ) + detail = state.run_detail() + assert detail["official_success"] is True + assert detail["timeline"][0]["official_success"] is True diff --git a/tests/unit_tests/rpent/memory/test_manager_contracts.py b/tests/unit_tests/rpent/memory/test_manager_contracts.py index 09f33c807..672abfc83 100644 --- a/tests/unit_tests/rpent/memory/test_manager_contracts.py +++ b/tests/unit_tests/rpent/memory/test_manager_contracts.py @@ -18,6 +18,7 @@ from pathlib import Path from typing import Any +import pytest import yaml from rpent.memory import MemoryManager @@ -211,6 +212,51 @@ def test_memory_manager_does_not_publish_unsolved_task_artifacts( assert not (memory_dir / "task_only" / f"{cell}_recipe.jsonl").exists() +@pytest.mark.parametrize("solved", [False, True]) +def test_memory_manager_publishes_generic_command_recipe_only_when_solved( + tmp_path: Path, solved: bool +) -> None: + from rpent.session import EnvState, write_command_recipe_from_states + + memory_dir = tmp_path / "memory" + output_dir = tmp_path / "run" + cell = "turning_on_radio_s0" + output_dir.mkdir() + audit = output_dir / f"{cell}.json" + audit.write_text(json.dumps({"raw_done": {"success": solved}})) + # Exercise the real shared writer, not a duplicated filename convention. + state = EnvState(output_dir) + with state.record_step( + state={}, command={"action": "press", "hand": "left"}, result={} + ): + pass + recipe = Path(write_command_recipe_from_states(state, cell)) + result = MemoryManager(memory_dir).merge_memory( + cell_tag=cell, run_state_dir=output_dir, solved=solved + ) + assert result["task"] == int(solved) + target = memory_dir / "task_only" / f"{cell}_recipe.jsonl" + assert target.exists() is solved + if solved: + assert target.read_bytes() == recipe.read_bytes() + assert (target.parent / audit.name).read_bytes() == audit.read_bytes() + + +def test_memory_manager_prefers_existing_sibling_recipe_name(tmp_path: Path) -> None: + memory_dir, output_dir = tmp_path / "memory", tmp_path / "run" + cell = "10_task_t2_s0" + _write_task_pair(output_dir, cell, solved=True) + canonical = output_dir / f"{cell}_recipe.jsonl" + (output_dir / f"recipe_{cell}.jsonl").write_text('{"action":"wrong_source"}\n') + result = MemoryManager(memory_dir).merge_memory( + cell_tag=cell, run_state_dir=output_dir, solved=True + ) + assert result["task"] == 1 + assert ( + memory_dir / "task_only" / canonical.name + ).read_bytes() == canonical.read_bytes() + + def test_memory_manager_skips_invalid_draft_without_archiving_its_inbox( tmp_path: Path, ) -> None: diff --git a/tests/unit_tests/rpent/planner/test_codex_contracts.py b/tests/unit_tests/rpent/planner/test_codex_contracts.py index 3f8a189b3..011c93a89 100644 --- a/tests/unit_tests/rpent/planner/test_codex_contracts.py +++ b/tests/unit_tests/rpent/planner/test_codex_contracts.py @@ -383,6 +383,132 @@ def test_successful_fake_codex_lifecycle_uses_fake_mcp_and_accounts_events( assert any(isinstance(event, UsageEvent) for event in sink.events) +def test_finish_downgrades_recorder_only_stream_error( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + install_fake_backend(monkeypatch) + FakeCodex.events = [ + { + "method": "error", + "payload": { + "error": { + "message": "Reconnecting... 2/5", + "codex_error_info": {"response_stream_disconnected": {}}, + }, + "will_retry": True, + }, + }, + { + "method": "item/completed", + "payload": { + "item": { + "type": "mcpToolCall", + "tool": "mcp__rpent__finish", + "status": "completed", + "arguments": {"status": "success", "summary": "done"}, + "result": "accepted", + } + }, + }, + { + "method": "turn/completed", + "payload": {"turn": {"status": "completed", "duration_ms": 1000}}, + }, + ] + + result = make_planner(tmp_path, RecordingSink()).solve( + system_prompt="", + user_message="task", + toolkit=FakeToolkit(), + max_turns=3, + ) + + assert result.finish_result == { + "_finish": True, + "status": "success", + "summary": "done", + } + assert result.error is None + assert result.stats["nonfatal_recorder_errors"] == 1 + + +def test_recorder_stream_error_remains_fatal_without_finish( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + install_fake_backend(monkeypatch) + FakeCodex.events = [ + { + "method": "error", + "payload": { + "error": { + "message": "Reconnecting... 2/5", + "codex_error_info": {"response_stream_disconnected": {}}, + }, + "will_retry": True, + }, + }, + { + "method": "turn/completed", + "payload": {"turn": {"status": "completed", "duration_ms": 1000}}, + }, + ] + + result = make_planner(tmp_path, RecordingSink()).solve( + system_prompt="", + user_message="task", + toolkit=FakeToolkit(), + max_turns=3, + ) + + assert result.finish_result is None + assert result.error is not None + assert "Reconnecting... 2/5" in result.error + assert result.stats["nonfatal_recorder_errors"] == 0 + + +def test_failed_turn_error_is_not_suppressed_after_finish( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + install_fake_backend(monkeypatch) + FakeCodex.events = [ + { + "method": "item/completed", + "payload": { + "item": { + "type": "mcpToolCall", + "tool": "mcp__rpent__finish", + "status": "completed", + "arguments": {"status": "success", "summary": "done"}, + "result": "accepted", + } + }, + }, + { + "method": "turn/completed", + "payload": { + "turn": { + "status": "failed", + "error": {"message": "Codex turn failed"}, + } + }, + }, + ] + + result = make_planner(tmp_path, RecordingSink()).solve( + system_prompt="", + user_message="task", + toolkit=FakeToolkit(), + max_turns=3, + ) + + assert result.finish_result is not None + assert result.error == "Codex turn failed" + assert result.stats["nonfatal_recorder_errors"] == 0 + + def test_rejected_finish_item_is_not_promoted() -> None: from rpent.planner.codex import _Recorder diff --git a/tests/unit_tests/rpent/robots/test_config_contracts.py b/tests/unit_tests/rpent/robots/test_config_contracts.py index e7e9656e0..4dee19c08 100644 --- a/tests/unit_tests/rpent/robots/test_config_contracts.py +++ b/tests/unit_tests/rpent/robots/test_config_contracts.py @@ -36,6 +36,11 @@ def _parser(robot_name: str, *, dashboard: bool = False) -> argparse.ArgumentPar @pytest.mark.parametrize( ("robot_name", "required_args", "identity_fields"), [ + ( + "behavior", + ["--task-name", "turning_on_radio", "--public-seed", "1"], + ("task_name", "public_seed"), + ), ("libero", ["--suite", "libero_object_task", "--task", "2"], ("suite", "task")), ("robocasa", ["--task-name", "OpenDrawer"], ("task_name",)), ( @@ -62,6 +67,7 @@ def test_robot_arguments_are_required_on_cli_but_deferred_for_dashboard( @pytest.mark.parametrize( ("robot_name", "message"), [ + ("behavior", "--task-name is required"), ("libero", "--suite is required"), ("robocasa", "--task-name is required"), ("robotwin", "--task-name is required"), @@ -106,6 +112,45 @@ def test_libero_default_evaluation_config(tmp_path: Path) -> None: } +def test_behavior_uses_only_the_official_local_memory_profile(tmp_path: Path) -> None: + memory_dir = tmp_path / "behavior-memory" + args = _parser("behavior").parse_args( + [ + "--task-name", + "turning_on_radio", + "--public-seed", + "1", + "--memory-dir", + str(memory_dir), + "--output-dir", + str(tmp_path / "output"), + ] + ) + + config = get_robot_spec("behavior").parse_config(args) + + assert args.memory_profile == "local" + assert args.memory_dir == str(memory_dir.resolve()) + assert config.prompt_vars["memory_profile"] == "local" + assert config.prompt_vars["memory_dir"] == str(memory_dir.resolve()) + assert config.prompt_vars["memory_inbox"].endswith( + "_internal/inbox/turning_on_radio_s1" + ) + + invalid = _parser("behavior").parse_args( + [ + "--task-name", + "turning_on_radio", + "--public-seed", + "1", + "--memory-profile", + "hf", + ] + ) + with pytest.raises(ValueError, match="requires --memory-profile local"): + get_robot_spec("behavior").parse_config(invalid) + + def test_libero_exploration_uses_local_memory_and_session_metadata( tmp_path: Path, ) -> None: diff --git a/tests/unit_tests/rpent/robots/test_registry_contracts.py b/tests/unit_tests/rpent/robots/test_registry_contracts.py index 8aa6d8714..95b305213 100644 --- a/tests/unit_tests/rpent/robots/test_registry_contracts.py +++ b/tests/unit_tests/rpent/robots/test_registry_contracts.py @@ -14,10 +14,12 @@ from __future__ import annotations +import argparse from dataclasses import FrozenInstanceError from pathlib import Path from string import Formatter +import numpy as np import pytest from robots.robotwin.robot_spec import ( @@ -29,9 +31,25 @@ from rpent.robots import enumerate_robots, get_robot_spec from rpent.robots.robot_spec import RobotSpec, RunConfig -EXPECTED_ROBOTS = ("libero", "robocasa", "robotwin") +EXPECTED_ROBOTS = ("behavior", "libero", "robocasa", "robotwin") PROMPT_VARIABLES = { + "behavior": { + "behavior_mode": "eval", + "task_name": "turning_on_radio", + "task": 0, + "task_language": "Turn on the radio.", + "task_instruction": "Turn on the radio.", + "public_seed": 1, + "recipe_tag": "turning_on_radio_s1", + "output_dir": Path("/output"), + "max_episode_steps": 43200, + "wall_clock_seconds": 7200, + "public_capabilities": ["observe", "pi0_nav_pick", "finish"], + "memory_dir": "/memory", + "memory_profile": "local", + "memory_inbox": "/memory/_inbox/turning_on_radio_s1", + }, "libero": { "suite": "libero_object_task", "task": 2, @@ -83,6 +101,10 @@ def test_registry_discovers_exactly_the_source_checkout_robots() -> None: assert callable(spec.add_cli_args) assert callable(spec.parse_config) assert callable(spec.init_runtime) + if name == "behavior": + assert callable(spec.on_explore_session) + else: + assert spec.on_explore_session is None @pytest.mark.parametrize("robot_name", EXPECTED_ROBOTS) @@ -178,3 +200,207 @@ def test_robotwin_runtime_contracts_contain_execution_critical_metadata() -> Non )["action_layouts"] ) assert "mutated" not in vla_runtime_contract()["camera_order"] + + +def test_behavior_uses_the_shared_pi05_registry_and_wire_contract() -> None: + from rpent.robots.components.pi05_vla_client import Pi05VLAClient + from rpent.robots.components.pi05_vla_server import ( + PI05_EMBODIMENTS, + build_model_cfg, + ) + + assert PI05_EMBODIMENTS["behavior"]["openpi"]["config_name"] == "pi05_behavior" + assert PI05_EMBODIMENTS["behavior"]["openpi"]["action_chunk"] == 32 + assert PI05_EMBODIMENTS["behavior"]["openpi"]["action_env_dim"] == 23 + assert PI05_EMBODIMENTS["behavior"]["openpi_data"]["norm_stats_path"] == ( + "assets/behavior-1k/2025-challenge-demos/norm_stats.json" + ) + assert "assets" not in PI05_EMBODIMENTS["behavior"]["openpi_data"] + cfg = build_model_cfg("/tmp/pi05-behavior", PI05_EMBODIMENTS["behavior"]) + assert cfg.openpi_data.norm_stats_path == ( + "/tmp/pi05-behavior/assets/behavior-1k/2025-challenge-demos/norm_stats.json" + ) + + class FakeRpcClient: + def call(self, method, *, args, timeout_s): + assert method == "vla.predict" + observation, options = args + assert observation["main_images"].shape == (1, 720, 720, 3) + assert observation["wrist_images"].shape == (1, 2, 480, 480, 3) + assert observation["states"].shape == (1, 256) + assert observation["task_descriptions"] == ["turn on the radio"] + assert options == {"mode": "eval"} + assert timeout_s == 600.0 + return np.zeros((1, 32, 23), dtype=np.float32) + + model = Pi05VLAClient(FakeRpcClient(), embodiment="behavior") + action = model.predict( + { + "main_images": np.zeros((720, 720, 3), dtype=np.uint8), + "wrist_images": np.zeros((2, 480, 480, 3), dtype=np.uint8), + "states": np.zeros(256, dtype=np.float32), + "task_descriptions": "turn on the radio", + }, + options={"mode": "eval"}, + ) + assert action.shape == (32, 23) + assert action.dtype == np.float32 + + +def test_behavior_resolves_the_official_challenge_instance_layout( + tmp_path: Path, +) -> None: + from robots.behavior.rlinf_env import ( + _bootstrap_template_path, + _resolve_activity_instance_dir, + ) + + dataset_root = tmp_path / "2025-challenge-task-instances" + json_dir = dataset_root / "scenes" / "house" / "json" + instance_dir = json_dir / "house_task_demo_instances" + instance_dir.mkdir(parents=True) + bootstrap = json_dir / "house_task_demo_0_0_template.json" + bootstrap.write_text("{}") + selected = instance_dir / "house_task_demo_0_242_template-tro_state.json" + selected.write_text("{}") + + resolved = _resolve_activity_instance_dir( + dataset_root, + scene_model="house", + task_name="demo", + activity_definition_id=0, + activity_instance_id=242, + ) + + assert resolved == instance_dir + assert ( + _bootstrap_template_path( + resolved, + scene_model="house", + task_name="demo", + activity_definition_id=0, + ) + == bootstrap + ) + + +@pytest.mark.parametrize( + ("mode", "task_name", "public_seed", "role_title"), + [ + ("eval", "turning_on_radio", 1, "ROLE AND EVALUATION"), + ("explore", "picking_up_trash", 0, "ROLE AND MODE"), + ], +) +def test_behavior_prompts_strictly_render_real_run_config( + tmp_path: Path, + mode: str, + task_name: str, + public_seed: int, + role_title: str, +) -> None: + spec = get_robot_spec("behavior") + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir") + parser.add_argument("--memory-profile", choices=["hf", "local"], default=None) + parser.add_argument("--memory-dir") + spec.add_cli_args(parser, use_dashboard=False) + args = parser.parse_args( + [ + "--task-name", + task_name, + "--public-seed", + str(public_seed), + "--behavior-mode", + mode, + "--output-dir", + str(tmp_path / "output"), + "--memory-dir", + str(tmp_path / "memory"), + ] + ) + variables = spec.parse_config(args).prompt_vars + + system = spec.prompts.render("system", variables=variables) + user = spec.prompts.render("user", variables=variables) + + assert "{{" not in system + assert "{{" not in user + for value in ( + task_name, + variables["task_instruction"], + str(public_seed), + mode, + variables["recipe_tag"], + variables["output_dir"], + str(variables["max_episode_steps"]), + str(variables["wall_clock_seconds"]), + str(variables["public_capabilities"]), + variables["memory_profile"], + variables["memory_dir"], + variables["memory_inbox"], + ): + assert value in system or value in user + + ordered_sections = [ + role_title, + "CURRENT INVOCATION", + "INVOCATION MODEL", + "RUNTIME", + "YOUR GOAL", + "MEMORY", + "EVIDENCE", + "PLANNER TOOLS", + "TERMINATION", + "OUTPUT DISCIPLINE", + ] + positions = [system.index(title) for title in ordered_sections] + assert positions == sorted(positions) + if mode == "explore": + assert f"{variables['output_dir']}/{variables['recipe_tag']}.json" in system + assert f"{variables['memory_inbox']}/wip/" in system + assert "Only after official success" in system + assert "Do not invoke merge yourself" in system + for tool in ( + "pi0_nav_pick", + "observe", + "pixel_to_world", + "navigate_to", + "move_to", + "rotate_wrist", + "close", + "open", + "press", + ): + assert f"`{tool}`" in system + assert [user.index(title) for title in ("CELL", "MODE", "BEGIN")] == sorted( + user.index(title) for title in ("CELL", "MODE", "BEGIN") + ) + + required = { + "behavior_mode", + "task_name", + "task_instruction", + "public_seed", + "recipe_tag", + "output_dir", + "max_episode_steps", + "wall_clock_seconds", + "public_capabilities", + "memory_profile", + "memory_dir", + "memory_inbox", + } + for key in required: + incomplete = {name: value for name, value in variables.items() if name != key} + with pytest.raises(KeyError, match=key): + spec.prompts.render("system", variables=incomplete) + + literal = {**variables, "task_instruction": "literal {{do_not_expand}}"} + assert "literal {{do_not_expand}}" in spec.prompts.render( + "system", variables=literal + ) + + with pytest.raises(ValueError, match="unsupported BEHAVIOR prompt mode"): + spec.prompts.render( + "system", variables={**variables, "behavior_mode": "unknown"} + ) diff --git a/tests/unit_tests/rpent/session/test_state_contracts.py b/tests/unit_tests/rpent/session/test_state_contracts.py index 5b7316e55..09267937d 100644 --- a/tests/unit_tests/rpent/session/test_state_contracts.py +++ b/tests/unit_tests/rpent/session/test_state_contracts.py @@ -18,6 +18,7 @@ from pathlib import Path from typing import Any +import imageio.v2 as imageio import numpy as np import pytest @@ -208,6 +209,56 @@ def test_env_state_run_level_artifact_round_trips(tmp_path: Path) -> None: assert env_state.load_bytes("episode.mp4", step=None) == b"offline-mp4-bytes" +def test_env_state_streams_video_artifact_without_frame_buffer(tmp_path: Path) -> None: + env_state = EnvState(tmp_path) + writer = env_state.open_video_writer( + "episode.mp4", + step=None, + fps=5, + max_frames=3, + ) + frame = np.zeros((16, 16, 3), dtype=np.uint8) + + assert not hasattr(writer, "_frames") + assert writer.append(frame) is True + assert writer.append(frame + 20) is True + assert writer.append(frame + 40) is True + assert writer.append(frame + 60) is False + assert writer.frames_written == 3 + assert writer.frames_dropped == 1 + assert writer.close() == "episode.mp4" + assert writer.close() == "episode.mp4" + + path = env_state.artifact_path("episode.mp4", step=None) + assert path.stat().st_size > 0 + reader = imageio.get_reader(path) + try: + frames = [np.asarray(item) for item in reader] + finally: + reader.close() + assert len(frames) == 3 + assert json.loads((tmp_path / "states.json").read_text())["run_artifacts"] == [ + "episode.mp4" + ] + assert list(tmp_path.glob(".*.tmp*.mp4")) == [] + + +def test_env_state_video_writer_abort_and_empty_close_do_not_publish( + tmp_path: Path, +) -> None: + env_state = EnvState(tmp_path) + empty = env_state.open_video_writer("empty.mp4", step=None) + assert empty.close() is None + assert not env_state.exists("empty.mp4", step=None) + + aborted = env_state.open_video_writer("aborted.mp4", step=None) + aborted.append(np.zeros((16, 16, 3), dtype=np.uint8)) + aborted.abort() + aborted.abort() + assert not env_state.exists("aborted.mp4", step=None) + assert list(tmp_path.glob(".*.tmp*.mp4")) == [] + + def test_env_state_per_step_paths_save_load_and_exists(tmp_path: Path) -> None: env_state = EnvState(tmp_path) diff --git a/tests/unit_tests/rpent/tools/test_toolkit_contracts.py b/tests/unit_tests/rpent/tools/test_toolkit_contracts.py index dee85228e..72e2e5071 100644 --- a/tests/unit_tests/rpent/tools/test_toolkit_contracts.py +++ b/tests/unit_tests/rpent/tools/test_toolkit_contracts.py @@ -93,9 +93,14 @@ def solved(self) -> bool: def test_tool_result_builds_text_and_images_without_mutating_result() -> None: image_payloads = { "_image_bytes": b"main", + "_depth_image_bytes": b"main-depth", "_image_cam_bytes": b"camera", "_image_nav_bytes": b"navigation", "_image_wrist_bytes": b"wrist", + "_image_left_wrist_bytes": b"left-wrist", + "_depth_left_wrist_bytes": b"left-wrist-depth", + "_image_right_wrist_bytes": b"right-wrist", + "_depth_right_wrist_bytes": b"right-wrist-depth", } result = {"status": "ok", "count": 2, **image_payloads} original = copy.deepcopy(result) @@ -115,6 +120,11 @@ def test_tool_result_builds_text_and_images_without_mutating_result() -> None: "image", "image", "image", + "image", + "image", + "image", + "image", + "image", ] assert [ base64.b64decode(block["source"]["data"]) diff --git a/tests/unit_tests/rpent/utils/test_daemon_lifecycle.py b/tests/unit_tests/rpent/utils/test_daemon_lifecycle.py index 2d95f8080..85722d3b1 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 subprocess import sys import time from collections.abc import Callable @@ -122,3 +123,40 @@ 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) + + +@pytest.mark.parametrize("close_stdin", [False, True]) +def test_parent_watch_does_not_abort_interpreter_shutdown(close_stdin: bool) -> None: + script = """ +import sys, threading, time +from rpent.utils.daemon import watch_parent_death +done = threading.Event() +watch_parent_death(lambda: (print('eof', flush=True), done.set())) +print('ready', flush=True) +if CLOSE_STDIN: + assert done.wait(5) +else: + time.sleep(0.1) # Leave the watcher blocked while the interpreter exits. +""".replace("CLOSE_STDIN", repr(close_stdin)) + proc = subprocess.Popen( + [sys.executable, "-c", script], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + try: + assert proc.stdout.readline() == b"ready\n" + if close_stdin: + proc.stdin.write(b"not EOF yet") + proc.stdin.flush() + assert proc.poll() is None + proc.stdin.close() + assert proc.wait(timeout=10) == 0 + assert proc.stdout.read() == (b"eof\n" if close_stdin else b"") + assert b"_enter_buffered_busy" not in proc.stderr.read() + finally: + if proc.poll() is None: + proc.kill() + proc.wait() + for stream in (proc.stdin, proc.stdout, proc.stderr): + stream.close()