fix link and generate

2026-07-31 13:45:20 +08:00
parent 0d2afbf2bb
commit f40874bc90
3 changed files with 41 additions and 15 deletions
--- 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
--- 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)
--- 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.")
+        return NotImplementedError("CTS rollout storage does not support RNNs yet.")