Added optimizers to policy

* Removed save_obs param from run_model
* Updated nsra.py
* Fixed _get_pos bug in hrl envs

Author: sash-a

configs/ns.json

@@ -1,6 +1,6 @@
 {
   "env": {
-    "name": "HopperBulletEnv-v0",
+    "name": "AntMaze-v0",
     "max_steps": 2000
   },
   "noise": {
@@ -10,17 +10,24 @@
     "std_decay": 1
   },
   "policy": {
-    "ac_std": 0.01,
+    "layer_sizes": [
+      256,
+      256,
+      256
+    ],
+    "ac_std": 0.05,
+    "ac_std_decay": 0.99,
     "l2coeff": 0.005,
     "lr": 0.01,
     "lr_limit": 0.001,
     "lr_decay": 1,
-    "save_obs_chance": 0.01
+    "save_obs_chance": 0.01,
+    "ob_clip": 5
   },
   "general": {
     "name": "ns",
     "gens": 5000,
-    "policies_per_gen": 9600,
+    "policies_per_gen": 4800,
     "eps_per_policy": 1,
     "n_policies": 5,
     "batch_size": 500,

multi_agent.py

@@ -14,7 +14,7 @@
 from src.gym.unity import UnityGymWrapper
 from src.nn.nn import FeedForward
 from src.nn.obstat import ObStat
-from src.nn.optimizers import Adam, Optimizer
+from src.nn.optimizers import Adam
 from src.utils import utils
 from src.utils.rankers import CenteredRanker
 from src.utils.reporters import LoggerReporter, ReporterSet, StdoutReporter, MLFlowReporter
@@ -92,29 +92,28 @@ def custom_test_params(n: int, policies: List[Policy], fit_fn, obstats: List[ObS
     # initializing obstat, policy, optimizer, noise and ranker
     obstats: List[ObStat] = [ObStat(env.observation_space[i].shape, 1e-2) for i in range(2)]
     neuralnets = [FeedForward(cfg.policy.layer_sizes, torch.nn.Tanh(), env, cfg.policy.ac_std, cfg.policy.ob_clip)]
-    policies: List[Policy] = [Policy(nn, cfg.noise.std) for nn in neuralnets]
-    optims: List[Optimizer] = [Adam(policy, cfg.policy.lr) for policy in policies]
+    policies: List[Policy] = [Policy(nn, cfg, Adam) for nn in neuralnets]
     nt: NoiseTable = NoiseTable.create_shared(comm, cfg.noise.tbl_size, len(policies[0]), None, cfg.general.seed)
     ranker = CenteredRanker()
 
 
     def r_fn(models: List[torch.nn.Module], use_ac_noise=True) -> TrainingResult:
         save_obs = rs.random() < cfg.policy.save_obs_chance
-        rews, behv, obs, stps = gym_runner.multi_agent_gym_runner(models,
-                                                                  env,
-                                                                  cfg.env.max_steps,
-                                                                  rs if use_ac_noise else None,
-                                                                  save_obs)
-        return MultiAgentTrainingResult(rews, behv, obs, stps)
+        rews, behv, obs, steps = gym_runner.multi_agent_gym_runner(models,
+                                                                   env,
+                                                                   cfg.env.max_steps,
+                                                                   rs if use_ac_noise else None)
+        return MultiAgentTrainingResult(rews, behv,
+                                        obs if save_obs else np.array([np.zeros(env.observation_space.shape)]), steps)
 
 
     for gen in range(cfg.general.gens):
         reporter.start_gen()
         gen_obstats = [ObStat(env.observation_space[i].shape, 0) for i in range(2)]
         results = custom_test_params(eps_per_proc, policies, r_fn, gen_obstats)
-        for (pos_res, neg_res, inds, steps), policy, optim in zip(results, policies, optims):
+        for (pos_res, neg_res, inds, steps), policy in zip(results, policies):
             ranker.rank(pos_res, neg_res, inds)
-            es.approx_grad(ranker, nt, policy.flat_params, optim, cfg.general.batch_size, cfg.policy.l2coeff)
+            es.approx_grad(policy, ranker, nt, policy.flat_params, cfg.general.batch_size, cfg.policy.l2coeff)
             noiseless_result = RewardResult([0], [0], np.empty(1), 0)
             reporter.log_gen(ranker.fits, noiseless_result, policy, steps)
 

nsra.py

@@ -15,8 +15,7 @@
 from src.gym import gym_runner
 from src.gym.training_result import NSRResult, NSResult
 from src.nn.nn import FeedForward
-from src.nn.obstat import ObStat
-from src.nn.optimizers import Adam, Optimizer
+from src.nn.optimizers import Adam
 from src.utils import utils
 from src.utils.novelty import update_archive, novelty
 from src.utils.rankers import CenteredRanker, MultiObjectiveRanker
@@ -89,23 +88,21 @@ def main(cfg: Munch):
 
     archive: Optional[np.ndarray] = None
 
-    def ns_fn(model: torch.nn.Module) -> NSRResult:
+    def ns_fn(model: torch.nn.Module, use_ac_noise=True) -> NSRResult:
         """Reward function"""
-        rews, behv, obs, steps = gym_runner.run_model(model, env, cfg.env.max_steps, rs)
-        return NSRResult(rews, behv, obs, steps, archive, cfg.novelty.k)
+        save_obs = rs.random() < cfg.policy.save_obs_chance
+        rews, behv, obs, steps = gym_runner.run_model(model, env, cfg.env.max_steps, rs if use_ac_noise else None)
+        return NSRResult(rews, behv, obs if save_obs else np.array([np.zeros(env.observation_space.shape)]), steps,
+                         archive, cfg.novelty.k)
 
     # init population
     population = []
     nns = []
     for _ in range(cfg.general.n_policies):
         nns.append(FeedForward(cfg.policy.layer_sizes, torch.nn.Tanh(), env, cfg.policy.ac_std, cfg.policy.ob_clip))
-        population.append(Policy(nns[-1], cfg.noise.std))
+        population.append(Policy(nns[-1], cfg, Adam))
     # init optimizer and noise table
-    optims: List[Optimizer] = [Adam(policy, cfg.policy.lr) for policy in population]
     nt: NoiseTable = NoiseTable.create_shared(comm, cfg.noise.tbl_size, len(population[0]), reporter, cfg.general.seed)
-
-    obstat: ObStat = ObStat(env.observation_space.shape, 1e-2)  # eps to prevent dividing by zero at the beginning
-
     policies_best_rewards = [-np.inf] * cfg.general.n_policies
     time_since_best = [0 for _ in range(cfg.general.n_policies)]  # TODO should this be per individual?
     obj_weight = [cfg.nsr.initial_w for _ in range(cfg.general.n_policies)]
@@ -120,7 +117,6 @@ def ns_fn(model: torch.nn.Module) -> NSRResult:
         idx = random.choices(list(range(len(policies_novelties))), weights=policies_novelties, k=1)[0]
         if cfg.nsr.progressive: idx = gen % cfg.general.n_policies
         idx = comm.scatter([idx] * comm.size)
-        nns[idx].set_ob_mean_std(obstat.mean, obstat.std)
         ranker = MultiObjectiveRanker(CenteredRanker(), obj_weight[idx])
         # reporting
         if cfg.general.mlflow: mlflow_reporter.set_active_run(idx)
@@ -129,11 +125,12 @@ def ns_fn(model: torch.nn.Module) -> NSRResult:
         reporter.log({'w': obj_weight[idx]})
         reporter.log({'time since best': time_since_best[idx]})
         # running es
-        tr, gen_obstat = es.step(cfg, comm, population[idx], optims[idx], nt, env, ns_fn, rs, ranker, reporter)
+        tr, gen_obstat = es.step(cfg, comm, population[idx], nt, env, ns_fn, rs, ranker, reporter)
+        for policy in population:
+            policy.update_obstat(gen_obstat)  # shared obstat
+
         # sharing result and obstat
         tr = comm.scatter([tr] * comm.size)
-        gen_obstat.mpi_inc(comm)
-        obstat += gen_obstat
         # updating the weighting for choosing the next policy to be evaluated
         behv = comm.scatter([mean_behv(population[idx], ns_fn, cfg.novelty.rollouts)] * comm.size)
         nov = comm.scatter([novelty(behv, archive, cfg.novelty.k)] * comm.size)

obj.py

@@ -12,7 +12,7 @@
 from src.gym import gym_runner
 from src.gym.training_result import TrainingResult, RewardResult
 from src.nn.nn import FeedForward, BaseNet
-from src.nn.optimizers import Adam, Optimizer
+from src.nn.optimizers import Adam
 from src.utils import utils
 from src.utils.rankers import CenteredRanker, EliteRanker
 from src.utils.reporters import LoggerReporter, ReporterSet, StdoutReporter, MLFlowReporter
@@ -42,9 +42,9 @@ def main(cfg):
         nn: BaseNet = policy._module
     else:
         nn: BaseNet = FeedForward(cfg.policy.layer_sizes, torch.nn.Tanh(), env, cfg.policy.ac_std, cfg.policy.ob_clip)
-        policy: Policy = Policy(nn, cfg.noise.std)
+        policy: Policy = Policy(nn, cfg, Adam)
+    # optim: Optimizer = Adam(policy, cfg.policy.lr)
 
-    optim: Optimizer = Adam(policy, cfg.policy.lr)
     nt: NoiseTable = NoiseTable.create_shared(comm, cfg.noise.tbl_size, len(policy), reporter, cfg.general.seed)
 
     ranker = CenteredRanker()
@@ -59,12 +59,12 @@ def r_fn(model: torch.nn.Module, use_ac_noise=True) -> TrainingResult:
         save_obs = rs.random() < cfg.policy.save_obs_chance
         rews = np.zeros(cfg.env.max_steps)
         for _ in range(max(1, cfg.general.eps_per_policy)):
-            rew, behv, obs, steps = gym_runner.run_model(model, env, cfg.env.max_steps,
-                                                         rs if use_ac_noise else None, save_obs)
+            rew, behv, obs, steps = gym_runner.run_model(model, env, cfg.env.max_steps, rs if use_ac_noise else None)
             rews[:len(rew)] += np.array(rew)
 
         rews /= max(1, cfg.general.eps_per_policy)
-        return RewardResult(rews.tolist(), behv, obs, steps)
+        return RewardResult(rews.tolist(), behv, obs if save_obs else np.array([np.zeros(env.observation_space.shape)]),
+                            steps)
 
     time_since_best = 0
     noise_std_inc = 0.08
@@ -75,16 +75,16 @@ def r_fn(model: torch.nn.Module, use_ac_noise=True) -> TrainingResult:
         if cfg.noise.std_decay != 1:
             reporter.log({'noise std': policy.std})
         if cfg.policy.lr_decay != 1:
-            reporter.log({'lr': optim.lr})
+            reporter.log({'lr': policy.optim.lr})
         if cfg.policy.ac_std_decay != 1:
             reporter.log({'ac std': nn._action_std})
 
-        tr, gen_obstat = es.step(cfg, comm, policy, optim, nt, env, r_fn, rs, ranker, reporter)
+        tr, gen_obstat = es.step(cfg, comm, policy, nt, env, r_fn, rs, ranker, reporter)
         policy.update_obstat(gen_obstat)
 
         cfg.policy.ac_std = nn._action_std = nn._action_std * cfg.policy.ac_std_decay
         cfg.noise.std = policy.std = max(cfg.noise.std * cfg.noise.std_decay, cfg.noise.std_limit)
-        cfg.policy.lr = optim.lr = max(cfg.policy.lr * cfg.policy.lr_decay, cfg.policy.lr_limit)
+        cfg.policy.lr = policy.optim.lr = max(cfg.policy.lr * cfg.policy.lr_decay, cfg.policy.lr_limit)
 
         reporter.log({'obs recorded': policy.obstat.count})
 

simple_example.py

@@ -10,7 +10,7 @@
 from src.gym.training_result import TrainingResult, RewardResult
 from src.nn.nn import FeedForward
 from src.nn.obstat import ObStat
-from src.nn.optimizers import Adam, Optimizer
+from src.nn.optimizers import Adam
 from src.utils import utils
 from src.utils.rankers import CenteredRanker
 from src.utils.utils import generate_seed
@@ -28,34 +28,31 @@
     rs = utils.seed(comm, cfg.general.seed, env)
 
     # initializing obstat, policy, optimizer, noise and ranker
-    obstat: ObStat = ObStat(env.observation_space.shape, 1e-2)  # eps to prevent dividing by zero at the beginning
     nn = FeedForward(cfg.policy.layer_sizes, torch.nn.Tanh(), env, cfg.policy.ac_std, cfg.policy.ob_clip)
-    policy: Policy = Policy(nn, cfg.noise.std)
-    optim: Optimizer = Adam(policy, cfg.policy.lr)
+    policy: Policy = Policy(nn, cfg, Adam)
     nt: NoiseTable = NoiseTable.create_shared(comm, cfg.noise.tbl_size, len(policy), None, cfg.general.seed)
     ranker = CenteredRanker()
 
 
     def r_fn(model: torch.nn.Module) -> TrainingResult:
         save_obs = (rs.random() if rs is not None else np.random.random()) < cfg.policy.save_obs_chance
-        rews, behv, obs, steps = gym_runner.run_model(model, env, 10000, rs, save_obs)
-        return RewardResult(rews, behv, obs, steps)
+        rews, behv, obs, steps = gym_runner.run_model(model, env, 10000, rs)
+        return RewardResult(rews, behv, obs if save_obs else np.array([np.zeros(env.observation_space.shape)]), steps)
 
 
     assert cfg.general.policies_per_gen % comm.size == 0 and (cfg.general.policies_per_gen / comm.size) % 2 == 0
     eps_per_proc = int((cfg.general.policies_per_gen / comm.size) / 2)
     for gen in range(cfg.general.gens):  # main loop
         if comm.rank == 0: print(f'Generation:{gen}')  # only print on one process
-        nn.set_ob_mean_std(obstat.mean, obstat.std)  # for normalizing the observation space
 
         # the block below is encapsulated in es.step(...), but this is more flexible. Example use can be seen in obj.py
         gen_obstat = ObStat(env.observation_space.shape, 0)  # for normalizing the observation space
         # obtaining the fitnesses from many perturbed policies
         pos_fits, neg_fits, inds, steps = es.test_params(comm, eps_per_proc, policy, nt, gen_obstat, r_fn, rs)
-        obstat += gen_obstat  # adding the new observations to the global obstat
+        policy.update_obstat(gen_obstat)
         ranker.rank(pos_fits, neg_fits, inds)  # ranking the fitnesses between -1 and 1
         # approximating the gradient and updating policy.flat_params (pseudo backprop)
-        es.approx_grad(ranker, nt, policy.flat_params, optim, cfg.general.batch_size, cfg.policy.l2coeff)
+        es.approx_grad(policy, ranker, nt, policy.flat_params, cfg.general.batch_size, cfg.policy.l2coeff)
 
         if comm.rank == 0: print(f'avg fitness:{np.mean(np.concatenate((pos_fits, neg_fits)))}\n\n')
         if gen % 10 and comm.rank == 0:  # save policy every 10 generations

src/core/es.py

@@ -16,7 +16,6 @@
 from src.core.policy import Policy
 from src.gym.training_result import TrainingResult
 from src.nn.obstat import ObStat
-from src.nn.optimizers import Optimizer
 from src.utils.rankers import Ranker, CenteredRanker
 from src.utils.reporters import StdoutReporter, Reporter
 from src.utils.utils import scale_noise
@@ -26,7 +25,6 @@
 def step(cfg,
          comm: MPI.Comm,
          policy: Policy,
-         optim: Optimizer,
          nt: NoiseTable,
          env: gym.Env,
          fit_fn: Callable[[Module], TrainingResult],
@@ -48,7 +46,7 @@ def step(cfg,
     reporter.print(f'n dupes: {len(inds) - len(set(inds))}')
 
     ranker.rank(pos_res, neg_res, inds)
-    approx_grad(ranker, nt, policy.flat_params, optim, cfg.general.batch_size, cfg.policy.l2coeff)
+    approx_grad(policy, ranker, nt, policy.flat_params, cfg.general.batch_size, cfg.policy.l2coeff)
     noiseless_result = fit_fn(policy.pheno(np.zeros(len(policy))), False)
     reporter.log_gen(ranker.fits, noiseless_result, policy, steps)
 
@@ -99,7 +97,7 @@ def _share_results(comm: MPI.Comm,
     return results.reshape((-1, 1 + 2 * objectives))  # flattening the process dim
 
 
-def approx_grad(ranker: Ranker, nt: NoiseTable, params: ndarray, optim: Optimizer, batch_size: int, l2coeff: float):
+def approx_grad(policy: Policy, ranker: Ranker, nt: NoiseTable, params: ndarray, batch_size: int, l2coeff: float):
     """Approximating gradient and update policy params"""
     grad = scale_noise(ranker.ranked_fits, ranker.noise_inds, nt, batch_size) / ranker.n_fits_ranked
-    optim.step(l2coeff * params - grad)
+    policy.optim.step(l2coeff * params - grad)

src/core/policy.py

@@ -2,29 +2,32 @@
 
 import os
 import pickle
+from typing import Type
 
 import numpy as np
 import torch
+from munch import Munch
 
 from src.nn.nn import BaseNet
 from src.nn.obstat import ObStat
+from src.nn.optimizers import Optimizer
 
 
 def init_normal(m):
     if type(m) == torch.nn.Linear:
         torch.nn.init.kaiming_normal_(m.weight)
 
 
-class Policy(torch.nn.Module):
-    def __init__(self, module: BaseNet, std: float):
-        super().__init__()
+class Policy:
+    def __init__(self, module: BaseNet, cfg: Munch, OptimType: Type[Optimizer]):
         module.apply(init_normal)
 
         self._module: BaseNet = module
-        self.std = std
+        self.std = cfg.noise.std
 
         self.flat_params: np.ndarray = Policy.get_flat(module)
         self.obstat: ObStat = ObStat(module._obmean.shape, 1e-2)
+        self.optim = OptimType(self, cfg.policy.lr)
 
     def __len__(self):
         return len(self.flat_params)
@@ -68,6 +71,3 @@ def pheno(self, noise: np.ndarray = None) -> torch.nn.Module:
     def update_obstat(self, obstat: ObStat):
         self.obstat += obstat  # adding the new observations to the global obstat
         self._module.set_ob_mean_std(self.obstat.mean, self.obstat.std)
-
-    def forward(self, inp):
-        self._module.forward(inp)

src/gym/gym_runner.py

@@ -15,7 +15,6 @@ def run_model(model: torch.nn.Module,
               env: gym.Env,
               max_steps: int,
               rs: np.random.RandomState = None,
-              save_obs: bool = False,
               render: bool = False) -> Tuple[List[float], List[float], np.ndarray, int]:
     """
     Evaluates model on the provided env
@@ -32,10 +31,9 @@ def run_model(model: torch.nn.Module,
 
             action = model(ob, rs=rs)
             ob, rew, done, _ = env.step(action.numpy())
-            if save_obs:
-                obs.append(ob)
 
             rews += [rew]
+            obs.append(ob)
             behv.extend(_get_pos(env.unwrapped))
 
             if render:
@@ -44,9 +42,6 @@ def run_model(model: torch.nn.Module,
             if done:
                 break
 
-    if not save_obs:
-        obs.append(np.zeros(ob.shape))
-
     behv += behv[-3:] * (max_steps - int(len(behv) / 3))  # extending the behaviour vector to have `max_steps` elements
     return rews, behv, np.array(obs), step
 
@@ -96,8 +91,8 @@ def multi_agent_gym_runner(policies: List[torch.nn.Module],
 
 
 def _get_pos(env):
-    if env.spec.id[:-3] in ["AntMaze", "AntPush", "AntFall"]:
-        return env._robot_x, env._robot_y, 0
+    if env.spec.id[:-3] in ["AntMaze", "AntPush", "AntFall"]:  # hrl ant env
+        return env.wrapped_env.get_body_com("torso")[:3]
 
     if BULLET_ENV_SUFFIX in env.spec.id:  # bullet env
         return env.robot_body.get_pose()[:3]