diff --git a/UPDATE.md b/UPDATE.md index 9ae2b2e..8619c89 100644 --- a/UPDATE.md +++ b/UPDATE.md @@ -1,3 +1,6 @@ +# 20260107 +## v0.1.6 +1. 加入`go2_rem_cts`, student使用MoE结构, teacher使用普通CTS, 使用非共享权重和全goal输入 # 20260106 ## v0.1.5 1. 加入`go2_ac_moe_cts`, 参考MoELoco将MoE加载Actor-Critic上, 使用非共享权重和全goal输入 diff --git a/legged_gym/envs/__init__.py b/legged_gym/envs/__init__.py index dc158b1..52bf28b 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, GO2CfgDualMoECTS +from legged_gym.envs.go2.go2_config import GO2Cfg, GO2CfgPPO, GO2CfgCTS, GO2CfgMoECTS, GO2CfgMCPCTS, GO2CfgACMoECTS, GO2CfgDualMoECTS, GO2CfgREMCTS from .base.legged_robot import LeggedRobot from legged_gym.utils.task_registry import task_registry @@ -12,3 +12,4 @@ 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()) +task_registry.register("go2_rem_cts", Go2Robot, GO2Cfg(), GO2CfgREMCTS()) diff --git a/legged_gym/envs/base/legged_robot_config.py b/legged_gym/envs/base/legged_robot_config.py index 94032ca..768cb4c 100644 --- a/legged_gym/envs/base/legged_robot_config.py +++ b/legged_gym/envs/base/legged_robot_config.py @@ -384,4 +384,13 @@ class LeggedRobotCfgDualMoECTS(LeggedRobotCfgCTS): class runner(LeggedRobotCfgCTS.runner): policy_class_name = 'ActorCriticDualMoECTS' - algorithm_class_name = 'DualMoECTS' \ No newline at end of file + algorithm_class_name = 'DualMoECTS' + +class LeggedRobotCfgREMCTS(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 = 'ActorCriticREMCTS' + algorithm_class_name = 'REMCTS' \ 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 bd971db..feee3d2 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, LeggedRobotCfgDualMoECTS +from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMCPCTS, LeggedRobotCfgACMoECTS, LeggedRobotCfgDualMoECTS, LeggedRobotCfgREMCTS class GO2Cfg(LeggedRobotCfg): class init_state(LeggedRobotCfg.init_state): @@ -295,3 +295,13 @@ class GO2CfgDualMoECTS(LeggedRobotCfgDualMoECTS): experiment_name = 'go2_dual_moe_cts' max_iterations = 150000 save_interval = 500 + +class GO2CfgREMCTS(LeggedRobotCfgREMCTS): + class policy(LeggedRobotCfgREMCTS.policy): + expert_num = 8 # number of experts in the student model + + class runner(LeggedRobotCfgREMCTS.runner): + run_name = '' + experiment_name = 'go2_rem_cts' + max_iterations = 150000 + save_interval = 500 diff --git a/legged_gym/utils/exporter.py b/legged_gym/utils/exporter.py index 48ad95b..e73c9e8 100644 --- a/legged_gym/utils/exporter.py +++ b/legged_gym/utils/exporter.py @@ -82,6 +82,8 @@ 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 not hasattr(policy, "obs_no_goal_mask"): + self.forward = self.forward_rem_cts 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() @@ -139,6 +141,13 @@ 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_rem_cts(self, x): # x is single observations + x = self.normalizer(x) + self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1) + latent, weights = self.student_moe_encoder(self.history.flatten(1)) + 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) diff --git a/rsl_rl/rsl_rl/algorithms/__init__.py b/rsl_rl/rsl_rl/algorithms/__init__.py index 269d4fe..b22e669 100644 --- a/rsl_rl/rsl_rl/algorithms/__init__.py +++ b/rsl_rl/rsl_rl/algorithms/__init__.py @@ -33,4 +33,5 @@ from .cts import CTS from .moe_cts import MoECTS from .mcp_cts import MCPCTS from .ac_moe_cts import ACMoECTS -from .dual_moe_cts import DualMoECTS \ No newline at end of file +from .dual_moe_cts import DualMoECTS +from .rem_cts import REMCTS \ No newline at end of file diff --git a/rsl_rl/rsl_rl/algorithms/rem_cts.py b/rsl_rl/rsl_rl/algorithms/rem_cts.py new file mode 100644 index 0000000..a916c2f --- /dev/null +++ b/rsl_rl/rsl_rl/algorithms/rem_cts.py @@ -0,0 +1,234 @@ +# 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 ActorCriticMoECTS +from rsl_rl.storage import RolloutStorageCTS +from rsl_rl.algorithms.cts import CTS + +class REMCTS(CTS): + model: ActorCriticMoECTS + 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.parameters()}, + {"params": self.model.actor.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 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 = 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, 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]) + load_balance_loss = torch.mean((mean_usage - target_usage).pow(2)) + # load_balance_loss = torch.sum(mean_usage.pow(2)) * gating_weights.shape[1] # Switch Transformer style + + student_loss = latent_loss + self.load_balance_coef * 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_load_balance_loss += 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_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 ed1b698..4a2e9a2 100644 --- a/rsl_rl/rsl_rl/modules/__init__.py +++ b/rsl_rl/rsl_rl/modules/__init__.py @@ -34,4 +34,5 @@ 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 -from .actor_critic_dual_moe_cts import ActorCriticDualMoECTS \ No newline at end of file +from .actor_critic_dual_moe_cts import ActorCriticDualMoECTS +from .actor_critic_rem_cts import ActorCriticREMCTS \ No newline at end of file diff --git a/rsl_rl/rsl_rl/modules/actor_critic_rem_cts.py b/rsl_rl/rsl_rl/modules/actor_critic_rem_cts.py new file mode 100644 index 0000000..6f0309d --- /dev/null +++ b/rsl_rl/rsl_rl/modules/actor_critic_rem_cts.py @@ -0,0 +1,141 @@ +# -*- 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 : Mixture of Experts Concurrent Teacher Student Network +@Refer : CTS https://arxiv.org/abs/2405.10830, Switch Transformers https://arxiv.org/abs/2101.03961 +''' +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 L2Norm, SimNorm, StudentMoEEncoder, MLP + +class ActorCriticREMCTS(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], + 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_a = latent_dim + num_obs + mlp_input_dim_c = latent_dim + num_critic_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=activation), + L2Norm() if norm_type == 'l2norm' else SimNorm() + ) + + # Student MoE encoder + self.student_moe_encoder = StudentMoEEncoder( + expert_num=expert_num, + input_dim=mlp_input_dim_s, + hidden_dims=student_encoder_hidden_dims, + output_dim=latent_dim, + activation=activation, + norm_type=norm_type, + ) + + # Policy + self.actor = MLP([mlp_input_dim_a, *actor_hidden_dims, num_actions], activation=activation) + + # Value function + self.critic = MLP([mlp_input_dim_c, *critic_hidden_dims, 1], activation=activation) + + print(f"Actor MLP: {self.actor}") + print(f"Critic MLP: {self.critic}") + print(f"Teacher Encoder: {self.teacher_encoder}") + print(f"Student MoE Encoder: {self.student_moe_encoder}") + + # Action noise + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args = False + + @staticmethod + # not used at the moment + def init_weights(sequential, scales): + [torch.nn.init.orthogonal_(module.weight, gain=scales[idx]) for idx, module in + enumerate(mod for mod in sequential if isinstance(mod, nn.Linear))] + + + 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, latent_and_obs): + mean = self.actor(latent_and_obs) + 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) + actions_mean = self.actor(x) + 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_moe_encoder(history) + x = torch.cat([latent.detach(), privileged_obs], dim=1) + value = self.critic(x) + return value 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 bb1c7fa..035e70b 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, DualMoECTS -from rsl_rl.modules import ActorCriticCTS, ActorCriticMoECTS, ActorCriticMCPCTS, ActorCriticACMoECTS, ActorCriticDualMoECTS +from rsl_rl.algorithms import CTS, MoECTS, MCPCTS, ACMoECTS, DualMoECTS, REMCTS +from rsl_rl.modules import ActorCriticCTS, ActorCriticMoECTS, ActorCriticMCPCTS, ActorCriticACMoECTS, ActorCriticDualMoECTS, ActorCriticREMCTS 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, ActorCriticDualMoECTS] = actor_critic_class( + model: Union[ActorCriticCTS, ActorCriticMoECTS, ActorCriticMCPCTS, ActorCriticACMoECTS, ActorCriticDualMoECTS, ActorCriticREMCTS] = 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, DualMoECTS] = alg_class(model, self.env.num_envs, history_length, device=self.device, **self.alg_cfg) + self.alg: Union[CTS, MoECTS, MCPCTS, ACMoECTS, DualMoECTS, REMCTS] = 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"] @@ -183,7 +183,7 @@ class OnPolicyRunnerCTS: 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"] in ["MoECTS", "ACMoECTS"]: + elif self.cfg["algorithm_class_name"] in ["MoECTS", "ACMoECTS", "REMCTS"]: 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()