From 93fa055f79ec67dae10a9fd39aeaeb5bc1879654 Mon Sep 17 00:00:00 2001 From: wertyuilife Date: Wed, 8 Apr 2026 22:43:50 +0800 Subject: [PATCH] add env without action delay (up to test actuator level action delay). --- .../robot_lab/robot_lab/tasks/go2/__init__.py | 2 +- .../robot_lab/tasks/go2/env/go2_env.py | 20 ++++++-- .../tasks/go2/manager/action_manager.py | 48 ++++++++++++++++++- 3 files changed, 63 insertions(+), 7 deletions(-) diff --git a/source/robot_lab/robot_lab/tasks/go2/__init__.py b/source/robot_lab/robot_lab/tasks/go2/__init__.py index 52f6333..99335d7 100644 --- a/source/robot_lab/robot_lab/tasks/go2/__init__.py +++ b/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={ diff --git a/source/robot_lab/robot_lab/tasks/go2/env/go2_env.py b/source/robot_lab/robot_lab/tasks/go2/env/go2_env.py index 3317f9a..db07a06 100644 --- a/source/robot_lab/robot_lab/tasks/go2/env/go2_env.py +++ b/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 + + + diff --git a/source/robot_lab/robot_lab/tasks/go2/manager/action_manager.py b/source/robot_lab/robot_lab/tasks/go2/manager/action_manager.py index c5e5896..667fd9a 100644 --- a/source/robot_lab/robot_lab/tasks/go2/manager/action_manager.py +++ b/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)