Add mcp-cts

UPDATE.md

@@ -1,6 +1,7 @@
 # 20251230
 ## v0.1.1
 1. 给cts算法加入robogauge异步评估
+2. 加入MCP-CTS
 Fix Bug: 修复MoE中专家使用了共享权重的问题, 换成Conv1D
 # 20251221
 1. 修改最大地形速度限制, y在所有地形上最大为1.0, z只有平地最大为2.0, x最大为2.0

deploy/deploy_mujoco/deploy_go2_lab.py

@@ -1,126 +0,0 @@
-import time
-
-import mujoco.viewer
-import mujoco
-import numpy as np
-from legged_gym import LEGGED_GYM_ROOT_DIR
-import torch
-import yaml
-
-
-def get_gravity_orientation(quaternion):
-    qw = quaternion[0]
-    qx = quaternion[1]
-    qy = quaternion[2]
-    qz = quaternion[3]
-
-    gravity_orientation = np.zeros(3)
-
-    gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
-    gravity_orientation[1] = -2 * (qz * qy + qw * qx)
-    gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
-
-    return gravity_orientation
-
-
-def pd_control(target_q, q, kp, target_dq, dq, kd):
-    """Calculates torques from position commands"""
-    return (target_q - q) * kp + (target_dq - dq) * kd
-
-
-if __name__ == "__main__":
-    # get config file name from command line
-    import argparse
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("config_file", type=str, help="config file name in the config folder")
-    args = parser.parse_args()
-    config_file = args.config_file
-    with open(f"{LEGGED_GYM_ROOT_DIR}/deploy/deploy_mujoco/configs/{config_file}", "r") as f:
-        config = yaml.load(f, Loader=yaml.FullLoader)
-        policy_path = config["policy_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
-        xml_path = config["xml_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
-
-        simulation_duration = config["simulation_duration"]
-        simulation_dt = config["simulation_dt"]
-        control_decimation = config["control_decimation"]
-
-        kps = np.array(config["kps"], dtype=np.float32)
-        kds = np.array(config["kds"], dtype=np.float32)
-
-        default_angles = np.array(config["default_angles"], dtype=np.float32)
-
-        joint_ids_map = config["joint_ids_map"]
-        ang_vel_scale = config["ang_vel_scale"]
-        dof_pos_scale = config["dof_pos_scale"]
-        dof_vel_scale = config["dof_vel_scale"]
-        action_scale = config["action_scale"]
-        cmd_scale = np.array(config["cmd_scale"], dtype=np.float32)
-
-        num_actions = config["num_actions"]
-        num_obs = config["num_obs"]
-
-        cmd = np.array(config["cmd_init"], dtype=np.float32)
-
-    # define context variables
-    action = np.zeros(num_actions, dtype=np.float32)
-    target_dof_pos = default_angles.copy()
-    obs = np.zeros(num_obs, dtype=np.float32)
-
-    counter = 0
-
-    # Load robot model
-    m = mujoco.MjModel.from_xml_path(xml_path)
-    d = mujoco.MjData(m)
-    m.opt.timestep = simulation_dt
-
-    # load policy
-    policy = torch.jit.load(policy_path)
-
-    with mujoco.viewer.launch_passive(m, d) as viewer:
-        # Close the viewer automatically after simulation_duration wall-seconds.
-        start = time.time()
-        while viewer.is_running() and time.time() - start < simulation_duration:
-            step_start = time.time()
-            temp = target_dof_pos[[0,4,8,1,5,9,2,6,10,3,7,11]]
-            tau = pd_control(temp, d.qpos[7:], kps, np.zeros_like(kds), d.qvel[6:], kds)
-            d.ctrl[:] = tau
-            # mj_step can be replaced with code that also evaluates
-            # a policy and applies a control signal before stepping the physics.
-            mujoco.mj_step(m, d)
-
-            counter += 1
-            if counter % control_decimation == 0:
-                # Apply control signal here.
-
-                # create observation
-                qj = d.qpos[7:]
-                dqj = d.qvel[6:]
-                quat = d.qpos[3:7]
-                ang_vel = d.qvel[3:6]
-
-                qj = (qj - default_angles) * dof_pos_scale
-
-                dqj = dqj * dof_vel_scale
-                gravity_orientation = get_gravity_orientation(quat)
-                ang_vel = ang_vel * ang_vel_scale
-
-                obs[:3] = ang_vel
-                obs[3:6] = gravity_orientation
-                obs[6:9] = cmd * cmd_scale
-                obs[9 : 9 + num_actions] = qj[joint_ids_map]
-                obs[9 + num_actions : 9 + 2 * num_actions] = dqj[joint_ids_map]
-                obs[9 + 2 * num_actions : 9 + 3 * num_actions] = action
-                obs_tensor = torch.from_numpy(obs).unsqueeze(0)
-                # policy inference
-                action = policy(obs_tensor).detach().numpy().squeeze()
-                # transform action to target_dof_pos
-                target_dof_pos = action * action_scale + default_angles[joint_ids_map]
-
-            # Pick up changes to the physics state, apply perturbations, update options from GUI.
-            viewer.sync()
-
-            # Rudimentary time keeping, will drift relative to wall clock.
-            time_until_next_step = m.opt.timestep - (time.time() - step_start)
-            if time_until_next_step > 0:
-                time.sleep(time_until_next_step)

legged_gym/envs/__init__.py

@@ -1,7 +1,7 @@
 from legged_gym import LEGGED_GYM_ROOT_DIR, LEGGED_GYM_ENVS_DIR
 
 from legged_gym.envs.go2.go2_env import Go2Robot
-from legged_gym.envs.go2.go2_config import GO2Cfg, GO2CfgPPO, GO2CfgCTS, GO2CfgMoECTS
+from legged_gym.envs.go2.go2_config import GO2Cfg, GO2CfgPPO, GO2CfgCTS, GO2CfgMoECTS, GO2CfgMCPCTS
 from .base.legged_robot import LeggedRobot
 
 from legged_gym.utils.task_registry import task_registry
@@ -9,3 +9,4 @@ from legged_gym.utils.task_registry import task_registry
 task_registry.register("go2", Go2Robot, GO2Cfg(), GO2CfgPPO())
 task_registry.register("go2_cts", Go2Robot, GO2Cfg(), GO2CfgCTS())
 task_registry.register("go2_moe_cts", Go2Robot, GO2Cfg(), GO2CfgMoECTS())
+task_registry.register("go2_mcp_cts", Go2Robot, GO2Cfg(), GO2CfgMCPCTS())

legged_gym/envs/base/legged_robot_config.py

@@ -359,3 +359,12 @@ class LeggedRobotCfgMoECTS(LeggedRobotCfgCTS):
     class runner(LeggedRobotCfgCTS.runner):
         policy_class_name = 'ActorCriticMoECTS'
         algorithm_class_name = 'MoECTS'
+
+class LeggedRobotCfgMCPCTS(LeggedRobotCfgCTS):
+    class policy(LeggedRobotCfgCTS.policy):
+        obs_no_goal_mask = None # mask for observation without goal inputs
+        student_expert_num = 8 # number of experts in the student model
+
+    class runner(LeggedRobotCfgCTS.runner):
+        policy_class_name = 'ActorCriticMCPCTS'
+        algorithm_class_name = 'MCPCTS'

legged_gym/envs/go2/go2_config.py

@@ -1,5 +1,5 @@
 import math
-from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMoECTS
+from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMCPCTS
 
 class GO2Cfg(LeggedRobotCfg):
     class init_state(LeggedRobotCfg.init_state):
@@ -264,3 +264,14 @@ class GO2CfgMoECTS(LeggedRobotCfgMoECTS):
         experiment_name = 'go2_moe_cts'
         max_iterations = 150000
         save_interval = 500
+
+class GO2CfgMCPCTS(LeggedRobotCfgMCPCTS):
+    class policy(LeggedRobotCfgMCPCTS.policy):
+        obs_no_goal_mask = [True] * 6 + [False] * 3 + [True] * 36  # mask for obs without command info
+        student_expert_num = 8 # number of experts in the student model
+
+    class runner(LeggedRobotCfgMCPCTS.runner):
+        run_name = ''
+        experiment_name = 'go2_mcp_cts'
+        max_iterations = 150000
+        save_interval = 500

legged_gym/utils/exporter.py

@@ -81,7 +81,11 @@ class _TorchPolicyExporter(torch.nn.Module):
             self.history_length = policy.history.shape[1]
             self.history = torch.zeros([1, policy.history.shape[1], policy.history.shape[2]], device='cpu')
             self.forward = self.forward_moe_cts
-        if hasattr(policy, "actor"):
+        if hasattr(policy, "actor_mcp"):
+            self.actor = copy.deepcopy(policy.actor_mcp)
+            self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
+            self.forward = self.forward_mcp_cts
+        elif hasattr(policy, "actor"):
             self.actor = copy.deepcopy(policy.actor)
             if self.is_recurrent:
                 self.rnn = copy.deepcopy(policy.memory_a.rnn)
@@ -130,6 +134,16 @@ class _TorchPolicyExporter(torch.nn.Module):
         x = torch.cat([latent, x], dim=1)
         return self.actor(x), (weights, latent)
     
+    def forward_mcp_cts(self, x):  # x is single observations
+        x = self.normalizer(x)
+        self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
+        x_no_goal = x[:, self.obs_no_goal_mask]
+        latent = self.student_encoder(self.history.flatten(1))
+        x = torch.cat([latent, x], dim=1)
+        x_no_goal = torch.cat([latent, x_no_goal], dim=1)
+        mean_action, _, weights = self.actor(x, x_no_goal)
+        return mean_action, weights
+
     @torch.jit.export
     def reset(self):
         if hasattr(self, 'history'):
@@ -176,6 +190,11 @@ class _OnnxPolicyExporter(torch.nn.Module):
                 self.rnn = copy.deepcopy(policy.memory_a.rnn)
             if self.input_dim is None:
                  self.input_dim = self.actor[0].in_features
+        elif hasattr(policy, "actor_mcp"):
+            self.actor = copy.deepcopy(policy.actor_mcp)
+            self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
+            self.history_length = policy.history.shape[1]
+            self.forward = self.forward_mcp_cts 
         else:
             raise ValueError("Policy does not have an actor/student module.")
 
@@ -246,6 +265,33 @@ class _OnnxPolicyExporter(torch.nn.Module):
 
         return self.actor(x), weights, latent
     
+    def forward_mcp_cts(self, x):
+        x = self.normalizer(x)
+        term_dims = [3, 3, 3, self.num_actions, self.num_actions, self.num_actions]
+        obs_dim = sum(term_dims)
+        frames = x.shape[1] // obs_dim
+        
+        split_sizes = [dim * frames for dim in term_dims]
+        term_chunks = torch.split(x, split_sizes, dim=1)
+        frame_terms_reshaped = [chunk.view(-1, frames, dim) for chunk, dim in zip(term_chunks, term_dims)]
+        history_by_frame = []
+        for i in range(frames):
+            terms_for_this_frame = [ftr[:, i, :] for ftr in frame_terms_reshaped]
+            history_by_frame.append(torch.cat(terms_for_this_frame, dim=1))
+        history = torch.cat(history_by_frame, dim=1)
+
+        last_obs = history[:, -obs_dim:]
+        
+        obs_no_goal = last_obs[:, self.obs_no_goal_mask]
+        
+        latent = self.student_encoder(history)
+        
+        x_in = torch.cat([latent, last_obs], dim=1)
+        x_no_goal_in = torch.cat([latent, obs_no_goal], dim=1)
+        
+        mean_action, _, weights = self.actor(x_in, x_no_goal_in)
+        return mean_action, weights
+
     def export(self, path, filename):
         self.to("cpu")
         obs = torch.zeros(1, self.input_dim)
@@ -254,6 +300,8 @@ class _OnnxPolicyExporter(torch.nn.Module):
         if self.forward == self.forward_moe_cts:
             output_names.append("weights")
             output_names.append("latent")
+        if self.forward == self.forward_mcp_cts:
+            output_names.append("weights")
 
         torch.onnx.export(
             self,

rsl_rl/rsl_rl/algorithms/__init__.py

@@ -31,3 +31,4 @@
 from .ppo import PPO
 from .cts import CTS
 from .moe_cts import MoECTS
+from .mcp_cts import MCPCTS

rsl_rl/rsl_rl/algorithms/mcp_cts.py

@@ -0,0 +1,220 @@
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# 
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# Copyright (c) 2021 ETH Zurich, Nikita Rudin
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+import itertools
+from rsl_rl.modules import ActorCriticMCPCTS
+from rsl_rl.storage import RolloutStorageCTS
+from rsl_rl.algorithms.cts import CTS
+
+class MCPCTS(CTS):
+    model: ActorCriticMCPCTS
+    def __init__(self,
+                model,
+                num_envs,
+                history_length,
+                num_learning_epochs=1,
+                num_mini_batches=1,
+                clip_param=0.2,
+                gamma=0.998,
+                lam=0.95,
+                value_loss_coef=1.0,
+                entropy_coef=0.0,
+                learning_rate=1e-3,
+                student_encoder_learning_rate=1e-3,
+                max_grad_norm=1.0,
+                use_clipped_value_loss=True,
+                schedule="fixed",
+                desired_kl=0.01,
+                teacher_env_ratio=0.75,
+                device='cpu',
+                ):
+
+        self.device = device
+
+        self.desired_kl = desired_kl
+        self.schedule = schedule
+        self.learning_rate = learning_rate
+        self.history_length = history_length
+
+        # CTS components
+        self.model = model
+        self.model.to(self.device)
+        self.storage = None # initialized later
+        params1 = [
+            {"params": self.model.teacher_encoder.parameters()},
+            {"params": self.model.critic.parameters()},
+            {"params": self.model.actor_mcp.parameters()},
+        ]
+        self.optimizer1 = optim.Adam(params1, lr=learning_rate)
+        self.optimizer2 = optim.Adam(self.model.student_encoder.parameters(), lr=student_encoder_learning_rate)
+        self.transition = RolloutStorageCTS.Transition()
+
+        # CTS parameters
+        self.clip_param = clip_param
+        self.num_learning_epochs = num_learning_epochs
+        self.num_mini_batches = num_mini_batches
+        self.value_loss_coef = value_loss_coef
+        self.entropy_coef = entropy_coef
+        self.gamma = gamma
+        self.lam = lam
+        self.max_grad_norm = max_grad_norm
+        self.use_clipped_value_loss = use_clipped_value_loss
+        self.teacher_num_envs = max(int(num_envs * teacher_env_ratio), 1)
+        self.student_num_envs = num_envs - self.teacher_num_envs
+        student_env_ratio = 1 - teacher_env_ratio
+        self.teacher_env_idxs = torch.tensor([i for i in range(num_envs) if i % int(1/student_env_ratio) != 0], device=self.device)
+        self.student_env_idxs = torch.tensor([i for i in range(num_envs) if i % int(1/student_env_ratio) == 0], device=self.device)
+        assert len(self.teacher_env_idxs) == self.teacher_num_envs, f"{len(self.teacher_env_idxs)=} != {self.teacher_num_envs=}"
+        assert len(self.student_env_idxs) == self.student_num_envs, f"{len(self.student_env_idxs)=} != {self.student_num_envs=}"
+
+    def update(self):
+        mean_value_loss = 0
+        mean_surrogate_loss = 0
+        mean_entropy_loss = 0
+        mean_latent_loss = 0
+        assert not self.model.is_recurrent
+        data = list(self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs))
+        teacher_samples = self.teacher_num_envs * self.storage.num_transitions_per_env // self.num_mini_batches
+        student_samples = self.student_num_envs * self.storage.num_transitions_per_env // self.num_mini_batches
+        for sample in data:
+            (
+                obs_batch, privileged_obs_batch, actions_batch, history_batch,
+                target_values_batch, advantages_batch, returns_batch,
+                old_actions_log_prob_batch, old_mu_batch, old_sigma_batch,
+                hid_states_batch, masks_batch
+            ) = sample
+            def get_results(start, end, is_teacher):
+                self.model.act(obs_batch[start:end], privileged_obs_batch[start:end], history_batch[start:end], is_teacher)
+                actions_log_prob = self.model.get_actions_log_prob(actions_batch[start:end])
+                value = self.model.evaluate(privileged_obs_batch[start:end], history_batch[start:end], is_teacher)
+                mu = self.model.action_mean
+                sigma = self.model.action_std
+                entropy = self.model.entropy
+                return actions_log_prob, value, mu, sigma, entropy
+            teacher_results = get_results(0, teacher_samples, True)
+            student_results = get_results(teacher_samples, teacher_samples + student_samples, False)
+            results = []
+            for x1, x2 in zip(teacher_results, student_results):
+                results.append(torch.cat([x1, x2], dim=0))
+            actions_log_prob_batch = results[0]
+            value_batch = results[1]
+            mu_batch = results[2]
+            sigma_batch = results[3]
+            entropy_batch = results[4]
+
+            # KL
+            if self.desired_kl != None and self.schedule == 'adaptive':
+                with torch.inference_mode():
+                    kl = torch.sum(
+                        torch.log(
+                            sigma_batch / old_sigma_batch + 1.e-5) + (
+                                torch.square(old_sigma_batch) +
+                                torch.square(old_mu_batch - mu_batch)
+                            ) / (2.0 * torch.square(sigma_batch)) - 0.5, axis=-1)
+                    kl_mean = torch.mean(kl)
+
+                    if kl_mean > self.desired_kl * 2.0:
+                        self.learning_rate = max(1e-5, self.learning_rate / 1.5)
+                    elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
+                        self.learning_rate = min(1e-2, self.learning_rate * 1.5)
+                    
+                    for param_group in self.optimizer1.param_groups:
+                        param_group['lr'] = self.learning_rate
+
+
+            # Surrogate loss
+            ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch))
+            surrogate = -torch.squeeze(advantages_batch) * ratio
+            surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp(ratio, 1.0 - self.clip_param,
+                                                                            1.0 + self.clip_param)
+            surrogate_losses = torch.max(surrogate, surrogate_clipped)
+            teacher_surrogate_loss = surrogate_losses[:teacher_samples].mean()
+            student_surrogate_loss = surrogate_losses[teacher_samples:].mean()
+            surrogate_loss = teacher_surrogate_loss + student_surrogate_loss
+            # surrogate_loss = teacher_surrogate_loss
+
+            # Value function loss
+            if self.use_clipped_value_loss:
+                value_clipped = target_values_batch + (value_batch - target_values_batch).clamp(-self.clip_param,
+                                                                                                self.clip_param)
+                value_losses = (value_batch - returns_batch).pow(2)
+                value_losses_clipped = (value_clipped - returns_batch).pow(2)
+                value_loss = torch.max(value_losses, value_losses_clipped).mean()
+            else:
+                value_loss = (returns_batch - value_batch).pow(2).mean()
+            # teacher_value_loss = value_losses[:teacher_samples].mean()
+            # student_value_loss = value_losses[teacher_samples:].mean()
+            # value_loss = teacher_value_loss  # + student_value_loss
+
+            loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean()
+
+            # Gradient step
+            self.optimizer1.zero_grad()
+            loss.backward()
+            params_to_clip = itertools.chain.from_iterable(g['params'] for g in self.optimizer1.param_groups)
+            nn.utils.clip_grad_norm_(params_to_clip, self.max_grad_norm)
+            self.optimizer1.step()
+
+            mean_value_loss += value_loss.item()
+            mean_surrogate_loss += surrogate_loss.item()
+            mean_entropy_loss += entropy_batch.mean().item()
+        
+        for sample in data:
+            (
+                obs_batch, privileged_obs_batch, actions_batch, history_batch,
+                target_values_batch, advantages_batch, returns_batch,
+                old_actions_log_prob_batch, old_mu_batch, old_sigma_batch,
+                hid_states_batch, masks_batch
+            ) = sample
+            # Student encoder update
+            student_latent = self.model.student_encoder(history_batch[teacher_samples:])
+            with torch.no_grad():
+                teacher_latent = self.model.teacher_encoder(privileged_obs_batch[teacher_samples:])
+            latent_loss = (teacher_latent - student_latent).pow(2).mean()
+
+            self.optimizer2.zero_grad()
+            latent_loss.backward()
+            nn.utils.clip_grad_norm_(self.model.student_encoder.parameters(), self.max_grad_norm)
+            self.optimizer2.step()
+
+            mean_latent_loss += latent_loss.item()
+
+        num_updates = self.num_learning_epochs * self.num_mini_batches
+        mean_value_loss /= num_updates
+        mean_surrogate_loss /= num_updates
+        mean_entropy_loss /= num_updates
+        mean_latent_loss /= num_updates
+        self.storage.clear()
+
+        return mean_value_loss, mean_surrogate_loss, mean_entropy_loss, mean_latent_loss

rsl_rl/rsl_rl/modules/__init__.py

@@ -32,3 +32,4 @@ from .actor_critic import ActorCritic
 from .actor_critic_recurrent import ActorCriticRecurrent
 from .actor_critic_cts import ActorCriticCTS
 from .actor_critic_moe_cts import ActorCriticMoECTS
+from .actor_critic_mcp_cts import ActorCriticMCPCTS

rsl_rl/rsl_rl/modules/actor_critic_mcp_cts.py

@@ -0,0 +1,292 @@
+# -*- coding: utf-8 -*-
+'''
+@File    : actor_critic_moe_cts.py
+@Time    : 2025/12/30 21:06:46
+@Author  : wty-yy
+@Version : 1.0
+@Blog    : https://wty-yy.github.io/
+@Desc    : Multiplicative Compositional Policies Concurrent Teacher Student Network
+@Refer   : CTS https://arxiv.org/abs/2405.10830, MCP https://arxiv.org/abs/1905.09808
+'''
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Normal
+
+class ActorCriticMCPCTS(nn.Module):
+    is_recurrent = False
+    def __init__(self,  num_obs,
+                        num_critic_obs,
+                        num_actions,
+                        num_envs,
+                        history_length,
+                        obs_no_goal_mask,
+                        actor_hidden_dims=[512, 256],
+                        critic_hidden_dims=[512, 256, 128],
+                        teacher_encoder_hidden_dims=[512, 256],
+                        student_encoder_hidden_dims=[512, 256],
+                        student_expert_num=8,
+                        activation='elu',
+                        latent_dim=32,
+                        norm_type='l2norm',
+                        **kwargs):
+        if kwargs:
+            print("ActorCritic.__init__ got unexpected arguments, which will be ignored: " + str([key for key in kwargs.keys()]))
+        assert norm_type in ['l2norm', 'simnorm'], f"Normalization type {norm_type} not supported!"
+        super().__init__()
+        self.num_actions = num_actions
+        self.history_length = history_length
+        self.register_buffer("obs_no_goal_mask", torch.tensor(obs_no_goal_mask, dtype=torch.bool), persistent=False)
+        self.num_obs_no_goal = torch.sum(self.obs_no_goal_mask).item()
+
+        activation_str = activation
+        activation = get_activation(activation)
+
+        mlp_input_dim_t = num_critic_obs
+        mlp_input_dim_s = num_obs * history_length
+        mlp_input_dim_c = latent_dim + num_critic_obs
+        actor_input_dim_g = latent_dim + num_obs
+        actor_input_dim_p = latent_dim + self.num_obs_no_goal
+
+        # History
+        self.register_buffer("history", torch.zeros((num_envs, history_length, num_obs)), persistent=False)
+
+        # Teacher encoder
+        encoder_layers = []
+        encoder_layers.append(nn.Linear(mlp_input_dim_t, teacher_encoder_hidden_dims[0]))
+        encoder_layers.append(activation)
+        for l in range(len(teacher_encoder_hidden_dims)):
+            if l == len(teacher_encoder_hidden_dims) - 1:
+                encoder_layers.append(nn.Linear(teacher_encoder_hidden_dims[l], latent_dim))
+                if norm_type == 'l2norm':
+                    encoder_layers.append(L2Norm())
+                elif norm_type == 'simnorm':
+                    encoder_layers.append(SimNorm())
+            else:
+                encoder_layers.append(nn.Linear(teacher_encoder_hidden_dims[l], teacher_encoder_hidden_dims[l + 1]))
+                encoder_layers.append(activation)
+        self.teacher_encoder = nn.Sequential(*encoder_layers)
+
+        # Student encoder
+        encoder_layers = []
+        encoder_layers.append(nn.Linear(mlp_input_dim_s, student_encoder_hidden_dims[0]))
+        encoder_layers.append(activation)
+        for l in range(len(student_encoder_hidden_dims)):
+            if l == len(student_encoder_hidden_dims) - 1:
+                encoder_layers.append(nn.Linear(student_encoder_hidden_dims[l], latent_dim))
+                if norm_type == 'l2norm':
+                    encoder_layers.append(L2Norm())
+                elif norm_type == 'simnorm':
+                    encoder_layers.append(SimNorm())
+            else:
+                encoder_layers.append(nn.Linear(student_encoder_hidden_dims[l], student_encoder_hidden_dims[l + 1]))
+                encoder_layers.append(activation)
+        self.student_encoder = nn.Sequential(*encoder_layers)
+
+        # MCP Actor
+        self.actor_mcp = ActorMCP(
+             input_dim=actor_input_dim_g,
+             input_dim_no_goal=actor_input_dim_p,
+             action_dim=num_actions,
+             hidden_dims=actor_hidden_dims,
+             expert_num=student_expert_num,
+             activation=activation_str,
+        )
+
+        # Value function
+        critic_layers = []
+        critic_layers.append(nn.Linear(mlp_input_dim_c, critic_hidden_dims[0]))
+        critic_layers.append(activation)
+        for l in range(len(critic_hidden_dims)):
+            if l == len(critic_hidden_dims) - 1:
+                critic_layers.append(nn.Linear(critic_hidden_dims[l], 1))
+            else:
+                critic_layers.append(nn.Linear(critic_hidden_dims[l], critic_hidden_dims[l + 1]))
+                critic_layers.append(activation)
+        self.critic = nn.Sequential(*critic_layers)
+
+        print(f"Actor MCP: {self.actor_mcp}")
+        print(f"Critic MLP: {self.critic}")
+        print(f"Teacher Encoder: {self.teacher_encoder}")
+        print(f"Student Encoder: {self.student_encoder}")
+
+        self.distribution = None
+        # disable args validation for speedup
+        Normal.set_default_validate_args = False
+
+    def reset(self, dones=None):
+        self.history[dones > 0] = 0.0
+
+    def forward(self):
+        raise NotImplementedError
+    
+    @property
+    def action_mean(self):
+        return self.distribution.mean
+
+    @property
+    def action_std(self):
+        return self.distribution.stddev
+    
+    @property
+    def entropy(self):
+        return self.distribution.entropy().sum(dim=-1)
+
+    def update_distribution(self, x, x_no_goal):
+        mean, std, _ = self.actor_mcp(x, x_no_goal)
+        self.distribution = Normal(mean, std)
+
+    def act(self, obs, privileged_obs, history, is_teacher, **kwargs):
+        if is_teacher:
+            latent = self.teacher_encoder(privileged_obs)
+        else:
+            with torch.no_grad():
+                latent = self.student_encoder(history)
+        x = torch.cat([latent, obs], dim=1)
+        obs_no_goal = obs[:, self.obs_no_goal_mask]
+        x_no_goal = torch.cat([latent, obs_no_goal], dim=1)
+        self.update_distribution(x, x_no_goal)
+        return self.distribution.sample()
+    
+    def get_actions_log_prob(self, actions):
+        return self.distribution.log_prob(actions).sum(dim=-1)
+
+    def act_inference(self, obs):
+        self.history = torch.cat([self.history[:, 1:], obs.unsqueeze(1)], dim=1)
+        latent = self.student_encoder(self.history.flatten(1))
+        x = torch.cat([latent, obs], dim=1)
+        obs_no_goal = obs[:, self.obs_no_goal_mask]
+        x_no_goal = torch.cat([latent, obs_no_goal], dim=1)
+        actions_mean, _, _ = self.actor_mcp(x, x_no_goal)
+        return actions_mean
+
+    def evaluate(self, privileged_obs, history, is_teacher, **kwargs):
+        if is_teacher:
+            latent = self.teacher_encoder(privileged_obs)
+        else:
+            latent = self.student_encoder(history)
+        x = torch.cat([latent.detach(), privileged_obs], dim=1)
+        value = self.critic(x)
+        return value
+
+class ActorMCP(nn.Module):
+    def __init__(
+        self,
+        input_dim,          # latent + full obs
+        input_dim_no_goal,  # latent + obs without goal
+        action_dim,
+        hidden_dims=[512, 256],
+        expert_num=8,
+        expert_hidden_dim=256,
+        activation='elu',
+    ):
+        super().__init__()
+        self.expert_num = expert_num
+        self.action_dim = action_dim
+        activation = get_activation(activation)
+
+        # Gating network
+        gating_layers = []
+        last_dim = input_dim
+        for l in hidden_dims:
+            gating_layers.append(nn.Linear(last_dim, l))
+            gating_layers.append(activation)
+            last_dim = l
+        gating_layers.append(nn.Linear(last_dim, expert_num))
+        gating_layers.append(nn.Sigmoid())
+        self.gating_network = nn.Sequential(*gating_layers)
+
+        # Expert networks
+        expert_layers = []
+        last_dim = input_dim_no_goal
+        for l in hidden_dims:
+            expert_layers.append(nn.Linear(last_dim, l))
+            expert_layers.append(activation)
+            last_dim = l
+        self.experts_backbone = nn.Sequential(*expert_layers)
+        self.experts_hidden = nn.Sequential(
+            nn.Linear(last_dim, expert_num * expert_hidden_dim),
+            activation
+        )
+        self.experts_out = nn.Conv1d(
+             in_channels=expert_num*expert_hidden_dim,
+             out_channels=expert_num*action_dim*2,
+             kernel_size=1,
+             groups=expert_num
+        )
+    
+    def forward(self, x, x_no_goal):
+        """
+        x: latent + full goal
+        x_no_goal: latent + obs without goal
+        """
+        B = x.shape[0]
+        weights = self.gating_network(x).unsqueeze(-1)  # (batch, expert_num, 1)
+        shared_features = self.experts_backbone(x_no_goal)
+        expert_hidden = self.experts_hidden(shared_features)
+        expert_hidden = expert_hidden.unsqueeze(-1)  # (batch, channels, 1)
+        expert_out = self.experts_out(expert_hidden)  # (batch, expert_num * action_dim * 2, 1)
+        expert_out = expert_out.view(B, self.expert_num, self.action_dim * 2)
+        mu, log_std = torch.chunk(expert_out, 2, dim=-1)  # (batch, expert_num, action_dim)
+        log_std = torch.clamp(log_std, -5.0, 2.0)
+        var = torch.exp(2 * log_std) + 1e-9
+
+        # MCP Composition
+        # Formula: var_total = 1 / sum(w_i / var_i)
+        #          mu_total = var_total * sum(w_i * mu_i / var_i)
+        weighted_sum = torch.sum(weights / var, dim=1) + 1e-9  # (batch, action_dim)
+        var_total = 1.0 / weighted_sum
+        sigma_total = torch.sqrt(var_total)
+
+        mu_weighted_sum = torch.sum(weights * mu / var, dim=1)  # (batch, action_dim)
+        mu_total = var_total * mu_weighted_sum
+
+        return mu_total, sigma_total, weights.squeeze(-1)
+
+def get_activation(act_name):
+    if act_name == "elu":
+        return nn.ELU()
+    elif act_name == "selu":
+        return nn.SELU()
+    elif act_name == "relu":
+        return nn.ReLU()
+    elif act_name == "crelu":
+        return nn.ReLU()
+    elif act_name == "lrelu":
+        return nn.LeakyReLU()
+    elif act_name == "tanh":
+        return nn.Tanh()
+    elif act_name == "sigmoid":
+        return nn.Sigmoid()
+    else:
+        print("invalid activation function!")
+        return None
+
+class L2Norm(nn.Module):
+    
+	def __init__(self):
+		super().__init__()
+
+	def forward(self, x):
+		return F.normalize(x, p=2.0, dim=-1)
+
+class SimNorm(nn.Module):
+	"""
+	Simplicial normalization.
+	Adapted from https://arxiv.org/abs/2204.00616.
+	"""
+
+	def __init__(self):
+		super().__init__()
+		self.dim = 8  # for latent dim 512
+
+	def forward(self, x):
+		shp = x.shape
+		x = x.view(*shp[:-1], -1, self.dim)
+		x = F.softmax(x, dim=-1)
+		return x.view(*shp)
+
+	def __repr__(self):
+		return f"SimNorm(dim={self.dim})"

rsl_rl/rsl_rl/runners/on_policy_runner_cts.py

@@ -36,8 +36,8 @@ import statistics
 from torch.utils.tensorboard import SummaryWriter
 import torch
 
-from rsl_rl.algorithms import CTS, MoECTS
+from rsl_rl.algorithms import CTS, MoECTS, MCPCTS
-from rsl_rl.modules import ActorCriticCTS, ActorCriticMoECTS
+from rsl_rl.modules import ActorCriticCTS, ActorCriticMoECTS, ActorCriticMCPCTS
 from rsl_rl.env import VecEnv
 
 import yaml
@@ -80,7 +80,7 @@ class OnPolicyRunnerCTS:
             num_critic_obs = self.env.num_obs
         history_length = train_cfg["history_length"]
         actor_critic_class = eval(self.cfg["policy_class_name"])
-        model: Union[ActorCriticCTS, ActorCriticMoECTS] = actor_critic_class(
+        model: Union[ActorCriticCTS, ActorCriticMoECTS, ActorCriticMCPCTS] = actor_critic_class(
             self.env.num_obs,
             num_critic_obs,
             self.env.num_actions,
@@ -88,7 +88,7 @@ class OnPolicyRunnerCTS:
             history_length,
             **self.policy_cfg).to(self.device)
         alg_class = eval(self.cfg["algorithm_class_name"])
-        self.alg: Union[CTS, MoECTS] = alg_class(model, self.env.num_envs, history_length, device=self.device, **self.alg_cfg)
+        self.alg: Union[CTS, MoECTS, MCPCTS] = alg_class(model, self.env.num_envs, history_length, device=self.device, **self.alg_cfg)
         self.num_steps_per_env = self.cfg["num_steps_per_env"]
         self.save_interval = self.cfg["save_interval"]
 
@@ -175,7 +175,7 @@ class OnPolicyRunnerCTS:
                 start = stop
                 self.alg.compute_returns(privileged_obs, self.history.flatten(1))
             
-            if self.cfg["algorithm_class_name"] == "CTS":
+            if self.cfg["algorithm_class_name"] in ["CTS", "MCPCTS"]:
                 mean_value_loss, mean_surrogate_loss, mean_entropy_loss, mean_latent_loss = self.alg.update()
             elif self.cfg["algorithm_class_name"] == "MoECTS":
                 mean_value_loss, mean_surrogate_loss, mean_entropy_loss, mean_latent_loss, mean_load_balance_loss = self.alg.update()
@@ -212,6 +212,7 @@ class OnPolicyRunnerCTS:
                 else:
                     self.writer.add_scalar('Episode/' + key, value, locs['it'])
                 ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n"""
+        if 'mcp' not in self.cfg["algorithm_class_name"].lower():
             mean_std = self.alg.model.std.mean()
         fps = int(self.num_steps_per_env * self.env.num_envs / (locs['collection_time'] + locs['learn_time']))
 
@@ -222,6 +223,7 @@ class OnPolicyRunnerCTS:
         if 'mean_load_balance_loss' in locs:
             self.writer.add_scalar('Loss/load_balance', locs['mean_load_balance_loss'], locs['it'])
         self.writer.add_scalar('Loss/learning_rate', self.alg.learning_rate, locs['it'])
+        if 'mcp' not in self.cfg["algorithm_class_name"].lower():
             self.writer.add_scalar('Policy/mean_noise_std', mean_std.item(), locs['it'])
         self.writer.add_scalar('Perf/total_fps', fps, locs['it'])
         self.writer.add_scalar('Perf/collection time', locs['collection_time'], locs['it'])
@@ -249,6 +251,7 @@ class OnPolicyRunnerCTS:
                       f"""{'Latent loss:':>{pad}} {locs['mean_latent_loss']:.4f}\n""")
         if 'mean_load_balance_loss' in locs:
             log_string += f"""{'Load balance loss:':>{pad}} {locs['mean_load_balance_loss']:.4f}\n"""
+        if 'mcp' not in self.cfg["algorithm_class_name"].lower():
             log_string += f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n"""
         if len(locs['teacher_rewbuffer']):
             log_string += (f"""{'Mean teacher reward:':>{pad}} {statistics.mean(locs['teacher_rewbuffer']):.2f}\n"""
@@ -274,19 +277,22 @@ class OnPolicyRunnerCTS:
             'iter': self.current_learning_iteration,
             'infos': infos,
             }, path)
-        self.update_robogauge(path, it)
+        self.update_robogauge(it)
     
-    def update_robogauge(self, model_path, it):
+    def update_robogauge(self, it):
         if it % 500 == 0:
             # export jit model
             jit_dir = os.path.join(self.log_dir, 'jit_models')
             jit_path = os.path.join(jit_dir, f'policy_jit_{it}.pt')
             export_policy_as_jit(self.alg.model, jit_dir, filename=f'policy_jit_{it}.pt')
             # upload to robogauge
+            task_name = 'go2'
+            if 'moe' in self.cfg["algorithm_class_name"].lower() or 'mcp' in self.cfg["algorithm_class_name"].lower():
+                task_name = 'go2_moe'
             self.robogauge_client.submit_task(
                 model_path=jit_path,
                 step=it,
-                task_name='go2_moe' if 'moe' in self.cfg["algorithm_class_name"].lower() else 'go2',
+                task_name=task_name,
                 experiment_name=self.cfg["experiment_name"]
             )
         self.robogauge_client.monitor_tasks()