Reformat script arguments (#17)

* Reformat script arguments
Add support for loading network/opitimizer

* update readme

Author: benbellheron

README.md

@@ -66,15 +66,15 @@ The log directory contains the tensorboard file, saved models, and other metadat
 ```
 # Local Mode (A2C)
 # We recommend 4GB+ GPU memory, 8GB+ RAM, 4+ Cores
-python -m adept.scripts.local --env-id BeamRiderNoFrameskip-v4
+python -m adept.scripts.local ActorCritic --env-id BeamRiderNoFrameskip-v4
 
 # Towered Mode (A3C Variant, requires mpi4py)
 # We recommend 2+ GPUs, 8GB+ GPU memory, 32GB+ RAM, 4+ Cores
-python -m adept.scripts.towered --env-id BeamRiderNoFrameskip-v4
+python -m adept.scripts.towered ActorCritic --env-id BeamRiderNoFrameskip-v4
 
 # IMPALA (requires mpi4py and is resource intensive)
 # We recommend 2+ GPUs, 8GB+ GPU memory, 32GB+ RAM, 4+ Cores
-mpiexec -n 3 python -m adept.scripts.impala --env-id BeamRiderNoFrameskip-v4
+mpiexec -n 3 python -m adept.scripts.impala ActorCriticVtrace --env-id BeamRiderNoFrameskip-v4
 
 # StarCraft 2 (IMPALA not supported yet)
 # Warning: much more resource intensive than Atari

adept/agents/__init__.py

@@ -14,6 +14,7 @@
 You should have received a copy of the GNU General Public License
 along with this program.  If not, see <http://www.gnu.org/licenses/>.
 """
+
 from .actor_critic import ActorCritic
 from .impala import ActorCriticVtrace
 

adept/agents/actor_critic.py

@@ -14,14 +14,15 @@
 You should have received a copy of the GNU General Public License
 along with this program.  If not, see <http://www.gnu.org/licenses/>.
 """
+from argparse import ArgumentParser
 from collections import OrderedDict
 
 import torch
 from adept.environments import Engines
 from torch.nn import functional as F
 
 from adept.expcaches.rollout import RolloutCache
-from adept.utils.util import listd_to_dlist
+from adept.utils.util import listd_to_dlist, parse_bool
 from ._base import Agent
 
 
@@ -66,6 +67,47 @@ def from_args(cls, network, device, reward_normalizer, gpu_preprocessor, engine,
             args.nb_env, args.exp_length, args.discount, args.generalized_advantage_estimation, args.tau, args.normalize_advantage
         )
 
+    @classmethod
+    def add_args(cls, parser: ArgumentParser):
+        parser.add_argument(
+            '-ae',
+            '--exp-length',
+            type=int,
+            default=20,
+            help='Experience length (default: 20)'
+        )
+        parser.add_argument(
+            '-ag',
+            '--generalized-advantage-estimation',
+            type=parse_bool,
+            nargs='?',
+            const=True,
+            default=True,
+            help='Use generalized advantage estimation for the policy loss. (default: True)'
+        )
+        parser.add_argument(
+            '-at',
+            '--tau',
+            type=float,
+            default=1.00,
+            help='parameter for GAE (default: 1.00)'
+        )
+        parser.add_argument(
+            '--entropy-weight',
+            type=float,
+            default=0.01,
+            help='Entropy penalty (default: 0.01)'
+        )
+        parser.add_argument(
+            '--normalize-advantage',
+            type=parse_bool,
+            nargs='?',
+            const=True,
+            default=False,
+            help='Normalize the advantage when calculating policy loss. (default: False)'
+        )
+
+
     @property
     def exp_cache(self):
         return self._exp_cache

adept/agents/impala/actor_critic_vtrace.py

@@ -15,6 +15,7 @@
 along with this program.  If not, see <http://www.gnu.org/licenses/>.
 """
 # Use https://github.com/deepmind/scalable_agent/blob/master/vtrace.py for reference
+from argparse import ArgumentParser
 from collections import OrderedDict
 import torch
 from torch.nn import functional as F
@@ -54,6 +55,16 @@ def from_args(cls, network, device, reward_normalizer, gpu_preprocessor, engine,
             args.nb_env, args.exp_length, args.discount
         )
 
+    @classmethod
+    def add_args(cls, parser):
+        parser.add_argument(
+            '-ae',
+            '--exp-length',
+            type=int,
+            default=20,
+            help='Experience length (default: 20)'
+        )
+
     @property
     def exp_cache(self):
         return self._exp_cache
@@ -440,3 +451,5 @@ def _vtrace_returns(log_diff_behavior_vs_current, discount_terminal_mask, reward
 
         weighted_advantage = clamped_importance_pg * advantage
         return v_s, weighted_advantage, importance
+
+

adept/containers/impala.py

@@ -311,6 +311,7 @@ def summary_writer(self):
         return self._summary_writer
 
     def run(self, initial_count=0):
+        self.local_step_count = initial_count
         next_obs = self.environment.reset()
         self._starting_internals = self.agent.internals
         while not self.should_stop():

adept/containers/towered.py

@@ -301,6 +301,7 @@ def nb_env(self):
         return self._nb_env
 
     def run(self, initial_count=0):
+        self.local_step_count = initial_count
         next_obs = self.environment.reset()
         self.start_time = time.time()
         while not self.should_stop():
@@ -344,7 +345,7 @@ def submit(self):
         if host_info is not None:
             self.global_step = host_info
         else:
-            self.global_step = 0
+            self.global_step = self.local_step_count
         # host decides when it wants pytorch buffers
         if self.mpi_buffer_request.test()[0]:
             buffer_list = [x.cpu().numpy() for x in self.network._all_buffers()]

adept/environments/atari.py

@@ -27,14 +27,14 @@
 import numpy as np
 
 
-def make_atari_env(env_id, skip_rate, max_ep_length, do_zscore_norm, do_frame_stack, seed):
+def make_atari_env(env_id, skip_rate, max_ep_length, do_frame_stack, seed):
     def _f():
-        env = atari_env(env_id, skip_rate, max_ep_length, do_zscore_norm, do_frame_stack, seed)
+        env = atari_env(env_id, skip_rate, max_ep_length, do_frame_stack, seed)
         return env
     return _f
 
 
-def atari_env(env_id, skip_rate, max_ep_length, do_zscore_norm, do_frame_stack, seed):
+def atari_env(env_id, skip_rate, max_ep_length, do_frame_stack, seed):
     env = gym.make(env_id)
     if hasattr(env.unwrapped, 'ale'):
         if 'FIRE' in env.unwrapped.get_action_meanings():

adept/scripts/benchmark_atari.py

@@ -54,7 +54,7 @@
         help='Number of eval steps allowed to run per second decreasing this amount can improve training speed. 0 to disable (default: 0)'
     )
 
-    args = parser.parse_args()
+    args = parser.add_args()
 
     args.mode_name = 'Local'
 

adept/scripts/impala.py

@@ -16,6 +16,7 @@
 along with this program.  If not, see <http://www.gnu.org/licenses/>.
 """
 import os
+from copy import deepcopy
 from mpi4py import MPI as mpi
 import torch
 from absl import flags
@@ -41,7 +42,7 @@ def main(args):
     if rank == 0:
         timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
         log_id = make_log_id_from_timestamp(args.tag, args.mode_name, args.agent,
-                                            args.vision_network + args.network_body,
+                                            args.network_vision + args.network_body,
                                             timestamp)
         log_id_dir = os.path.join(args.log_dir, args.env_id, log_id)
         os.makedirs(log_id_dir)
@@ -53,34 +54,54 @@ def main(args):
 
     if rank != 0:
         log_id = make_log_id_from_timestamp(args.tag, args.mode_name, args.agent,
-                                            args.vision_network + args.network_body,
+                                            args.network_vision + args.network_body,
                                             timestamp)
         log_id_dir = os.path.join(args.log_dir, args.env_id, log_id)
 
     comm.Barrier()
 
     # construct env
-    seed = args.seed if rank == 0 else args.seed + (args.nb_env * (rank - 1))  # unique seed per process
-    env = make_env(args, seed)
+    # unique seed per process
+    seed = args.seed if rank == 0 else args.seed + args.nb_env * (rank - 1)
+    # don't make a ton of envs if host
+    if rank == 0:
+        env_args = deepcopy(args)
+        env_args.nb_env = 1
+        env = make_env(env_args, seed)
+    else:
+        env = make_env(args, seed)
 
     # construct network
     torch.manual_seed(args.seed)
     network_head_shapes = get_head_shapes(env.action_space, args.agent)
     network = make_network(env.observation_space, network_head_shapes, args)
 
-    # sync network params
-    if rank == 0:
-        for v in network.parameters():
-            comm.Bcast(v.detach().cpu().numpy(), root=0)
-        print('Root variables synced')
+    # possibly load network
+    initial_step_count = 0
+    if args.load_network:
+        network.load_state_dict(
+            torch.load(
+                args.load_network, map_location=lambda storage, loc: storage
+            )
+        )
+        # get step count from network file
+        epoch_dir = os.path.split(args.load_network)[0]
+        initial_step_count = int(os.path.split(epoch_dir)[-1])
+        print('Reloaded network from {}'.format(args.load_network))
+    # only sync network params if not loading
     else:
-        # can just use the numpy buffers
-        variables = [v.detach().cpu().numpy() for v in network.parameters()]
-        for v in variables:
-            comm.Bcast(v, root=0)
-        for shared_v, model_v in zip(variables, network.parameters()):
-            model_v.data.copy_(torch.from_numpy(shared_v), non_blocking=True)
-        print('{} variables synced'.format(rank))
+        if rank == 0:
+            for v in network.parameters():
+               comm.Bcast(v.detach().cpu().numpy(), root=0)
+            print('Root variables synced')
+        else:
+            # can just use the numpy buffers
+            variables = [v.detach().cpu().numpy() for v in network.parameters()]
+            for v in variables:
+                comm.Bcast(v, root=0)
+            for shared_v, model_v in zip(variables, network.parameters()):
+                model_v.data.copy_(torch.from_numpy(shared_v), non_blocking=True)
+            print('{} variables synced'.format(rank))
 
     # construct agent
     # host is always the first gpu, workers are distributed evenly across the rest
@@ -120,7 +141,7 @@ def main(args):
             profiler.stop()
             print(profiler.output_text(unicode=True, color=True))
         else:
-            container.run()
+            container.run(initial_step_count)
         env.close()
     # host
     else:
@@ -136,6 +157,12 @@ def main(args):
         # Construct the optimizer
         def make_optimizer(params):
             opt = torch.optim.RMSprop(params, lr=args.learning_rate, eps=1e-5, alpha=0.99)
+            if args.load_optimizer:
+                opt.load_state_dict(
+                    torch.load(
+                        args.load_optimizer, map_location=lambda storage, loc: storage
+                    )
+                )
             return opt
 
         container = ImpalaHost(agent, comm, make_optimizer, summary_writer, args.summary_frequency, saver,
@@ -169,59 +196,41 @@ def make_optimizer(params):
     import argparse
     from adept.utils.script_helpers import add_base_args, parse_bool
 
-    parser = argparse.ArgumentParser(description='AdeptRL IMPALA Mode')
-    parser = add_base_args(parser)
-    parser.add_argument('--gpu-id', type=int, nargs='+', default=[0],
-                        help='Which GPU to use for training. The host will always be the first gpu, workers are distributed evenly across the rest (default: [0])')
-    parser.add_argument(
-        '-vn', '--vision-network', default='Nature',
-        help='name of preset network (default: Nature)'
-    )
-    parser.add_argument(
-        '-dn', '--discrete-network', default='Identity',
-    )
-    parser.add_argument(
-        '-nb', '--network-body', default='LSTM',
-    )
-    parser.add_argument(
-        '--agent', default='ActorCriticVtrace',
-        help='name of preset agent (default: ActorCriticVtrace)'
-    )
-    parser.add_argument(
-        '--profile', type=parse_bool, nargs='?', const=True, default=False,
-        help='displays profiling tree after 10e3 steps (default: False)'
-    )
-    parser.add_argument(
-        '--debug', type=parse_bool, nargs='?', const=True, default=False,
-        help='debug mode sends the logs to /tmp/ and overrides number of workers to 3 (default: False)'
-    )
-    parser.add_argument(
-        '--max-queue-length', type=int, default=(size - 1) * 2,
-        help='Maximum rollout queue length. If above the max, workers will wait to append (default: (size - 1) * 2)'
-    )
-    parser.add_argument(
-        '--num-rollouts-in-batch', type=int, default=(size - 1),
-        help='The batch size in rollouts (so total batch is this number * nb_env * seq_len). '
-             + 'Not compatible with --dynamic-batch (default: (size - 1))'
-    )
-    parser.add_argument(
-        '--max-dynamic-batch', type=int, default=0,
-        help='When > 0 uses dynamic batching (disables cudnn and --num-rollouts-in-batch). '
-             + 'Limits the maximum rollouts in the batch to limit GPU memory usage. (default: 0 (False))'
-    )
-    parser.add_argument(
-        '--min-dynamic-batch', type=int, default=0,
-        help='Guarantees a minimum number of rollouts in the batch when using dynamic batching. (default: 0)'
-    )
-    parser.add_argument(
-        '--host-training-info-interval', type=int, default=100,
-        help='The number of training steps before the host writes an info summary. (default: 100)'
-    )
-    parser.add_argument(
-        '--use-local-buffers', type=parse_bool, nargs='?', const=True, default=False,
-        help='If true all workers use their local network buffers (for batch norm: mean & var are not shared) (default: False)'
-    )
-    args = parser.parse_args()
+    base_parser = argparse.ArgumentParser(description='AdeptRL IMPALA Mode')
+
+    def add_args(parser):
+        parser = parser.add_argument_group('IMPALA Mode Args')
+        parser.add_argument('--gpu-id', type=int, nargs='+', default=[0],
+                            help='Which GPU to use for training. The host will always be the first gpu, workers are distributed evenly across the rest (default: [0])')
+        parser.add_argument(
+            '--max-queue-length', type=int, default=(size - 1) * 2,
+            help='Maximum rollout queue length. If above the max, workers will wait to append (default: (size - 1) * 2)'
+        )
+        parser.add_argument(
+            '--num-rollouts-in-batch', type=int, default=(size - 1),
+            help='The batch size in rollouts (so total batch is this number * nb_env * seq_len). '
+                 + 'Not compatible with --dynamic-batch (default: (size - 1))'
+        )
+        parser.add_argument(
+            '--max-dynamic-batch', type=int, default=0,
+            help='When > 0 uses dynamic batching (disables cudnn and --num-rollouts-in-batch). '
+                 + 'Limits the maximum rollouts in the batch to limit GPU memory usage. (default: 0 (False))'
+        )
+        parser.add_argument(
+            '--min-dynamic-batch', type=int, default=0,
+            help='Guarantees a minimum number of rollouts in the batch when using dynamic batching. (default: 0)'
+        )
+        parser.add_argument(
+            '--host-training-info-interval', type=int, default=100,
+            help='The number of training steps before the host writes an info summary. (default: 100)'
+        )
+        parser.add_argument(
+            '--use-local-buffers', type=parse_bool, nargs='?', const=True, default=False,
+            help='If true all workers use their local network buffers (for batch norm: mean & var are not shared) (default: False)'
+        )
+
+    add_base_args(base_parser, add_args)
+    args = base_parser.parse_args()
 
     if args.debug:
         args.nb_env = 3