From f40874bc90c72264dbf46989d90e15a9715c1393 Mon Sep 17 00:00:00 2001 From: "youyuan.chen" Date: Fri, 31 Jul 2026 13:45:20 +0800 Subject: [PATCH] fix link and generate --- .../robot_lab/robot_lab/tasks/go1/env_cfg.py | 12 ++++---- source/rsl_rl/rsl_rl/algorithms/moe_cts.py | 16 +++++++---- .../rsl_rl/storage/rollout_storage_cts.py | 28 ++++++++++++++++--- 3 files changed, 41 insertions(+), 15 deletions(-) diff --git a/source/robot_lab/robot_lab/tasks/go1/env_cfg.py b/source/robot_lab/robot_lab/tasks/go1/env_cfg.py index 7b67ecf..0dbf406 100644 --- a/source/robot_lab/robot_lab/tasks/go1/env_cfg.py +++ b/source/robot_lab/robot_lab/tasks/go1/env_cfg.py @@ -27,7 +27,7 @@ JOINT_NAMES = [ "RL_hip_joint", "RL_thigh_joint", "RL_calf_joint", ] -BASE_LINK_NAME = "trunk" +BASE_LINK_NAME = "base" FOOT_LINK_NAME = ".*_foot" BASE_HEIGHT_TARGET = 0.30 # Go1 nominal trunk height used by the original Gym task. @@ -63,7 +63,7 @@ class Go1SceneCfg(InteractiveSceneCfg): robot: ArticulationCfg = GO1_CFG_UNITREE.replace(prim_path="{ENV_REGEX_NS}/Robot") height_scanner = RayCasterCfg( - prim_path="{ENV_REGEX_NS}/Robot/trunk", + prim_path="{ENV_REGEX_NS}/Robot/base", offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)), ray_alignment="yaw", pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]), @@ -71,7 +71,7 @@ class Go1SceneCfg(InteractiveSceneCfg): mesh_prim_paths=["/World/ground"], ) height_scanner_small = RayCasterCfg( - prim_path="{ENV_REGEX_NS}/Robot/trunk", + prim_path="{ENV_REGEX_NS}/Robot/base", offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)), ray_alignment="yaw", pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[0.4, 0.3]), @@ -260,7 +260,7 @@ class EventCfg: func=mdp.randomize_rigid_body_mass, mode="startup", params={ - "asset_cfg": SceneEntityCfg("robot", body_names="^(?!.*trunk).*"), + "asset_cfg": SceneEntityCfg("robot", body_names="^(?!.*base).*"), "mass_distribution_params": (0.9, 1.1), "operation": "scale", "recompute_inertia": True, @@ -466,7 +466,7 @@ class Go1EnvCfg(ManagerBasedRLEnvCfg): """Merged configuration for the Go1 robot on rough terrain.""" # Scene settings - scene: Go1SceneCfg = Go1SceneCfg(num_envs=16384, env_spacing=0.5) + scene: Go1SceneCfg = Go1SceneCfg(num_envs=4096, env_spacing=0.5) # Basic settings observations: ObservationsCfg = ObservationsCfg() actions: ActionsCfg = ActionsCfg() @@ -489,7 +489,7 @@ class Go1EnvCfg(ManagerBasedRLEnvCfg): # Physics material settings from subclass self.sim.physics_material = self.scene.terrain.physics_material self.sim.physx.gpu_max_rigid_patch_count = int(1 * 1024 * 1024) # 1 million - self.sim.physx.gpu_collision_stack_size = int(512 * 1024 * 1024) # 128 MB + self.sim.physx.gpu_collision_stack_size = int(128 * 1024 * 1024) # 128 MB self.sim.physx.enable_external_forces_every_iteration = True # Update sensor periods diff --git a/source/rsl_rl/rsl_rl/algorithms/moe_cts.py b/source/rsl_rl/rsl_rl/algorithms/moe_cts.py index 7a42c06..8730d03 100644 --- a/source/rsl_rl/rsl_rl/algorithms/moe_cts.py +++ b/source/rsl_rl/rsl_rl/algorithms/moe_cts.py @@ -233,9 +233,12 @@ class MoECTS: # RND loss mean_rnd_loss = 0 if self.rnd else None - # Get mini batch generator - generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - data = list(generator) + # Reuse the exact same shuffled samples for both update phases without + # materializing every epoch and mini-batch on the GPU at once. + batch_indices = self.storage.generate_mini_batch_indices() + generator = self.storage.mini_batch_generator( + self.num_mini_batches, self.num_learning_epochs, batch_indices=batch_indices + ) # Iterate over batches teacher_samples = self.teacher_num_envs * self.storage.num_transitions_per_env // self.num_mini_batches @@ -251,7 +254,7 @@ class MoECTS: old_sigma_batch, hidden_states_batch, masks_batch, - ) in data: + ) in generator: original_batch_size = obs_batch.batch_size[0] # Check if we should normalize advantages per mini batch @@ -377,6 +380,9 @@ class MoECTS: if mean_rnd_loss is not None: mean_rnd_loss += rnd_loss.item() + generator = self.storage.mini_batch_generator( + self.num_mini_batches, self.num_learning_epochs, batch_indices=batch_indices + ) for ( obs_batch, actions_batch, @@ -388,7 +394,7 @@ class MoECTS: old_sigma_batch, hidden_states_batch, masks_batch, - ) in data: + ) in generator: # Student encoder loss obs_a_batch = self.policy.get_actor_obs(obs_batch) obs_a_batch = self.policy.actor_obs_normalizer(obs_a_batch) diff --git a/source/rsl_rl/rsl_rl/storage/rollout_storage_cts.py b/source/rsl_rl/rsl_rl/storage/rollout_storage_cts.py index 4efda08..e3c3162 100644 --- a/source/rsl_rl/rsl_rl/storage/rollout_storage_cts.py +++ b/source/rsl_rl/rsl_rl/storage/rollout_storage_cts.py @@ -127,7 +127,25 @@ class RolloutStorageCTS: yield self.observations[i], self.actions[i], self.privileged_actions[i], self.dones[i] # For reinforcement learning with feedforward networks - def mini_batch_generator(self, num_mini_batches: int, num_epochs: int = 8) -> Generator: + def generate_mini_batch_indices(self) -> tuple[torch.Tensor, torch.Tensor]: + """Generate the teacher and student permutations shared by both update phases.""" + if self.training_type != "rl": + raise ValueError("This function is only available for reinforcement learning training.") + + teacher_samples_num = self.teacher_num_envs * self.num_transitions_per_env + student_samples_num = self.student_num_envs * self.num_transitions_per_env + teacher_indices = torch.randperm(teacher_samples_num, requires_grad=False, device=self.device) + student_indices = teacher_samples_num + torch.randperm( + student_samples_num, requires_grad=False, device=self.device + ) + return teacher_indices, student_indices + + def mini_batch_generator( + self, + num_mini_batches: int, + num_epochs: int = 8, + batch_indices: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> Generator: if self.training_type != "rl": raise ValueError("This function is only available for reinforcement learning training.") @@ -136,8 +154,10 @@ class RolloutStorageCTS: student_samples_num = self.student_num_envs * self.num_transitions_per_env teacher_mini_batch_size = teacher_samples_num // num_mini_batches student_mini_batch_size = student_samples_num // num_mini_batches - teacher_indices = torch.randperm(teacher_samples_num, requires_grad=False, device=self.device) - student_indices = teacher_samples_num + torch.randperm(student_samples_num, requires_grad=False, device=self.device) + if batch_indices is None: + teacher_indices, student_indices = self.generate_mini_batch_indices() + else: + teacher_indices, student_indices = batch_indices # Core observations = self.observations.transpose(0, 1).flatten(0, 1) @@ -204,4 +224,4 @@ class RolloutStorageCTS: return NotImplementedError("CTS rollout storage does not support RNNs yet.") def _save_hidden_states(self, hidden_states: tuple[HiddenState, HiddenState]) -> None: - return NotImplementedError("CTS rollout storage does not support RNNs yet.") \ No newline at end of file + return NotImplementedError("CTS rollout storage does not support RNNs yet.")