fix(cusrl): fix cusrl train and play scripts; update cusrl distillation configuration

Author: chengruiz

scripts/reinforcement_learning/cusrl/play.py

@@ -11,10 +11,6 @@
 
 from isaaclab.app import AppLauncher
 
-# local imports
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
-from rl_utils import camera_follow
-
 # add argparse arguments
 parser = argparse.ArgumentParser(description="Evaluate an RL agent with CusRL.")
 parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
@@ -56,7 +52,6 @@
 
 import cusrl
 import gymnasium as gym
-import robot_lab.tasks  # noqa: F401
 import torch
 from cusrl.environment.isaaclab import TrainerCfg
 
@@ -73,6 +68,12 @@
 
 from isaaclab_tasks.utils.hydra import hydra_task_config  # noqa: F401
 
+import robot_lab.tasks  # noqa: F401  # isort: skip
+
+# local imports
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+from rl_utils import camera_follow
+
 
 class CameraFollowPlayerHook(cusrl.Player.Hook):
     def step(self, step: int, transition: dict, metrics: dict):

scripts/reinforcement_learning/cusrl/train.py

@@ -58,7 +58,6 @@
 
 import cusrl
 import gymnasium as gym
-import robot_lab.tasks  # noqa: F401
 import torch
 from cusrl.environment.isaaclab import TrainerCfg
 
@@ -73,6 +72,8 @@
 
 from isaaclab_tasks.utils.hydra import hydra_task_config  # noqa: F401
 
+import robot_lab.tasks  # noqa: F401  # isort: skip
+
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 torch.backends.cudnn.deterministic = False
@@ -128,7 +129,7 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
         agent_factory=agent_cfg.agent_factory.override(
             device=args_cli.device, autocast=args_cli.autocast, compile=args_cli.compile
         ),
-        logger_factory=cusrl.make_logger_factory(args_cli.logger, log_dir, add_datetime_prefix=False),
+        logger_factory=cusrl.make_logger_factory(args_cli.logger, log_dir, name=None),
         num_iterations=agent_cfg.max_iterations,
         save_interval=agent_cfg.save_interval,
         checkpoint_path=args_cli.checkpoint,

source/robot_lab/robot_lab/tasks/manager_based/locomotion/velocity/config/quadruped/anymal_d/agents/cusrl_distillation_cfg.py

@@ -28,7 +28,7 @@ class AnymalDFlatDistillationTrainerCfg(TrainerCfg):
         hooks=[
             cusrl.hook.ModuleInitialization(init_actor=False, init_critic=False, distribution_std=0.1),
             cusrl.hook.OnPolicyPreparation(),
-            cusrl.hook.PolicyDistillationLoss(""),
+            cusrl.hook.PolicyDistillation(""),
             cusrl.hook.GradientClipping(1.0),
         ],
     )