Exporting RNN policy that is not an LSTM (RSL RL)

[WT-MM]

What

RSL RL provides a recurrent actor critic with LSTM or GRU. However, the exporter in isaaclab_rl/rsl_rl/exporter.py assumes all recurrent networks are LSTMs. (see code snippets)

@configclass
class RslRlPpoActorCriticRecurrentCfg(RslRlPpoActorCriticCfg):
    """Configuration for the PPO actor-critic networks with recurrent layers."""

    class_name: str = "ActorCriticRecurrent"
    """The policy class name. Default is ActorCriticRecurrent."""

    rnn_type: str = MISSING
    """The type of RNN to use. Either "lstm" or "gru"."""

    rnn_hidden_dim: int = MISSING
    """The dimension of the RNN layers."""

    rnn_num_layers: int = MISSING
    """The number of RNN layers."""

class _TorchPolicyExporter(torch.nn.Module):
    """Exporter of actor-critic into JIT file."""

    def __init__(self, policy, normalizer=None):
        super().__init__()
        self.is_recurrent = policy.is_recurrent
        # copy policy parameters
        if hasattr(policy, "actor"):
            self.actor = copy.deepcopy(policy.actor)
            if self.is_recurrent:
                self.rnn = copy.deepcopy(policy.memory_a.rnn)
        elif hasattr(policy, "student"):
            self.actor = copy.deepcopy(policy.student)
            if self.is_recurrent:
                self.rnn = copy.deepcopy(policy.memory_s.rnn)
        else:
            raise ValueError("Policy does not have an actor/student module.")
        # set up recurrent network
        if self.is_recurrent:
            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))
            self.forward = self.forward_lstm
            self.reset = self.reset_memory
        # copy normalizer if exists
        if normalizer:
            self.normalizer = copy.deepcopy(normalizer)
        else:
            self.normalizer = torch.nn.Identity()

    def forward_lstm(self, x):
        x = self.normalizer(x)
        x, (h, c) = self.rnn(x.unsqueeze(0), (self.hidden_state, self.cell_state))
        self.hidden_state[:] = h
        self.cell_state[:] = c
        x = x.squeeze(0)
        return self.actor(x)

    def forward(self, x):
        return self.actor(self.normalizer(x))

    @torch.jit.export
    def reset(self):
        pass

    def reset_memory(self):
        self.hidden_state[:] = 0.0
        self.cell_state[:] = 0.0

    def export(self, path, filename):
        os.makedirs(path, exist_ok=True)
        path = os.path.join(path, filename)
        self.to("cpu")
        traced_script_module = torch.jit.script(self)
        traced_script_module.save(path)

Steps to reproduce

I'm working off of a fork for K-Scale Labs' K-Bot here -- https://github.com/kscalelabs/IsaacLab/blob/main/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/kbot/agents/rsl_rl_ppo_cfg.py

I added a new agent with the following config:

@configclass
class KBotRoughRNNPPORunnerCfg(RslRlOnPolicyRunnerCfg):
    num_steps_per_env = 24
    max_iterations = 3000
    save_interval = 50
    experiment_name = "kbot_rough_rnn"
    empirical_normalization = False
    policy = RslRlPpoActorCriticRecurrentCfg(
        init_noise_std=1.0,
        actor_hidden_dims=[512, 256, 128],
        critic_hidden_dims=[512, 256, 128],
        activation="elu",
        rnn_type="gru",
        rnn_hidden_dim=512,
        rnn_num_layers=2,
    )
    algorithm = RslRlPpoAlgorithmCfg(
        value_loss_coef=1.0,
        use_clipped_value_loss=True,
        clip_param=0.2,
        entropy_coef=0.008,
        num_learning_epochs=5,
        num_mini_batches=4,
        learning_rate=1.0e-3,
        schedule="adaptive",
        gamma=0.99,
        lam=0.95,
        desired_kl=0.01,
        max_grad_norm=1.0,
    )

And registered the new tasks like so:

gym.register(
    id="Isaac-Velocity-Rough-Kbot-RNN-v0",
    entry_point="isaaclab.envs:ManagerBasedRLEnv",
    disable_env_checker=True,
    kwargs={
        "env_cfg_entry_point": f"{__name__}.rough_env_cfg:KBotRoughEnvCfg",
        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:KBotRoughRNNPPORunnerCfg",
        "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml",
    },
)

gym.register(
    id="Isaac-Velocity-Rough-Kbot-RNN-v0-Play",
    entry_point="isaaclab.envs:ManagerBasedRLEnv",
    disable_env_checker=True,
    kwargs={
        "env_cfg_entry_point": f"{__name__}.rough_env_cfg:KBotRoughEnvCfg_PLAY",
        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:KBotRoughRNNPPORunnerCfg",
        "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml",
    },
)

Then I ran the train script to get an initial .pt checkpoint (python scripts/reinforcement_learning/rsl_rl/train.py --task=Isaac-Velocity-Rough-Kbot-RNN-v0)

and then tried to play the checkpoint: python scripts/reinforcement_learning/rsl_rl/play.py --task=Isaac-Velocity-Rough-Kbot-RNN-v0 --num_envs=4

I received this error message:

[INFO]: Loading model checkpoint from: /home/wesley/Github/IsaacLab/logs/rsl_rl/kbot_rough_rnn/2025-07-23_11-44-46/model_100.pt
ActorCriticRecurrent.__init__ got unexpected arguments, which will be ignored: dict_keys(['noise_std_type'])
Actor MLP: Sequential(
  (0): Linear(in_features=512, out_features=512, bias=True)
  (1): ELU(alpha=1.0)
  (2): Linear(in_features=512, out_features=256, bias=True)
  (3): ELU(alpha=1.0)
  (4): Linear(in_features=256, out_features=128, bias=True)
  (5): ELU(alpha=1.0)
  (6): Linear(in_features=128, out_features=20, bias=True)
)
Critic MLP: Sequential(
  (0): Linear(in_features=512, out_features=512, bias=True)
  (1): ELU(alpha=1.0)
  (2): Linear(in_features=512, out_features=256, bias=True)
  (3): ELU(alpha=1.0)
  (4): Linear(in_features=256, out_features=128, bias=True)
  (5): ELU(alpha=1.0)
  (6): Linear(in_features=128, out_features=1, bias=True)
)
Actor RNN: Memory(
  (rnn): GRU(69, 512, num_layers=2)
)
Critic RNN: Memory(
  (rnn): GRU(328, 512, num_layers=2)
)
Traceback (most recent call last):
  File "/home/wesley/Github/IsaacLab/scripts/reinforcement_learning/rsl_rl/play.py", line 169, in <module>
    main()
  File "/home/wesley/Github/IsaacLab/scripts/reinforcement_learning/rsl_rl/play.py", line 133, in main
    export_policy_as_jit(policy_nn, ppo_runner.obs_normalizer, path=export_model_dir, filename="policy.pt")
  File "/home/wesley/Github/IsaacLab/source/isaaclab_rl/isaaclab_rl/rsl_rl/exporter.py", line 21, in export_policy_as_jit
    policy_exporter.export(path, filename)
  File "/home/wesley/Github/IsaacLab/source/isaaclab_rl/isaaclab_rl/rsl_rl/exporter.py", line 100, in export
    traced_script_module = torch.jit.script(self)
  File "/home/wesley/.conda/envs/isaaclab/lib/python3.10/site-packages/torch/jit/_script.py", line 1429, in script
    ret = _script_impl(
  File "/home/wesley/.conda/envs/isaaclab/lib/python3.10/site-packages/torch/jit/_script.py", line 1147, in _script_impl
    return torch.jit._recursive.create_script_module(
  File "/home/wesley/.conda/envs/isaaclab/lib/python3.10/site-packages/torch/jit/_recursive.py", line 557, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/home/wesley/.conda/envs/isaaclab/lib/python3.10/site-packages/torch/jit/_recursive.py", line 634, in create_script_module_impl
    create_methods_and_properties_from_stubs(
  File "/home/wesley/.conda/envs/isaaclab/lib/python3.10/site-packages/torch/jit/_recursive.py", line 466, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(
RuntimeError: 
Arguments for call are not valid.
The following variants are available:
  
  forward__0(__torch__.torch.nn.modules.rnn.GRU self, Tensor input, Tensor? hx=None) -> ((Tensor, Tensor)):
  Expected a value of type 'Optional[Tensor]' for argument 'hx' but instead found type 'Tuple[Tensor, Tensor]'.
  
  forward__1(__torch__.torch.nn.modules.rnn.GRU self, __torch__.torch.nn.utils.rnn.PackedSequence input, Tensor? hx=None) -> ((__torch__.torch.nn.utils.rnn.PackedSequence, Tensor)):
  Expected a value of type '__torch__.torch.nn.utils.rnn.PackedSequence' for argument 'input' but instead found type 'Tensor'.

The original call is:
  File "/home/wesley/Github/IsaacLab/source/isaaclab_rl/isaaclab_rl/rsl_rl/exporter.py", line 79
    def forward_lstm(self, x):
        x = self.normalizer(x)
        x, (h, c) = self.rnn(x.unsqueeze(0), (self.hidden_state, self.cell_state))
                    ~~~~~~~~ <--- HERE
        self.hidden_state[:] = h
        self.cell_state[:] = c

Which only occurs because the play.py script also exports the policy as TorchScript and ONNX:

Info

• Commit: kscalelabs@d128018

This is probably a simple fix that I can patch in if there is not enough bandwidth by other people.

[WT-MM]

Actually I've drafted a fix here #3009

[RandomOakForest]

Thanks for posting this. It seems this will require a closer look. But we welcome any PRs we could review on this subject.

[WT-MM]

Thanks for posting this. It seems this will require a closer look. But we welcome any PRs we could review on this subject.

Yep! Opened a fix here -- #3009