diff --git a/UPDATE.md b/UPDATE.md index 7234555..3bc7e2c 100644 --- a/UPDATE.md +++ b/UPDATE.md @@ -1,3 +1,6 @@ +# 20260106 +## v0.1.5 +1. 加入`go2_ac_moe_cts`, 参考MoELoco将MoE加载Actor-Critic上, 使用非共享权重和全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 fb92b50..eee91b5 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 +from legged_gym.envs.go2.go2_config import GO2Cfg, GO2CfgPPO, GO2CfgCTS, GO2CfgMoECTS, GO2CfgMCPCTS, GO2CfgACMoECTS from .base.legged_robot import LeggedRobot from legged_gym.utils.task_registry import task_registry @@ -10,3 +10,4 @@ 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()) +task_registry.register("go2_ac_moe_cts", Go2Robot, GO2Cfg(), GO2CfgACMoECTS()) diff --git a/legged_gym/envs/base/legged_robot_config.py b/legged_gym/envs/base/legged_robot_config.py index aeb2887..fa412a7 100644 --- a/legged_gym/envs/base/legged_robot_config.py +++ b/legged_gym/envs/base/legged_robot_config.py @@ -368,3 +368,11 @@ class LeggedRobotCfgMCPCTS(LeggedRobotCfgCTS): class runner(LeggedRobotCfgCTS.runner): policy_class_name = 'ActorCriticMCPCTS' algorithm_class_name = 'MCPCTS' + +class LeggedRobotCfgACMoECTS(LeggedRobotCfgCTS): + class policy(LeggedRobotCfgCTS.policy): + expert_num = 8 # number of experts in the student model + + class runner(LeggedRobotCfgCTS.runner): + policy_class_name = 'ActorCriticACMoECTS' + algorithm_class_name = 'ACMoECTS' \ 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 a967b9d..e3c63a2 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 +from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMCPCTS, LeggedRobotCfgACMoECTS class GO2Cfg(LeggedRobotCfg): class init_state(LeggedRobotCfg.init_state): @@ -275,3 +275,13 @@ class GO2CfgMCPCTS(LeggedRobotCfgMCPCTS): experiment_name = 'go2_mcp_cts' max_iterations = 150000 save_interval = 500 + +class GO2CfgACMoECTS(LeggedRobotCfgACMoECTS): + class policy(LeggedRobotCfgACMoECTS.policy): + expert_num = 8 # number of experts in the student model + + class runner(LeggedRobotCfgACMoECTS.runner): + run_name = '' + experiment_name = 'go2_ac_moe_cts' + max_iterations = 150000 + save_interval = 500 diff --git a/legged_gym/utils/exporter.py b/legged_gym/utils/exporter.py index 4905a3c..602c1ef 100644 --- a/legged_gym/utils/exporter.py +++ b/legged_gym/utils/exporter.py @@ -85,6 +85,9 @@ class _TorchPolicyExporter(torch.nn.Module): 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_moe"): + self.actor = copy.deepcopy(policy.actor_moe) + self.forward = self.forward_ac_moe elif hasattr(policy, "actor"): self.actor = copy.deepcopy(policy.actor) if self.is_recurrent: @@ -144,6 +147,14 @@ class _TorchPolicyExporter(torch.nn.Module): mean_action, _, weights = self.actor(x, x_no_goal) return mean_action, (weights, latent) + def forward_ac_moe(self, x): # x is single observations + x = self.normalizer(x) + self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1) + latent = self.student_encoder(self.history.flatten(1)) + x = torch.cat([latent, x], dim=1) + mean, weights = self.actor(x) + return mean, (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 a9fd6f7..506a47b 100644 --- a/rsl_rl/rsl_rl/algorithms/__init__.py +++ b/rsl_rl/rsl_rl/algorithms/__init__.py @@ -31,4 +31,5 @@ from .ppo import PPO from .cts import CTS from .moe_cts import MoECTS -from .mcp_cts import MCPCTS \ No newline at end of file +from .mcp_cts import MCPCTS +from .ac_moe_cts import ACMoECTS \ No newline at end of file diff --git a/rsl_rl/rsl_rl/algorithms/ac_moe_cts.py b/rsl_rl/rsl_rl/algorithms/ac_moe_cts.py new file mode 100644 index 0000000..052b547 --- /dev/null +++ b/rsl_rl/rsl_rl/algorithms/ac_moe_cts.py @@ -0,0 +1,277 @@ +# 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 ActorCriticACMoECTS +from rsl_rl.storage import RolloutStorageCTS +from rsl_rl.algorithms.cts import CTS + +class ACMoECTS(CTS): + model: ActorCriticACMoECTS + 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_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_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]) + 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 * 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_load_balance_loss += 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 = 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 + mean_load_balance_loss /= num_updates + self.storage.clear() + + return mean_value_loss, mean_surrogate_loss, mean_entropy_loss, mean_latent_loss, mean_load_balance_loss diff --git a/rsl_rl/rsl_rl/modules/__init__.py b/rsl_rl/rsl_rl/modules/__init__.py index 9a65ca9..fe54d23 100644 --- a/rsl_rl/rsl_rl/modules/__init__.py +++ b/rsl_rl/rsl_rl/modules/__init__.py @@ -32,4 +32,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 -from .actor_critic_mcp_cts import ActorCriticMCPCTS \ No newline at end of file +from .actor_critic_mcp_cts import ActorCriticMCPCTS +from .actor_critic_ac_moe_cts import ActorCriticACMoECTS \ No newline at end of file diff --git a/rsl_rl/rsl_rl/modules/actor_critic_ac_moe_cts.py b/rsl_rl/rsl_rl/modules/actor_critic_ac_moe_cts.py new file mode 100644 index 0000000..a772208 --- /dev/null +++ b/rsl_rl/rsl_rl/modules/actor_critic_ac_moe_cts.py @@ -0,0 +1,146 @@ +# -*- 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 +import torch.nn.functional as F +from torch.distributions import Normal +from rsl_rl.modules.utils import get_activation, MLP, MoE, Experts, L2Norm, SimNorm + +class ActorCriticACMoECTS(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], + 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_encoder = nn.Sequential( + MLP([mlp_input_dim_s, *student_encoder_hidden_dims, latent_dim], activation), + L2Norm() if norm_type == 'l2norm' else SimNorm() + ) + + # 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 Encoder: {self.student_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_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_encoder(self.history.flatten(1)) + x = torch.cat([latent, obs], dim=1) + mean, weights = 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_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/modules/actor_critic_mcp_cts.py b/rsl_rl/rsl_rl/modules/actor_critic_mcp_cts.py index fa9fa14..9ed7af7 100644 --- a/rsl_rl/rsl_rl/modules/actor_critic_mcp_cts.py +++ b/rsl_rl/rsl_rl/modules/actor_critic_mcp_cts.py @@ -198,7 +198,7 @@ class ActorMCP(nn.Module): gating_layers.append(nn.Sigmoid()) self.gating_network = nn.Sequential(*gating_layers) - # Expert networks + # Expert networks (Share backbone version) expert_layers = [] last_dim = input_dim_no_goal for l in hidden_dims: diff --git a/rsl_rl/rsl_rl/modules/utils.py b/rsl_rl/rsl_rl/modules/utils.py new file mode 100644 index 0000000..b1f359d --- /dev/null +++ b/rsl_rl/rsl_rl/modules/utils.py @@ -0,0 +1,126 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Experts(nn.Module): + def __init__(self, + expert_num, + input_dim, + backbone_hidden_dims, + expert_hidden_dim, + output_dim, + activation='elu', + ): + super().__init__() + self.expert_num = expert_num + self.output_dim = output_dim + + self.backbone = MLP([input_dim, *backbone_hidden_dims, expert_num * expert_hidden_dim], activation, last_activation=True) + self.experts = nn.Conv1d( + in_channels=expert_num*expert_hidden_dim, + out_channels=expert_num*output_dim, + kernel_size=1, + groups=expert_num, + ) + + def forward(self, x): + shared_features = self.backbone(x).unsqueeze(-1) # (B, expert_num * expert_hidden_dim, 1) + expert_outs = self.experts(shared_features).squeeze(-1) # (B, expert_num * output_dim) + expert_outs = expert_outs.reshape(-1, self.expert_num, self.output_dim) + return expert_outs + +class MoE(nn.Module): + def __init__(self, + expert_num, + input_dim, + hidden_dims, + output_dim, + activation='elu', + ): + super().__init__() + + # Expert networks + self.experts = Experts( + expert_num=expert_num, + input_dim=input_dim, + backbone_hidden_dims=hidden_dims[:-1], + expert_hidden_dim=hidden_dims[-1], + output_dim=output_dim, + activation=activation, + ) + + # Gating network + self.gating_network = nn.Sequential( + MLP([input_dim, *hidden_dims, expert_num], activation), + nn.Softmax(dim=-1) + ) + + def forward(self, x): + weights = self.gating_network(x) # (B, expert_num) + expert_outs = self.experts(x) # (B, expert_num, output_dim) + output = torch.sum(weights.unsqueeze(-1) * expert_outs, dim=1) # (B, output_dim) + return output, weights + +class MLP(nn.Module): + def __init__(self, dims, activation='elu', last_activation=False): + super().__init__() + activation = get_activation(activation) + layers = [] + last_dim = dims[0] + for h_dim in dims[1:-1]: + layers.append(nn.Linear(last_dim, h_dim)) + layers.append(activation) + last_dim = h_dim + layers.append(nn.Linear(last_dim, dims[-1])) + if last_activation: + layers.append(activation) + self.network = nn.Sequential(*layers) + + def forward(self, x): + return self.network(x) + +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 0c8de75..c5e51ce 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 -from rsl_rl.modules import ActorCriticCTS, ActorCriticMoECTS, ActorCriticMCPCTS +from rsl_rl.algorithms import CTS, MoECTS, MCPCTS, ACMoECTS +from rsl_rl.modules import ActorCriticCTS, ActorCriticMoECTS, ActorCriticMCPCTS, ActorCriticACMoECTS 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] = actor_critic_class( + model: Union[ActorCriticCTS, ActorCriticMoECTS, ActorCriticMCPCTS, ActorCriticACMoECTS] = 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] = alg_class(model, self.env.num_envs, history_length, device=self.device, **self.alg_cfg) + self.alg: Union[CTS, MoECTS, MCPCTS, ACMoECTS] = 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"] @@ -176,11 +176,14 @@ class OnPolicyRunnerCTS: # Learning step start = stop - self.alg.compute_returns(privileged_obs, self.history.flatten(1)) + if self.cfg["algorithm_class_name"] == "ACMoECTS": + self.alg.compute_returns(obs, privileged_obs, self.history.flatten(1)) + else: + self.alg.compute_returns(privileged_obs, self.history.flatten(1)) 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": + 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() stop = time.time() learn_time = stop - start diff --git a/setup.py b/setup.py index c0ff322..2e10c31 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ from setuptools import find_packages from distutils.core import setup setup(name='go2_rl_gym', - version='0.1.4', + version='0.1.5', author='Wu Tianyang', license="MIT", packages=find_packages(), @@ -13,7 +13,8 @@ setup(name='go2_rl_gym', 'rsl-rl', 'matplotlib', 'numpy==1.20', - 'tensorboard', + 'tensorboard==2.14.0', + 'google-auth==2.45.0', 'mujoco==3.2.3', 'pyyaml', 'onnx==1.17.0',