Add RL-friendly velocity and effort action commands

docs/source/overview/gym/action_functors.md

@@ -16,6 +16,51 @@ This page lists all available action terms that can be used with the Action Mana
 **Using an AI coding agent?** Use the **`/add-functor`** skill to scaffold a new action term with the correct class structure, `ActionTermCfg` registration, and module placement in `actions.py`.
 ````
 
+## Policy and Command Contract
+
+The Action Manager exposes one flat {class}`gymnasium.spaces.Box` to the
+policy. Each ``pre`` term owns a contiguous slice in configuration order. The
+manager processes those slices into typed ``qpos``, ``qvel``, or ``qf``
+commands and applies each command to that term's selected joints.
+
+By default, every term's policy range is ``[-1, 1]``. Use ``scale`` to map that
+normalized value to a useful physical magnitude, or use
+{class}`~actions.QposDenormalizedTerm` to map the full normalized range to
+joint position limits. Physical commands are clipped to the robot's qpos,
+qvel, or qf limits by default.
+
+The following parameters are common to pre-processing terms:
+
+- ``joint_ids``: Static active-joint indices controlled by the term.
+- ``control_part``: Named robot control part; use this instead of
+  ``joint_ids``.
+- ``action_range``: Two finite policy-space bounds. The default is
+  ``[-1, 1]``. ``QposDenormalizedTerm`` uses its existing ``range`` parameter.
+- ``clip``: Clip the processed physical command to robot limits. Defaults to
+  ``true``.
+- ``allow_overlap``: Permit two terms to address the same joint only when it
+  is explicitly ``true`` on both terms. Disjoint groups are recommended
+  because overlapping position, velocity, and effort semantics are otherwise
+  ambiguous.
+
+Flat tensors are the standard RL interface. A mapping may also address terms
+by their configuration names, which is useful for scripted controllers:
+
+```python
+env.step({
+    "arm_velocity": arm_velocity_action,
+    "gripper_effort": gripper_effort_action,
+})
+```
+
+Without an Action Manager, ``env.step(tensor)`` remains a qpos command for
+backward compatibility. Direct typed commands are also accepted:
+
+```python
+env.step({"qvel": target_velocity})
+env.step({"qf": target_effort})
+```
+
 ## Joint Position Control
 
 ```{list-table} Joint Position Action Terms
@@ -46,7 +91,7 @@ This page lists all available action terms that can be used with the Action Mana
   - Normalize action from qpos limits -> [range[0], range[1]]. Maps joint positions to a normalized range based on joint limits. Typically used for post-processing action outputs.
 
     ```json
-    {"func": "QposNormalizedTerm", "params": {"range": [0.0, 1.0]}}
+    {"func": "QposNormalizedTerm", "mode": "post", "params": {"range": [0.0, 1.0]}}
     ```
 ```
 
@@ -75,13 +120,13 @@ This page lists all available action terms that can be used with the Action Mana
 * - Action Term
   - Description
 * - {class}`~actions.QvelTerm`
-  - Joint velocity action: scale * action -> qvel. The policy outputs target joint velocities.
+  - Joint velocity action: scale * action -> qvel. The policy outputs target joint velocities. Configure zero position stiffness on these joints when a position drive should not oppose the velocity target.
 
     ```json
     {"func": "QvelTerm", "params": {"scale": 1.0}}
     ```
 * - {class}`~actions.QfTerm`
-  - Joint force/torque action: scale * action -> qf. The policy outputs target joint torques/forces.
+  - Joint force/torque action: scale * action -> qf. The policy outputs target joint torques/forces. The command is reapplied before every physics substep so it is held across control decimation.
 
     ```json
     {"func": "QfTerm", "params": {"scale": 1.0}}
@@ -117,6 +162,7 @@ actions = {
 actions = {
     "normalize_qpos": ActionTermCfg(
         func="QposNormalizedTerm",
+        mode="post",
         params={
             "range": [0.0, 1.0],  # Normalize to [0, 1] range
         },
@@ -133,11 +179,46 @@ actions = {
         },
     ),
 }
+
+# Example: one flat RL action controlling disjoint joint groups
+actions = {
+    "arm_velocity": ActionTermCfg(
+        func="QvelTerm",
+        params={
+            "joint_ids": [0, 1, 2, 3, 4, 5],
+            "scale": 1.5,
+        },
+    ),
+    "gripper_effort": ActionTermCfg(
+        func="QfTerm",
+        params={
+            "control_part": "gripper",
+            "scale": 20.0,
+        },
+    ),
+}
 ```
 
+For the mixed example, the policy outputs
+``[arm_velocity..., gripper_effort...]``. The environment retains that exact
+flat action in RL and trajectory buffers while routing the processed commands
+to ``set_qvel`` and ``set_qf`` respectively.
+
+````{attention}
+Velocity and effort commands do not automatically change the robot's drive
+configuration. A qvel term on joints with non-zero stiffness may fight the
+position drive. A qf term with non-zero stiffness or damping is additive to
+the active drive rather than pure torque control. The Action Manager emits a
+warning for these combinations; configure the robot drive properties to match
+the intended control mode.
+````
+
 ## Action Term Properties
 
 All action terms provide the following properties:
 
 - ``action_dim``: The dimension of the action space (number of values the policy should output)
+- ``action_space``: The per-term policy-space bounds
+- ``joint_ids``: The resolved robot joints controlled by the term
+- ``command_key``: The physical output type (``qpos``, ``qvel``, or ``qf``)
 - ``process_action(action)``: Method to convert raw policy output to robot control format

docs/source/overview/gym/env.md

@@ -291,6 +291,8 @@ The dataset manager is called automatically during {meth}`~envs.Env.step()`, ens
 For RL tasks, EmbodiChain uses the **Action Manager** integrated into {class}`~envs.EmbodiedEnv`:
 
 * **Action Preprocessing**: Configurable via ``actions`` in {class}`~envs.EmbodiedEnvCfg`. Supports DeltaQposTerm, QposTerm, QposDenormalizedTerm, EefPoseTerm, QvelTerm, QfTerm. For a complete list of available action terms, please refer to {doc}`action_functors`.
+* **Flat RL Interface**: The Action Manager concatenates all ``pre`` terms into one flat ``Box`` policy action space, then routes the slices to typed qpos, qvel, or qf commands on their selected joints.
+* **Command Safety**: Processed commands are checked for batch shape and finite values and are clipped to robot limits by default. Effort commands are held across every physics substep.
 * **Standardized Info Structure**: {class}`~envs.EmbodiedEnv` provides ``compute_task_state``, ``get_info``, and ``evaluate`` for task-specific success/failure and metrics.
 * **Episode Management**: Configurable episode length and truncation logic.
 
@@ -325,6 +327,29 @@ In a gym config file, use the ``actions`` section:
 }
 ```
 
+Multiple terms may control disjoint joint groups. Their dimensions are
+concatenated in configuration order, so standard continuous-control policies
+still produce one tensor:
+
+```json
+"actions": {
+    "arm_velocity": {
+        "func": "QvelTerm",
+        "params": {"joint_ids": [0, 1, 2, 3, 4, 5], "scale": 1.5}
+    },
+    "gripper_effort": {
+        "func": "QfTerm",
+        "params": {"control_part": "gripper", "scale": 20.0}
+    }
+}
+```
+
+When no Action Manager is configured, a bare tensor passed to ``step`` remains
+a qpos command. A direct mapping such as ``{"qvel": value}`` or
+``{"qf": value}`` selects another physical command explicitly. For RL and
+trajectory recording, prefer the Action Manager because its flat action space
+and raw-action ordering are stable.
+
 
 ## Creating a Custom Task
 

docs/source/overview/rl/models.md

@@ -15,10 +15,12 @@ This module contains RL policy networks and related model implementations, suppo
 
 ### ActorCritic
 - Typical actor-critic policy, includes actor (action distribution) and critic (value function). Used with PPO.
+- Supports ``squash_actions`` for a tanh-bounded Gaussian. The sampled and reevaluated log probabilities include the tanh Jacobian correction required by PPO.
 
 ### ActorOnly
 - Actor-only policy without Critic. Used with GRPO (Group Relative Policy Optimization), which estimates advantages via group-level return comparison instead of a value function.
 - Supports Gaussian action distributions, learnable log_std, suitable for continuous action spaces.
+- Supports the same corrected ``squash_actions`` path as ``ActorCritic``.
 - Key methods:
     - `forward`: Actor network outputs mean, samples action, and writes policy outputs into a `TensorDict`.
     - `evaluate_actions`: Used for loss calculation in PPO/GRPO algorithms.
@@ -52,6 +54,12 @@ log_prob = step_td["sample_log_prob"]
 value = step_td["value"]
 ```
 
+For simulator RL, training enables ``squash_actions`` automatically when an
+Action Manager exposes the standard ``[-1, 1]`` action range. Set it explicitly
+to ``false`` to retain an unbounded Gaussian. Custom Action Manager bounds do
+not enable automatic tanh squashing; use a policy distribution whose support
+matches those bounds.
+
 ## Extension and Customization
 - Supports custom network architectures (e.g., CNN, Transformer) by implementing the Policy interface.
 - Can extend to multi-head policies, distributional actors, hybrid action spaces, etc.

docs/source/tutorial/rl.rst

@@ -140,6 +140,7 @@ The ``policy`` section defines the neural network policy:
 
 - **name**: Policy name (e.g., "actor_critic", "vla")
 - **action_dim**: Optional policy output action dimension. If omitted, it is inferred from ``env.action_space``.
+- **squash_actions**: Optional tanh bounding for built-in Gaussian policies. Simulator training enables it by default when the Action Manager uses ``[-1, 1]`` bounds; log probabilities include the tanh Jacobian correction.
 - **actor**: Actor network configuration (required for actor_critic)
 - **critic**: Critic network configuration (required for actor_critic)
 

embodichain/lab/gym/envs/base_env.py

@@ -20,7 +20,7 @@
 import numpy as np
 import gymnasium as gym
 
-from typing import Dict, List, Union, Tuple, Any, Sequence
+from typing import Any, Callable, Dict, List, Sequence, Tuple, Union
 from functools import cached_property
 from tensordict import TensorDict
 
@@ -589,6 +589,31 @@ def _postprocess_action(self, action: EnvAction) -> EnvAction:
         """
         return action
 
+    def _before_sim_step(self, substep_index: int) -> None:
+        """Hook invoked immediately before every physics substep.
+
+        Args:
+            substep_index: Zero-based substep within the current environment
+                control step.
+
+        .. tip::
+            Override this hook for commands, such as generalized efforts, that
+            must be reapplied throughout action decimation.
+        """
+        del substep_index
+
+    def _get_before_sim_step_callback(self) -> Callable[[int], None] | None:
+        """Return an optional callback for physics-substep control updates.
+
+        Returns:
+            The overridden :meth:`_before_sim_step` hook, or ``None`` when the
+            hook is unchanged so ordinary environments do not pay a Python
+            callback cost during action decimation.
+        """
+        if type(self)._before_sim_step is BaseEnv._before_sim_step:
+            return None
+        return self._before_sim_step
+
     def _step_action(self, action: EnvAction) -> EnvAction:
         """Set action control command into simulation.
 
@@ -666,7 +691,11 @@ def step(
                 action = self._step_action(action=action)
 
             with self._profiler.section("sim_update"):
-                self.sim.update(self.sim_cfg.physics_dt, self.cfg.sim_steps_per_control)
+                self.sim.update(
+                    self.sim_cfg.physics_dt,
+                    self.cfg.sim_steps_per_control,
+                    before_step_callback=self._get_before_sim_step_callback(),
+                )
             with self._profiler.section("update_sim_state"):
                 self._update_sim_state(**kwargs)