Merge commit from research-vgsm branch.

Author: TolgaOk

examples/mujoco/a2c_torch.py

@@ -1,5 +1,6 @@
 from typing import Any, Dict
 from argparse import ArgumentParser
+import numpy as np
 
 from modular_baselines.algorithms.a2c import A2C, A2CArgs
 from modular_baselines.algorithms.a2c.torch_agent import TorchA2CAgent
@@ -9,8 +10,8 @@
 from torch_setup import MujocoTorchConfig, setup, parallel_run, add_arguments
 
 
-def a2c_setup(env_name: str, config: Dict[str, Any], seed: int):
-    return setup(A2C, TorchA2CAgent, SeparateFeatureNetwork, env_name, config, seed)
+def a2c_setup(env_name: str, config: Dict[str, Any], experiment_name, device: str):
+    return setup(A2C, TorchA2CAgent, SeparateFeatureNetwork, experiment_name, env_name, config, device=device)
 
 
 a2c_mujoco_config = MujocoTorchConfig(
@@ -24,15 +25,22 @@ def a2c_setup(env_name: str, config: Dict[str, Any], seed: int):
         max_grad_norm=1.0,
         normalize_advantage=True,
     ),
+    name="default",
     n_envs=16,
     total_timesteps=5_000_000,
     log_interval=256,
-    device="cpu",
+    use_vec_normalization=True,
+    record_video=False,
+    seed=np.random.randint(2**10, 2**30),
 )
 
 if __name__ == "__main__":
     parser = ArgumentParser("A2C Mujoco")
     add_arguments(parser)
     cli_args = parser.parse_args()
-    parallel_run(a2c_setup, a2c_mujoco_config, n_procs=cli_args.n_procs,
-                 env_names=cli_args.env_names, n_seeds=cli_args.n_seeds)
+    parallel_run(a2c_setup,
+                 a2c_mujoco_config,
+                 n_procs=cli_args.n_procs,
+                 env_names=cli_args.env_names,
+                 experiment_name=cli_args.experiment_name,
+                 cuda_devices=cli_args.cuda_devices)

examples/mujoco/lstm_ppo_torch.py

@@ -1,5 +1,6 @@
 from typing import Any, Dict
 from argparse import ArgumentParser
+import multiprocessing as mp
 
 from modular_baselines.algorithms.ppo.ppo import LstmPPO, LstmPPOArgs
 from modular_baselines.algorithms.ppo.torch_lstm_agent import TorchLstmPPOAgent
@@ -9,11 +10,11 @@
 from torch_setup import MujocoTorchConfig, setup, parallel_run, add_arguments
 
 
-def ppo_setup(env_name: str, config: Dict[str, Any], seed: int):
-    return setup(LstmPPO, TorchLstmPPOAgent, LSTMSeparateNetwork, env_name, config, seed)
+def lstm_ppo_setup(experiment_name: str, env_name: str, config: MujocoTorchConfig, seed: int, device:str):
+    return setup(LstmPPO, TorchLstmPPOAgent, LSTMSeparateNetwork, experiment_name, env_name, config, seed, device)
 
 
-lstm_ppo_mujoco_config = MujocoTorchConfig(
+lstm_ppo_mujoco_config = [MujocoTorchConfig(
     args=LstmPPOArgs(
         rollout_len=2048,
         ent_coef=1e-4,
@@ -23,21 +24,23 @@ def ppo_setup(env_name: str, config: Dict[str, Any], seed: int):
         epochs=10,
         lr=LinearAnnealing(3e-4, 0.0, 5_000_000 // (2048 * 16)),
         clip_value=LinearAnnealing(0.2, 0.2, 5_000_000 // (2048 * 16)),
-        batch_size=8,
+        batch_size=64 // n_step,
         max_grad_norm=1.0,
         normalize_advantage=True,
-        mini_rollout_size=8,
+        mini_rollout_size=n_step,
         use_sampled_hidden=False,
     ),
+    name=f"{n_step}_step",
     n_envs=16,
     total_timesteps=5_000_000,
     log_interval=1,
-    device="cpu",
-)
+) for n_step in (1, 2, 4, 8, 16, 32, 64)]
 
 if __name__ == "__main__":
+    mp.set_start_method("spawn")
     parser = ArgumentParser("PPO Mujoco")
     add_arguments(parser)
     cli_args = parser.parse_args()
-    parallel_run(ppo_setup, lstm_ppo_mujoco_config, n_procs=cli_args.n_procs,
-                 env_names=cli_args.env_names, n_seeds=cli_args.n_seeds)
+    parallel_run(lstm_ppo_setup, lstm_ppo_mujoco_config, n_procs=cli_args.n_procs,
+                 env_names=cli_args.env_names, experiment_name=cli_args.experiment_name,
+                 n_seeds=cli_args.n_seeds, cuda_devices=cli_args.cuda_devices)

examples/mujoco/torch_setup.py

@@ -1,64 +1,112 @@
-from typing import List, Any, Dict, Union, Optional, Tuple, Callable, Type
+from typing import List, Any, Dict, Union, Optional, Tuple, Callable, Type, Iterable
 import torch
 import os
 import numpy as np
 import sys
 from multiprocessing import Process, Queue
 from dataclasses import dataclass
 import argparse
+import time
+import gym
+from datetime import datetime
 
 from stable_baselines3.common.env_util import make_vec_env
+from stable_baselines3.common.vec_env.base_vec_env import VecEnv
 from stable_baselines3.common.vec_env.subproc_vec_env import SubprocVecEnv
 from stable_baselines3.common.vec_env.vec_normalize import VecNormalize
 from stable_baselines3.common.logger import HumanOutputFormat, CSVOutputFormat, JSONOutputFormat
+from stable_baselines3.common.vec_env.vec_video_recorder import VecVideoRecorder
+from stable_baselines3.common.running_mean_std import RunningMeanStd
 
 from modular_baselines.algorithms.algorithm import BaseAlgorithm
 from modular_baselines.algorithms.agent import BaseAgent
-from modular_baselines.loggers.writers import ScalarWriter, DictWriter
+from modular_baselines.loggers.writers import ScalarWriter, DictWriter, BaseWriter, SaveModelParametersWriter, LogConfigs
 from modular_baselines.loggers.data_logger import DataLogger
 
 
 @dataclass(frozen=True)
 class MujocoTorchConfig():
     args: Any
+    name: str
     n_envs: int
     total_timesteps: int
     log_interval: int
-    device: str
-
-
-def setup(algorithm_cls: Type[BaseAlgorithm],
-          agent_cls: Type[BaseAgent],
-          network: Type[torch.nn.Module],
-          env_name: str,
-          config: MujocoTorchConfig,
-          seed: int
-          ) -> BaseAlgorithm:
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-
-    log_dir = f"logs/{algorithm_cls.__name__}-{env_name.lower()}/{seed}"
+    record_video: bool
+    seed: int
+
+
+def pre_setup(experiment_name: str,
+              env: Union[gym.Env, str],
+              config: MujocoTorchConfig,
+              ) -> Tuple[DataLogger, List[BaseWriter], VecEnv]:
+    """ Prepare loggers and vectorized environment
+
+    Args:
+        experiment_name (str): Name of the experiment
+        env (Union[gym.Env, str]): Name of the environment or the environment
+        config (MujocoTorchConfig): Torch Mujoco configuration
+
+    Returns:
+        Tuple[DataLogger, List[BaseWriter], VecEnv]: Data logger, writers list and vectorized 
+            environment
+    """
+    np.random.seed(config.seed)
+    torch.manual_seed(config.seed)
+    env_name = env if isinstance(env, str) else env.__class__.__name__
+    date_time = datetime.now().strftime("%Y-%m-%dT%H:%M:%S")
+
+    log_dir = f"logs/{experiment_name}-{env_name.lower()}/{config.name}/{date_time}"
     data_logger = DataLogger()
     os.makedirs(log_dir, exist_ok=True)
     sb3_writers = [HumanOutputFormat(sys.stdout),
                    CSVOutputFormat(os.path.join(log_dir, "progress.csv")),
                    JSONOutputFormat(os.path.join(log_dir, "progress.json"))]
     logger_callbacks = [
         ScalarWriter(interval=config.log_interval, dir_path=log_dir, writers=sb3_writers),
-        DictWriter(interval=config.log_interval, dir_path=log_dir)
+        DictWriter(interval=config.log_interval, dir_path=log_dir),
+        SaveModelParametersWriter(interval=config.log_interval * 1, dir_path=log_dir)
     ]
 
     vecenv = make_vec_env(
-        env_name,
+        env,
         n_envs=config.n_envs,
-        seed=seed,
+        seed=config.seed,
         wrapper_class=None,
         vec_env_cls=SubprocVecEnv)
-    vecenv = VecNormalize(vecenv, training=True, gamma=config.args.gamma)
+    if config.args.use_vec_normalization:
+        vecenv = VecNormalize(
+            vecenv,
+            training=True,
+            gamma=config.args.gamma,
+            **config.args.vec_norm_info)
+        if config.args.vec_norm_info["norm_obs"] is False:
+            vecenv.obs_rms = RunningMeanStd(shape=vecenv.observation_space.shape)
+    if config.record_video:
+        vecenv = VecVideoRecorder(
+            vecenv,
+            f"{log_dir}/videos",
+            record_video_trigger=lambda x: x % 25000 == 0, video_length=1000
+        )
+    LogConfigs(config=config, dir_path=log_dir)
+
+    return data_logger, logger_callbacks, vecenv
+
+
+def setup(algorithm_cls: Type[BaseAlgorithm],
+          agent_cls: Type[BaseAgent],
+          network: Type[torch.nn.Module],
+          experiment_name: str,
+          env_name: str,
+          config: MujocoTorchConfig,
+          device: str
+          ) -> BaseAlgorithm:
+
+    experiment_name = "-".join([experiment_name, algorithm_cls.__name__])
+    data_logger, logger_callbacks, vecenv = pre_setup(experiment_name, env_name, config)
 
     policy = network(observation_space=vecenv.observation_space,
                      action_space=vecenv.action_space)
-    policy.to(config.device)
+    policy.to(device)
     optimizer = torch.optim.Adam(policy.parameters(), eps=1e-5)
     agent = agent_cls(policy,
                       optimizer,
@@ -80,39 +128,48 @@ def setup(algorithm_cls: Type[BaseAlgorithm],
 
 
 def add_arguments(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--experiment-name", type=str, default="",
+                        help="Prefix of the experiment name")
     parser.add_argument("--n-procs", type=int, default=1,
                         help="Number of parallelized processes for experiments")
-    parser.add_argument("--n-seeds", type=int, default=1,
-                        help="Number of seeds/runs per environment")
     parser.add_argument("--env-names", nargs='+', type=str, required=True,
                         help="Gym environment names")
+    parser.add_argument("--cuda-devices", nargs='+', type=int, required=False,
+                        help="Available cuda devices")
 
 
-def worker(setup_fn, argument_queue: Queue, rank: int) -> None:
+def worker(setup_fn, argument_queue: Queue, rank: int, cuda_devices) -> None:
+    device = "cpu" if cuda_devices is None else f"cuda:{cuda_devices[rank % len(cuda_devices)]}"
+    print(f"Worker-{rank} use device: {device}")
     while not argument_queue.empty():
         kwargs = argument_queue.get()
-        setup_fn(**kwargs)
+        setup_fn(device=device, **kwargs)
 
 
 def parallel_run(setup_fn: Callable[[str, MujocoTorchConfig, int], BaseAlgorithm],
-                 config: MujocoTorchConfig,
+                 configs: Union[MujocoTorchConfig, Iterable[MujocoTorchConfig]],
+                 experiment_name: str,
                  n_procs: int,
                  env_names: Tuple[str],
-                 n_seeds: int
+                 cuda_devices: Tuple[int],
                  ) -> None:
 
-    arguments = [dict(env_name=env_name, seed=seed, config=config)
+    if not isinstance(configs, Iterable):
+        configs = [configs]
+
+    arguments = [dict(env_name=env_name, config=config, experiment_name=experiment_name)
                  for env_name in env_names
-                 for seed in np.random.randint(2 ** 10, 2 ** 30, size=n_seeds).tolist()]
+                 for config in configs]
 
     argument_queue = Queue()
     for arg in arguments:
         argument_queue.put(arg)
 
-    processes = [Process(target=worker, args=(setup_fn, argument_queue, rank))
+    processes = [Process(target=worker, args=(setup_fn, argument_queue, rank, cuda_devices))
                  for rank in range(n_procs)]
 
     for proc in processes:
+        time.sleep(1.5)  # To avoid having the same log name
         proc.start()
 
     for proc in processes: