From e4aa714eab5d9f9dce7b8e93c02f98d6ff1a8e3d Mon Sep 17 00:00:00 2001 From: wty-yy <993660140@qq.com> Date: Wed, 7 Jan 2026 00:21:59 +0800 Subject: [PATCH] v0.1.5 prev1; Add DualMoE --- UPDATE.md | 1 + legged_gym/envs/__init__.py | 3 +- legged_gym/envs/base/legged_robot_config.py | 11 +- legged_gym/envs/go2/go2_config.py | 12 +- legged_gym/utils/exporter.py | 13 +- rsl_rl/rsl_rl/algorithms/__init__.py | 3 +- rsl_rl/rsl_rl/algorithms/dual_moe_cts.py | 287 ++++++++++++++++++ rsl_rl/rsl_rl/modules/__init__.py | 3 +- .../modules/actor_critic_dual_moe_cts.py | 148 +++++++++ rsl_rl/rsl_rl/runners/on_policy_runner_cts.py | 18 +- 10 files changed, 487 insertions(+), 12 deletions(-) create mode 100644 rsl_rl/rsl_rl/algorithms/dual_moe_cts.py create mode 100644 rsl_rl/rsl_rl/modules/actor_critic_dual_moe_cts.py diff --git a/UPDATE.md b/UPDATE.md index 3bc7e2c..9ae2b2e 100644 --- a/UPDATE.md +++ b/UPDATE.md @@ -1,6 +1,7 @@ # 20260106 ## v0.1.5 1. 加入`go2_ac_moe_cts`, 参考MoELoco将MoE加载Actor-Critic上, 使用非共享权重和全goal输入 +2. 加入`go2_dual_moe_cts`, student和actor都使用MoE结构, 使用非共享权重和全goal输入 # 20260105 ## v0.1.4 1. `legged_gym/utils/terrain.py`加入地形难度选择默认`IS_HARD=True` diff --git a/legged_gym/envs/__init__.py b/legged_gym/envs/__init__.py index eee91b5..dc158b1 100644 --- a/legged_gym/envs/__init__.py +++ b/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, GO2CfgMCPCTS, GO2CfgACMoECTS +from legged_gym.envs.go2.go2_config import GO2Cfg, GO2CfgPPO, GO2CfgCTS, GO2CfgMoECTS, GO2CfgMCPCTS, GO2CfgACMoECTS, GO2CfgDualMoECTS from .base.legged_robot import LeggedRobot from legged_gym.utils.task_registry import task_registry @@ -11,3 +11,4 @@ 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()) task_registry.register("go2_ac_moe_cts", Go2Robot, GO2Cfg(), GO2CfgACMoECTS()) +task_registry.register("go2_dual_moe_cts", Go2Robot, GO2Cfg(), GO2CfgDualMoECTS()) diff --git a/legged_gym/envs/base/legged_robot_config.py b/legged_gym/envs/base/legged_robot_config.py index fa412a7..94032ca 100644 --- a/legged_gym/envs/base/legged_robot_config.py +++ b/legged_gym/envs/base/legged_robot_config.py @@ -375,4 +375,13 @@ class LeggedRobotCfgACMoECTS(LeggedRobotCfgCTS): class runner(LeggedRobotCfgCTS.runner): policy_class_name = 'ActorCriticACMoECTS' - algorithm_class_name = 'ACMoECTS' \ No newline at end of file + algorithm_class_name = 'ACMoECTS' + +class LeggedRobotCfgDualMoECTS(LeggedRobotCfgCTS): + class policy(LeggedRobotCfgCTS.policy): + expert_num = 8 # number of experts in the student model + student_encoder_hidden_dims = [512, 256, 128] + + class runner(LeggedRobotCfgCTS.runner): + policy_class_name = 'ActorCriticDualMoECTS' + algorithm_class_name = 'DualMoECTS' \ No newline at end of file diff --git a/legged_gym/envs/go2/go2_config.py b/legged_gym/envs/go2/go2_config.py index e3c63a2..bd971db 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, LeggedRobotCfgMCPCTS, LeggedRobotCfgACMoECTS +from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMCPCTS, LeggedRobotCfgACMoECTS, LeggedRobotCfgDualMoECTS class GO2Cfg(LeggedRobotCfg): class init_state(LeggedRobotCfg.init_state): @@ -285,3 +285,13 @@ class GO2CfgACMoECTS(LeggedRobotCfgACMoECTS): experiment_name = 'go2_ac_moe_cts' max_iterations = 150000 save_interval = 500 + +class GO2CfgDualMoECTS(LeggedRobotCfgDualMoECTS): + class policy(LeggedRobotCfgDualMoECTS.policy): + expert_num = 8 # number of experts in the student model + + class runner(LeggedRobotCfgDualMoECTS.runner): + run_name = '' + experiment_name = 'go2_dual_moe_cts' + max_iterations = 150000 + save_interval = 500 diff --git a/legged_gym/utils/exporter.py b/legged_gym/utils/exporter.py index 602c1ef..48ad95b 100644 --- a/legged_gym/utils/exporter.py +++ b/legged_gym/utils/exporter.py @@ -77,7 +77,8 @@ class _TorchPolicyExporter(torch.nn.Module): self.forward = self.forward_cts if hasattr(policy, "student_moe_encoder"): self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder).cpu() - self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu() + if hasattr(policy, "obs_no_goal_mask"): + self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu() 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 @@ -98,6 +99,8 @@ class _TorchPolicyExporter(torch.nn.Module): self.rnn = copy.deepcopy(policy.memory_s.rnn) else: raise ValueError("Policy does not have an actor/student module.") + if hasattr(policy, "student_moe_encoder") and hasattr(policy, "actor_moe"): + self.forward = self.forward_dual_moe_cts # set up recurrent network if self.is_recurrent: self.rnn.cpu() @@ -155,6 +158,14 @@ class _TorchPolicyExporter(torch.nn.Module): mean, weights = self.actor(x) return mean, (weights, latent) + def forward_dual_moe_cts(self, x): # x is single observations + x = self.normalizer(x) + self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1) + latent, student_weights = self.student_moe_encoder(self.history.flatten(1)) + x = torch.cat([latent, x], dim=1) + mean, actor_weights = self.actor(x) + return mean, (student_weights, actor_weights, latent) + @torch.jit.export def reset(self): if hasattr(self, 'history'): diff --git a/rsl_rl/rsl_rl/algorithms/__init__.py b/rsl_rl/rsl_rl/algorithms/__init__.py index 506a47b..269d4fe 100644 --- a/rsl_rl/rsl_rl/algorithms/__init__.py +++ b/rsl_rl/rsl_rl/algorithms/__init__.py @@ -32,4 +32,5 @@ from .ppo import PPO from .cts import CTS from .moe_cts import MoECTS from .mcp_cts import MCPCTS -from .ac_moe_cts import ACMoECTS \ No newline at end of file +from .ac_moe_cts import ACMoECTS +from .dual_moe_cts import DualMoECTS \ No newline at end of file diff --git a/rsl_rl/rsl_rl/algorithms/dual_moe_cts.py b/rsl_rl/rsl_rl/algorithms/dual_moe_cts.py new file mode 100644 index 0000000..2179249 --- /dev/null +++ b/rsl_rl/rsl_rl/algorithms/dual_moe_cts.py @@ -0,0 +1,287 @@ +# 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 ActorCriticDualMoECTS +from rsl_rl.storage import RolloutStorageCTS +from rsl_rl.algorithms.cts import CTS + +class DualMoECTS(CTS): + model: ActorCriticDualMoECTS + 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, + load_balance_coef=0.01, + 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_experts.parameters()}, + {"params": self.model.actor_moe.parameters()}, + {"params": self.model.std} + ] + self.optimizer1 = optim.Adam(params1, lr=learning_rate) + self.optimizer2 = optim.Adam(self.model.student_moe_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.load_balance_coef = load_balance_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 act(self, obs, privileged_obs, history): + history = history.clone() + def get_results(obs, privileged_obs, history, is_teacher): + actions = self.model.act(obs, privileged_obs, history, is_teacher).detach() + return ( + actions, + self.model.evaluate(obs, privileged_obs, history, is_teacher)[0].detach(), + self.model.get_actions_log_prob(actions).detach(), + self.model.action_mean.detach(), + self.model.action_std.detach(), + ) + ti, si = self.teacher_env_idxs, self.student_env_idxs + teacher_results = get_results(obs[ti], privileged_obs[ti], history[ti], True) + student_results = get_results(obs[si], privileged_obs[si], history[si], False) + results = [] + for x1, x2 in zip(teacher_results, student_results): + results.append(torch.cat([x1, x2], dim=0)) + # Compute the actions and values + self.transition.actions = results[0] + self.transition.values = results[1] + self.transition.actions_log_prob = results[2] + self.transition.action_mean = results[3] + self.transition.action_sigma = results[4] + # need to record obs and critic_obs before env.step() + self.transition.history = torch.cat([history[ti], history[si]], dim=0) + self.transition.observations = torch.cat([obs[ti], obs[si]], dim=0) + self.transition.critic_observations = torch.cat([privileged_obs[ti], privileged_obs[si]], dim=0) + real_actions = torch.zeros_like(self.transition.actions) + real_actions[ti] = self.transition.actions[:self.teacher_num_envs] + real_actions[si] = self.transition.actions[self.teacher_num_envs:] + return real_actions + + def compute_returns(self, last_obs, last_privileged_obs, last_history): + ti, si = self.teacher_env_idxs, self.student_env_idxs + last_values = torch.cat([ + self.model.evaluate(last_obs[ti], last_privileged_obs[ti], last_history[ti], True)[0].detach(), + self.model.evaluate(last_obs[si], last_privileged_obs[si], last_history[si], False)[0].detach(), + ], dim=0) + self.storage.compute_returns(last_values, self.gamma, self.lam) + + def update(self): + mean_value_loss = 0 + mean_surrogate_loss = 0 + mean_entropy_loss = 0 + mean_latent_loss = 0 + mean_student_load_balance_loss = 0 + mean_actor_load_balance_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, weights = self.model.evaluate(obs_batch[start:end], 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, weights + 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] + ac_weights = results[5] + + # 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 + + # Load balance loss + mean_usage = torch.mean(ac_weights, dim=0) + target_usage = torch.full_like(mean_usage, 1.0 / ac_weights.shape[1]) + actor_load_balance_loss = torch.mean((mean_usage - target_usage).pow(2)) + + loss = ( + surrogate_loss + + self.value_loss_coef * value_loss - + self.entropy_coef * entropy_batch.mean() + + self.load_balance_coef * actor_load_balance_loss + ) + + # 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() + mean_actor_load_balance_loss += actor_load_balance_loss.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, gating_weights = self.model.student_moe_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() + + # Load balance loss + mean_usage = torch.mean(gating_weights, dim=0) + target_usage = torch.full_like(mean_usage, 1.0 / gating_weights.shape[1]) + student_load_balance_loss = torch.mean((mean_usage - target_usage).pow(2)) + + student_loss = latent_loss + self.load_balance_coef * student_load_balance_loss + + self.optimizer2.zero_grad() + student_loss.backward() + nn.utils.clip_grad_norm_(self.model.student_moe_encoder.parameters(), self.max_grad_norm) + self.optimizer2.step() + + mean_latent_loss += latent_loss.item() + mean_student_load_balance_loss += student_load_balance_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 + mean_student_load_balance_loss /= num_updates + mean_actor_load_balance_loss /= num_updates + self.storage.clear() + + return mean_value_loss, mean_surrogate_loss, mean_entropy_loss, mean_latent_loss, mean_student_load_balance_loss, mean_actor_load_balance_loss diff --git a/rsl_rl/rsl_rl/modules/__init__.py b/rsl_rl/rsl_rl/modules/__init__.py index fe54d23..ed1b698 100644 --- a/rsl_rl/rsl_rl/modules/__init__.py +++ b/rsl_rl/rsl_rl/modules/__init__.py @@ -33,4 +33,5 @@ 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 -from .actor_critic_ac_moe_cts import ActorCriticACMoECTS \ No newline at end of file +from .actor_critic_ac_moe_cts import ActorCriticACMoECTS +from .actor_critic_dual_moe_cts import ActorCriticDualMoECTS \ No newline at end of file diff --git a/rsl_rl/rsl_rl/modules/actor_critic_dual_moe_cts.py b/rsl_rl/rsl_rl/modules/actor_critic_dual_moe_cts.py new file mode 100644 index 0000000..5b667ba --- /dev/null +++ b/rsl_rl/rsl_rl/modules/actor_critic_dual_moe_cts.py @@ -0,0 +1,148 @@ +# -*- 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, + Switch Transformers (Load Balance) https://arxiv.org/abs/2101.03961 + MoE-Loco (AC MoE) http://arxiv.org/abs/2503.08564 +''' +import numpy as np + +import torch +import torch.nn as nn +from torch.distributions import Normal +from rsl_rl.modules.utils import MLP, MoE, Experts, L2Norm, SimNorm + +class ActorCriticDualMoECTS(nn.Module): + is_recurrent = False + def __init__(self, num_obs, + num_critic_obs, + num_actions, + num_envs, + history_length, + actor_hidden_dims=[512, 256, 128], + critic_hidden_dims=[512, 256, 128], + teacher_encoder_hidden_dims=[512, 256], + student_encoder_hidden_dims=[512, 256, 128], # last dim is expert hidden dim + expert_num=8, + activation='elu', + init_noise_std=1.0, + 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 + + mlp_input_dim_t = num_critic_obs + mlp_input_dim_s = num_obs * history_length + mlp_input_dim_c = latent_dim + num_critic_obs + mlp_input_dim_a = latent_dim + num_obs + + # History + self.register_buffer("history", torch.zeros((num_envs, history_length, num_obs)), persistent=False) + + # Teacher encoder + self.teacher_encoder = nn.Sequential( + MLP([mlp_input_dim_t, *teacher_encoder_hidden_dims, latent_dim], activation), + L2Norm() if norm_type == 'l2norm' else SimNorm() + ) + + # Student encoder + self.student_moe_encoder = MoE( + expert_num=expert_num, + input_dim=mlp_input_dim_s, + hidden_dims=student_encoder_hidden_dims, + output_dim=latent_dim, + activation=activation, + ) + + # MCP Actor + self.actor_moe = MoE( + expert_num=expert_num, + input_dim=mlp_input_dim_a, + hidden_dims=actor_hidden_dims, + output_dim=num_actions, + activation=activation, + ) + + # Value function + self.critic_experts = Experts( + expert_num=expert_num, + input_dim=mlp_input_dim_c, + backbone_hidden_dims=critic_hidden_dims[:-1], + expert_hidden_dim=critic_hidden_dims[-1], + output_dim=1, + activation=activation, + ) + + print(f"Actor MoE: {self.actor_moe}") + print(f"Critic Experts: {self.critic_experts}") + print(f"Teacher Encoder: {self.teacher_encoder}") + print(f"Student MoE Encoder: {self.student_moe_encoder}") + + self.distribution = None + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + # 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): + mean, _ = self.actor_moe(x) + self.distribution = Normal(mean, mean*0. + self.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_moe_encoder(history) + x = torch.cat([latent, obs], dim=1) + self.update_distribution(x) + 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_moe_encoder(self.history.flatten(1)) + x = torch.cat([latent, obs], dim=1) + mean, _ = self.actor_moe(x) + return mean + + def evaluate(self, obs, privileged_obs, history, is_teacher, **kwargs): + if is_teacher: + latent = self.teacher_encoder(privileged_obs) + else: + latent, _ = self.student_moe_encoder(history) + x_actor = torch.cat([latent, obs], dim=1) + weights = self.actor_moe.gating_network(x_actor) # (B, expert_num) + x_critic = torch.cat([latent.detach(), privileged_obs], dim=1) + experts_value = self.critic_experts(x_critic) + value = torch.sum(weights.unsqueeze(-1) * experts_value, dim=1) + return value, weights 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 c5e51ce..bb1c7fa 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, MCPCTS, ACMoECTS -from rsl_rl.modules import ActorCriticCTS, ActorCriticMoECTS, ActorCriticMCPCTS, ActorCriticACMoECTS +from rsl_rl.algorithms import CTS, MoECTS, MCPCTS, ACMoECTS, DualMoECTS +from rsl_rl.modules import ActorCriticCTS, ActorCriticMoECTS, ActorCriticMCPCTS, ActorCriticACMoECTS, ActorCriticDualMoECTS from rsl_rl.env import VecEnv import yaml @@ -79,7 +79,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, ActorCriticMCPCTS, ActorCriticACMoECTS] = actor_critic_class( + model: Union[ActorCriticCTS, ActorCriticMoECTS, ActorCriticMCPCTS, ActorCriticACMoECTS, ActorCriticDualMoECTS] = actor_critic_class( self.env.num_obs, num_critic_obs, self.env.num_actions, @@ -87,7 +87,7 @@ class OnPolicyRunnerCTS: history_length, **self.policy_cfg).to(self.device) alg_class = eval(self.cfg["algorithm_class_name"]) - self.alg: Union[CTS, MoECTS, MCPCTS, ACMoECTS] = alg_class(model, self.env.num_envs, history_length, device=self.device, **self.alg_cfg) + self.alg: Union[CTS, MoECTS, MCPCTS, ACMoECTS, DualMoECTS] = 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"] @@ -113,7 +113,7 @@ class OnPolicyRunnerCTS: # robogauge client try: from robogauge.scripts.client import RoboGaugeClient - self.robogauge_client = RoboGaugeClient() + self.robogauge_client = RoboGaugeClient("http://127.0.0.1:9973") # Change PORT to your server port if needed, default is 9973 except: self.robogauge_client = None @@ -176,7 +176,7 @@ class OnPolicyRunnerCTS: # Learning step start = stop - if self.cfg["algorithm_class_name"] == "ACMoECTS": + if self.cfg["algorithm_class_name"] in ["ACMoECTS", "DualMoECTS"]: self.alg.compute_returns(obs, privileged_obs, self.history.flatten(1)) else: self.alg.compute_returns(privileged_obs, self.history.flatten(1)) @@ -185,6 +185,8 @@ class OnPolicyRunnerCTS: mean_value_loss, mean_surrogate_loss, mean_entropy_loss, mean_latent_loss = self.alg.update() elif self.cfg["algorithm_class_name"] in ["MoECTS", "ACMoECTS"]: mean_value_loss, mean_surrogate_loss, mean_entropy_loss, mean_latent_loss, mean_load_balance_loss = self.alg.update() + elif self.cfg["algorithm_class_name"] == "DualMoECTS": + mean_value_loss, mean_surrogate_loss, mean_entropy_loss, mean_latent_loss, mean_load_balance_loss, mean_actor_load_balance_loss = self.alg.update() stop = time.time() learn_time = stop - start self.current_learning_iteration += 1 @@ -228,6 +230,8 @@ class OnPolicyRunnerCTS: self.writer.add_scalar('Loss/latent', locs['mean_latent_loss'], locs['it']) if 'mean_load_balance_loss' in locs: self.writer.add_scalar('Loss/load_balance', locs['mean_load_balance_loss'], locs['it']) + if 'mean_actor_load_balance_loss' in locs: + self.writer.add_scalar('Loss/actor_load_balance', locs['mean_actor_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']) @@ -257,6 +261,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""" + if 'mean_actor_load_balance_loss' in locs: + log_string += f"""{'Actor load balance loss:':>{pad}} {locs['mean_actor_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']):