Integrating hydra with DQN

examples/ddpg/ddpg.py

@@ -10,7 +10,6 @@
 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
@@ -71,9 +70,6 @@
 def main(cfg: "DictConfig"):
     from torch.utils.tensorboard import SummaryWriter
 
-    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(cfg.reward_scaling, float):
@@ -171,7 +167,7 @@ def main(cfg: "DictConfig"):
 
     # remove video recorder from recorder to have matching state_dict keys
     if cfg.record_video:
-        recorder_rm = TransformedEnv(recorder.env)
+        recorder_rm = TransformedEnv(recorder.base_env)
         for transform in recorder.transform:
             if not isinstance(transform, VideoRecorder):
                 recorder_rm.append_transform(transform)

examples/dqn/dqn.py

@@ -3,77 +3,68 @@
 # 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.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 torchrl.envs import ParallelEnv, EnvCreator
 from torchrl.envs.transforms import RewardScaling, TransformedEnv
 from torchrl.modules import EGreedyWrapper
+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_dqn_loss, parser_loss_args
+from torchrl.trainers.helpers.losses import make_dqn_loss, LossConfig
 from torchrl.trainers.helpers.models import (
     make_dqn_actor,
-    parser_model_args_discrete,
+    DiscreteModelConfig,
 )
-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="DQN")
-    parser_model_args_discrete(parser)
-    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,
+        DiscreteModelConfig,
+        RecorderConfig,
+        ReplayArgsConfig,
+    )
+    for config_field in dataclasses.fields(config_cls)
+]
+Config = dataclasses.make_dataclass(cls_name="Config", fields=config_fields)
+cs = ConfigStore.instance()
+cs.store(name="config", node=Config)
 
-parser = make_args()
 
+@hydra.main(version_base=None, config_path=None, config_name="config")
+def main(cfg: "DictConfig"):
 
-def main(args):
     from torch.utils.tensorboard import SummaryWriter
 
-    args = correct_for_frame_skip(args)
+    cfg = correct_for_frame_skip(cfg)
 
-    if not isinstance(args.reward_scaling, float):
-        args.reward_scaling = 1.0
+    if not isinstance(cfg.reward_scaling, float):
+        cfg.reward_scaling = 1.0
 
     device = (
         torch.device("cpu")
@@ -84,41 +75,42 @@ def main(args):
     exp_name = "_".join(
         [
             "DQN",
-            args.exp_name,
+            cfg.exp_name,
             str(uuid.uuid4())[:8],
             datetime.now().strftime("%y_%m_%d-%H_%M_%S"),
         ]
     )
     writer = SummaryWriter(f"dqn_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_dqn_actor(
         proof_environment=proof_env,
-        args=args,
+        cfg=cfg,
         device=device,
     )
 
-    loss_module, target_net_updater = make_dqn_loss(model, args)
-    model_explore = EGreedyWrapper(model, annealing_num_steps=args.annealing_frames).to(
+    loss_module, target_net_updater = make_dqn_loss(model, cfg)
+    model_explore = EGreedyWrapper(model, annealing_num_steps=cfg.annealing_frames).to(
         device
     )
+
     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,
@@ -127,26 +119,26 @@ def main(args):
     collector = make_collector_offpolicy(
         make_env=create_env_fn,
         actor_model_explore=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:
-        recorder_rm = TransformedEnv(recorder.env)
+    if cfg.record_video:
+        recorder_rm = TransformedEnv(recorder.base_env)
         for transform in recorder.transform:
             if not isinstance(transform, VideoRecorder):
                 recorder_rm.append_transform(transform)
@@ -171,10 +163,10 @@ def main(args):
         loss_module,
         recorder,
         target_net_updater,
-        model_explore,
+        model,
         replay_buffer,
         writer,
-        args,
+        cfg,
     )
 
     def select_keys(batch):
@@ -191,13 +183,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()

examples/ppo/ppo.py

@@ -10,7 +10,6 @@
 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
@@ -56,10 +55,6 @@
 def main(cfg: "DictConfig"):
     from torch.utils.tensorboard import SummaryWriter
 
-    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(cfg.reward_scaling, float):
@@ -142,7 +137,7 @@ def main(cfg: "DictConfig"):
 
     # remove video recorder from recorder to have matching state_dict keys
     if cfg.record_video:
-        recorder_rm = TransformedEnv(recorder.env)
+        recorder_rm = TransformedEnv(recorder.base_env)
         for transform in recorder.transform:
             if not isinstance(transform, VideoRecorder):
                 recorder_rm.append_transform(transform)

examples/redq/redq.py

@@ -10,7 +10,6 @@
 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
@@ -72,9 +71,6 @@
 def main(cfg: "DictConfig"):
     from torch.utils.tensorboard import SummaryWriter  # avoid loading on each process
 
-    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(cfg.reward_scaling, float):
@@ -171,7 +167,7 @@ def main(cfg: "DictConfig"):
 
     # remove video recorder from recorder to have matching state_dict keys
     if cfg.record_video:
-        recorder_rm = TransformedEnv(recorder.env)
+        recorder_rm = TransformedEnv(recorder.base_env)
         for transform in recorder.transform:
             if not isinstance(transform, VideoRecorder):
                 recorder_rm.append_transform(transform)

examples/sac/sac.py

@@ -10,7 +10,6 @@
 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
@@ -72,10 +71,6 @@
 def main(cfg: "DictConfig"):
     from torch.utils.tensorboard import SummaryWriter
 
-    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(cfg.reward_scaling, float):
@@ -168,7 +163,7 @@ def main(cfg: "DictConfig"):
 
     # remove video recorder from recorder to have matching state_dict keys
     if cfg.record_video:
-        recorder_rm = TransformedEnv(recorder.env)
+        recorder_rm = TransformedEnv(recorder.base_env)
         for transform in recorder.transform:
             if not isinstance(transform, VideoRecorder):
                 recorder_rm.append_transform(transform)

torchrl/data/tensor_specs.py

@@ -1003,7 +1003,11 @@ def project(self, val: _TensorDict) -> _TensorDict:
 
     def rand(self, shape=torch.Size([])):
         return TensorDict(
-            {key: value.rand(shape) for key, value in self._specs.items()},
+            {
+                key: value.rand(shape)
+                for key, value in self._specs.items()
+                if value is not None
+            },
             batch_size=shape,
         )