From 9ed3f0e1440c76eb560599df53370f98b2931091 Mon Sep 17 00:00:00 2001 From: wty-yy <993660140@qq.com> Date: Wed, 31 Dec 2025 01:16:33 +0800 Subject: [PATCH] Add mcp-cts --- UPDATE.md | 1 + deploy/deploy_mujoco/deploy_go2_lab.py | 126 -------- legged_gym/envs/__init__.py | 9 +- legged_gym/envs/base/legged_robot_config.py | 9 + legged_gym/envs/go2/go2_config.py | 13 +- legged_gym/utils/exporter.py | 50 ++- rsl_rl/rsl_rl/algorithms/__init__.py | 3 +- rsl_rl/rsl_rl/algorithms/mcp_cts.py | 220 +++++++++++++ rsl_rl/rsl_rl/modules/__init__.py | 3 +- rsl_rl/rsl_rl/modules/actor_critic_mcp_cts.py | 292 ++++++++++++++++++ rsl_rl/rsl_rl/runners/on_policy_runner_cts.py | 28 +- 11 files changed, 609 insertions(+), 145 deletions(-) delete mode 100644 deploy/deploy_mujoco/deploy_go2_lab.py create mode 100644 rsl_rl/rsl_rl/algorithms/mcp_cts.py create mode 100644 rsl_rl/rsl_rl/modules/actor_critic_mcp_cts.py diff --git a/UPDATE.md b/UPDATE.md index 6753b9d..2b0acb9 100644 --- a/UPDATE.md +++ b/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 diff --git a/deploy/deploy_mujoco/deploy_go2_lab.py b/deploy/deploy_mujoco/deploy_go2_lab.py deleted file mode 100644 index 19f45ea..0000000 --- a/deploy/deploy_mujoco/deploy_go2_lab.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/legged_gym/envs/__init__.py b/legged_gym/envs/__init__.py index ebe4a29..fb92b50 100644 --- a/legged_gym/envs/__init__.py +++ b/legged_gym/envs/__init__.py @@ -1,11 +1,12 @@ 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 -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", 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()) diff --git a/legged_gym/envs/base/legged_robot_config.py b/legged_gym/envs/base/legged_robot_config.py index 3c55afc..aeb2887 100644 --- a/legged_gym/envs/base/legged_robot_config.py +++ b/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' diff --git a/legged_gym/envs/go2/go2_config.py b/legged_gym/envs/go2/go2_config.py index 1baada9..a967b9d 100644 --- a/legged_gym/envs/go2/go2_config.py +++ b/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 diff --git a/legged_gym/utils/exporter.py b/legged_gym/utils/exporter.py index 12af729..04339ff 100644 --- a/legged_gym/utils/exporter.py +++ b/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) @@ -129,6 +133,16 @@ class _TorchPolicyExporter(torch.nn.Module): latent, weights = self.student_moe_encoder(self.history.flatten(1), history_no_goal) 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): @@ -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.") @@ -245,6 +264,33 @@ class _OnnxPolicyExporter(torch.nn.Module): x = torch.cat([latent, last_obs], dim=1) 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") @@ -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, diff --git a/rsl_rl/rsl_rl/algorithms/__init__.py b/rsl_rl/rsl_rl/algorithms/__init__.py index ba9dadf..a9fd6f7 100644 --- a/rsl_rl/rsl_rl/algorithms/__init__.py +++ b/rsl_rl/rsl_rl/algorithms/__init__.py @@ -30,4 +30,5 @@ from .ppo import PPO from .cts import CTS -from .moe_cts import MoECTS \ No newline at end of file +from .moe_cts import MoECTS +from .mcp_cts import MCPCTS \ No newline at end of file diff --git a/rsl_rl/rsl_rl/algorithms/mcp_cts.py b/rsl_rl/rsl_rl/algorithms/mcp_cts.py new file mode 100644 index 0000000..77c4580 --- /dev/null +++ b/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 diff --git a/rsl_rl/rsl_rl/modules/__init__.py b/rsl_rl/rsl_rl/modules/__init__.py index b234fe3..9a65ca9 100644 --- a/rsl_rl/rsl_rl/modules/__init__.py +++ b/rsl_rl/rsl_rl/modules/__init__.py @@ -31,4 +31,5 @@ from .actor_critic import ActorCritic from .actor_critic_recurrent import ActorCriticRecurrent from .actor_critic_cts import ActorCriticCTS -from .actor_critic_moe_cts import ActorCriticMoECTS \ No newline at end of file +from .actor_critic_moe_cts import ActorCriticMoECTS +from .actor_critic_mcp_cts import ActorCriticMCPCTS \ No newline at end of file diff --git a/rsl_rl/rsl_rl/modules/actor_critic_mcp_cts.py b/rsl_rl/rsl_rl/modules/actor_critic_mcp_cts.py new file mode 100644 index 0000000..fa9fa14 --- /dev/null +++ b/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})" diff --git a/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py b/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py index 407e3df..bf9c473 100644 --- a/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py +++ b/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.modules import ActorCriticCTS, ActorCriticMoECTS +from rsl_rl.algorithms import CTS, MoECTS, MCPCTS +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,7 +212,8 @@ class OnPolicyRunnerCTS: else: self.writer.add_scalar('Episode/' + key, value, locs['it']) ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" - mean_std = self.alg.model.std.mean() + 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'])) self.writer.add_scalar('Loss/value_function', locs['mean_value_loss'], locs['it']) @@ -222,7 +223,8 @@ 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']) - self.writer.add_scalar('Policy/mean_noise_std', mean_std.item(), 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']) self.writer.add_scalar('Perf/learning_time', locs['learn_time'], locs['it']) @@ -249,7 +251,8 @@ 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""" - log_string += f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\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""" f"""{'Mean teacher episode length:':>{pad}} {statistics.mean(locs['teacher_lenbuffer']):.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()