Update evolutionary algorithms example

Author: m-wojnar

examples/cart-pole-evo/main.py

@@ -1,16 +1,15 @@
 from argparse import ArgumentParser
 
+import evosax.algorithms
 import gymnasium as gym
-import jax
 import numpy as np
 from chex import Array
-from evosax.algorithms import PGPE
 from flax import linen as nn
 
 from reinforced_lib import RLib
 from reinforced_lib.agents.neuro import Evosax
 from reinforced_lib.exts import GymnasiumVectorized
-from reinforced_lib.logs import StdoutLogger, TensorboardLogger, WeightsAndBiasesLogger
+from reinforced_lib.logs import CsvLogger, StdoutLogger
 
 
 class Network(nn.Module):
@@ -20,17 +19,19 @@ def __call__(self, x: Array) -> Array:
         x = nn.tanh(x)
         x = nn.Dense(64)(x)
         x = nn.tanh(x)
-        logits = nn.Dense(2)(x)
-        action = jax.random.categorical(self.make_rng('rlib'), logits)
+        x = nn.Dense(1)(x)
+        action = 2 * nn.tanh(x)
         return action
 
 
-def run(num_epochs: int, population_size: int, seed: int) -> None:
+def run(evo_alg: type, num_epochs: int, population_size: int, seed: int) -> None:
     """
-    Run ``num_epochs`` cart-pole Gymnasium environments in parallel using an evolutionary strategy to optimize the policy.
+    Run ``num_envs`` Pendulum Gymnasium environments in parallel using an evolutionary strategy to optimize the policy.
 
     Parameters
     ----------
+    evo_alg : type
+        Evolutionary strategy to use (from evosax).
     num_epochs : int
         Number of simulation steps to perform.
     population_size : int
@@ -43,44 +44,46 @@ def run(num_epochs: int, population_size: int, seed: int) -> None:
         agent_type=Evosax,
         agent_params={
             'network': Network(),
-            'evo_strategy': PGPE,
+            'evo_strategy': evo_alg,
             'evo_strategy_default_params': {'std_init': 0.1},
             'population_size': population_size
         },
         ext_type=GymnasiumVectorized,
-        ext_params={'env_id': 'CartPole-v1', 'num_envs': population_size},
-        logger_types=[StdoutLogger, TensorboardLogger, WeightsAndBiasesLogger]
+        ext_params={'env_id': 'Pendulum-v1', 'num_envs': population_size},
+        logger_types=[CsvLogger, StdoutLogger],
+        logger_params={'csv_path': f'pendulum-{evo_alg.__name__}-evo-{seed}.csv'}
     )
 
     def make_env():
-        return gym.make('CartPole-v1', render_mode='no')
+        return gym.make('Pendulum-v1')
 
-    for step in range(num_epochs):
-        env = gym.vector.SyncVectorEnv([make_env for _ in range(population_size)])
+    env = gym.vector.SyncVectorEnv([make_env for _ in range(population_size)])
 
-        _, _ = env.reset(seed=seed + step)
+    for epoch in range(num_epochs):
+        _, _ = env.reset(seed=seed + epoch)
         actions = env.action_space.sample()
+        return_pop = np.zeros(population_size, dtype=float)
 
-        terminal = np.array([False] * population_size)
-        max_epoch_len = 0
-
-        while not np.all(terminal):
+        for _ in range(env.envs[0].spec.max_episode_steps):
             env_states = env.step(np.asarray(actions))
             actions = rl.sample(*env_states)
+            return_pop += env_states[1]
 
-            terminal = terminal | env_states[2] | env_states[3]
-            max_epoch_len += 1
-
-        rl.log('max_epoch_len', max_epoch_len)
+        rl.log('mean_return', return_pop.mean())
+        rl.log('max_return', return_pop.max())
+        rl.log('epoch', epoch + 1)
 
 
 if __name__ == '__main__':
     args = ArgumentParser()
 
+    args.add_argument('--evo_alg', type=str, required=True)
     args.add_argument('--num_epochs', default=300, type=int)
     args.add_argument('--population_size', default=64, type=int)
     args.add_argument('--seed', default=42, type=int)
 
     args = args.parse_args()
 
-    run(**(vars(args)))
+    args = vars(args)
+    args['evo_alg'] = getattr(evosax.algorithms, args['evo_alg'])
+    run(**args)

examples/cart-pole-evo/run_all.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+evo_algs=("CMA_ES" "PGPE" "SimpleGA")
+seeds=(1 2 3 4 5 6 7 8 9 10)
+
+for alg in "${evo_algs[@]}"; do
+    for s in "${seeds[@]}"; do
+        echo "Running with $n environments and seed $s"
+        python main.py --evo_alg $alg --seed $s
+    done
+done

examples/cart-pole-vectorized/main.py

@@ -49,7 +49,8 @@ def __call__(self, x: Array) -> tuple[Array, Array]:
 
 def run(time_limit: float, num_envs: int, seed: int) -> None:
     """
-    Run ``num_steps`` cart-pole Gymnasium steps.
+    Run ``num_envs`` CartPole Gymnasium environments in parallel using PPO to optimize the policy.
+    The experiment runs for a maximum of ``time_limit`` seconds.
 
     Parameters
     ----------
@@ -86,7 +87,7 @@ def run(time_limit: float, num_envs: int, seed: int) -> None:
     )
 
     def make_env():
-        return gym.make('CartPole-v1', render_mode='no')
+        return gym.make('CartPole-v1')
 
     env = gym.vector.SyncVectorEnv([make_env for _ in range(num_envs)])
     _, _ = env.reset(seed=seed)
@@ -112,7 +113,7 @@ def make_env():
 if __name__ == '__main__':
     args = ArgumentParser()
 
-    args.add_argument('--time_limit', default=120, type=float)
+    args.add_argument('--time_limit', default=85, type=float)
     args.add_argument('--num_envs', default=64, type=int)
     args.add_argument('--seed', default=42, type=int)
 

examples/cart-pole-vectorized/run_all.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-num_envs=(2 4 8 16 32 64 128)
+num_envs=(2 4 8 16 32 64)
 seeds=(1 2 3 4 5 6 7 8 9 10)
 
 for n in "${num_envs[@]}"; do