evolutionary-policy-optimization 0.0.36__tar.gz → 0.0.38__tar.gz

{evolutionary_policy_optimization-0.0.36 → evolutionary_policy_optimization-0.0.38}/PKG-INFO

@@ -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.36 → evolutionary_policy_optimization-0.0.38}/README.md

@@ -8,8 +8,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.36 → evolutionary_policy_optimization-0.0.38}/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-0.0.36 → evolutionary_policy_optimization-0.0.38}/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-0.0.38/evolutionary_policy_optimization/experimental.py

@@ -0,0 +1,47 @@
+import torch
+from einops import rearrange
+
+def crossover_weights(w1, w2, transpose = False):
+    assert w2.shape == w2.shape
+
+    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 = tuple(rearrange(t, 'b i j -> b j i') for t in (w1, w2))
+
+    rank = min(w2.shape[1:])
+    assert rank >= 2
+
+    batch = w1.shape[0]
+
+    u1, s1, v1 = torch.svd(w1)
+    u2, s2, v2 = torch.svd(w2)
+
+    batch_randperm = torch.randn((batch, rank), device = w1.device).argsort(dim = -1)
+    mask = batch_randperm < (rank // 2)
+
+    u = torch.where(mask[:, None, :], u1, u2)
+    s = torch.where(mask, s1, s2)
+    v = torch.where(mask[:, None, :], v1, v2)
+
+    out = u @ torch.diag_embed(s) @ v.mT
+
+    if transpose:
+        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-0.0.36 → evolutionary_policy_optimization-0.0.38}/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 → evolutionary_policy_optimization-0.0.38}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.0.36"
+version = "0.0.38"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{evolutionary_policy_optimization-0.0.36 → evolutionary_policy_optimization-0.0.38}/tests/test_epo.py

@@ -16,8 +16,7 @@ def test_readme(
 
     latent_pool = LatentGenePool(
         num_latents = 128,
-        dim_latent = 32,
-        dim_state = 512,
+        dim_latent = 32,        
         num_islands = num_islands,
     )
 

evolutionary_policy_optimization-0.0.36/evolutionary_policy_optimization/experimental.py

@@ -1,27 +0,0 @@
-import torch
-
-def crossover_weights(w1, w2, transpose = False):
-    assert w2.shape == w2.shape
-    assert w1.ndim == 2
-
-    if transpose:
-        w1, w2 = w1.t(), w2.t()
-
-    rank = min(w2.shape)
-    assert rank >= 2
-
-    u1, s1, v1 = torch.svd(w1)
-    u2, s2, v2 = torch.svd(w2)
-
-    mask = torch.randperm(rank) < (rank // 2)
-
-    u = torch.where(mask[None, :], u1, u2)
-    s = torch.where(mask, s1, s2)
-    v = torch.where(mask[None, :], v1, v2)
-
-    out = u @ torch.diag_embed(s) @ v.mT
-
-    if transpose:
-        out = out.t()
-
-    return out