move push_function random_set to np_env

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):

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)

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

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

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

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)

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: