evolutionary-policy-optimization 0.0.18__py3-none-any.whl → 0.0.22__py3-none-any.whl

evolutionary_policy_optimization/epo.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from collections import namedtuple
 
 import torch
-from torch import nn, cat
+from torch import nn, cat, is_tensor, tensor
 import torch.nn.functional as F
 from torch.nn import Linear, Module, ModuleList
 from torch.utils.data import TensorDataset, DataLoader
@@ -176,6 +176,8 @@ class MLP(Module):
             if latent.ndim == 1:
                 latent = repeat(latent, 'd -> b d', b = batch)
 
+            assert latent.shape[0] == x.shape[0], f'received state with batch size {x.shape[0]} but latent ids received had batch size {latent_id.shape[0]}'
+
             x = cat((x, latent), dim = -1)
 
         # layers
@@ -299,6 +301,7 @@ class LatentGenePool(Module):
         dim_latent,                      # gene dimension
         num_latent_sets = 1,             # allow for sets of latents / gene per individual, expression of a set controlled by the environment
         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,
         frac_tournaments = 0.25,         # fraction of genes to participate in tournament - the lower the value, the more chance a less fit gene could be selected
@@ -319,7 +322,7 @@ class LatentGenePool(Module):
 
         self.num_latents = num_latents
         self.needs_latent_gate = num_latent_sets > 1
-        self.latents = nn.Parameter(latents, requires_grad = False)
+        self.latents = nn.Parameter(latents, requires_grad = not frozen_latents)
 
         self.maybe_l2norm = maybe_l2norm
 
@@ -441,27 +444,42 @@ class LatentGenePool(Module):
         net: Module | None = None,
         **kwargs,
     ):
+        device = self.latents.device
 
         # if only 1 latent, assume doing ablation and get lone gene
 
         if not exists(latent_id) and self.num_latents == 1:
             latent_id = 0
 
-        assert 0 <= latent_id < self.num_latents
+        if not is_tensor(latent_id):
+            latent_id = tensor(latent_id, device = device)
+
+        assert (0 <= latent_id).all() and (latent_id < self.num_latents).all()
 
         # fetch latent
 
+        fetching_multiple_latents = latent_id.numel() > 1
+
         latent = self.latents[latent_id]
 
         if self.needs_latent_gate:
             assert exists(state), 'state must be passed in if greater than number of 1 latent set'
 
+            if not fetching_multiple_latents:
+                latent = repeat(latent, '... -> b ...', b = state.shape[0])
+
+            assert latent.shape[0] == state.shape[0]
+
             gates = self.to_latent_gate(state)
-            latent = einsum(latent, gates, 'n g, b n -> b g')
+            latent = einsum(latent, gates, 'b n g, b n -> b g')
+
+        elif fetching_multiple_latents:
+            latent = latent[:, 0]
         else:
-            assert latent.shape[0] == 1
             latent = latent[0]
 
+        latent = self.maybe_l2norm(latent)
+
         if not exists(net):
             return latent
 

{evolutionary_policy_optimization-0.0.18.dist-info → evolutionary_policy_optimization-0.0.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.18
+Version: 0.0.22
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -60,10 +60,15 @@ Besides their latent variable strategy, I'll also throw in some attempts with cr
 
 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
+$ pip install evolutionary-policy-optimization
+```
+
 ## Usage
 
 ```python
-
 import torch
 
 from evolutionary_policy_optimization import (

evolutionary_policy_optimization-0.0.22.dist-info/RECORD

@@ -0,0 +1,7 @@
+evolutionary_policy_optimization/__init__.py,sha256=Qavcia0n13jjaWIS_LPW7QrxSLT_BBeKujCjF9kQjbA,133
+evolutionary_policy_optimization/epo.py,sha256=TbUX2L-Wa2zIZ2b7iHmBtaym-qDSLAFrC7iU7xReX_k,18449
+evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
+evolutionary_policy_optimization-0.0.22.dist-info/METADATA,sha256=L3G-tesSEyhrc_SbTN6HuJQlXfogEUvr3W9SXPcnRVw,4931
+evolutionary_policy_optimization-0.0.22.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.22.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.22.dist-info/RECORD,,

evolutionary_policy_optimization-0.0.18.dist-info/RECORD

@@ -1,7 +0,0 @@
-evolutionary_policy_optimization/__init__.py,sha256=Qavcia0n13jjaWIS_LPW7QrxSLT_BBeKujCjF9kQjbA,133
-evolutionary_policy_optimization/epo.py,sha256=3pvYPwAEZdrxwwV95Ea1qG4CQjLnyaxAr40opk07LDw,17747
-evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
-evolutionary_policy_optimization-0.0.18.dist-info/METADATA,sha256=BIyCXw2IbMs-x2hDbFs9NR5s2dYEbfbeK_LadUeUc8Q,4860
-evolutionary_policy_optimization-0.0.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.18.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.18.dist-info/RECORD,,