evolutionary-policy-optimization 0.0.36__py3-none-any.whl → 0.0.38__py3-none-any.whl

evolutionary_policy_optimization/__init__.py

@@ -4,5 +4,8 @@ from evolutionary_policy_optimization.epo import (
     Critic,
     create_agent,
     Agent,
-    LatentGenePool
+    LatentGenePool,
+    EPO
 )
+
+from evolutionary_policy_optimization.mock_env import Env

evolutionary_policy_optimization/epo.py

@@ -5,13 +5,13 @@ from pathlib import Path
 from collections import namedtuple
 
 import torch
-from torch import nn, cat, is_tensor, tensor
+from torch import nn, cat, stack, is_tensor, tensor
 import torch.nn.functional as F
 from torch.nn import Linear, Module, ModuleList
 from torch.utils.data import TensorDataset, DataLoader
 
 import einx
-from einops import rearrange, repeat, einsum
+from einops import rearrange, repeat, einsum, pack
 from einops.layers.torch import Rearrange
 
 from assoc_scan import AssocScan
@@ -319,7 +319,6 @@ class LatentGenePool(Module):
         num_latents,                     # same as gene pool size
         dim_latent,                      # gene dimension
         num_islands = 1,                 # add the island strategy, which has been effectively used in a few recent works
-        dim_state = None,
         frozen_latents = True,
         crossover_random = True,         # random interp from parent1 to parent2 for crossover, set to `False` for averaging (0.5 constant value)
         l2norm_latent = False,           # whether to enforce latents on hypersphere,
@@ -384,7 +383,6 @@ class LatentGenePool(Module):
         fitness,
         beta0 = 2.,           # exploitation factor, moving fireflies of low light intensity to high
         gamma = 1.,           # controls light intensity decay over distance - setting this to zero will make firefly equivalent to vanilla PSO
-        alpha = 0.1,          # exploration factor
         inplace = True,
     ):
         islands = self.num_islands
@@ -555,7 +553,6 @@ class LatentGenePool(Module):
     def forward(
         self,
         *args,
-        state: Tensor | None = None,
         latent_id: int | None = None,
         net: Module | None = None,
         net_latent_kwarg_name = 'latent',
@@ -575,8 +572,6 @@ class LatentGenePool(Module):
 
         # fetch latent
 
-        fetching_multiple_latents = latent_id.numel() > 1
-
         latent = self.latents[latent_id]
 
         latent = self.maybe_l2norm(latent)
@@ -713,6 +708,7 @@ class Agent(Module):
         memories, next_value = memories_and_next_value
 
         (
+            _,
             states,
             latent_gene_ids,
             actions,
@@ -743,9 +739,11 @@ class Agent(Module):
                 old_values
             ) in dataloader:
 
+                latents = self.latent_gene_pool(latent_gene_ids)
+
                 # learn actor
 
-                logits = self.actor(states)
+                logits = self.actor(states, latents)
                 actor_loss = self.actor_loss(logits, log_probs, actions, advantages)
 
                 actor_loss.backward()
@@ -754,7 +752,12 @@ class Agent(Module):
 
                 # learn critic with maybe classification loss
 
-                critic_loss = self.critic(states, advantages + old_values)
+                critic_loss = self.critic(
+                    states,
+                    latents,
+                    targets = advantages + old_values
+                )
+
                 critic_loss.backward()
 
                 self.critic_optim.step()
@@ -778,7 +781,7 @@ def actor_loss(
 
     clipped_ratio = ratio.clamp(min = 1. - eps_clip, max = 1. + eps_clip)
 
-    actor_loss = -torch.min(clipped_ratio * advantage, ratio * advantage)
+    actor_loss = -torch.min(clipped_ratio * advantages, ratio * advantages)
 
     # add entropy loss for exploration
 
@@ -821,7 +824,6 @@ def create_agent(
     )
 
     latent_gene_pool = LatentGenePool(
-        dim_state = dim_state,
         num_latents = num_latents,
         dim_latent = dim_latent,
     )
@@ -832,6 +834,7 @@ def create_agent(
 # the tricky part is that the latent ids for each episode / trajectory needs to be tracked
 
 Memory = namedtuple('Memory', [
+    'episode_id',
     'state',
     'latent_gene_id',
     'action',
@@ -843,21 +846,33 @@ Memory = namedtuple('Memory', [
 
 MemoriesAndNextValue = namedtuple('MemoriesAndNextValue', [
     'memories',
-    'next_value'
+    'next_value',
+    'cumulative_rewards'
 ])
 
 class EPO(Module):
 
     def __init__(
         self,
-        agent: Agent
+        agent: Agent,
+        episodes_per_latent,
+        max_episode_length
     ):
         super().__init__()
         self.agent = agent
 
+        self.num_latents = agent.latent_gene_pool.num_latents
+        self.episodes_per_latent = episodes_per_latent
+        self.max_episode_length = max_episode_length
+
+    @torch.no_grad()
     def forward(
         self,
         env
     ) -> MemoriesAndNextValue:
 
+        self.agent.eval()
+
+        memories: list[Memory] = []
+
         raise NotImplementedError

evolutionary_policy_optimization/experimental.py

@@ -1,27 +1,47 @@
 import torch
+from einops import rearrange
 
 def crossover_weights(w1, w2, transpose = False):
     assert w2.shape == w2.shape
-    assert w1.ndim == 2
+
+    no_batch = w1.ndim == 2
+
+    if no_batch:
+        w1, w2 = tuple(rearrange(t, '... -> 1 ...') for t in (w1, w2))
+
+    assert w1.ndim == 3
 
     if transpose:
-        w1, w2 = w1.t(), w2.t()
+        w1, w2 = tuple(rearrange(t, 'b i j -> b j i') for t in (w1, w2))
 
-    rank = min(w2.shape)
+    rank = min(w2.shape[1:])
     assert rank >= 2
 
+    batch = w1.shape[0]
+
     u1, s1, v1 = torch.svd(w1)
     u2, s2, v2 = torch.svd(w2)
 
-    mask = torch.randperm(rank) < (rank // 2)
+    batch_randperm = torch.randn((batch, rank), device = w1.device).argsort(dim = -1)
+    mask = batch_randperm < (rank // 2)
 
-    u = torch.where(mask[None, :], u1, u2)
+    u = torch.where(mask[:, None, :], u1, u2)
     s = torch.where(mask, s1, s2)
-    v = torch.where(mask[None, :], v1, v2)
+    v = torch.where(mask[:, None, :], v1, v2)
 
     out = u @ torch.diag_embed(s) @ v.mT
 
     if transpose:
-        out = out.t()
+        out = rearrange(out, 'b j i -> b i j')
+
+    if no_batch:
+        out = rearrange(out, '1 ... -> ...')
 
     return out
+
+if __name__ == '__main__':
+    w1 = torch.randn(32, 16)
+    w2 = torch.randn(32, 16)
+    child = crossover_weights(w2, w2)
+
+    assert child.shape == w2.shape

evolutionary_policy_optimization/mock_env.py

@@ -4,15 +4,20 @@ import torch
 from torch import tensor, randn, randint
 from torch.nn import Module
 
+# functions
+
+def cast_tuple(v):
+    return v if isinstance(v, tuple) else v\
+
 # mock env
 
 class Env(Module):
     def __init__(
         self,
-        state_shape: tuple[int, ...]
+        state_shape: int | tuple[int, ...]
     ):
         super().__init__()
-        self.state_shape = state_shape
+        self.state_shape = cast_tuple(state_shape)
         self.register_buffer('dummy', tensor(0))
 
     @property
@@ -31,6 +36,6 @@ class Env(Module):
     ):
         state = randn(self.state_shape, device = self.device)
         reward = randint(0, 5, (), device = self.device).float()
-        done = zeros((), device = self.device, dtype = torch.bool)
+        done = torch.zeros((), device = self.device, dtype = torch.bool)
 
         return state, reward, done

{evolutionary_policy_optimization-0.0.36.dist-info → evolutionary_policy_optimization-0.0.38.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.36
+Version: 0.0.38
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -60,8 +60,6 @@ This paper stands out, as I have witnessed the positive effects first hand in an
 
 Besides their latent variable strategy, I'll also throw in some attempts with crossover in weight space
 
-Update: I see, mixing genetic algorithms with gradient based method is already a research field, under [Memetic algorithms](https://en.wikipedia.org/wiki/Memetic_algorithm)
-
 ## Install
 
 ```bash

evolutionary_policy_optimization-0.0.38.dist-info/RECORD

@@ -0,0 +1,8 @@
+evolutionary_policy_optimization/__init__.py,sha256=0q0aBuFgWi06MLMD8FiHzBYQ3_W4LYWrwmCtF3u5H2A,201
+evolutionary_policy_optimization/epo.py,sha256=z41p5LmvOHULq6o5aIj9Q6lpyka5DvkqsJ493-WL-EQ,26175
+evolutionary_policy_optimization/experimental.py,sha256=9FrJGviLESlYysHI3i83efT9g2ZB9ha4u3K9HXN98_w,1100
+evolutionary_policy_optimization/mock_env.py,sha256=QqVPZVJtrvQmSDcnYDTob_A5sDwiUzGj6_tmo6BII5c,918
+evolutionary_policy_optimization-0.0.38.dist-info/METADATA,sha256=Lofrc6waEB8qBD19pjBKQjbKBYMNUyjYZZrJCO1fji8,4818
+evolutionary_policy_optimization-0.0.38.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.38.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.38.dist-info/RECORD,,

evolutionary_policy_optimization-0.0.36.dist-info/RECORD

@@ -1,8 +0,0 @@
-evolutionary_policy_optimization/__init__.py,sha256=Qavcia0n13jjaWIS_LPW7QrxSLT_BBeKujCjF9kQjbA,133
-evolutionary_policy_optimization/epo.py,sha256=wRqGjoiksWY33BQc9jypJbKWroHm3i_aEPNx1twVjWk,25819
-evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
-evolutionary_policy_optimization/mock_env.py,sha256=3xrd-gwjZeVd_sEvxIyX0lppnMWcfQGOapO-XjKmExI,816
-evolutionary_policy_optimization-0.0.36.dist-info/METADATA,sha256=WQpJa1PuiQx1qANilbJ0E7tZoKHDm2wAvjMccQoPH5Q,4992
-evolutionary_policy_optimization-0.0.36.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.36.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.36.dist-info/RECORD,,