default dict not working like i think

adept/agents/agent_registry.py

@@ -59,5 +59,12 @@ def lookup_agent(self, agent_id):
         """
         return self._agent_class_by_id[agent_id]
 
-    def lookup_nb_output(self):
-
+    def lookup_output_shape(self, agent_id, action_space):
+        """
+        For a given agent_id, determine the shapes of the outputs.
+
+        :param agent_id: str
+        :param action_space:
+        :return:
+        """
+        return self._agent_class_by_id[agent_id].output_shape(action_space)

adept/environments/env_registry.py

@@ -74,7 +74,7 @@ class EnvPluginRegistry:
     def __init__(self):
         self._engine_ids_by_env_id_set = {}
         self._plugin_class_by_engine_id = {}
-        self._reward_norm_by_env_id = defaultdict(Clip)
+        self._reward_norm_by_env_id = defaultdict(Clip())
 
         from adept.environments.openai_gym import AdeptGymEnv
         self.register_env(Engines.GYM, AdeptGymEnv, ATARI_ENVS)

adept/scripts/local.py

@@ -81,7 +81,7 @@
     SimpleModelSaver
 )
 from adept.utils.script_helpers import (
-    make_network, get_head_shapes, count_parameters
+    make_network, count_parameters
 )
 
 # hack to use bypass pysc2 flags
@@ -148,8 +148,11 @@ def main(
 
     # 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)
+    network = make_network(
+        env.observation_space,
+        agent_registry.lookup_output_shape(args.agent, env.action_space),
+        args
+    )
     # possibly load network
     initial_step_count = 0
     if args.load_network:
@@ -211,7 +214,9 @@ def make_optimizer(params):
             eval_args, seed=args.seed + args.nb_env, registry=env_registry
         )
         eval_net = make_network(
-            eval_env.observation_space, network_head_shapes, eval_args
+            eval_env.observation_space,
+            agent_registry.lookup_output_shape(args.agent, env.action_space),
+            eval_args
         )
         eval_agent = agent_registry.lookup_agent(args.agent).from_args(
             eval_args,

adept/utils/script_helpers.py

@@ -14,11 +14,8 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 from argparse import ArgumentParser  # for type hinting
 
-from adept.agents import AGENTS
-from adept.environments._env import reward_normalizer_by_env_id
 from adept.networks import VISION_NETWORKS, DISCRETE_NETWORKS, NETWORK_BODIES
 from adept.networks._base import NetworkTrunk, ModularNetwork, NetworkHead
-from adept.utils.util import parse_bool
 
 
 def make_network(
@@ -61,185 +58,3 @@ def make_network(
 
 def count_parameters(model):
     return sum(p.numel() for p in model.parameters() if p.requires_grad)
-
-
-def get_head_shapes(action_space, agent_name):
-    Agent = AGENTS[agent_name]
-    return Agent.output_shape(action_space)
-
-
-def _add_common_agent_args(parser: ArgumentParser):
-    parser.add_argument(
-        '-al',
-        '--learning-rate',
-        type=float,
-        default=7e-4,
-        help='learning rate (default: 7e-4)'
-    )
-    parser.add_argument(
-        '-ad',
-        '--discount',
-        type=float,
-        default=0.99,
-        help='discount factor for rewards (default: 0.99)'
-    )
-
-
-def _add_agent_args(subparser: ArgumentParser):
-    agent_parsers = []
-    for agent_name, agent_class in AGENTS.items():
-        parser_agent = subparser.add_parser(agent_name)
-        agent_group = parser_agent.add_argument_group('Agent Args')
-        _add_common_agent_args(agent_group)
-        agent_class.add_args(agent_group)
-        agent_parsers.append(parser_agent)
-    return agent_parsers
-
-
-def _add_network_args(parser: ArgumentParser):
-    subparser = parser.add_argument_group('Network Args')
-    subparser.add_argument('-nv', '--network-net3d', default='FourConv')
-    subparser.add_argument('-nd', '--network-net1d', default='Identity')
-    subparser.add_argument('-nb', '--network-body', default='LSTM')
-    subparser.add_argument(
-        '--normalize',
-        type=parse_bool,
-        nargs='?',
-        const=True,
-        default=True,
-        help=
-        'Applies batch norm between linear/convolutional'
-        'layers and layer norm for LSTMs (default: True)'
-    )
-
-
-def _add_reload_args(parser: ArgumentParser):
-    subparser = parser.add_argument_group('Reload Args')
-    # Reload from save
-    subparser.add_argument(
-        '-ln',
-        '--load-network',
-        default='',
-        help='Load network from this path. Sets initial step count'
-    )
-    subparser.add_argument(
-        '-lo',
-        '--load-optimizer',
-        default='',
-        help='Load optimizer from this path'
-    )
-
-
-def _add_env_args(parser: ArgumentParser):
-    subparser = parser.add_argument_group('Environment Args')
-    subparser.add_argument(
-        '-e',
-        '--env-id',
-        default='PongNoFrameskip-v4',
-        help='environment to train on (default: PongNoFrameskip-v4)'
-    )
-    subparser.add_argument(
-        '-ne',
-        '--nb-env',
-        type=int,
-        default=32,
-        help='number of envs to run in parallel (default: 32)'
-    )
-    subparser.add_argument(
-        '-es',
-        '--skip-rate',
-        type=int,
-        default=4,
-        help='frame skip rate (default: 4)'
-    )
-    subparser.add_argument(
-        '-em',
-        '--max-episode-length',
-        type=int,
-        default=10000,
-        metavar='MEL',
-        help='maximum length of an episode (default: 10000)'
-    )
-
-
-def _add_common_args(parser: ArgumentParser):
-    _add_env_args(parser)
-    _add_network_args(parser)
-    _add_reload_args(parser)
-    """
-    Common Arguments
-    """
-    subparser = parser.add_argument_group('Common Args')
-    subparser.add_argument(
-        '-cl',
-        '--log-dir',
-        default='/tmp/adept_logs/',
-        help='Folder to save logs. (default: /tmp/adept_logs)'
-    )
-    subparser.add_argument(
-        '-ct',
-        '--tag',
-        default='',
-        help=
-        'Identify your experiment with a tag that'
-        'gets prepended to the experiment log directory'
-    )
-    subparser.add_argument(
-        '-cm',
-        '--max-train-steps',
-        type=int,
-        default=10e6,
-        help='Number of steps to train for (default: 10e6)'
-    )
-    subparser.add_argument(
-        '-ce',
-        '--epoch-len',
-        type=int,
-        default=1e6,
-        metavar='FREQ',
-        help='Save models every FREQ steps'
-    )
-    subparser.add_argument(
-        '-cf',
-        '--summary-frequency',
-        default=10,
-        help='Write tensorboard summaries every FREQ seconds'
-    )
-    subparser.add_argument(
-        '--seed', type=int, default=0, help='Random seed (default: 0)'
-    )
-    subparser.add_argument(
-        '--profile',
-        type=parse_bool,
-        nargs='?',
-        const=True,
-        default=False,
-        help='displays profiling tree after 10e3 steps (default: False)'
-    )
-    subparser.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)'
-    )
-
-
-def _add_args_to_parsers(arg_fn, parsers):
-    return [arg_fn(x) for x in parsers]
-
-
-def add_base_args(parser: ArgumentParser, additional_args_fn=None):
-    # TODO: there must be a better way of adding
-    #  args to subparsers while keeping the help message
-    # TODO: some agents may not run in certain modes not sure the best way to handle this
-    subparser_agent = parser.add_subparsers(title='Agents', dest='agent')
-    subparser_agent.required = True
-    agent_parsers = _add_agent_args(subparser_agent)
-    _add_args_to_parsers(_add_common_args, agent_parsers)
-    _add_args_to_parsers(additional_args_fn, agent_parsers)
-
-    return parser