Frlai/fix rnn leapp tracing

[frlai]

Description

Fixes leapp lstm policy export. Previously it hooked onto structures in rslrl=3.1.2. The update now makes it work for rslrl=5.0.1, same version as the one in the pyproject.toml. Added a regression test to test if rslrl rnns work new test generates a dummy policy with rnn and tests the export. Previous it was assumed that there was one policy with a pretrained checkpoint that had an rnn but turns out there isn't any that by default run rnns so this error slipped through.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist

• I have read and understood the contribution guidelines
• I have run the pre-commit checks with ./isaaclab.sh --format
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• I have updated the changelog and the corresponding version in the extension's config/extension.toml file
• I have added my name to the CONTRIBUTORS.md or my name already exists there

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot — PR #6000

Summary: This PR updates the LEAPP RNN tracing export to work with RSL-RL 5.x's RNNModel API (which exposes .rnn and .get_hidden_state()) instead of the old 3.1.2 API (memory_a/memory_s/get_hidden_states()). It also adds a well-structured regression test.

✅ Strengths

• Clean abstraction: The new helper functions (get_actor_memory_module, is_actor_recurrent_policy, get_actor_hidden_state, set_actor_hidden_state) provide a single point of adaptation for the RSL-RL recurrent API, making future version changes easier.
• Regression test included: The new test file properly mocks the RSL-RL 5.x module structure and validates the state-handling helpers end-to-end.
• Backward compatibility consideration: Using hasattr checks and getattr with defaults makes the code resilient to missing attributes.

⚠️ Findings

1. policy_nn = policy assignment (line 321): The old code used getattr(policy, "__self__", None) to unwrap the bound method returned by get_inference_policy(). The new code assigns policy_nn = policy directly. This works if RSL-RL 5.x's get_inference_policy() now returns the module itself (not a bound method), but if it ever returns a bound method, policy_nn.parameters() and hasattr(policy_nn, "rnn") would fail silently. A brief comment explaining why __self__ unwrapping is no longer needed would help future maintainers.

2. Redundant torch_module alias (line 217): The torch_module = torch + assert torch_module is not None pattern seems like a workaround for type-narrowing. Consider adding a comment or using a more idiomatic pattern like assert torch is not None directly.

3. get_actor_hidden_state has dual paths: It first checks get_hidden_state() (method on policy_nn), then falls back to memory.hidden_state. If get_hidden_state() exists but returns stale data while memory.hidden_state is the ground truth, this ordering could be surprising. Document which path is canonical for RSL-RL 5.x.

4. Test file import fragility: The test uses importlib to load the export script as a module with mocked dependencies. This is clever but fragile — if the export script adds new top-level imports, the test will break without clear errors. Consider adding a try/except ImportError wrapper with a descriptive skip message.

5. Missing CI check: The Check changelog fragments CI check is failing. Ensure a changelog fragment is added before merge.

📋 Verdict

The fix is correct and well-scoped for the RSL-RL 3.x → 5.x migration. The key logic changes are sound — the new accessor pattern properly handles LSTM tuple states and GRU single-tensor states. No breaking side effects for non-RNN LEAPP tracing (guarded by is_actor_recurrent_policy). Minor documentation suggestions above.

---

Update (388195d): New commits add documentation improvements only:

• Updated LEAPP version requirement to >=0.5.2
• Fixed broken doc link (now points to https://nvidia-isaac.github.io/leapp/)
• Added usage example for deploy.py with --viz kit option

✅ No code changes — previous inline comments on export.py and test file still apply. No new issues introduced.

---

Update (7601453): New commit adds changelog fragment (source/isaaclab_rl/changelog.d/frlai-fix_rnn_leapp_tracing.rst).

✅ Finding #5 (Missing CI check): Addressed — changelog fragment now included.

No code changes in this commit. Previous inline comments on export.py and the test file remain applicable.

[isaaclab-review-bot]

Observation: The old code used getattr(policy, "__self__", None) to extract the underlying nn.Module from the bound method returned by get_inference_policy(). Now policy_nn = policy directly.

This works if RSL-RL 5.x's get_inference_policy() returns the module object itself (which appears to be the case based on the .rnn / .is_recurrent access pattern), but a brief comment here explaining the API change would help future readers understand why the unwrapping was removed.

[isaaclab-review-bot]

Nit: get_actor_hidden_state prefers policy_nn.get_hidden_state() over memory.hidden_state. Since set_actor_hidden_state always writes to memory.hidden_state, there's a subtle contract: get_hidden_state() must read from memory.hidden_state internally for round-trip consistency. This is true for RSL-RL 5.x's RNNModel, but worth a brief docstring note for robustness.

[isaaclab-review-bot]

Minor: torch_module = torch + assert torch_module is not None appears to be a type-narrowing workaround for the module-level torch = None pattern. Consider a short comment (e.g., # narrow type after _load_runtime_dependencies()) to clarify intent.

[isaaclab-review-bot]

Suggestion: This dynamic import approach is clever but fragile. If the export script gains new unresolvable imports, this will fail with an opaque error. Consider wrapping the spec.loader.exec_module(module) call with a more descriptive error message, or adding a pytest.importorskip-style guard to produce a clear skip reason in CI.

[greptile-apps]

Greptile Summary

This PR fixes the LEAPP LSTM policy export to work with RSL-RL 5.0.1 (the version pinned in pyproject.toml) by replacing the old 3.x attribute lookups (memory_a, memory_s, get_hidden_states()) with the new 5.x structure (.rnn, get_hidden_state()). It also corrects policy_nn assignment to use the inference policy object directly rather than extracting a bound-method's __self__.

• New helper functions (is_actor_recurrent_policy, get_actor_hidden_state, set_actor_hidden_state) encapsulate version-specific RNN introspection and are tested by a new regression test that stubs out the Isaac Sim/isaaclab runtime dependencies.
• The regression test creates a minimal RSL-RL 5.x RNNModel mock and validates the full hidden-state round-trip (initialize → register → restore → verify).

Confidence Score: 4/5

The fix correctly adapts the RNN hidden-state plumbing from RSL-RL 3.x to 5.x; the core logic is straightforward and backed by a new regression test.

The assert torch_module is not None guard in ensure_actor_hidden_state_initialized would silently vanish under Python's -O flag, producing a cryptic downstream error rather than a clear message. The getter/setter asymmetry is not a current defect but leaves the code inconsistently structured. Both are non-blocking.

scripts/reinforcement_learning/leapp/rsl_rl/export.py — the assert guard and setter asymmetry in the new helper functions warrant a closer look.

Important Files Changed

Filename	Overview
scripts/reinforcement_learning/leapp/rsl_rl/export.py	Adapts RNN policy introspection from RSL-RL 3.x attributes (memory_a/memory_s) to RSL-RL 5.x structure (.rnn); introduces get/set/is_recurrent helpers and fixes policy_nn assignment to use the inference policy object directly.
source/isaaclab_rl/test/export/test_rsl_rl_export_recurrent_state.py	New regression test that loads the export script via importlib (mocking isaaclab deps) and validates recurrent state helpers with a minimal RSL-RL 5.x RNNModel mock.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[export_rsl_rl_agent] --> B[runner.get_inference_policy]
    B --> C[policy_nn = policy]
    C --> D{is_actor_recurrent_policy?}
    D -->|No| E[actions = policy obs]
    D -->|Yes| F[ensure_actor_hidden_state_initialized]
    F --> G{get_actor_hidden_state}
    G -->|has get_hidden_state method| H[policy_nn.get_hidden_state]
    G -->|no method| I[memory.hidden_state]
    H --> J[annotate.state_tensors]
    I --> J
    J --> K[set_actor_hidden_state]
    K --> E
    E --> L{is_actor_recurrent_policy?}
    L -->|Yes| M[get_actor_hidden_state after step]
    M --> N[annotate.update_state]
    L -->|No| O[env.step]
    N --> O

Loading

_{Reviews (1): Last reviewed commit: "removed bindings for rslrl 3.1.2" | Re-trigger Greptile}

[greptile-apps]

Using assert for a runtime guard is unsafe here: Python optimisations (-O or -OO) silently disable all assert statements, so the check would be skipped and torch_module would remain None, causing an obscure AttributeError on torch_module.zeros(...) rather than a clear error message. Raise an explicit RuntimeError instead.

Suggested change

	torch_module = torch
	assert torch_module is not None
	actor_state = get_actor_hidden_state(policy_nn)
	torch_module = torch
	if torch_module is None:
	raise RuntimeError("Runtime dependencies not loaded; call _load_runtime_dependencies() before this function.")
	actor_state = get_actor_hidden_state(policy_nn)

[greptile-apps]

Asymmetry between getter and setter accessor strategies

get_actor_hidden_state prefers the API method get_hidden_state() when present, then falls back to the raw attribute. set_actor_hidden_state only uses raw attribute access (memory.hidden_state = ...), so if a future RSL-RL version exposes a set_hidden_state() method with extra bookkeeping, the setter would silently bypass it. Consider applying the same API-first pattern for consistency: check for set_hidden_state() before falling back to direct attribute assignment.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[StafaH]

Left some minor comments to fix, otherwise LGTM

[StafaH]

This pattern seems odd, why do we need to do it this way? If torch is required just do the common:

try:
   import torch
except:
   raise ...

But even that might not be needed, if torch is not installed the error is quire obvious to the user to install torch

[StafaH]

Can we remove this line? Seems like a no-op, just use policy everywhere below?

[StafaH]

It might be better to just an add a RNN test to the previously existing test file, this would allow re-use and reduce code duplication, and it would fit in with the other set of tests