v0.1.5 prev1; Add ACMoE
This commit is contained in:
wty-yy
@@ -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
|
||||
from .actor_critic_mcp_cts import ActorCriticMCPCTS
|
||||
from .actor_critic_ac_moe_cts import ActorCriticACMoECTS
|
||||
146
rsl_rl/rsl_rl/modules/actor_critic_ac_moe_cts.py
Normal file
146
rsl_rl/rsl_rl/modules/actor_critic_ac_moe_cts.py
Normal file
@@ -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
|
||||
@@ -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:
|
||||
|
||||
126
rsl_rl/rsl_rl/modules/utils.py
Normal file
126
rsl_rl/rsl_rl/modules/utils.py
Normal file
@@ -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})"
|
||||
Reference in New Issue
Block a user