From 31b2e224e15a3006a35dd57fbfad140073c3d711 Mon Sep 17 00:00:00 2001 From: Kelly Guo Date: Mon, 8 Sep 2025 14:32:55 -0700 Subject: [PATCH 1/4] Passes logdir from train/play scripts to environments --- scripts/reinforcement_learning/rl_games/play.py | 3 +++ scripts/reinforcement_learning/rl_games/train.py | 3 +++ scripts/reinforcement_learning/rsl_rl/play.py | 3 +++ scripts/reinforcement_learning/rsl_rl/train.py | 3 +++ scripts/reinforcement_learning/sb3/play.py | 3 +++ scripts/reinforcement_learning/sb3/train.py | 3 +++ scripts/reinforcement_learning/skrl/play.py | 3 +++ scripts/reinforcement_learning/skrl/train.py | 3 +++ source/isaaclab/isaaclab/envs/direct_marl_env_cfg.py | 3 +++ source/isaaclab/isaaclab/envs/direct_rl_env_cfg.py | 3 +++ source/isaaclab/isaaclab/envs/manager_based_env_cfg.py | 3 +++ .../direct/shadow_hand/feature_extractor.py | 8 ++++++-- .../direct/shadow_hand/shadow_hand_vision_env.py | 3 ++- 13 files changed, 41 insertions(+), 3 deletions(-) diff --git a/scripts/reinforcement_learning/rl_games/play.py b/scripts/reinforcement_learning/rl_games/play.py index dd2185b82b0..135980e92c7 100644 --- a/scripts/reinforcement_learning/rl_games/play.py +++ b/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) diff --git a/scripts/reinforcement_learning/rl_games/train.py b/scripts/reinforcement_learning/rl_games/train.py index cc1e54b1756..a73886345fe 100644 --- a/scripts/reinforcement_learning/rl_games/train.py +++ b/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) diff --git a/scripts/reinforcement_learning/rsl_rl/play.py b/scripts/reinforcement_learning/rsl_rl/play.py index 9e89c6ff318..11ef7399462 100644 --- a/scripts/reinforcement_learning/rsl_rl/play.py +++ b/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) diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py index 33bfc9f63d4..ea630c5988d 100644 --- a/scripts/reinforcement_learning/rsl_rl/train.py +++ b/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) diff --git a/scripts/reinforcement_learning/sb3/play.py b/scripts/reinforcement_learning/sb3/play.py index 05c52390749..c803c1807ba 100644 --- a/scripts/reinforcement_learning/sb3/play.py +++ b/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) diff --git a/scripts/reinforcement_learning/sb3/train.py b/scripts/reinforcement_learning/sb3/train.py index ba45398f108..b2a9832a117 100644 --- a/scripts/reinforcement_learning/sb3/train.py +++ b/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) diff --git a/scripts/reinforcement_learning/skrl/play.py b/scripts/reinforcement_learning/skrl/play.py index 990ef5b558d..b4d52c39e8c 100644 --- a/scripts/reinforcement_learning/skrl/play.py +++ b/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) diff --git a/scripts/reinforcement_learning/skrl/train.py b/scripts/reinforcement_learning/skrl/train.py index e3399f204b5..c66efae0138 100644 --- a/scripts/reinforcement_learning/skrl/train.py +++ b/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) diff --git a/source/isaaclab/isaaclab/envs/direct_marl_env_cfg.py b/source/isaaclab/isaaclab/envs/direct_marl_env_cfg.py index 210b5139730..15f57cb4c03 100644 --- a/source/isaaclab/isaaclab/envs/direct_marl_env_cfg.py +++ b/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.""" diff --git a/source/isaaclab/isaaclab/envs/direct_rl_env_cfg.py b/source/isaaclab/isaaclab/envs/direct_rl_env_cfg.py index 6b26bdc7500..33297a228af 100644 --- a/source/isaaclab/isaaclab/envs/direct_rl_env_cfg.py +++ b/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.""" diff --git a/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py b/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py index a50e5336f9b..1a6bdacb795 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py +++ b/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.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py index 60a27649119..143dea88ffa 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py +++ b/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): """Initialize the feature extractor model. Args: cfg (FeatureExtractorCfg): Configuration for the feature extractor model. device (str): Device to run the model on. + log_dir (str, optional): Directory to save checkpoints. If None, uses local logs folder. """ 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) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py index 6cde7d06fc1..42e8c4f03c4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py +++ b/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 From 6a2785f5c68999ff585aff57c3eea4e4ef3bef8e Mon Sep 17 00:00:00 2001 From: Kelly Guo Date: Tue, 9 Sep 2025 16:47:54 -0700 Subject: [PATCH 2/4] Update source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py Co-authored-by: Mayank Mittal <12863862+Mayankm96@users.noreply.github.com> Signed-off-by: Kelly Guo --- .../isaaclab_tasks/direct/shadow_hand/feature_extractor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py index 143dea88ffa..090cc7a7ea9 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py @@ -73,7 +73,7 @@ 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, log_dir: str = None): + def __init__(self, cfg: FeatureExtractorCfg, device: str, log_dir: str | None = None): """Initialize the feature extractor model. Args: From 27fccfbefcee6c56f335b5903bfee972b903bea0 Mon Sep 17 00:00:00 2001 From: Kelly Guo Date: Tue, 9 Sep 2025 16:48:00 -0700 Subject: [PATCH 3/4] Update source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py Co-authored-by: Mayank Mittal <12863862+Mayankm96@users.noreply.github.com> Signed-off-by: Kelly Guo --- .../direct/shadow_hand/feature_extractor.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py index 090cc7a7ea9..f604052fb50 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py @@ -76,10 +76,10 @@ class FeatureExtractor: 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. - log_dir (str, optional): Directory to save checkpoints. If None, uses local logs folder. + Args: + 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 From bafab1433561a0e4ce3903c9e6b0d3d319411d25 Mon Sep 17 00:00:00 2001 From: Kelly Guo Date: Tue, 9 Sep 2025 16:49:13 -0700 Subject: [PATCH 4/4] Update feature_extractor.py Signed-off-by: Kelly Guo --- .../isaaclab_tasks/direct/shadow_hand/feature_extractor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py index f604052fb50..82d76ec7f1e 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/feature_extractor.py @@ -76,9 +76,9 @@ class FeatureExtractor: def __init__(self, cfg: FeatureExtractorCfg, device: str, log_dir: str | None = None): """Initialize the feature extractor model. - Args: - cfg: Configuration for the feature extractor model. - device: Device to run the model on. + Args: + 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. """