Integrating hydra with SAC

examples/sac/sac.py

@@ -3,69 +3,59 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+import dataclasses
 import uuid
 from datetime import datetime
 
-from torchrl.envs import ParallelEnv, EnvCreator
-from torchrl.envs.utils import set_exploration_mode
-from torchrl.record import VideoRecorder
-
-try:
-    import configargparse as argparse
-
-    _configargparse = True
-except ImportError:
-    import argparse
-
-    _configargparse = False
+import hydra
 import torch.cuda
+from hydra.core.config_store import ConfigStore
+from omegaconf import OmegaConf
+from torchrl.envs import ParallelEnv, EnvCreator
 from torchrl.envs.transforms import RewardScaling, TransformedEnv
+from torchrl.envs.utils import set_exploration_mode
 from torchrl.modules import OrnsteinUhlenbeckProcessWrapper
+from torchrl.record import VideoRecorder
 from torchrl.trainers.helpers.collectors import (
     make_collector_offpolicy,
-    parser_collector_args_offpolicy,
+    OffPolicyCollectorConfig,
 )
 from torchrl.trainers.helpers.envs import (
     correct_for_frame_skip,
     get_stats_random_rollout,
     parallel_env_constructor,
-    parser_env_args,
     transformed_env_constructor,
+    EnvConfig,
 )
-from torchrl.trainers.helpers.losses import make_sac_loss, parser_loss_args
+from torchrl.trainers.helpers.losses import make_sac_loss, LossConfig
 from torchrl.trainers.helpers.models import (
     make_sac_model,
-    parser_model_args_continuous,
+    SACModelConfig,
 )
-from torchrl.trainers.helpers.recorder import parser_recorder_args
+from torchrl.trainers.helpers.recorder import RecorderConfig
 from torchrl.trainers.helpers.replay_buffer import (
     make_replay_buffer,
-    parser_replay_args,
+    ReplayArgsConfig,
 )
-from torchrl.trainers.helpers.trainers import make_trainer, parser_trainer_args
-
-
-def make_args():
-    parser = argparse.ArgumentParser()
-    if _configargparse:
-        parser.add_argument(
-            "-c",
-            "--config",
-            required=True,
-            is_config_file=True,
-            help="config file path",
-        )
-    parser_trainer_args(parser)
-    parser_collector_args_offpolicy(parser)
-    parser_env_args(parser)
-    parser_loss_args(parser, algorithm="SAC")
-    parser_model_args_continuous(parser, "SAC")
-    parser_recorder_args(parser)
-    parser_replay_args(parser)
-    return parser
-
+from torchrl.trainers.helpers.trainers import make_trainer, TrainerConfig
+
+config_fields = [
+    (config_field.name, config_field.type, config_field)
+    for config_cls in (
+        TrainerConfig,
+        OffPolicyCollectorConfig,
+        EnvConfig,
+        LossConfig,
+        SACModelConfig,
+        RecorderConfig,
+        ReplayArgsConfig,
+    )
+    for config_field in dataclasses.fields(config_cls)
+]
 
-parser = make_args()
+Config = dataclasses.make_dataclass(cls_name="Config", fields=config_fields)
+cs = ConfigStore.instance()
+cs.store(name="config", node=Config)
 
 DEFAULT_REWARD_SCALING = {
     "Hopper-v1": 5,
@@ -78,13 +68,18 @@ def make_args():
 }
 
 
-def main(args):
+@hydra.main(version_base=None, config_path=None, config_name="config")
+def main(cfg: "DictConfig"):
     from torch.utils.tensorboard import SummaryWriter
 
-    args = correct_for_frame_skip(args)
+    if cfg.config_file is not None:
+        overriding_cfg = OmegaConf.load(cfg.config_file)
+        cfg = OmegaConf.merge(cfg, overriding_cfg)
+
+    cfg = correct_for_frame_skip(cfg)
 
-    if not isinstance(args.reward_scaling, float):
-        args.reward_scaling = DEFAULT_REWARD_SCALING.get(args.env_name, 5.0)
+    if not isinstance(cfg.reward_scaling, float):
+        cfg.reward_scaling = DEFAULT_REWARD_SCALING.get(cfg.env_name, 5.0)
 
     device = (
         torch.device("cpu")
@@ -95,47 +90,47 @@ def main(args):
     exp_name = "_".join(
         [
             "SAC",
-            args.exp_name,
+            cfg.exp_name,
             str(uuid.uuid4())[:8],
             datetime.now().strftime("%y_%m_%d-%H_%M_%S"),
         ]
     )
     writer = SummaryWriter(f"sac_logging/{exp_name}")
-    video_tag = exp_name if args.record_video else ""
+    video_tag = exp_name if cfg.record_video else ""
 
     stats = None
-    if not args.vecnorm and args.norm_stats:
-        proof_env = transformed_env_constructor(args=args, use_env_creator=False)()
+    if not cfg.vecnorm and cfg.norm_stats:
+        proof_env = transformed_env_constructor(cfg=cfg, use_env_creator=False)()
         stats = get_stats_random_rollout(
-            args, proof_env, key="next_pixels" if args.from_pixels else None
+            cfg, proof_env, key="next_pixels" if cfg.from_pixels else None
         )
         # make sure proof_env is closed
         proof_env.close()
-    elif args.from_pixels:
+    elif cfg.from_pixels:
         stats = {"loc": 0.5, "scale": 0.5}
     proof_env = transformed_env_constructor(
-        args=args, use_env_creator=False, stats=stats
+        cfg=cfg, use_env_creator=False, stats=stats
     )()
     model = make_sac_model(
         proof_env,
-        args=args,
+        cfg=cfg,
         device=device,
     )
-    loss_module, target_net_updater = make_sac_loss(model, args)
+    loss_module, target_net_updater = make_sac_loss(model, cfg)
 
     actor_model_explore = model[0]
-    if args.ou_exploration:
+    if cfg.ou_exploration:
         actor_model_explore = OrnsteinUhlenbeckProcessWrapper(
             actor_model_explore,
-            annealing_num_steps=args.annealing_frames,
-            sigma=args.ou_sigma,
-            theta=args.ou_theta,
+            annealing_num_steps=cfg.annealing_frames,
+            sigma=cfg.ou_sigma,
+            theta=cfg.ou_theta,
         ).to(device)
     if device == torch.device("cpu"):
         # mostly for debugging
         actor_model_explore.share_memory()
 
-    if args.gSDE:
+    if cfg.gSDE:
         with torch.no_grad(), set_exploration_mode("random"):
             # get dimensions to build the parallel env
             proof_td = actor_model_explore(proof_env.reset().to(device))
@@ -145,7 +140,7 @@ def main(args):
         action_dim_gsde, state_dim_gsde = None, None
     proof_env.close()
     create_env_fn = parallel_env_constructor(
-        args=args,
+        cfg=cfg,
         stats=stats,
         action_dim_gsde=action_dim_gsde,
         state_dim_gsde=state_dim_gsde,
@@ -154,25 +149,25 @@ def main(args):
     collector = make_collector_offpolicy(
         make_env=create_env_fn,
         actor_model_explore=actor_model_explore,
-        args=args,
+        cfg=cfg,
         # make_env_kwargs=[
         #     {"device": device} if device >= 0 else {}
         #     for device in args.env_rendering_devices
         # ],
     )
 
-    replay_buffer = make_replay_buffer(device, args)
+    replay_buffer = make_replay_buffer(device, cfg)
 
     recorder = transformed_env_constructor(
-        args,
+        cfg,
         video_tag=video_tag,
         norm_obs_only=True,
         stats=stats,
         writer=writer,
     )()
 
     # remove video recorder from recorder to have matching state_dict keys
-    if args.record_video:
+    if cfg.record_video:
         recorder_rm = TransformedEnv(recorder.env)
         for transform in recorder.transform:
             if not isinstance(transform, VideoRecorder):
@@ -202,7 +197,7 @@ def main(args):
         actor_model_explore,
         replay_buffer,
         writer,
-        args,
+        cfg,
     )
 
     def select_keys(batch):
@@ -219,13 +214,12 @@ def select_keys(batch):
 
     trainer.register_op("batch_process", select_keys)
 
-    final_seed = collector.set_seed(args.seed)
-    print(f"init seed: {args.seed}, final seed: {final_seed}")
+    final_seed = collector.set_seed(cfg.seed)
+    print(f"init seed: {cfg.seed}, final seed: {final_seed}")
 
     trainer.train()
     return (writer.log_dir, trainer._log_dict, trainer.state_dict())
 
 
 if __name__ == "__main__":
-    args = parser.parse_args()
-    main(args)
+    main()

torchrl/trainers/helpers/losses.py

@@ -14,7 +14,7 @@
     "make_redq_loss",
 ]
 
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Any
 
 from torchrl.modules import ActorValueOperator, ActorCriticOperator
 from torchrl.objectives import (
@@ -226,7 +226,7 @@ class LossConfig:
     # use two (or more!) different qvalue networks trained independently and choose the lowest value
     # predicted to predict the state action value. This can be disabled by using this flag.
     # REDQ uses an arbitrary number of Q-value functions to speed up learning in MF contexts.
-    target_entropy: float = None
+    target_entropy: Any = None
     # Target entropy for the policy distribution. Default is None (auto calculated as the `target_entropy = -action_dim`)