feat(env): manage env state centrally

docs/source-en/rst_source/development/add_primitive.rst

@@ -35,14 +35,17 @@ call. They differ only in how the method is implemented.
 Add a scripted primitive
 ------------------------
 
-Adding a scripted primitive usually involves three steps:
+Adding a scripted primitive usually involves two steps:
 
 1. **Add a method to the primitives.** Add the method to the
    current environment's primitives class, such as
    ``LiberoPrimitives`` or ``MyRobotPrimitives``. The method accepts
    the tool-call arguments, performs the work, usually through one or
    more ``self._env.step(...)`` calls, and returns a small log ``dict``.
 
+  Primitive methods capture and re-render state (``get_env_state``)
+  automatically after they run:
+
    .. code-block:: python
 
       def open_drawer(self, dx: float = 0.15) -> dict:
@@ -51,8 +54,12 @@ Adding a scripted primitive usually involves three steps:
               self._env.step(build_open_drawer_chunk(dx))
           return {"ok": True, "dx": dx}
 
+  You can mark read-only tools (``view_env_state``, ``back_project``, ``segment``,
+  ...) with :func:`~rpent.tools.toolkit.readonly` so the toolkit skips state
+  capture for them, improving performance.
+
 2. **Add the tool schema.** Add an entry to ``TOOLS_SPEC`` in
-   ``toolkit.py``:
+   ``robots/<env>/tools.py``:
 
    .. code-block:: python
 
@@ -67,13 +74,9 @@ Adding a scripted primitive usually involves three steps:
           },
       }
 
-3. **Register the tool in the toolkit.** Route it through the toolkit's
-   ``_step`` helper so that state is re-rendered after execution:
-
-   .. code-block:: python
-
-      self.add_tool("open_drawer", OPEN_DRAWER_SPEC,
-                    lambda **kw: self._step("open_drawer", **kw))
+Once both exist, the toolkit registers the tool automatically: it iterates
+``TOOLS_SPEC`` and binds each spec to the matching primitive-driver method
+(e.g. ``getattr(self._primitives, name)``).
 
 After these steps, the ``api``, ``claude_code``, and ``codex`` planners
 can all call the primitive without any other code changes.
@@ -167,8 +170,10 @@ Design principles for a new primitive
   the state dump reflecting the post-action world. Don't let the
   primitive return before the render finishes.
 - **Return small dicts.** Tool return values are fed back to the LLM
-  as text. Store larger content, such as images, depth data, and
-  ``states.json``, in the state dump instead.
+  as text. Save larger observations through ``EnvState.save``; ``EnvState``
+  automatically records each logical base name in its owned
+  ``StepRecord.artifacts`` set. Expose images through ``view_env_state`` and
+  geometry through environment tools rather than returning raw paths.
 - **Guardrails belong in env_server**, not in the toolkit. The LLM
   can and will call any tool with any arguments; workspace bounds
   and safety clamps must be enforced on the server side.

docs/source-en/rst_source/development/add_robot.rst

@@ -220,23 +220,29 @@ state needed for the current run. It exposes one method per primitive tool
 **Tool definitions and handlers** — a module-level ``TOOLS_SPEC`` list of
 Anthropic-style tool definitions (``name``, ``description``, ``input_schema``),
 plus any module-level functions referenced by the toolkit (e.g.
-``view_driver_state``, ``back_project``, ``finish``).
-
-**Per-step state dump** — ``dump_state(primitives, output_dir, step_idx, log)``
-serializes whatever state the agent will read back via the ``view_*`` tools
-(images, depths, JSON state, camera meta) into ``output_dir``.
+``view_env_state``, ``back_project``, ``finish``).
+
+**Per-step state dump** — ``dump_state(driver, env_state, log)`` opens
+``env_state.record_step(...)`` and receives the allocated step index; the
+``StepRecord`` is appended and committed immediately. Save large observations
+through ``env_state.save(...)`` — inside a ``record_step`` block the ``step``
+argument may be omitted (it defaults to the new step), pass an explicit
+``step=<int>`` to target a different step, and ``step=None`` for run-level
+artifacts. ``EnvState`` adds every successfully saved base name to the step's
+flat ``artifacts`` set automatically. Readers use the canonical artifact
+filenames rather than maintaining a parallel observation index.
 
 **Toolkit class** — subclass ``rpent.tools.toolkit.Toolkit``:
 
 - build the primitives in ``__init__`` through a custom initialization
-  helper (named ``init_primitives_clean`` in LIBERO; it wipes stale
-  ``images/`` etc., constructs the primitives, and dumps step 0),
+  helper (named ``init_primitives_clean`` in LIBERO; it calls
+  ``EnvState.reset()``, constructs the primitives, and dumps step 0),
 - register each tool with ``self.add_tool(name, spec, handler)`` — stateless
-  readers (``view_driver_state``, ``finish``, …) bind directly to module-level
+  readers (``view_env_state``, ``finish``, …) bind directly to module-level
   functions; primitive tools route through ``_step(name, **kwargs)`` which
   calls ``getattr(self._primitives, name)(**kwargs)`` and re-renders state,
-- override ``close()`` to write any remaining agent-side artifacts (e.g. the
-  LIBERO toolkit saves the agentview MP4 there).
+- override ``close()`` to save remaining agent-side artifacts through
+  ``EnvState`` (for example ``state.save("episode.mp4", frames, step=None)``).
 
 ``primitives_kwargs`` (forwarded from ``__init__.py:get_toolkit``) is the dict
 the toolkit passes verbatim to your primitives' ``__init__`` — typically
@@ -246,14 +252,15 @@ Conventions worth keeping
 -------------------------
 
 - ``output_dir`` is the working directory that the runner creates for each
-  run. Images, depths, ``states.json``, transcripts, ``episode.mp4``, and other
-  artifacts go there.
+  run. Environment observations are owned by ``EnvState``; callers use logical
+  base names and never construct storage paths. Transcripts and other
+  run-management outputs share the same run directory.
 - Tool definitions use the Anthropic format (``name`` / ``description`` /
   ``input_schema``). Every tool registered with ``self.add_tool(...)`` is
   exposed to all planners.
 - Server-side return values must be picklable and torch-free.
 - Each primitive tool dumps a fresh state snapshot after running so the next
-  ``view_driver_state`` call reflects the post-action world.
+  ``view_env_state`` call reflects the post-action world.
 - Treat ``dump_state`` as the source of truth for what the agent sees — any new
   modality (e.g. tactile, force) goes through it.
 

docs/source-en/rst_source/development/interfaces.rst

@@ -90,7 +90,7 @@ Subclass ``Toolkit`` in ``robots/<env>/toolkit.py`` and register env tools with
        ends; optional ``_image_bytes`` (etc.) to return camera images.
 
 The base class already registers common file tools; call ``super().__init__()`` then
-``add_tool`` for env tools. Per-step state and ``view_driver_state`` are in
+``add_tool`` for env tools. Per-step state and ``view_env_state`` are in
 :doc:`add_primitive`.
 
 Inter-process communication

docs/source-en/rst_source/quickstart.rst

@@ -103,14 +103,11 @@ A successful run:
    by the elapsed time, token usage, and path to the run record.
 3. With the Dashboard enabled, also streams agent output, camera views,
    the action timeline, and clip replays to the Dashboard.
-4. By default, artifacts are saved under
-   ``logs/<timestamp>_<suite>_t<task>_s<seed>/``. They include
-   ``transcript_*.json`` (run record), ``states.json`` (one record per
-   environment step), ``recipe_*.jsonl`` (action sequence), and
-   ``episode.mp4`` (episode video).
-
-After the run, inspect the final record in ``states.json``:
-``libero_terminated`` set to ``true`` means LIBERO judged the task complete.
-You can also open ``episode.mp4`` to review the run.
+4. By default, artifacts are saved under ``logs/<timestamp>_<suite>_t<task>_s<seed>/``. They include ``transcript_*.json`` (run record), ``states.json`` (the ``EnvState`` manifest), ``recipe_*.jsonl`` (action sequence), and ``episode.mp4`` (episode video). Each step artifact has a directory named after its logical artifact name; zero-padded step files live inside it, for example ``agentview_depth.npy/00.npy`` and ``agentview_depth.npy/01.npy``. Run-level artifacts remain at the output root.
+
+Inspect the final state through the Dashboard or
+``view_env_state(step=-1)``. Its top-level ``terminated`` value is the
+benchmark outcome. ``states.json`` is internal ``EnvState`` storage and should
+not be parsed by callers. You can also open ``episode.mp4`` to review the run.
 If something goes wrong, inspect the four log files described at the
 bottom of :doc:`installation`.

docs/source-en/rst_source/usage/libero.rst

@@ -116,7 +116,7 @@ What runs where
   transports (HTTP or socket). It returns only the top compressed PNG mask.
 - **toolkit** (``robots/libero/toolkit.py``) — defines the tools the
   LLM can call: ``pi0_pick`` (fed to Pi0.5), ``move_to``,
-  ``rotate_wrist``, ``back_project``, ``view_driver_state``,
+  ``rotate_wrist``, ``back_project``, ``view_env_state``,
   ``finish``, …
 
 Tools the planner can call
@@ -147,8 +147,10 @@ Physical action tools advance the environment and record new state and images.
   coordinates.
 - ``segment(prompt=... / point=..., ...)`` — use SAM3 to segment an existing
   image with a text or point prompt.
-- ``view_driver_state(step=None)`` — read an existing state and image record.
-- ``view_camera_meta(camera=..., step=None)`` — read existing camera metadata.
+- ``view_env_state(step=-1)`` — read a recorded state and its embedded
+  observation images. Step ``0`` is initial; ``-1`` is latest.
+- ``view_camera_meta(camera=..., step=-1)`` — read camera metadata for a
+  recorded step. Step ``-1`` is latest.
 - ``finish(status, summary)`` — end the current run.
 
 These tools do not advance the environment.

docs/source-zh/rst_source/development/add_primitive.rst

@@ -33,13 +33,16 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法
 添加一个脚本化原语
 ------------------
 
-添加脚本化原语通常需要以下三个步骤：
+添加脚本化原语通常需要以下两个步骤：
 
 1. **在 primitives 中添加方法。** 在当前环境的 primitives
    类（如 ``LiberoPrimitives``、``MyRobotPrimitives``）中添加
    一个方法。该方法接收工具调用的参数，执行一次或多次
    ``self._env.step(...)``，并返回一个简短的日志字典。
 
+     primitive 方法执行后默认会自动捕获并重新渲染状态
+     （``get_env_state``）：
+
    .. code-block:: python
 
       def open_drawer(self, dx: float = 0.15) -> dict:
@@ -48,7 +51,11 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法
               self._env.step(build_open_drawer_chunk(dx))
           return {"ok": True, "dx": dx}
 
-2. **添加工具定义。** 在 ``toolkit.py`` 的 ``TOOLS_SPEC`` 中新增一项：
+   只读工具（``view_env_state``、``back_project``、``segment`` 等）
+     可以使用 :func:`~rpent.tools.toolkit.readonly` 标记，toolkit 会跳过
+     它们的状态捕获，提升性能。
+
+2. **添加工具定义。** 在 ``robots/<env>/tools.py`` 的 ``TOOLS_SPEC`` 中新增一项：
 
    .. code-block:: python
 
@@ -63,13 +70,8 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法
           },
       }
 
-3. **在 toolkit 中注册工具。** 通过 toolkit 的 ``_step`` 辅助函数运行
-   该工具，使其在执行结束后自动重新渲染状态：
-
-   .. code-block:: python
-
-      self.add_tool("open_drawer", OPEN_DRAWER_SPEC,
-                    lambda **kw: self._step("open_drawer", **kw))
+两者就位后，toolkit 会自动注册该工具：它遍历 ``TOOLS_SPEC``，把每个定义
+绑定到对应的 primitive 方法（如 ``getattr(self._primitives, name)``）。
 
 完成以上步骤后，``api``、``claude_code`` 和 ``codex`` 三种 planner
 都可以调用该工具，无需修改其他代码。
@@ -152,7 +154,9 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法
 - **每个工具执行结束后都要保存新的状态快照。** 下一轮需要读取动作执行后的
   环境状态，因此原语不能在渲染完成前返回。
 - **工具只返回简短的字典。** 返回值会以文本形式提供给 LLM；图像、深度数据和
-  ``states.json`` 等较大的内容则通过状态快照提供。
+  其他大型观测应通过 ``EnvState.save`` 保存；``EnvState`` 会把每个逻辑基础
+  文件名自动加入其持有的 ``StepRecord.artifacts`` 集合。图像通过
+  ``view_env_state`` 提供，几何数据通过环境工具访问，不返回原始路径。
 - **安全限制由 ``env_server`` 强制执行。** LLM 可能使用任意参数调用工具，
   因此工作空间边界和安全限制不能只依赖 toolkit。
 

docs/source-zh/rst_source/development/add_robot.rst

@@ -207,23 +207,27 @@ toolkit 模块通常包含四部分：
 **工具定义和处理函数** 包括模块级的 ``TOOLS_SPEC`` 列表（列表元素采用
 Anthropic API 的工具定义格式，包含 ``name``、``description`` 和
 ``input_schema``），以及 toolkit 引用的模块级函数，例如
-``view_driver_state``、``back_project`` 和 ``finish``。
+``view_env_state``、``back_project`` 和 ``finish``。
 
-**每步状态 dump** —— ``dump_state(primitives, output_dir, step_idx, log)`` 把 agent
-之后会通过 ``view_*`` 工具读回的所有状态 (图像、深度、JSON 状态、camera meta)
-序列化到 ``output_dir``。
+**每步状态 dump** —— ``dump_state(driver, env_state, log)`` 通过
+``env_state.record_step(...)`` 创建由 ``EnvState`` 持有的步骤，并取得分配的
+step index；该 ``StepRecord`` 会被立即追加并提交。大型观测通过
+``env_state.save(...)`` 保存——在 ``record_step`` 块内可省略 ``step`` 参数
+（默认指向刚创建的步骤），传显式 ``step=<int>`` 可指定其它步骤，``step=None``
+用于运行级工件。每次保存成功后，``EnvState`` 会自动把基础文件名加入该
+``StepRecord`` 的扁平 ``artifacts`` 集合；读取方直接使用规范化的工件文件名。
 
 **Toolkit 类** 继承 ``rpent.tools.toolkit.Toolkit``：
 
 - 在 ``__init__`` 中通过自定义的初始化辅助方法构建 primitives（LIBERO
-  中的方法名为 ``init_primitives_clean``；它会清理过期的 ``images/`` 等目录、
-  构造原语并 dump 第 0 步）,
+  中的方法名为 ``init_primitives_clean``；它会调用 ``EnvState.reset()``、构造
+  原语并 dump 第 0 步）,
 - 用 ``self.add_tool(name, spec, handler)`` 注册每个工具。无状态的读取工具
-  （如 ``view_driver_state``、``finish``）直接绑定模块级函数；原语工具通过
+  （如 ``view_env_state``、``finish``）直接绑定模块级函数；原语工具通过
   ``_step(name, **kwargs)`` 调用。``_step`` 使用
-  ``getattr(self._primitives, name)(**kwargs)`` 调用 primitives 方法并重新渲染状态；
-- 重写 ``close()``，将 agent 侧生成的文件写入磁盘（例如 LIBERO toolkit
-  在这里保存 agentview MP4）。
+  ``getattr(self._primitives, name)(**kwargs)`` 调用 driver 方法并重新渲染状态；
+- 重写 ``close()``，通过 ``EnvState`` 保存 agent 侧剩余工件（例如
+  ``state.save("episode.mp4", frames, step=None)``）。
 
 ``primitives_kwargs`` 由 ``__init__.py:get_toolkit`` 转发给 toolkit，再原样传入
 primitives 的 ``__init__``。其中通常包含
@@ -232,14 +236,15 @@ primitives 的 ``__init__``。其中通常包含
 建议遵循的约定
 --------------
 
-- ``output_dir`` 是 runner 为单次运行创建的临时目录。图像、深度数据、
-  ``states.json``、transcript 和 ``episode.mp4`` 等工件都写入该目录。
+- ``output_dir`` 是 runner 为单次运行创建的工作目录。环境观测由
+  ``EnvState`` 管理；调用方只使用逻辑基础文件名，不自行拼接存储路径。
+  transcript 等运行管理输出与环境工件共享该目录。
 - 工具定义使用 Anthropic API 格式（``name`` / ``description`` /
   ``input_schema``）。
   每个用 ``self.add_tool(...)`` 注册的工具都会暴露给所有 planner。
 - 环境侧的返回值必须可 pickle，且不包含 torch 对象。
 - 每个原语工具执行后要 dump 一次新的状态快照, 这样下一次
-  ``view_driver_state`` 看到的是动作后的世界。
+  ``view_env_state`` 看到的是动作后的世界。
 - ``dump_state`` 是 Agent 获取环境状态的唯一数据来源；任何新的模态
   （例如触觉、力）都通过它提供。
 

docs/source-zh/rst_source/development/interfaces.rst

@@ -87,7 +87,7 @@ Planner
        需要回传相机图时可设 ``_image_bytes`` 等字段。
 
 基类已注册公共文件工具；子类 ``super().__init__()`` 后追加本环境工具即可。逐步状态与
-``view_driver_state`` 见 :doc:`add_primitive`。
+``view_env_state`` 见 :doc:`add_primitive`。
 
 进程间通信
 ----------

docs/source-zh/rst_source/quickstart.rst

@@ -95,7 +95,9 @@ LIBERO-PRO 仿真资源。下面以 LIBERO-PRO 和 ``claude_code`` planner
 1. 终端会先显示 ``env_server``、``vla_server`` 和 ``sam3_server`` 的启动信息。
 2. 智能体的逐轮输出和工具调用会显示在终端中；运行结束时还会显示耗时、token 用量和运行记录的路径。
 3. 启用 Dashboard 后，智能体的输出、相机视图、动作时间线和片段回放也会实时显示在 Dashboard 中。
-4. 默认输出目录为 ``logs/<timestamp>_<suite>_t<task>_s<seed>/``，其中包含 ``transcript_*.json``\ （运行记录）、``states.json``\ （每个环境步的记录）、``recipe_*.jsonl``\ （动作序列）和 ``episode.mp4``\ （回合录像）。
+4. 默认输出目录为 ``logs/<timestamp>_<suite>_t<task>_s<seed>/``，其中包含 ``transcript_*.json``\ （运行记录）、``states.json``\ （``EnvState`` 清单）、``recipe_*.jsonl``\ （动作序列）和 ``episode.mp4``\ （回合录像）。每种逐步工件使用一个与逻辑工件同名的目录，目录内按步骤保存零填充文件，例如 ``agentview_depth.npy/00.npy`` 和 ``agentview_depth.npy/01.npy``；运行级工件仍保存在输出目录根部。
 
-运行结束后，查看 ``states.json`` 的最后一条记录：``libero_terminated`` 为 ``true`` 表示 LIBERO 已判定任务完成；也可以打开 ``episode.mp4`` 复核运行过程。
+通过 Dashboard 或 ``view_env_state(step=-1)`` 查看最终状态；其顶层
+``terminated`` 即为基准任务结果。``states.json`` 是 ``EnvState`` 的内部
+存储，调用方不应直接解析。也可以打开 ``episode.mp4`` 复核运行过程。
 出问题时，参考 :doc:`installation` 页底部提到的四份日志文件。

docs/source-zh/rst_source/usage/libero.rst

@@ -112,7 +112,7 @@ LIBERO-PRO 核心套件一览
   排名第一的压缩 PNG mask。
 - **toolkit（工具集）** （``robots/libero/toolkit.py``）—— 定义 LLM
   能调用的工具：``pi0_pick``（交给 Pi0.5）、``move_to``、``rotate_wrist``、
-  ``back_project``、``view_driver_state``、``finish``…
+  ``back_project``、``view_env_state``、``finish``…
 
 Planner 能调用的工具
 --------------------
@@ -141,8 +141,10 @@ LIBERO 工具分为物理动作工具和只读工具。
 - ``back_project(row, col, ...)`` —— 将图像像素反投影到世界坐标。
 - ``segment(prompt=... / point=..., ...)`` —— 通过 SAM3 对已有图像进行文本或
   点提示分割。
-- ``view_driver_state(step=None)`` —— 读取已有的状态和图像记录。
-- ``view_camera_meta(camera=..., step=None)`` —— 读取已有的相机元数据。
+- ``view_env_state(step=-1)`` —— 读取已记录的状态和内嵌观测图像；第 0 步为
+  初始状态，``-1`` 表示最新状态。
+- ``view_camera_meta(camera=..., step=-1)`` —— 读取指定步骤的相机元数据；
+  ``-1`` 表示最新状态。
 - ``finish(status, summary)`` —— 结束当前运行。
 
 这些工具不会推进环境。

pyproject.toml

@@ -58,6 +58,7 @@ openpi = [
 ]
 libero = [
     "rlinf-libero",
+    "h5py",
 ]
 libero-pro = [
     "rpent[libero]",

robots/libero/__init__.py

@@ -45,15 +45,13 @@ def get_toolkit(
     *,
     primitives_kwargs: dict[str, Any],
     dashboard_events: DashboardEventSink,
-    video_path: str | None = None,
 ):
     """Return the LIBERO toolkit (common tools + LIBERO primitives)."""
     from robots.libero.toolkit import LiberoToolkit
 
     return LiberoToolkit(
         primitives_kwargs=primitives_kwargs,
         dashboard_events=dashboard_events,
-        video_path=video_path,
     )