Add reusable RL lifecycle and checkpoint resource hooks

.agents/skills/add-functor/SKILL.md

@@ -17,7 +17,7 @@ Scaffold a new functor following EmbodiChain's Functor/FunctorCfg pattern.
 
 | Functor Type | Config Class | Module File | Manager | Signature |
 |-------------|-------------|-------------|---------|-----------|
-| Observation | `ObservationCfg` (extends `FunctorCfg`) | `managers/observations.py` | `ObservationManager` | `(env, obs, entity_cfg, ...) -> Tensor` |
+| Observation | `ObservationCfg` (extends `FunctorCfg`) | `managers/observations/` | `ObservationManager` | `(env, obs, entity_cfg, ...) -> Tensor` |
 | Condition | `ConditionCfg` | `managers/conditions.py` | `TerminationManager` | `(env, ...) -> ConditionResult` |
 | Reward | `RewardCfg` (extends `FunctorCfg`) | `managers/rewards.py` | `RewardManager` | `(env, obs, action, info, ...) -> Tensor` |
 | Event | `EventCfg` (extends `FunctorCfg`) | `managers/events.py` | `EventManager` | `(env, env_ids, ...) -> None` |
@@ -95,7 +95,7 @@ Place the functor in the existing module for its type:
 
 | Type | File |
 |------|------|
-| Observation | `embodichain/lab/gym/envs/managers/observations.py` |
+| Observation | `embodichain/lab/gym/envs/managers/observations/` |
 | Reward | `embodichain/lab/gym/envs/managers/rewards.py` |
 | Event | `embodichain/lab/gym/envs/managers/events.py` |
 | Action | `embodichain/lab/gym/envs/managers/actions.py` |
@@ -128,6 +128,8 @@ Key rules:
 ### 4. Update `__all__`
 
 Add the new functor to the module's `__all__` list. If no `__all__` exists, create one.
+For observation implementations, re-export the functor from the observation
+package so existing config-name resolution finds it.
 
 ### 5. Write a Test
 

agent_context/MAP.yaml

@@ -90,8 +90,8 @@ topics:
   status: active
 - id: manager-functor
   title: Manager & Functor Pattern
-  aliases: [manager functor, functor pattern, observation manager, reward manager, event manager, action
-      manager, dataset manager, 管理器]
+  aliases: [manager functor, functor pattern, observation manager, observation functor, obs functor,
+    reward manager, event manager, action manager, dataset manager, 观测 functor, 观测管理器, 管理器]
   keywords: [FunctorCfg, ObservationCfg, RewardCfg, ActionTermCfg, ObservationManager, RewardManager,
     EventManager, ActionManager, ActionDescriptor, ActionTermDescriptor, ParallelGripperAction,
     DatasetManager, SceneEntityCfg, TerminationManager, ConditionResult, ObservationDescriptor, ObservationSample, get_entity_state, observation consumers]
@@ -103,6 +103,7 @@ topics:
   - embodichain/lab/gym/envs/managers/manager_base.py
   - embodichain/lab/gym/envs/managers/cfg.py
   - embodichain/lab/gym/envs/managers/observation_manager.py
+  - embodichain/lab/gym/envs/managers/observations/
   - embodichain/lab/gym/envs/managers/observation_types.py
   - embodichain/lab/gym/envs/managers/reward_manager.py
   - embodichain/lab/gym/envs/managers/event_manager.py

agent_context/topics/env-framework/configuration.md

@@ -28,6 +28,12 @@ class MyTaskEnv(EmbodiedEnv):
    classes/types). A `RuntimeError` is raised otherwise.
 6. Use `override=True` to re-register an existing uid (useful in scripts/tests).
 
+In a source checkout, `embodichain_tasks/__init__.py` resolves the sibling
+`configs` package before eager registration. This keeps task code and resources
+in the same selected worktree even when an editable install maps configs to
+another checkout. Its temporary search path is removed before walking task
+modules; wheel installations retain their normal nested resource layout.
+
 ### Gym ID convention
 
 Format: `<TaskName>-v<N>` (e.g. `PourWater-v1`, `PushCubeRL`).

agent_context/topics/env-framework/env-framework.md

@@ -92,6 +92,13 @@ recorded boundary. Missing or foreign evidence cannot satisfy a predicate.
 
 ## Lifecycle and boundaries
 
+An owned `scene_preparation` functor prepares ordinary `EmbodiedEnvCfg` before
+World creation; registered defaults and explicit config use the same route.
+Explicit config wins. Preparation returns configuration and does not own a
+runtime environment. `resource_export` freezes generated task resources when
+saving native RL checkpoints; task-specific preparation/export stays in the
+task extension so registration can use ordinary `EmbodiedEnv`.
+
 Construction seeds before scene setup, acquires enabled sensors, prepares
 simulation and initializes dependent managers. Robot joint selections are
 instance-owned and preserve control-part order across repeated setup.

agent_context/topics/manager-functor/manager-functor.md

@@ -6,7 +6,7 @@
 |------|------|
 | Base classes (`ManagerBase`, `Functor`) | `embodichain/lab/gym/envs/managers/manager_base.py` |
 | All config classes (`FunctorCfg`, `EventCfg`, `ObservationCfg`, `RewardCfg`, `ActionTermCfg`, `DatasetFunctorCfg`, `SceneEntityCfg`) | `embodichain/lab/gym/envs/managers/cfg.py` |
-| Observation manager + built-in functors | `managers/observation_manager.py`, `managers/observations.py` |
+| Observation manager + built-in functors | `managers/observation_manager.py`, `managers/observations/` |
 | Reward manager + built-in functors | `managers/reward_manager.py`, `managers/rewards.py` |
 | Event manager + built-in functors | `managers/event_manager.py`, `managers/events.py` |
 | Episode camera recording events | `managers/record.py` |
@@ -16,6 +16,9 @@
 | Randomization functors (event sub-type) | `managers/randomization/` (spatial, visual, physics, geometry) |
 
 All paths relative to `embodichain/lab/gym/envs/`.
+Observation implementations are grouped in `managers/observations/` and
+re-exported by that package for configuration lookup and existing imports.
+The historical locomotion module forwards its existing observation name.
 
 ---
 

agent_context/topics/rl-learning/rl-learning.md

@@ -119,6 +119,27 @@ orbit/zoom and recorder view. Use
 camera changes; the [evaluation guide](../../../docs/source/guides/policy_evaluation.md)
 owns user controls.
 
+## Episode evaluation and saved resources
+
+Auto-reset preserves completed-row observations and IDs for timeout bootstrap;
+only true timeouts enter terminal value inference, with module training modes
+restored after inference. GRPO does not request those values. True termination
+suppresses bootstrap. Legacy environments retain their fallback.
+Evaluation defaults to first completion order; `balanced` allocates per-arena
+quotas before stepping. `on_episode` retains owned environment terminal reports;
+`on_scored_episode` emits one scalar record per scored completion after quota
+filtering. The native headless CLI retains both quota selection and terminal
+reports. Evaluation uses an independent environment and restores policy mode.
+
+Training snapshots freeze physical components relative to the declaring
+config before moving them into a run. `EmbodiedEnvCfg.resource_export` packages
+materialized task resources from the unwrapped environment and an owned frozen
+mapping. Explicit local physical-resource paths resolve from the saved config.
+Task Program snapshots retain embodiment skill metadata and bind physical and
+semantic resource references before relocation. Resource-packaging failures
+are synchronized across ranks before resume or training starts. Packaging does
+not redefine checkpoint/resume compatibility.
+
 ## Invariants and validation
 
 Rollout kind, collector, payload and trainer must agree. APG requires graph-

docs/source/api_reference/public_api.rst

@@ -5086,6 +5086,17 @@ qualified DriverOpeningBinding type must be rebuilt with its Atomic owner.
 .. autoclass:: DriverOpeningBinding
    :members:
 
+Shared locomotion observation functor
+------------------------------------
+
+The observation package exposes task-built actor or critic state and preserves the
+existing locomotion configuration name while implementations move into a package.
+
+.. currentmodule:: embodichain.lab.gym.envs.managers.observations
+
+.. autofunction:: velocity_locomotion_observation
+
+
 Stability execution ownership
 ----------------------------
 

embodichain/lab/gym/envs/base_env.py

@@ -1147,9 +1147,11 @@ def step(
                 rollout or segment boundaries.
 
         Returns:
-            The observation, reward, termination and truncation flags, and
-            info. Hosts that publish terminal reports include completed rows
-            under `terminal_outcomes`.
+            Observation, reward, terminated, truncated, and an info mapping.
+            Completed rows may carry environment-owned ``terminal_outcomes``.
+            Auto-reset also preserves owned ``terminal_observation`` rows,
+            ``terminal_env_ids`` and a full ``terminal_observation_mask`` for
+            time-limit bootstrapping; live observations contain reset states.
         """
 
         if recording_active_mask is not None:
@@ -1222,10 +1224,23 @@ def step(
             if not replay and getattr(self, "_demo_instruction_snapshot", None) is None:
                 reset_env_ids = dones.nonzero(as_tuple=False).squeeze(-1)
                 if len(reset_env_ids) > 0:
+                    if isinstance(info, TensorDict):
+                        info = info.to_dict()
                     self._seal_step_outcomes(
                         reset_env_ids, terminateds, truncateds, info
                     )
                     if not no_auto_reset:
+                        # Snapshot completed rows before reset reuses buffers.
+                        info["terminal_env_ids"] = reset_env_ids.clone()
+                        info["terminal_observation_mask"] = (
+                            truncateds & ~terminateds
+                        ).clone()
+                        if isinstance(obs, TensorDict):
+                            info["terminal_observation"] = obs[reset_env_ids]
+                        else:
+                            info["terminal_observation"] = torch.utils._pytree.tree_map(
+                                lambda value: value[reset_env_ids], obs
+                            )
                         with self._profiler.section("auto_reset"):
                             obs, reset_info = self.reset(
                                 options={"reset_ids": reset_env_ids}

embodichain/lab/gym/envs/embodied_env.py

@@ -430,6 +430,16 @@ class EnvLightCfg:
     ``trajectory_save_dir`` before an explicitly committed episode reset.
     Closing discards trajectories that were never committed."""
 
+    scene_preparation: FunctorCfg | None = None
+    """Optional preparation function ``func(cfg, **params) -> EmbodiedEnvCfg``.
+
+    Runs on an owned configuration before World creation. The result declares
+    ordinary scene entities and manager terms; it does not receive an environment.
+    """
+
+    resource_export: FunctorCfg | None = None
+    """Optional ``func(env, config, config_dir, **params)`` checkpoint packager."""
+
     rl_task: embodichain.lab.gym.envs.managers.RLTaskCfg | None = None
     """Optional measured RL success, failure and truncation conditions.
 
@@ -550,8 +560,35 @@ def __init__(
         cfg: EmbodiedEnvCfg,
         *,
         task_program_adapter_factory: TaskProgramAdapterFactory | None = None,
+        scene_preparation: dict[str, Any] | None = None,
+        resource_export: dict[str, Any] | None = None,
         **kwargs,
     ):
+        # Registered defaults and explicit config use the same pre-World path.
+        # Copy before preparation so one reusable configuration can build peers.
+        if (
+            scene_preparation is not None
+            or resource_export is not None
+            or cfg.scene_preparation is not None
+        ):
+            cfg = copy.deepcopy(cfg)
+            for name, default in (
+                ("scene_preparation", scene_preparation),
+                ("resource_export", resource_export),
+            ):
+                if getattr(cfg, name) is None and default is not None:
+                    setattr(cfg, name, FunctorCfg(**copy.deepcopy(default)))
+            if cfg.scene_preparation is not None:
+                operation = cfg.scene_preparation
+                operation.validate()
+                prepare = (
+                    string_to_callable(operation.func)
+                    if isinstance(operation.func, str)
+                    else operation.func
+                )
+                cfg = prepare(cfg, **operation.params)
+                if not isinstance(cfg, EmbodiedEnvCfg):
+                    raise TypeError("scene_preparation must return EmbodiedEnvCfg.")
         replay = getattr(cfg, "replay", None)
         if replay is not None:
             cfg = copy.deepcopy(cfg)

embodichain/lab/gym/envs/managers/actions.py

@@ -257,7 +257,12 @@ def apply_actions(self) -> None:
 
 
 class RelativeJointPositionAction(_JointAction):
-    """Add scaled relative actions to current selected positions."""
+    """Add scaled relative actions to current selected positions.
+
+    Optional ``target_limits`` contains a lower/upper pair per selected joint.
+    It limits the resulting position command without changing the raw policy
+    action used by observation histories.
+    """
 
     contract_id: ClassVar[str] = "joint_position.relative@1"
     command_type = "qpos"
@@ -274,6 +279,20 @@ def __init__(self, cfg: ActionTermCfg, env: EmbodiedEnv) -> None:
         )
         clip = cfg.params.get("clip")
         self._clip = None if clip is None else float(clip)
+        limits = cfg.params.get("target_limits")
+        self._target_limits = None
+        if limits is not None:
+            limits = torch.as_tensor(limits, dtype=torch.float32, device=env.device)
+            if (
+                tuple(limits.shape) != (self.action_dim, 2)
+                or not torch.isfinite(limits).all()
+                or not (limits[:, 0] < limits[:, 1]).all()
+            ):
+                raise ValueError(
+                    "target_limits must contain finite ordered lower/upper pairs "
+                    "for each selected joint."
+                )
+            self._target_limits = limits
 
     @property
     def action_space(self) -> gym.spaces.Box:
@@ -290,6 +309,10 @@ def process_actions(self, actions: torch.Tensor) -> None:
             self._raw_actions.clamp_(-self._clip, self._clip)
         current = self._env.robot.get_qpos()[:, list(self._joint_ids)]
         self._processed_actions.copy_(current + self._raw_actions * self._scale)
+        if self._target_limits is not None:
+            self._processed_actions.clamp_(
+                min=self._target_limits[:, 0], max=self._target_limits[:, 1]
+            )
 
     def apply_actions(self) -> None:
         self._env.robot.set_qpos(

embodichain/lab/gym/envs/managers/observations/__init__.py

@@ -0,0 +1,54 @@
+# ----------------------------------------------------------------------------
+# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
+#
+# 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
+#
+#     http://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.
+# ----------------------------------------------------------------------------
+"""Observation functors for environment managers.
+
+Existing imports and configuration names resolve from this package. The
+implementation modules group general and locomotion observations.
+"""
+
+from __future__ import annotations
+
+from ._common import *
+from ._common import (
+    _capture_proprioception,
+    _get_deformable_centroid_state,
+    _get_evaluation_state,
+    _get_segment_start_state,
+    _state_descriptor,
+)
+from ._locomotion import *
+
+__all__ = [
+    "get_entity_state",
+    "get_object_pose",
+    "get_rigid_object_pose",
+    "get_object_body_scale",
+    "get_object_uid",
+    "get_rigid_object_velocity",
+    "normalize_robot_joint_data",
+    "get_sensor_pose_in_robot_frame",
+    "get_sensor_intrinsics",
+    "compute_semantic_mask",
+    "get_robot_eef_pose",
+    "target_position",
+    "compute_exteroception",
+    "get_rigid_object_physics_attributes",
+    "get_articulation_joint_drive",
+    "get_fluid_particle_state",
+    "get_fluid_summary",
+    "get_fluid_volume_in_region",
+    "velocity_locomotion_observation",
+]

embodichain/lab/gym/envs/managers/observations.py → embodichain/lab/gym/envs/managers/observations/_common.py

@@ -30,7 +30,7 @@
     ObservationDescriptor,
     ObservationSample,
 )
-from .observation_manager import _current_capture
+from ..observation_manager import _current_capture
 from embodichain.lab.gym.envs.managers.cfg import SceneEntityCfg
 from embodichain.lab.gym.envs.managers.events import resolve_dict
 from embodichain.lab.gym.envs.managers import Functor, FunctorCfg

embodichain/lab/gym/envs/managers/observations/_locomotion.py

@@ -0,0 +1,57 @@
+# ----------------------------------------------------------------------------
+# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
+#
+# 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
+#
+#     http://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.
+# ----------------------------------------------------------------------------
+
+"""Observation functors for native locomotion environments."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from embodichain.lab.sim.types import EnvObs
+
+if TYPE_CHECKING:
+    from embodichain.lab.gym.envs import EmbodiedEnv
+
+__all__ = ["velocity_locomotion_observation"]
+
+
+def velocity_locomotion_observation(
+    env: EmbodiedEnv,
+    obs: EnvObs,
+    privileged: bool = False,
+    enable_corruption: bool = False,
+) -> torch.Tensor:
+    """Return the task-owned actor or privileged critic observation.
+
+    Args:
+        env: Environment providing the task state and robot.
+        obs: Observation-manager input; the task builds its own observation tensors.
+        privileged: Return the critic observation when true, otherwise the actor observation.
+        enable_corruption: Request the task-defined observation noise.
+
+    Returns:
+        The selected actor or privileged critic tensor.
+    """
+    build = getattr(env, "build_velocity_locomotion_observations", None)
+    if not callable(build):
+        raise TypeError(
+            "velocity_locomotion_observation requires an environment that "
+            "implements build_velocity_locomotion_observations()."
+        )
+    actor, critic = build(enable_corruption=enable_corruption)
+    return critic if privileged else actor