Adds logdir configuration parameter to environments (#3391)

# Description

Adds a logdir parameter for all environment configs to allow passing of
the log dir for each experiment from the training/inferencing scripts to
the environment object.

This allows environments to access the logdir for the run and store log
files in there, such as from the feature extractor.


## Type of change

<!-- As you go through the list, delete the ones that are not
applicable. -->

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


## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] 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
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

---------
Signed-off-by: Kelly Guo <kellyg@nvidia.com>
Signed-off-by: Kelly Guo <kellyguo123@hotmail.com>
Co-authored-by: Mayank Mittal <12863862+Mayankm96@users.noreply.github.com>

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

@@ -165,6 +165,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,6 +90,9 @@ class FeatureExtractor:
         self.feature_extractor.to(self.device)
 
         self.step_count = 0
+        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