Adds logdir configuration parameter to environments

scripts/reinforcement_learning/rl_games/play.py

@@ -130,6 +130,9 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
         resume_path = retrieve_file_path(args_cli.checkpoint)
     log_dir = os.path.dirname(os.path.dirname(resume_path))
 
+    # set the log directory for the environment (works for all environment types)
+    env_cfg.log_dir = log_dir
+
     # wrap around environment for rl-games
     rl_device = agent_cfg["params"]["config"]["device"]
     clip_obs = agent_cfg["params"]["env"].get("clip_observations", math.inf)

scripts/reinforcement_learning/rl_games/train.py

@@ -160,6 +160,9 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
             "IO descriptors are only supported for manager based RL environments. No IO descriptors will be exported."
         )
 
+    # set the log directory for the environment (works for all environment types)
+    env_cfg.log_dir = log_dir
+
     # create isaac environment
     env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
 

scripts/reinforcement_learning/rsl_rl/play.py

@@ -112,6 +112,9 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
 
     log_dir = os.path.dirname(resume_path)
 
+    # set the log directory for the environment (works for all environment types)
+    env_cfg.log_dir = log_dir
+
     # create isaac environment
     env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
 

scripts/reinforcement_learning/rsl_rl/train.py

@@ -150,6 +150,9 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
             "IO descriptors are only supported for manager based RL environments. No IO descriptors will be exported."
         )
 
+    # set the log directory for the environment (works for all environment types)
+    env_cfg.log_dir = log_dir
+
     # create isaac environment
     env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
 

scripts/reinforcement_learning/sb3/play.py

@@ -127,6 +127,9 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
         checkpoint_path = args_cli.checkpoint
     log_dir = os.path.dirname(checkpoint_path)
 
+    # set the log directory for the environment (works for all environment types)
+    env_cfg.log_dir = log_dir
+
     # create isaac environment
     env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
 

scripts/reinforcement_learning/sb3/train.py

@@ -152,6 +152,9 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
             "IO descriptors are only supported for manager based RL environments. No IO descriptors will be exported."
         )
 
+    # set the log directory for the environment (works for all environment types)
+    env_cfg.log_dir = log_dir
+
     # create isaac environment
     env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
 

scripts/reinforcement_learning/skrl/play.py

@@ -165,6 +165,9 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, expe
         )
     log_dir = os.path.dirname(os.path.dirname(resume_path))
 
+    # set the log directory for the environment (works for all environment types)
+    env_cfg.log_dir = log_dir
+
     # create isaac environment
     env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
 

scripts/reinforcement_learning/skrl/train.py

@@ -182,6 +182,9 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
             "IO descriptors are only supported for manager based RL environments. No IO descriptors will be exported."
         )
 
+    # set the log directory for the environment (works for all environment types)
+    env_cfg.log_dir = log_dir
+
     # create isaac environment
     env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
 

source/isaaclab/isaaclab/envs/direct_marl_env_cfg.py

@@ -225,3 +225,6 @@ class DirectMARLEnvCfg:
 
     xr: XrCfg | None = None
     """Configuration for viewing and interacting with the environment through an XR device."""
+
+    log_dir: str | None = None
+    """Directory for logging experiment artifacts. Defaults to None, in which case no specific log directory is set."""

source/isaaclab/isaaclab/envs/direct_rl_env_cfg.py

@@ -229,3 +229,6 @@ class DirectRLEnvCfg:
 
     xr: XrCfg | None = None
     """Configuration for viewing and interacting with the environment through an XR device."""
+
+    log_dir: str | None = None
+    """Directory for logging experiment artifacts. Defaults to None, in which case no specific log directory is set."""

source/isaaclab/isaaclab/envs/manager_based_env_cfg.py

@@ -131,3 +131,6 @@ class ManagerBasedEnvCfg:
 
     io_descriptors_output_dir: str | None = None
     """The directory to export the IO descriptors to. Defaults to None."""
+
+    log_dir: str | None = None
+    """Directory for logging experiment artifacts. Defaults to None, in which case no specific log directory is set."""

source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py

@@ -73,12 +73,13 @@ class FeatureExtractor:
     If the train flag is set to True, the CNN is trained during the rollout process.
     """
 
-    def __init__(self, cfg: FeatureExtractorCfg, device: str):
+    def __init__(self, cfg: FeatureExtractorCfg, device: str, log_dir: str | None = None):
         """Initialize the feature extractor model.
 
         Args:
-            cfg (FeatureExtractorCfg): Configuration for the feature extractor model.
-            device (str): Device to run the model on.
+            cfg: Configuration for the feature extractor model.
+            device: Device to run the model on.
+            log_dir: Directory to save checkpoints. If None, uses local "logs" folder resolved with respect to this file.
         """
 
         self.cfg = cfg
@@ -89,7 +90,10 @@ def __init__(self, cfg: FeatureExtractorCfg, device: str):
         self.feature_extractor.to(self.device)
 
         self.step_count = 0
-        self.log_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "logs")
+        if log_dir is not None:
+            self.log_dir = log_dir
+        else:
+            self.log_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "logs")
         if not os.path.exists(self.log_dir):
             os.makedirs(self.log_dir)
 

source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py

@@ -65,7 +65,8 @@ class ShadowHandVisionEnv(InHandManipulationEnv):
 
     def __init__(self, cfg: ShadowHandVisionEnvCfg, render_mode: str | None = None, **kwargs):
         super().__init__(cfg, render_mode, **kwargs)
-        self.feature_extractor = FeatureExtractor(self.cfg.feature_extractor, self.device)
+        # Use the log directory from the configuration
+        self.feature_extractor = FeatureExtractor(self.cfg.feature_extractor, self.device, self.cfg.log_dir)
         # hide goal cubes
         self.goal_pos[:, :] = torch.tensor([-0.2, 0.1, 0.6], device=self.device)
         # keypoints buffer