add env without action delay (up to test actuator level action delay).

source/robot_lab/robot_lab/tasks/go2/__init__.py

@@ -19,7 +19,7 @@ from isaaclab_tasks.utils import import_packages
 ##
 gym.register(
     id="RobotLab-Go2-v0",
-    # entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    # entry_point="robot_lab.tasks.go2.env.go2_env:Go2Env",
     entry_point="robot_lab.tasks.go2.env.go2_env:ActionDelayGo2Env",
     disable_env_checker=True,
     kwargs={

source/robot_lab/robot_lab/tasks/go2/env/go2_env.py

@@ -1,9 +1,16 @@
 from isaaclab.envs import ManagerBasedRLEnv, ManagerBasedRLEnvCfg, VecEnvStepReturn
-from robot_lab.tasks.go2.manager.action_manager import ActionManagerWithDelay
+from robot_lab.tasks.go2.manager.action_manager import ActionManagerGo2, ActionManagerGo2WithDelay
 
 import torch
 
-from isaaclab.ui.widgets import ManagerLiveVisualizer
+class Go2Env(ManagerBasedRLEnv):
+    cfg: ManagerBasedRLEnvCfg
+    
+    def load_managers(self):
+        super().load_managers()
+        # override action manager
+        self.action_manager = ActionManagerGo2(self.cfg.actions, self)
+        print("[Go2Env-INFO] Overriding action manager with ActionManagerGo2: ", self.action_manager)
 
 
 class ActionDelayGo2Env(ManagerBasedRLEnv):
@@ -21,7 +28,7 @@ class ActionDelayGo2Env(ManagerBasedRLEnv):
         # Call the parent class initializer
         super().__init__(cfg=cfg, render_mode=render_mode, **kwargs)
         print(
-            "[WARNING] You are using ActionDelayGo2Env; "
+            "[ActionDelayGo2Env-WARNING] You are using ActionDelayGo2Env; "
             "make sure all ActionTerms support multiple calls to process_actions() "
             "within a single step()."
         )
@@ -29,8 +36,8 @@ class ActionDelayGo2Env(ManagerBasedRLEnv):
     def load_managers(self):
         super().load_managers()
         # override action manager
-        self.action_manager = ActionManagerWithDelay(self.cfg.actions, self)
-        print("[INFO] Overriding action manager with ActionManagerWithDelay: ", self.action_manager)            
+        self.action_manager = ActionManagerGo2WithDelay(self.cfg.actions, self)
+        print("[ActionDelayGo2Env-INFO] Overriding action manager with ActionManagerGo2WithDelay: ", self.action_manager)            
         
     def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         """Execute one time-step of the environment's dynamics and reset terminated environments.
@@ -126,3 +133,6 @@ class ActionDelayGo2Env(ManagerBasedRLEnv):
 
         # return observations, rewards, resets and extras
         return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, self.extras
+
+
+

source/robot_lab/robot_lab/tasks/go2/manager/action_manager.py

@@ -2,7 +2,53 @@ from isaaclab.managers import ActionManager
 import torch
 from collections.abc import Sequence
 
-class ActionManagerWithDelay(ActionManager):
+# ActionManagerGo2 is a simple custom ActionManager that 
+# maintain _prev_prev_action for action smoothness reward computation.
+class ActionManagerGo2(ActionManager):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._prev_prev_action = torch.zeros_like(self._action)
+    
+    def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]:
+        super().reset(env_ids)
+        if env_ids is None:
+            self._prev_prev_action.zero_()
+        else:
+            self._prev_prev_action[env_ids] = 0.0
+        return {}
+    
+    def process_action(self, action: torch.Tensor):
+        """Processes the actions sent to the environment.
+
+        Note:
+            This function should be called once per environment step.
+
+        Args:
+            action: The actions to process.
+        """
+        # check if action dimension is valid
+        if self.total_action_dim != action.shape[1]:
+            raise ValueError(f"Invalid action shape, expected: {self.total_action_dim}, received: {action.shape[1]}.")
+        # store the input actions
+        self._prev_prev_action[:] = self._prev_action
+        self._prev_action[:] = self._action
+        self._action[:] = action.to(self.device)
+
+        # split the actions and apply to each tensor
+        idx = 0
+        for term in self._terms.values():
+            term_actions = action[:, idx : idx + term.action_dim]
+            term.process_actions(term_actions)
+            idx += term.action_dim
+    
+    @property
+    def prev_prev_action(self):
+        return self._prev_prev_action
+    
+# ActionManagerGo2WithDelay is a custom ActionManager that 
+# maintain _prev_prev_action for action smoothness reward computation.
+# and also do random action delay by process_action_with_delay() function.
+class ActionManagerGo2WithDelay(ActionManager):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._prev_prev_action = torch.zeros_like(self._action)