From d725754139a0c9d5df1f3e30687b9104cb9dc4a4 Mon Sep 17 00:00:00 2001 From: Dixuan jiang Date: Fri, 13 Mar 2026 19:33:09 +0800 Subject: [PATCH] move push_function random_set to np_env --- scripts/train_rsl_rl.py | 44 ++++++++++++++++++++++++++ unilab/envs/backend/motrix_backend.py | 4 +-- unilab/envs/base.py | 4 +++ unilab/envs/locomotion/go1/base.py | 6 +--- unilab/envs/locomotion/go2/base.py | 5 +++ unilab/envs/locomotion/go2/joystick.py | 15 +++++++++ unilab/envs/np_env.py | 18 ++++++++--- 7 files changed, 84 insertions(+), 12 deletions(-) diff --git a/scripts/train_rsl_rl.py b/scripts/train_rsl_rl.py index 2b8719bee..ae80f4dbc 100644 --- a/scripts/train_rsl_rl.py +++ b/scripts/train_rsl_rl.py @@ -106,6 +106,7 @@ def step(self, actions): if hasattr(state, "info") and "log" in state.info: infos["log"] = state.info["log"] + obs_dict = TensorDict( {"policy": obs}, batch_size=self.num_envs, @@ -143,7 +144,50 @@ def get_privileged_observations(self): obs = to_torch(self.env.state.obs, self.device) return obs +def RslRlAacVecEnvWrapper(RslRlVecEnvWrapper): #Asymmetric Actor-Critic + def step(self, actions): + # Convert actions to numpy (CPU) + if isinstance(actions, torch.Tensor): + actions_np = actions.detach().cpu().numpy() + else: + actions_np = actions + + # Step the environment + state = self.env.step(actions_np) + + # Convert output to torch tensors on target device + obs = to_torch(state.obs, self.device) + rewards = to_torch(state.reward, self.device) + dones = to_torch(state.done, self.device).bool() + + # Update logging info + self.episode_returns += rewards + self.episode_lengths += 1 + + infos = {} + # Check for dones + done_indices = torch.nonzero(dones).flatten() + if len(done_indices) > 0: + # Handle limits and timeouts (RSL-RL expects 'time_outs' in extras/infos) + if hasattr(state, "truncated"): + infos["time_outs"] = to_torch(state.truncated, self.device).bool() + + # Reset buffers for done envs + self.episode_returns[done_indices] = 0 + self.episode_lengths[done_indices] = 0 + # Pass per-step logs if available (gs_playground style) + # prioritizing 'log' over 'episode' allows per-step metric logging + if hasattr(state, "info") and "log" in state.info: + infos["log"] = state.info["log"] + + obs_dict = TensorDict( + {"policy": obs}, + batch_size=self.num_envs, + device=self.device + ) + + return obs_dict, rewards, dones, infos def play_rsl_rl(args, cfg, device): diff --git a/unilab/envs/backend/motrix_backend.py b/unilab/envs/backend/motrix_backend.py index e42c9304e..c83d0705c 100644 --- a/unilab/envs/backend/motrix_backend.py +++ b/unilab/envs/backend/motrix_backend.py @@ -49,9 +49,7 @@ def push_robots(self, force_range): ex_force[:, 0] *= force_range[0] ex_force[:, 1] *= force_range[1] ex_force[:, 2] *= force_range[2] - self._body_link.add_external_force(self._data, ex_force, local=True) - print(ex_force) - + self._body_link.add_external_force(self._data, ex_force, local=True) def step(self, ctrl: np.ndarray, nsteps: int = 1) -> None: self._data.actuator_ctrls = np.ascontiguousarray(ctrl) diff --git a/unilab/envs/base.py b/unilab/envs/base.py index d6a52ac5f..abd207542 100644 --- a/unilab/envs/base.py +++ b/unilab/envs/base.py @@ -41,6 +41,10 @@ def validate(self): if self.sim_dt > self.ctrl_dt: raise ValueError("sim_dt must be less than or equal to ctrl_dt") +@dataclass +class obs_cfg: + obs_dict = {'vel': 3, 'gyro': 3, 'gravity': 3, 'diff': 12, + 'dof_vel': 12, 'action': 12, 'cmd': 3} # 'obs_name': dim class ABEnv(abc.ABC): @property diff --git a/unilab/envs/locomotion/go1/base.py b/unilab/envs/locomotion/go1/base.py index 2de5acc16..205e70081 100644 --- a/unilab/envs/locomotion/go1/base.py +++ b/unilab/envs/locomotion/go1/base.py @@ -78,11 +78,7 @@ def __init__(self, cfg: Go1BaseCfg, backend: SimBackend, num_envs=1): self._init_action_space() self._num_action = self._action_space.shape[0] self._init_buffers() - self.push_robots_flag = False - if self._backend.backend_type == 'motrix': - self._backend._process_rigid_body_props(cfg) - if self._cfg.domain_rand.push_robots == True: - self.push_robots_flag = True + def _init_action_space(self): model = self._backend.model diff --git a/unilab/envs/locomotion/go2/base.py b/unilab/envs/locomotion/go2/base.py index e6d1c7303..11e1a4e06 100644 --- a/unilab/envs/locomotion/go2/base.py +++ b/unilab/envs/locomotion/go2/base.py @@ -64,6 +64,11 @@ def __init__(self, cfg: Go2BaseCfg, backend: SimBackend, num_envs=1): self._init_action_space() self._num_action = self._action_space.shape[0] self._init_buffers() + self.push_robots_flag = False + if self._backend.backend_type == 'motrix': + self._backend._process_rigid_body_props(cfg) + if self._cfg.domain_rand.push_robots == True: + self.push_robots_flag = True def _init_action_space(self): model = self._backend.model diff --git a/unilab/envs/locomotion/go2/joystick.py b/unilab/envs/locomotion/go2/joystick.py index ec6ba561f..445b18aad 100644 --- a/unilab/envs/locomotion/go2/joystick.py +++ b/unilab/envs/locomotion/go2/joystick.py @@ -25,6 +25,19 @@ class Commands: [0.5, 0.0, 0.0], ] +@dataclass +class Domain_Rand: + # randomize_friction = True + # friction_range = [0.5, 1.25] + randomize_base_mass = True + added_mass_range = [-1.5, 1.5] + + random_com = True + com_offset_x = [-0.05, 0.05] + + push_robots = True + push_interval = 750 #step + max_force = [1, 1, 0.5] @dataclass class RewardConfig: @@ -56,9 +69,11 @@ class Go2JoystickCfg(Go2BaseCfg): init_state: InitState = field(default_factory=InitState) commands: Commands = field(default_factory=Commands) reward_config: RewardConfig = field(default_factory=RewardConfig) + domain_rand: Domain_Rand = field(default_factory=Domain_Rand) @registry.env("Go2JoystickFlatTerrain", sim_backend="mujoco") +@registry.env("Go2JoystickFlatTerrain", sim_backend="motrix") class Go2WalkTask(Go2BaseEnv): def __init__(self, cfg: Go2JoystickCfg, num_envs=1, backend_type="mujoco"): backend = create_backend(backend_type, cfg.model_file, num_envs, cfg.sim_dt, body_name=cfg.asset.body_name) diff --git a/unilab/envs/np_env.py b/unilab/envs/np_env.py index 859cac9bf..24bf7a21b 100644 --- a/unilab/envs/np_env.py +++ b/unilab/envs/np_env.py @@ -4,6 +4,7 @@ from typing import Tuple import numpy as np import gymnasium as gym +from typing import Optional from unilab.envs.base import ABEnv, EnvCfg from unilab.envs.backend import SimBackend @@ -17,6 +18,8 @@ class NpEnvState: terminated: np.ndarray truncated: np.ndarray info: dict + critic_obs: Optional[np.ndarray] = None + @property def done(self) -> np.ndarray: @@ -35,6 +38,11 @@ def __init__(self, cfg: EnvCfg, backend: SimBackend, num_envs: int): self._num_envs = num_envs self._state = None self.step_counter = 0 + self.push_robots_flag = False + if self._backend.backend_type == 'motrix': + self._backend._process_rigid_body_props(cfg) + if self._cfg.domain_rand.push_robots == True: + self.push_robots_flag = True @property def cfg(self) -> EnvCfg: @@ -55,7 +63,7 @@ def init_state(self) -> NpEnvState: terminated = np.ones((self._num_envs,), dtype=bool) truncated = np.zeros((self._num_envs,), dtype=bool) info = {"steps": np.zeros((self._num_envs,), dtype=np.uint32)} - + self._state = NpEnvState(obs, reward, terminated, truncated, info) self._reset_done_envs() return self._state @@ -128,14 +136,16 @@ def _reset_done_envs(self): self._state.info[key] = value elif isinstance(value, np.ndarray): self._state.info[key][env_indices] = value + + def push_robots(self): + if self.push_robots_flag == True: + if self.step_counter % self._cfg.domain_rand.push_interval == 0: + self._backend.push_robots(self._cfg.domain_rand.max_force) @abc.abstractmethod def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: """子类实现：action → ctrl""" - @abc.abstractmethod - def push_robots(self) -> None: - """子类实现""" @abc.abstractmethod def update_state(self, state: NpEnvState) -> NpEnvState: