make freeze optional

Author: pascal-roth

exts/nav_tasks/nav_tasks/mdp/actions/navigation_actions.py

@@ -30,7 +30,9 @@ def __init__(self, cfg: NavigationSE2ActionCfg, env: ManagerBasedRLEnv):
         # load policies
         file_bytes = read_file(self.cfg.low_level_policy_file)
         self.low_level_policy = torch.jit.load(file_bytes, map_location=self.device)
-        self.low_level_policy = torch.jit.freeze(self.low_level_policy.eval())
+        self.low_level_policy.eval()
+        if self.cfg.freeze_low_level_policy:
+            self.low_level_policy = torch.jit.freeze(self.low_level_policy)
 
         # prepare joint position actions
         if not isinstance(self.cfg.low_level_action, list):

exts/nav_tasks/nav_tasks/mdp/actions/navigation_actions_cfg.py

@@ -28,6 +28,11 @@ class NavigationSE2ActionCfg(ActionTermCfg):
     low_level_policy_file: str = MISSING
     """Path to the low level policy file."""
 
+    freeze_low_level_policy: bool = True
+    """Whether to freeze the low level policy.
+
+    Can improve performance but will also eliminate possible functions such as `reset`."""
+
     low_level_obs_group: str = "low_level_policy"
     """Observation group of the low level policy."""