Support pre and post mode for action functor

AGENTS.md

@@ -5,7 +5,6 @@
 **IMPORTANT**: The Python package name is `embodichain` (all lowercase, one word).
 - Repository folder: `EmbodiChain` (PascalCase)
 - Python package: `embodichain` (lowercase)
-- NEVER use: `embodiedchain`, `embodyichain`, or any other variant
 
 ## Project Structure
 

embodichain/lab/gym/envs/base_env.py

@@ -550,6 +550,21 @@ def _preprocess_action(self, action: EnvAction) -> EnvAction:
         """
         return action
 
+    def _postprocess_action(self, action: EnvAction) -> EnvAction:
+        """Postprocess action after applying to robot.
+
+        Post processing is usually used to modify the action after it has been applied to the robot,
+        performing normalization, noise addition, or any other modifications that need to be applied
+        for policy learning or evaluation purposes.
+
+        Args:
+            action: Action after preprocessing and robot control command generation
+
+        Returns:
+            Final action to be applied in the simulation
+        """
+        return action
+
     def _step_action(self, action: EnvAction) -> EnvAction:
         """Set action control command into simulation.
 
@@ -608,8 +623,10 @@ def step(
         Returns:
             A tuple contraining the observation, reward, terminated, truncated, and info dictionary.
         """
+
         action = self._preprocess_action(action=action)
         action = self._step_action(action=action)
+
         self.sim.update(self.sim_cfg.physics_dt, self.cfg.sim_steps_per_control)
         self._update_sim_state(**kwargs)
 
@@ -620,6 +637,9 @@ def step(
             rewards=rewards, obs=obs, action=action, info=info
         )
 
+        # Apply postprocessing to the action after all computations are done.
+        action = self._postprocess_action(action=action)
+
         self._elapsed_steps += 1
 
         terminateds = torch.logical_or(

embodichain/lab/gym/envs/embodied_env.py

@@ -707,6 +707,12 @@ def _preprocess_action(self, action: EnvAction) -> EnvAction:
             return self.action_manager.process_action(action)
         return super()._preprocess_action(action)
 
+    def _postprocess_action(self, action):
+        if self.action_manager is not None:
+            pass
+            # return self.action_manager.postprocess_action(action)
+        return super()._postprocess_action(action)
+
     def _setup_robot(self, **kwargs) -> Robot:
         """Setup the robot in the environment.
 

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

@@ -27,15 +27,6 @@
 from .event_manager import EventManager
 from .observation_manager import ObservationManager
 from .reward_manager import RewardManager
-from .action_manager import (
-    ActionManager,
-    ActionTerm,
-    DeltaQposTerm,
-    QposTerm,
-    QposNormalizedTerm,
-    EefPoseTerm,
-    QvelTerm,
-    QfTerm,
-)
+from .action_manager import *
+from .actions import *
 from .dataset_manager import DatasetManager
-from .datasets import *