Init

rsl_rl/rsl_rl/modules/__init__.py

@@ -0,0 +1,34 @@
+# 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
+
+from .actor_critic import ActorCritic
+from .actor_critic_recurrent import ActorCriticRecurrent
+from .actor_critic_cts import ActorCriticCTS
+from .actor_critic_moe_cts import ActorCriticMoECTS

rsl_rl/rsl_rl/modules/actor_critic.py

@@ -0,0 +1,155 @@
+# 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 numpy as np
+
+import torch
+import torch.nn as nn
+from torch.distributions import Normal
+from torch.nn.modules import rnn
+
+class ActorCritic(nn.Module):
+    is_recurrent = False
+    def __init__(self,  num_actor_obs,
+                        num_critic_obs,
+                        num_actions,
+                        actor_hidden_dims=[256, 256, 256],
+                        critic_hidden_dims=[256, 256, 256],
+                        activation='elu',
+                        init_noise_std=1.0,
+                        **kwargs):
+        if kwargs:
+            print("ActorCritic.__init__ got unexpected arguments, which will be ignored: " + str([key for key in kwargs.keys()]))
+        super(ActorCritic, self).__init__()
+
+        activation = get_activation(activation)
+
+        mlp_input_dim_a = num_actor_obs
+        mlp_input_dim_c = num_critic_obs
+
+        # Policy
+        actor_layers = []
+        actor_layers.append(nn.Linear(mlp_input_dim_a, actor_hidden_dims[0]))
+        actor_layers.append(activation)
+        for l in range(len(actor_hidden_dims)):
+            if l == len(actor_hidden_dims) - 1:
+                actor_layers.append(nn.Linear(actor_hidden_dims[l], num_actions))
+            else:
+                actor_layers.append(nn.Linear(actor_hidden_dims[l], actor_hidden_dims[l + 1]))
+                actor_layers.append(activation)
+        self.actor = nn.Sequential(*actor_layers)
+
+        # 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 MLP: {self.actor}")
+        print(f"Critic MLP: {self.critic}")
+
+        # 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
+        
+        # seems that we get better performance without init
+        # self.init_memory_weights(self.memory_a, 0.001, 0.)
+        # self.init_memory_weights(self.memory_c, 0.001, 0.)
+
+    @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):
+        pass
+
+    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, observations):
+        mean = self.actor(observations)
+        self.distribution = Normal(mean, mean*0. + self.std)
+
+    def act(self, observations, **kwargs):
+        self.update_distribution(observations)
+        return self.distribution.sample()
+    
+    def get_actions_log_prob(self, actions):
+        return self.distribution.log_prob(actions).sum(dim=-1)
+
+    def act_inference(self, observations):
+        actions_mean = self.actor(observations)
+        return actions_mean
+
+    def evaluate(self, critic_observations, **kwargs):
+        value = self.critic(critic_observations)
+        return value
+
+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

rsl_rl/rsl_rl/modules/actor_critic_cts.py

@@ -0,0 +1,243 @@
+# 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 numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Normal
+
+class ActorCriticCTS(nn.Module):
+    is_recurrent = False
+    def __init__(self,  num_actor_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],
+                        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(ActorCriticCTS, self).__init__()
+        self.num_actions = num_actions
+
+        activation = get_activation(activation)
+
+        mlp_input_dim_t = num_critic_obs
+        mlp_input_dim_s = num_actor_obs * history_length
+        mlp_input_dim_a = latent_dim + num_actor_obs
+        mlp_input_dim_c = latent_dim + num_critic_obs
+
+        # History
+        self.history = torch.zeros((num_envs, history_length, num_actor_obs), device='cuda')
+
+        # 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)
+
+        # Policy
+        actor_layers = []
+        actor_layers.append(nn.Linear(mlp_input_dim_a, actor_hidden_dims[0]))
+        actor_layers.append(activation)
+        for l in range(len(actor_hidden_dims)):
+            if l == len(actor_hidden_dims) - 1:
+                actor_layers.append(nn.Linear(actor_hidden_dims[l], num_actions))
+            else:
+                actor_layers.append(nn.Linear(actor_hidden_dims[l], actor_hidden_dims[l + 1]))
+                actor_layers.append(activation)
+        self.actor = nn.Sequential(*actor_layers)
+
+        # 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 MLP: {self.actor}")
+        print(f"Critic MLP: {self.critic}")
+        print(f"Teacher Encoder MLP: {self.teacher_encoder}")
+        print(f"Student Encoder MLP: {self.student_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
+        
+        # seems that we get better performance without init
+        # self.init_memory_weights(self.memory_a, 0.001, 0.)
+        # self.init_memory_weights(self.memory_c, 0.001, 0.)
+
+    @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:
+            latent = self.student_encoder(history).detach()
+        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)
+        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_encoder(history)
+        x = torch.cat([latent.detach(), privileged_obs], dim=1)
+        value = self.critic(x)
+        return value
+
+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
+        target_shape = list(shp[:-1]) + [-1, self.dim]
+        x = x.view(target_shape)
+        x = F.softmax(x, dim=-1)
+        return x.view(shp)
+
+    def __repr__(self):
+        return f"SimNorm(dim={self.dim})"

rsl_rl/rsl_rl/modules/actor_critic_moe_cts.py

@@ -0,0 +1,301 @@
+# 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 numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Normal
+
+class ActorCriticMoECTS(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, 128],
+                        critic_hidden_dims=[512, 256, 128],
+                        teacher_encoder_hidden_dims=[512, 256],
+                        student_encoder_hidden_dims=[512, 256],
+                        student_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(ActorCriticMoECTS, self).__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)
+
+        activation_str = activation
+        activation = get_activation(activation)
+
+        mlp_input_dim_t = num_critic_obs
+        mlp_input_dim_e = torch.sum(self.obs_no_goal_mask).item() * history_length  # exclude command inputs for expert
+        mlp_input_dim_g = num_obs * history_length  # all obs for gating
+        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
+        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 MoE encoder
+        self.student_moe_encoder = StudentMoEEncoder(
+            expert_dim=mlp_input_dim_e,
+            gating_dim=mlp_input_dim_g,
+            hidden_dims=student_encoder_hidden_dims,
+            expert_num=student_expert_num,
+            latent_dim=latent_dim,
+            activation=activation_str
+        )
+
+        # Policy
+        actor_layers = []
+        actor_layers.append(nn.Linear(mlp_input_dim_a, actor_hidden_dims[0]))
+        actor_layers.append(activation)
+        for l in range(len(actor_hidden_dims)):
+            if l == len(actor_hidden_dims) - 1:
+                actor_layers.append(nn.Linear(actor_hidden_dims[l], num_actions))
+            else:
+                actor_layers.append(nn.Linear(actor_hidden_dims[l], actor_hidden_dims[l + 1]))
+                actor_layers.append(activation)
+        self.actor = nn.Sequential(*actor_layers)
+
+        # 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 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
+        
+        # seems that we get better performance without init
+        # self.init_memory_weights(self.memory_a, 0.001, 0.)
+        # self.init_memory_weights(self.memory_c, 0.001, 0.)
+
+    @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.get_student_latent_and_weights(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.get_student_latent_and_weights(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.get_student_latent_and_weights(history)
+        x = torch.cat([latent.detach(), privileged_obs], dim=1)
+        value = self.critic(x)
+        return value
+    
+    def get_student_latent_and_weights(self, history):
+        B = history.shape[0]
+        history_no_goal = history.reshape(B, self.history_length, -1)[:, :, self.obs_no_goal_mask].reshape(B, -1)
+        return self.student_moe_encoder(history, history_no_goal)
+
+class StudentMoEEncoder(nn.Module):
+    def __init__(
+        self,
+        expert_dim,
+        gating_dim,
+        hidden_dims=[512, 256],
+        expert_num=8,
+        expert_hidden_dim=128,
+        latent_dim=32,
+        activation='elu',
+        norm_type='l2norm',
+    ):
+        super().__init__()
+        self.expert_num = expert_num
+        self.latent_dim = latent_dim
+        self.norm_layer = L2Norm() if norm_type == 'l2norm' else SimNorm()
+        activation = get_activation(activation)
+
+        # Expert networks
+        experts_layers = []
+        last_dim = expert_dim
+        for l in hidden_dims:
+            experts_layers.append(nn.Linear(last_dim, l))
+            experts_layers.append(activation)
+            last_dim = l
+        self.experts_backbone = nn.Sequential(*experts_layers)
+        self.experts_hidden = nn.Sequential(
+            nn.Linear(last_dim, expert_num * expert_hidden_dim),
+            activation
+        )
+        self.experts_out = nn.Linear(expert_hidden_dim, latent_dim)
+
+        # Gating network
+        gating_layers = []
+        last_dim = gating_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.Softmax(dim=-1))
+        self.gating_network = nn.Sequential(*gating_layers)
+    
+    def forward(self, obs, obs_no_goal):
+        weights = self.gating_network(obs)  # (batch, expert_num)
+        shared_features = self.experts_backbone(obs_no_goal)
+        expert_hidden = self.experts_hidden(shared_features)
+        expert_hidden = expert_hidden.view(-1, self.expert_num, expert_hidden.shape[-1] // self.expert_num)
+        expert_latent = self.experts_out(expert_hidden)  # (batch, expert_num, latent_dim)
+        latent = torch.sum(weights.unsqueeze(-1) * expert_latent, dim=1)  # (batch, latent_dim)
+        latent = self.norm_layer(latent)
+        return latent, weights
+
+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})"

rsl_rl/rsl_rl/modules/actor_critic_recurrent.py

@@ -0,0 +1,116 @@
+# 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 numpy as np
+
+import torch
+import torch.nn as nn
+from torch.distributions import Normal
+from torch.nn.modules import rnn
+from .actor_critic import ActorCritic, get_activation
+from rsl_rl.utils import unpad_trajectories
+
+class ActorCriticRecurrent(ActorCritic):
+    is_recurrent = True
+    def __init__(self,  num_actor_obs,
+                        num_critic_obs,
+                        num_actions,
+                        actor_hidden_dims=[256, 256, 256],
+                        critic_hidden_dims=[256, 256, 256],
+                        activation='elu',
+                        rnn_type='lstm',
+                        rnn_hidden_size=256,
+                        rnn_num_layers=1,
+                        init_noise_std=1.0,
+                        **kwargs):
+        if kwargs:
+            print("ActorCriticRecurrent.__init__ got unexpected arguments, which will be ignored: " + str(kwargs.keys()),)
+
+        super().__init__(num_actor_obs=rnn_hidden_size,
+                         num_critic_obs=rnn_hidden_size,
+                         num_actions=num_actions,
+                         actor_hidden_dims=actor_hidden_dims,
+                         critic_hidden_dims=critic_hidden_dims,
+                         activation=activation,
+                         init_noise_std=init_noise_std)
+
+        activation = get_activation(activation)
+
+        self.memory_a = Memory(num_actor_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_size)
+        self.memory_c = Memory(num_critic_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_size)
+
+        print(f"Actor RNN: {self.memory_a}")
+        print(f"Critic RNN: {self.memory_c}")
+
+    def reset(self, dones=None):
+        self.memory_a.reset(dones)
+        self.memory_c.reset(dones)
+
+    def act(self, observations, masks=None, hidden_states=None):
+        input_a = self.memory_a(observations, masks, hidden_states)
+        return super().act(input_a.squeeze(0))
+
+    def act_inference(self, observations):
+        input_a = self.memory_a(observations)
+        return super().act_inference(input_a.squeeze(0))
+
+    def evaluate(self, critic_observations, masks=None, hidden_states=None):
+        input_c = self.memory_c(critic_observations, masks, hidden_states)
+        return super().evaluate(input_c.squeeze(0))
+    
+    def get_hidden_states(self):
+        return self.memory_a.hidden_states, self.memory_c.hidden_states
+
+
+class Memory(torch.nn.Module):
+    def __init__(self, input_size, type='lstm', num_layers=1, hidden_size=256):
+        super().__init__()
+        # RNN
+        rnn_cls = nn.GRU if type.lower() == 'gru' else nn.LSTM
+        self.rnn = rnn_cls(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers)
+        self.hidden_states = None
+    
+    def forward(self, input, masks=None, hidden_states=None):
+        batch_mode = masks is not None
+        if batch_mode:
+            # batch mode (policy update): need saved hidden states
+            if hidden_states is None:
+                raise ValueError("Hidden states not passed to memory module during policy update")
+            out, _ = self.rnn(input, hidden_states)
+            out = unpad_trajectories(out, masks)
+        else:
+            # inference mode (collection): use hidden states of last step
+            out, self.hidden_states = self.rnn(input.unsqueeze(0), self.hidden_states)
+        return out
+
+    def reset(self, dones=None):
+        # When the RNN is an LSTM, self.hidden_states_a is a list with hidden_state and cell_state
+        for hidden_state in self.hidden_states:
+            hidden_state[..., dones, :] = 0.0