Add configs and adapt exporter for RSL-RL distillation

CONTRIBUTORS.md

@@ -46,6 +46,7 @@ Guidelines for modifications:
 * Calvin Yu
 * Cheng-Rong Lai
 * Chenyu Yang
+* Clemens Schwarke
 * CY (Chien-Ying) Chen
 * David Yang
 * Dorsa Rohani

scripts/benchmarks/benchmark_rsl_rl.py

@@ -190,7 +190,7 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
     # write git state to logs
     runner.add_git_repo_to_log(__file__)
     # save resume path before creating a new log_dir
-    if agent_cfg.resume:
+    if agent_cfg.load_run is not None:
         # get path to previous checkpoint
         resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
         print(f"[INFO]: Loading model checkpoint from: {resume_path}")

scripts/reinforcement_learning/rsl_rl/cli_args.py

@@ -27,7 +27,6 @@ def add_rsl_rl_args(parser: argparse.ArgumentParser):
     )
     arg_group.add_argument("--run_name", type=str, default=None, help="Run name suffix to the log directory.")
     # -- load arguments
-    arg_group.add_argument("--resume", type=bool, default=None, help="Whether to resume from a checkpoint.")
     arg_group.add_argument("--load_run", type=str, default=None, help="Name of the run folder to resume from.")
     arg_group.add_argument("--checkpoint", type=str, default=None, help="Checkpoint file to resume from.")
     # -- logger arguments
@@ -73,8 +72,6 @@ def update_rsl_rl_cfg(agent_cfg: RslRlOnPolicyRunnerCfg, args_cli: argparse.Name
         if args_cli.seed == -1:
             args_cli.seed = random.randint(0, 10000)
         agent_cfg.seed = args_cli.seed
-    if args_cli.resume is not None:
-        agent_cfg.resume = args_cli.resume
     if args_cli.load_run is not None:
         agent_cfg.load_run = args_cli.load_run
     if args_cli.checkpoint is not None:

scripts/reinforcement_learning/rsl_rl/train.py

@@ -30,7 +30,6 @@
 # local imports
 import cli_args  # isort: skip
 
-
 # add argparse arguments
 parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.")
 parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
@@ -138,7 +137,7 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
         env = multi_agent_to_single_agent(env)
 
     # save resume path before creating a new log_dir
-    if agent_cfg.resume:
+    if agent_cfg.load_run is not None:
         resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
 
     # wrap for video recording
@@ -161,7 +160,7 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
     # write git state to logs
     runner.add_git_repo_to_log(__file__)
     # load the checkpoint
-    if agent_cfg.resume:
+    if agent_cfg.load_run is not None:
         print(f"[INFO]: Loading model checkpoint from: {resume_path}")
         # load previously trained model
         runner.load(resume_path)

source/isaaclab_rl/isaaclab_rl/rsl_rl/__init__.py

@@ -16,7 +16,7 @@
 """
 
 from .exporter import export_policy_as_jit, export_policy_as_onnx
-from .rl_cfg import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg
+from .rl_cfg import *
 from .rnd_cfg import RslRlRndCfg
 from .symmetry_cfg import RslRlSymmetryCfg
 from .vecenv_wrapper import RslRlVecEnvWrapper

source/isaaclab_rl/isaaclab_rl/rsl_rl/exporter.py

@@ -8,34 +8,34 @@
 import torch
 
 
-def export_policy_as_jit(actor_critic: object, normalizer: object | None, path: str, filename="policy.pt"):
+def export_policy_as_jit(policy: object, normalizer: object | None, path: str, filename="policy.pt"):
     """Export policy into a Torch JIT file.
 
     Args:
-        actor_critic: The actor-critic torch module.
+        policy: The policy torch module.
         normalizer: The empirical normalizer module. If None, Identity is used.
         path: The path to the saving directory.
         filename: The name of exported JIT file. Defaults to "policy.pt".
     """
-    policy_exporter = _TorchPolicyExporter(actor_critic, normalizer)
+    policy_exporter = _TorchPolicyExporter(policy, normalizer)
     policy_exporter.export(path, filename)
 
 
 def export_policy_as_onnx(
-    actor_critic: object, path: str, normalizer: object | None = None, filename="policy.onnx", verbose=False
+    policy: object, path: str, normalizer: object | None = None, filename="policy.onnx", verbose=False
 ):
     """Export policy into a Torch ONNX file.
 
     Args:
-        actor_critic: The actor-critic torch module.
+        policy: The policy torch module.
         normalizer: The empirical normalizer module. If None, Identity is used.
         path: The path to the saving directory.
         filename: The name of exported ONNX file. Defaults to "policy.onnx".
         verbose: Whether to print the model summary. Defaults to False.
     """
     if not os.path.exists(path):
         os.makedirs(path, exist_ok=True)
-    policy_exporter = _OnnxPolicyExporter(actor_critic, normalizer, verbose)
+    policy_exporter = _OnnxPolicyExporter(policy, normalizer, verbose)
     policy_exporter.export(path, filename)
 
 
@@ -47,12 +47,23 @@ def export_policy_as_onnx(
 class _TorchPolicyExporter(torch.nn.Module):
     """Exporter of actor-critic into JIT file."""
 
-    def __init__(self, actor_critic, normalizer=None):
+    def __init__(self, policy, normalizer=None):
         super().__init__()
-        self.actor = copy.deepcopy(actor_critic.actor)
-        self.is_recurrent = actor_critic.is_recurrent
+        # infer policy type
+        if hasattr(policy, "actor"):
+            self.actor = copy.deepcopy(policy.actor)
+        elif hasattr(policy, "student"):
+            self.actor = copy.deepcopy(policy.student)
+        else:
+            raise ValueError("Policy does not have an actor/student module.")
+        self.is_recurrent = policy.is_recurrent
         if self.is_recurrent:
-            self.rnn = copy.deepcopy(actor_critic.memory_a.rnn)
+            if hasattr(policy, "memory_a"):
+                self.rnn = copy.deepcopy(policy.memory_a.rnn)
+            elif hasattr(policy, "memory_s"):
+                self.rnn = copy.deepcopy(policy.memory_s.rnn)
+            else:
+                raise ValueError("Policy does not have a memory module.")
             self.rnn.cpu()
             self.register_buffer("hidden_state", torch.zeros(self.rnn.num_layers, 1, self.rnn.hidden_size))
             self.register_buffer("cell_state", torch.zeros(self.rnn.num_layers, 1, self.rnn.hidden_size))
@@ -94,13 +105,24 @@ def export(self, path, filename):
 class _OnnxPolicyExporter(torch.nn.Module):
     """Exporter of actor-critic into ONNX file."""
 
-    def __init__(self, actor_critic, normalizer=None, verbose=False):
+    def __init__(self, policy, normalizer=None, verbose=False):
         super().__init__()
         self.verbose = verbose
-        self.actor = copy.deepcopy(actor_critic.actor)
-        self.is_recurrent = actor_critic.is_recurrent
+        # infer policy type
+        if hasattr(policy, "actor"):
+            self.actor = copy.deepcopy(policy.actor)
+        elif hasattr(policy, "student"):
+            self.actor = copy.deepcopy(policy.student)
+        else:
+            raise ValueError("Policy does not have an actor/student module.")
+        self.is_recurrent = policy.is_recurrent
         if self.is_recurrent:
-            self.rnn = copy.deepcopy(actor_critic.memory_a.rnn)
+            if hasattr(policy, "memory_a"):
+                self.rnn = copy.deepcopy(policy.memory_a.rnn)
+            elif hasattr(policy, "memory_s"):
+                self.rnn = copy.deepcopy(policy.memory_s.rnn)
+            else:
+                raise ValueError("Policy does not have a memory module.")
             self.rnn.cpu()
             self.forward = self.forward_lstm
         # copy normalizer if exists

source/isaaclab_rl/isaaclab_rl/rsl_rl/rl_cfg.py

@@ -3,6 +3,8 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
+from __future__ import annotations
+
 from dataclasses import MISSING
 from typing import Literal
 
@@ -11,6 +13,10 @@
 from .rnd_cfg import RslRlRndCfg
 from .symmetry_cfg import RslRlSymmetryCfg
 
+#########################
+# Policy configurations #
+#########################
+
 
 @configclass
 class RslRlPpoActorCriticCfg:
@@ -36,23 +42,76 @@ class RslRlPpoActorCriticCfg:
 
 
 @configclass
-class RslRlPpoAlgorithmCfg:
-    """Configuration for the PPO algorithm."""
+class RslRlPpoActorCriticRecurrentCfg(RslRlPpoActorCriticCfg):
+    """Configuration for the PPO actor-critic networks with recurrent layers."""
 
-    class_name: str = "PPO"
-    """The algorithm class name. Default is PPO."""
+    class_name: str = "ActorCriticRecurrent"
+    """The policy class name. Default is ActorCriticRecurrent."""
 
-    value_loss_coef: float = MISSING
-    """The coefficient for the value loss."""
+    rnn_type: str = MISSING
+    """The type of RNN to use. Either "lstm" or "gru"."""
 
-    use_clipped_value_loss: bool = MISSING
-    """Whether to use clipped value loss."""
+    rnn_hidden_dim: int = MISSING
+    """The dimension of the RNN layers."""
 
-    clip_param: float = MISSING
-    """The clipping parameter for the policy."""
+    rnn_num_layers: int = MISSING
+    """The number of RNN layers."""
 
-    entropy_coef: float = MISSING
-    """The coefficient for the entropy loss."""
+
+@configclass
+class RslRlDistillationStudentTeacherCfg:
+    """Configuration for the distillation student-teacher networks."""
+
+    class_name: str = "StudentTeacher"
+    """The policy class name. Default is StudentTeacher."""
+
+    init_noise_std: float = MISSING
+    """The initial noise standard deviation for the student policy."""
+
+    noise_std_type: Literal["scalar", "log"] = "scalar"
+    """The type of noise standard deviation for the policy. Default is scalar."""
+
+    student_hidden_dims: list[int] = MISSING
+    """The hidden dimensions of the student network."""
+
+    teacher_hidden_dims: list[int] = MISSING
+    """The hidden dimensions of the teacher network."""
+
+    activation: str = MISSING
+    """The activation function for the student and teacher networks."""
+
+
+@configclass
+class RslRlDistillationStudentTeacherRecurrentCfg(RslRlDistillationStudentTeacherCfg):
+    """Configuration for the distillation student-teacher recurrent networks."""
+
+    class_name: str = "StudentTeacherRecurrent"
+    """The policy class name. Default is StudentTeacherRecurrent."""
+
+    rnn_type: str = MISSING
+    """The type of the RNN network. Either "lstm" or "gru"."""
+
+    rnn_hidden_dim: int = MISSING
+    """The hidden dimension of the RNN network."""
+
+    rnn_num_layers: int = MISSING
+    """The number of layers of the RNN network."""
+
+    teacher_recurrent: bool = MISSING
+    """Whether the teacher network is recurrent too."""
+
+
+############################
+# Algorithm configurations #
+############################
+
+
+@configclass
+class RslRlPpoAlgorithmCfg:
+    """Configuration for the PPO algorithm."""
+
+    class_name: str = "PPO"
+    """The algorithm class name. Default is PPO."""
 
     num_learning_epochs: int = MISSING
     """The number of learning epochs per update."""
@@ -72,12 +131,24 @@ class RslRlPpoAlgorithmCfg:
     lam: float = MISSING
     """The lambda parameter for Generalized Advantage Estimation (GAE)."""
 
+    entropy_coef: float = MISSING
+    """The coefficient for the entropy loss."""
+
     desired_kl: float = MISSING
     """The desired KL divergence."""
 
     max_grad_norm: float = MISSING
     """The maximum gradient norm."""
 
+    value_loss_coef: float = MISSING
+    """The coefficient for the value loss."""
+
+    use_clipped_value_loss: bool = MISSING
+    """Whether to use clipped value loss."""
+
+    clip_param: float = MISSING
+    """The clipping parameter for the policy."""
+
     normalize_advantage_per_mini_batch: bool = False
     """Whether to normalize the advantage per mini-batch. Default is False.
 
@@ -93,6 +164,27 @@ class RslRlPpoAlgorithmCfg:
     in which case RND is not used.
     """
 
+@configclass
+class RslRlDistillationAlgorithmCfg:
+    """Configuration for the distillation algorithm."""
+
+    class_name: str = "Distillation"
+    """The algorithm class name. Default is Distillation."""
+
+    num_learning_epochs: int = MISSING
+    """The number of updates performed with each sample."""
+
+    learning_rate: float = MISSING
+    """The learning rate for the student policy."""
+
+    gradient_length: float = MISSING
+    """The number of environment steps the gradient flows back."""
+
+
+#########################
+# Runner configurations #
+#########################
+
 
 @configclass
 class RslRlOnPolicyRunnerCfg:
@@ -113,10 +205,10 @@ class RslRlOnPolicyRunnerCfg:
     empirical_normalization: bool = MISSING
     """Whether to use empirical normalization."""
 
-    policy: RslRlPpoActorCriticCfg = MISSING
+    policy: RslRlPpoActorCriticCfg | RslRlDistillationStudentTeacherCfg = MISSING
     """The policy configuration."""
 
-    algorithm: RslRlPpoAlgorithmCfg = MISSING
+    algorithm: RslRlPpoAlgorithmCfg | RslRlDistillationAlgorithmCfg = MISSING
     """The algorithm configuration."""
 
     clip_actions: float | None = None
@@ -161,13 +253,10 @@ class RslRlOnPolicyRunnerCfg:
     # Loading parameters
     ##
 
-    resume: bool = False
-    """Whether to resume. Default is False."""
-
-    load_run: str = ".*"
-    """The run directory to load. Default is ".*" (all).
+    load_run: None | str = None
+    """The run directory to load. Default is None.
 
-    If regex expression, the latest (alphabetical order) matching run will be loaded.
+    If regex expression, the latest (alphabetical order) matching run will be loaded, e.g, ".*" (all).
     """
 
     load_checkpoint: str = "model_.*.pt"

source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/agents/rsl_rl_ppo_cfg.py

@@ -15,7 +15,6 @@ class FrankaReachPPORunnerCfg(RslRlOnPolicyRunnerCfg):
     save_interval = 50
     experiment_name = "franka_reach"
     run_name = ""
-    resume = False
     empirical_normalization = False
     policy = RslRlPpoActorCriticCfg(
         init_noise_std=1.0,

source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/ur_10/agents/rsl_rl_ppo_cfg.py

@@ -15,7 +15,6 @@ class UR10ReachPPORunnerCfg(RslRlOnPolicyRunnerCfg):
     save_interval = 50
     experiment_name = "reach_ur10"
     run_name = ""
-    resume = False
     empirical_normalization = False
     policy = RslRlPpoActorCriticCfg(
         init_noise_std=1.0,

source/isaaclab_tasks/isaaclab_tasks/utils/parse_cfg.py

@@ -137,7 +137,7 @@ def parse_env_cfg(
 
 
 def get_checkpoint_path(
-    log_path: str, run_dir: str = ".*", checkpoint: str = ".*", other_dirs: list[str] = None, sort_alpha: bool = True
+    log_path: str, run_dir: str | None, checkpoint: str = ".*", other_dirs: list[str] = None, sort_alpha: bool = True
 ) -> str:
     """Get path to the model checkpoint in input directory.
 
@@ -164,8 +164,12 @@ def get_checkpoint_path(
     Raises:
         ValueError: When no runs are found in the input directory.
         ValueError: When no checkpoints are found in the input directory.
-
     """
+
+    # select the most recent run directory if run_dir is none
+    if run_dir is None:
+        run_dir = ".*"
+
     # check if runs present in directory
     try:
         # find all runs in the directory that math the regex expression