evolutionary-policy-optimization 0.0.69__py3-none-any.whl → 0.0.71__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -8,8 +8,10 @@ from functools import partial, wraps
 from collections import namedtuple
 from random import randrange
 
+import numpy as np
+
 import torch
-from torch import nn, cat, stack, is_tensor, tensor, Tensor
+from torch import nn, cat, stack, is_tensor, tensor, from_numpy, Tensor
 import torch.nn.functional as F
 import torch.distributed as dist
 from torch.nn import Linear, Module, ModuleList
@@ -58,6 +60,21 @@ def divisible_by(num, den):
 def to_device(inp, device):
     return tree_map(lambda t: t.to(device) if is_tensor(t) else t, inp)
 
+def interface_torch_numpy(fn, device):
+    # for a given function, move all inputs from torch tensor to numpy, and all outputs from numpy to torch tensor
+
+    @wraps(fn)
+    def decorated_fn(*args, **kwargs):
+
+        args, kwargs = tree_map(lambda t: t.cpu().numpy() if isinstance(t, Tensor) else t, (args, kwargs))
+
+        out = fn(*args, **kwargs)
+
+        out = tree_map(lambda t: from_numpy(t).to(device) if isinstance(t, np.ndarray) else t, out)
+        return out
+
+    return decorated_fn
+
 # tensor helpers
 
 def l2norm(t):
@@ -416,6 +433,8 @@ class LatentGenePool(Module):
         self.num_natural_selected = int(frac_natural_selected * latents_per_island)
         self.num_tournament_participants = int(frac_tournaments * self.num_natural_selected)
 
+        assert self.num_tournament_participants >= 2
+
         self.crossover_random  = crossover_random
 
         self.mutation_strength = mutation_strength
@@ -845,7 +864,7 @@ class Agent(Module):
 
         return self.latent_gene_pool.genetic_algorithm_step(fitnesses)
 
-    def forward(
+    def learn_from(
         self,
         memories_and_cumulative_rewards: MemoriesAndCumulativeRewards,
         epochs = 2
@@ -853,11 +872,13 @@ class Agent(Module):
     ):
         memories_and_cumulative_rewards = to_device(memories_and_cumulative_rewards, self.device)
 
-        memories, rewards_per_latent_episode = memories_and_cumulative_rewards
+        memories_list, rewards_per_latent_episode = memories_and_cumulative_rewards
 
         # stack memories
 
-        memories = map(stack, zip(*memories))
+        memories = map(stack, zip(*memories_list))
+
+        memories_list.clear()
 
         maybe_barrier()
 
@@ -979,7 +1000,6 @@ class Agent(Module):
         # apply evolution
 
         if self.has_latent_genes:
-
             self.latent_gene_pool.genetic_algorithm_step(fitness_scores)
 
 # reinforcement learning related - ppo
@@ -1159,9 +1179,10 @@ class EPO(Module):
             yield latent_id, episode_id, maybe_seed
 
     @torch.no_grad()
-    def forward(
+    def gather_experience_from(
         self,
         env,
+        memories: list[Memory] | None = None,
         fix_environ_across_latents = None
     ) -> MemoriesAndCumulativeRewards:
 
@@ -1171,9 +1192,10 @@ class EPO(Module):
 
         invalid_episode = tensor(-1) # will use `episode_id` value of `-1` for the `next_value`, needed for not discarding last reward for generalized advantage estimate
 
-        memories: list[Memory] = []
+        if not exists(memories):
+            memories = []
 
-        rewards_per_latent_episode = torch.zeros((self.num_latents, self.episodes_per_latent))
+        rewards_per_latent_episode = torch.zeros((self.num_latents, self.episodes_per_latent), device = self.device)
 
         rollout_gen = self.rollouts_for_machine(fix_environ_across_latents)
 
@@ -1188,7 +1210,7 @@ class EPO(Module):
             if fix_environ_across_latents:
                 reset_kwargs.update(seed = maybe_seed)
 
-            state = env.reset(**reset_kwargs)
+            state = interface_torch_numpy(env.reset, device = self.device)(**reset_kwargs)
 
             # get latent from pool
 
@@ -1210,7 +1232,7 @@ class EPO(Module):
 
                 # get the next state, action, and reward
 
-                state, reward, truncated, terminated = env(action)
+                state, reward, truncated, terminated = interface_torch_numpy(env.forward, device = self.device)(action)
 
                 done = truncated or terminated
 
@@ -1254,3 +1276,18 @@ class EPO(Module):
             memories = memories,
             cumulative_rewards = rewards_per_latent_episode
         )
+
+    def forward(
+        self,
+        agent: Agent,
+        env,
+        num_learning_cycles
+    ):
+
+        for _ in tqdm(range(num_learning_cycles), desc = 'learning cycle'):
+
+            memories = self.gather_experience_from(env)
+
+            agent.learn_from(memories)
+
+        print(f'training complete')

evolutionary_policy_optimization/mock_env.py

@@ -34,7 +34,7 @@ class Env(Module):
     ):
         state = randn(self.state_shape, device = self.device)
         self.step.zero_()
-        return state
+        return state.numpy()
 
     def forward(
         self,
@@ -51,4 +51,5 @@ class Env(Module):
 
         self.step.add_(1)
 
-        return state, reward, truncated, terminated
+        out = state, reward, truncated, terminated
+        return tuple(t.numpy() for t in out)

{evolutionary_policy_optimization-0.0.69.dist-info → evolutionary_policy_optimization-0.0.71.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.69
+Version: 0.0.71
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -114,11 +114,14 @@ from evolutionary_policy_optimization import (
 
 agent = create_agent(
     dim_state = 512,
-    num_latents = 8,
+    num_latents = 16,
     dim_latent = 32,
     actor_num_actions = 5,
     actor_dim_hiddens = (256, 128),
-    critic_dim_hiddens = (256, 128, 64)
+    critic_dim_hiddens = (256, 128, 64),
+    latent_gene_pool_kwargs = dict(
+        frac_natural_selected = 0.5
+    )
 )
 
 epo = EPO(
@@ -130,9 +133,7 @@ epo = EPO(
 
 env = Env((512,))
 
-memories = epo(env)
-
-agent(memories)
+epo(agent, env, num_learning_cycles = 5)
 
 # saving and loading
 

evolutionary_policy_optimization-0.0.71.dist-info/RECORD

@@ -0,0 +1,9 @@
+evolutionary_policy_optimization/__init__.py,sha256=0q0aBuFgWi06MLMD8FiHzBYQ3_W4LYWrwmCtF3u5H2A,201
+evolutionary_policy_optimization/distributed.py,sha256=7KgZdeS_wxBHo_du9XZFB1Cu318J-Bp66Xdr6Log_20,2423
+evolutionary_policy_optimization/epo.py,sha256=7_SFYcOcQA2TQcquIYswIJc-42HgwHj4hKYU9wGWZB0,39333
+evolutionary_policy_optimization/experimental.py,sha256=-IgqjJ_Wk_CMB1y9YYWpoYqTG9GZHAS6kbRdTluVevg,1563
+evolutionary_policy_optimization/mock_env.py,sha256=QhPYIazU3uFPKTMZyFw70KhXiDUBr5aU3v1idotfVFI,1391
+evolutionary_policy_optimization-0.0.71.dist-info/METADATA,sha256=veUHctceDr5Xvym2fCDND3fIjYb6ivsCsIfNnOtkdCg,6304
+evolutionary_policy_optimization-0.0.71.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.71.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.71.dist-info/RECORD,,

evolutionary_policy_optimization-0.0.69.dist-info/RECORD

@@ -1,9 +0,0 @@
-evolutionary_policy_optimization/__init__.py,sha256=0q0aBuFgWi06MLMD8FiHzBYQ3_W4LYWrwmCtF3u5H2A,201
-evolutionary_policy_optimization/distributed.py,sha256=7KgZdeS_wxBHo_du9XZFB1Cu318J-Bp66Xdr6Log_20,2423
-evolutionary_policy_optimization/epo.py,sha256=e83tghTNXfCW0zhhb4nIjvfbzDvzWRxgTlm3vKJd4rM,38189
-evolutionary_policy_optimization/experimental.py,sha256=-IgqjJ_Wk_CMB1y9YYWpoYqTG9GZHAS6kbRdTluVevg,1563
-evolutionary_policy_optimization/mock_env.py,sha256=Bv9ONFRbma8wpjUurc9aCk19A6ceiWitRnS3nwrIR64,1339
-evolutionary_policy_optimization-0.0.69.dist-info/METADATA,sha256=UZEaCY5lfTRMkuyEQs5PLA1AZzSOcsRzXey9kgdd9i0,6220
-evolutionary_policy_optimization-0.0.69.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.69.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.69.dist-info/RECORD,,