remove cli_utils wildcard imports from examples

Author: jakegrigsby

examples/00_meta_frozen_lake.py

@@ -1,8 +1,10 @@
+from argparse import ArgumentParser
+
 import amago
 from amago.envs.builtin.toy_gym import MetaFrozenLake
 from amago.envs import AMAGOEnv
 from amago.loading import DiskTrajDataset
-from amago.cli_utils import *
+from amago import cli_utils
 
 
 def add_cli(parser):
@@ -35,27 +37,27 @@ def add_cli(parser):
 
     config = {}
     # configure trajectory encoder (seq2seq memory model)
-    traj_encoder_type = switch_traj_encoder(
+    traj_encoder_type = cli_utils.switch_traj_encoder(
         config,
         arch=args.seq_model,
         memory_size=128,
         layers=3,
     )
     # configure timestep encoder
-    tstep_encoder_type = switch_tstep_encoder(
+    tstep_encoder_type = cli_utils.switch_tstep_encoder(
         config, arch="ff", n_layers=1, d_hidden=128, d_output=64, normalize_inputs=False
     )
 
     # we're using the default exploration strategy but being overly verbose about it for the example
-    exploration_wrapper_type = switch_exploration(
+    exploration_wrapper_type = cli_utils.switch_exploration(
         config,
         strategy="egreedy",
         eps_start=1.0,
         eps_end=0.05,
         steps_anneal=1_000_000,
         randomize_eps=True,
     )
-    use_config(config)
+    cli_utils.use_config(config)
 
     group_name = f"{args.run_name}_{args.seq_model}"
     for trial in range(args.trials):

examples/01_basic_gym.py

@@ -5,7 +5,7 @@
 
 import amago
 from amago.envs import AMAGOEnv
-from amago.cli_utils import *
+from amago import cli_utils
 
 
 def add_cli(parser):
@@ -26,7 +26,7 @@ def add_cli(parser):
 
 if __name__ == "__main__":
     parser = ArgumentParser()
-    add_common_cli(parser)
+    cli_utils.add_common_cli(parser)
     add_cli(parser)
     args = parser.parse_args()
 
@@ -43,24 +43,24 @@ def add_cli(parser):
         "amago.nets.policy_dists.Discrete.clip_prob_low": 1e-6,
     }
     # switch sequence model
-    traj_encoder_type = switch_traj_encoder(
+    traj_encoder_type = cli_utils.switch_traj_encoder(
         config,
         arch=args.traj_encoder,
         memory_size=args.memory_size,
         layers=args.memory_layers,
     )
     # switch agent
-    agent_type = switch_agent(
+    agent_type = cli_utils.switch_agent(
         config,
         args.agent_type,
         reward_multiplier=1.0,
     )
-    use_config(config, args.configs)
+    cli_utils.use_config(config, args.configs)
 
     group_name = f"{args.run_name}_{env_name}"
     for trial in range(args.trials):
         run_name = group_name + f"_trial_{trial}"
-        experiment = create_experiment_from_cli(
+        experiment = cli_utils.create_experiment_from_cli(
             args,
             make_train_env=make_train_env,
             make_val_env=make_train_env,
@@ -73,7 +73,7 @@ def add_cli(parser):
             group_name=group_name,
             val_timesteps_per_epoch=args.eval_timesteps,
         )
-        experiment = switch_async_mode(experiment, args.mode)
+        experiment = cli_utils.switch_async_mode(experiment, args.mode)
         experiment.start()
         if args.ckpt is not None:
             experiment.load_checkpoint(args.ckpt)

examples/02_gymnax.py

@@ -1,6 +1,6 @@
 """
-Support for gymnax is experimental and mainly meant to test the already_vectorized 
-env API used by XLand MiniGrid (an unsolved environment) with classic gym envs. 
+Support for gymnax is experimental and mainly meant to test the already_vectorized
+env API used by XLand MiniGrid (an unsolved environment) with classic gym envs.
 Many of the gymnax envs appear to be broken by recent versions of jax.
 There are a couple memory/meta-RL bsuite envs where AMAGO+Transformer
 is significantly better than the gymnax reference scores though.
@@ -24,7 +24,7 @@
 from amago.envs import AMAGOEnv
 from amago.envs.builtin.gymnax_envs import GymnaxCompatibility
 from amago.nets.cnn import GridworldCNN
-from amago.cli_utils import *
+from amago import cli_utils
 
 
 def add_cli(parser):
@@ -52,7 +52,7 @@ def guess_tstep_encoder(config, obs_shape):
     if len(obs_shape) == 3:
         print(f"Guessing CNN for observation of shape {obs_shape}")
         channels_first = np.argmin(obs_shape).item() == 0
-        return switch_tstep_encoder(
+        return cli_utils.switch_tstep_encoder(
             config,
             "cnn",
             cnn_type=GridworldCNN,
@@ -61,7 +61,7 @@ def guess_tstep_encoder(config, obs_shape):
     else:
         print(f"Guessing MLP for observation of shape {obs_shape}")
         dim = math.prod(obs_shape)  # FFTstepEncoder will flatten the obs on input
-        return switch_tstep_encoder(
+        return cli_utils.switch_tstep_encoder(
             config,
             "ff",
             d_hidden=max(dim // 3, 128),
@@ -72,7 +72,7 @@ def guess_tstep_encoder(config, obs_shape):
 
 if __name__ == "__main__":
     parser = ArgumentParser()
-    add_common_cli(parser)
+    cli_utils.add_common_cli(parser)
     add_cli(parser)
     args = parser.parse_args()
 
@@ -81,7 +81,7 @@ def guess_tstep_encoder(config, obs_shape):
 
     # config
     config = {}
-    traj_encoder_type = switch_traj_encoder(
+    traj_encoder_type = cli_utils.switch_traj_encoder(
         config,
         arch=args.traj_encoder,
         memory_size=args.memory_size,
@@ -91,16 +91,16 @@ def guess_tstep_encoder(config, obs_shape):
         test_env, env_params = gymnax.make(args.env)
         test_obs_shape = test_env.observation_space(env_params).shape
     tstep_encoder_type = guess_tstep_encoder(config, test_obs_shape)
-    agent_type = switch_agent(config, args.agent_type)
+    agent_type = cli_utils.switch_agent(config, args.agent_type)
 
-    use_config(config, args.configs)
+    cli_utils.use_config(config, args.configs)
     make_env = partial(
         make_gymnax_amago, env_name=args.env, parallel_envs=args.parallel_actors
     )
     group_name = f"{args.run_name}_{args.env}"
     for trial in range(args.trials):
         run_name = group_name + f"_trial_{trial}"
-        experiment = create_experiment_from_cli(
+        experiment = cli_utils.create_experiment_from_cli(
             args,
             make_train_env=make_env,
             make_val_env=make_env,
@@ -116,7 +116,7 @@ def guess_tstep_encoder(config, obs_shape):
             l2_coeff=1e-4,
             save_trajs_as="npz-compressed",
         )
-        experiment = switch_async_mode(experiment, args.mode)
+        experiment = cli_utils.switch_async_mode(experiment, args.mode)
         amago_device = experiment.DEVICE.index or torch.cuda.current_device()
         env_device = jax.devices("gpu")[amago_device]
         with jax.default_device(env_device):

examples/03_popgym_suite.py

@@ -2,8 +2,9 @@
 
 import wandb
 
+import amago
 from amago.envs.builtin.popgym_envs import POPGymAMAGO, MultiDomainPOPGymAMAGO
-from amago.cli_utils import *
+from amago import cli_utils
 
 
 def add_cli(parser):
@@ -19,7 +20,7 @@ def add_cli(parser):
 
 if __name__ == "__main__":
     parser = ArgumentParser()
-    add_common_cli(parser)
+    cli_utils.add_common_cli(parser)
     add_cli(parser)
     args = parser.parse_args()
 
@@ -43,32 +44,32 @@ def add_cli(parser):
     }
     # fmt: on
 
-    traj_encoder_type = switch_traj_encoder(
+    traj_encoder_type = cli_utils.switch_traj_encoder(
         config,
         arch=args.traj_encoder,
         memory_size=args.memory_size,  # paper: 256
         layers=args.memory_layers,  # paper: 3
     )
-    tstep_encoder_type = switch_tstep_encoder(
+    tstep_encoder_type = cli_utils.switch_tstep_encoder(
         config,
         arch="ff",
         n_layers=2,
         d_hidden=512,
         d_output=200,
     )
-    agent_type = switch_agent(
+    agent_type = cli_utils.switch_agent(
         config,
         args.agent_type,
         reward_multiplier=200.0 if args.multidomain else 100.0,
         tau=0.0025,
     )
     # steps_anneal can safely be set much lower (<500k) in most tasks. More sweeps needed.
-    exploration_type = switch_exploration(
+    exploration_type = cli_utils.switch_exploration(
         config,
         "egreedy",
         steps_anneal=1_000_000,
     )
-    use_config(config, args.configs)
+    cli_utils.use_config(config, args.configs)
 
     group_name = f"{args.run_name}_{args.env}"
     for trial in range(args.trials):
@@ -81,7 +82,7 @@ def add_cli(parser):
             make_train_env = lambda: POPGymAMAGO(
                 f"popgym-{args.env}-v0", truncated_is_done=True
             )
-        experiment = create_experiment_from_cli(
+        experiment = cli_utils.create_experiment_from_cli(
             args,
             make_train_env=make_train_env,
             make_val_env=make_train_env,
@@ -98,7 +99,7 @@ def add_cli(parser):
             grad_clip=1.0,
             lr_warmup_steps=2000,
         )
-        experiment = switch_async_mode(experiment, args.mode)
+        experiment = cli_utils.switch_async_mode(experiment, args.mode)
         experiment.start()
         if args.ckpt is not None:
             experiment.load_checkpoint(args.ckpt)

examples/04_tmaze.py

@@ -1,11 +1,12 @@
 from argparse import ArgumentParser
 
 import wandb
+import gin
 
 from amago.envs import AMAGOEnv
 from amago.envs.builtin.tmaze import TMazeAltPassive, TMazeAltActive
 from amago.envs.exploration import EpsilonGreedy
-from amago.cli_utils import *
+from amago import cli_utils
 
 
 def add_cli(parser):
@@ -60,31 +61,31 @@ def current_eps(self, local_step: int):
 
 if __name__ == "__main__":
     parser = ArgumentParser()
-    add_common_cli(parser)
+    cli_utils.add_common_cli(parser)
     add_cli(parser)
     args = parser.parse_args()
 
     config = {
         "TMazeExploration.horizon": args.horizon,
     }
-    traj_encoder_type = switch_traj_encoder(
+    traj_encoder_type = cli_utils.switch_traj_encoder(
         config,
         arch=args.traj_encoder,
         memory_size=args.memory_size,
         layers=args.memory_layers,
     )
-    tstep_encoder_type = switch_tstep_encoder(
+    tstep_encoder_type = cli_utils.switch_tstep_encoder(
         config,
         arch="ff",
         n_layers=2,
         d_hidden=128,
         d_output=128,
         normalize_inputs=False,
     )
-    agent_type = switch_agent(
+    agent_type = cli_utils.switch_agent(
         config, args.agent_type, reward_multiplier=100.0, gamma=0.9999
     )
-    use_config(config, args.configs)
+    cli_utils.use_config(config, args.configs)
 
     group_name = f"{args.run_name}_TMazePassive_H{args.horizon}"
     for trial in range(args.trials):
@@ -95,7 +96,7 @@ def current_eps(self, local_step: int):
             ),
             env_name=f"TMazePassive-H{args.horizon}",
         )
-        experiment = create_experiment_from_cli(
+        experiment = cli_utils.create_experiment_from_cli(
             args,
             make_train_env=make_env,
             make_val_env=make_env,
@@ -110,7 +111,7 @@ def current_eps(self, local_step: int):
             sample_actions=False,  # even softmax prob .999 isn't good enough for this env...
             exploration_wrapper_type=TMazeExploration,
         )
-        switch_async_mode(experiment, args.mode)
+        experiment = cli_utils.switch_async_mode(experiment, args.mode)
         experiment.start()
         if args.ckpt is not None:
             experiment.load_checkpoint(args.ckpt)

examples/05_dark_key_door.py

@@ -4,7 +4,7 @@
 
 from amago.envs.builtin.toy_gym import RoomKeyDoor
 from amago.envs import AMAGOEnv
-from amago.cli_utils import *
+from amago import cli_utils
 
 
 def add_cli(parser):
@@ -41,28 +41,30 @@ def add_cli(parser):
 
 if __name__ == "__main__":
     parser = ArgumentParser()
-    add_common_cli(parser)
+    cli_utils.add_common_cli(parser)
     add_cli(parser)
     args = parser.parse_args()
 
     config = {}
-    tstep_encoder_type = switch_tstep_encoder(
+    tstep_encoder_type = cli_utils.switch_tstep_encoder(
         config, arch="ff", n_layers=2, d_hidden=128, d_output=64
     )
-    traj_encoder_type = switch_traj_encoder(
+    traj_encoder_type = cli_utils.switch_traj_encoder(
         config,
         arch=args.traj_encoder,
         memory_size=args.memory_size,
         layers=args.memory_layers,
     )
-    agent_type = switch_agent(config, args.agent_type, reward_multiplier=100.0)
+    agent_type = cli_utils.switch_agent(
+        config, args.agent_type, reward_multiplier=100.0
+    )
     # the fancier exploration schedule mentioned in the appendix can help
     # when the domain is a true meta-RL problem and the "horizon" time limit
     # (above) is actually relevant for resetting the task.
-    exploration_type = switch_exploration(
+    exploration_type = cli_utils.switch_exploration(
         config, "bilevel", steps_anneal=500_000, rollout_horizon=args.meta_horizon
     )
-    use_config(config, args.configs)
+    cli_utils.use_config(config, args.configs)
 
     group_name = f"{args.run_name}_dark_key_door"
     for trial in range(args.trials):
@@ -77,7 +79,7 @@ def add_cli(parser):
             ),
             env_name=f"Dark-Key-To-Door-{args.room_size}x{args.room_size}",
         )
-        experiment = create_experiment_from_cli(
+        experiment = cli_utils.create_experiment_from_cli(
             args,
             agent_type=agent_type,
             tstep_encoder_type=tstep_encoder_type,
@@ -91,7 +93,7 @@ def add_cli(parser):
             val_timesteps_per_epoch=args.meta_horizon * 4,
             exploration_wrapper_type=exploration_type,
         )
-        switch_async_mode(experiment, args.mode)
+        experiment = cli_utils.switch_async_mode(experiment, args.mode)
         experiment.start()
         if args.ckpt is not None:
             experiment.load_checkpoint(args.ckpt)