minor cleanup adding docstrigs

Author: frlai

scripts/reinforcement_learning/rsl_rl/export.py

@@ -77,12 +77,6 @@
     default=False,
     help="Disable LEAPP graph visualization during compile_graph().",
 )
-parser.add_argument(
-    "--disable_automatic_module_annotation",
-    action="store_true",
-    default=False,
-    help="Disables automatic detection and annotation of modules that have internal states",
-)
 
 # append RSL-RL cli arguments
 cli_args.add_rsl_rl_args(parser)
@@ -123,6 +117,7 @@
 
 
 def get_actor_memory_module(policy_nn):
+    """Return the actor-side recurrent memory module when the policy exposes one."""
     if hasattr(policy_nn, "memory_a"):
         return policy_nn.memory_a
     if hasattr(policy_nn, "memory_s"):
@@ -131,6 +126,7 @@ def get_actor_memory_module(policy_nn):
 
 
 def ensure_actor_hidden_state_initialized(policy_nn, batch_size: int, device: torch.device, dtype: torch.dtype):
+    """Initialize and return the actor hidden state when a recurrent policy has not created it yet."""
     actor_state, _ = policy_nn.get_hidden_states()
     if actor_state is not None:
         return actor_state
@@ -151,6 +147,7 @@ def ensure_actor_hidden_state_initialized(policy_nn, batch_size: int, device: to
 
 
 def state_dict_from_actor_hidden(actor_hidden):
+    """Convert the actor hidden state into the named tensor mapping expected by LEAPP state APIs."""
     if actor_hidden is None:
         return {}
     if isinstance(actor_hidden, tuple):
@@ -159,6 +156,7 @@ def state_dict_from_actor_hidden(actor_hidden):
 
 
 def actor_hidden_from_registered(registered_state, original_hidden):
+    """Restore the registered LEAPP state to the hidden-state structure expected by the actor memory module."""
     if isinstance(original_hidden, tuple):
         if isinstance(registered_state, tuple):
             return registered_state

source/isaaclab/isaaclab/utils/leapp/export_annotator.py

@@ -107,7 +107,7 @@ def __init__(self, export_method: str, required_obs_groups: set[str] | None = No
         self._action_term_scene_keys: dict[str, str] = {}
 
     def setup(self, env):
-        """Patch observation and action managers on the unwrapped env."""
+        """Patch the unwrapped env in place for LEAPP-aware observation and action export."""
         unwrapped = env.env.unwrapped
         self.task_name = unwrapped.spec.id
 
@@ -551,9 +551,12 @@ def patch_env_for_export(
       reads **and** ``Articulation`` write methods.
 
     Data properties are resolved lazily through proxies — no hardcoded
-    class list is required. Properties with ``_leapp_semantics`` produce
-    rich annotations; properties without it are still traced so that no
-    tensor is silently baked as a constant.
+    class list is required. To produce LEAPP input annotations, the
+    accessed data property getter must carry ``_leapp_semantics``.
+    Likewise, action-side write methods must be annotated to produce
+    semantic LEAPP outputs. Undecorated reads and writes are forwarded
+    as normal runtime access, but they do not gain semantic annotation
+    metadata through this patching path.
 
     State reads are deduplicated across observation and action paths via a
     shared cache, so a property like ``joint_pos`` that is read by both an