evolutionary-policy-optimization 0.0.18__py3-none-any.whl → 0.0.20__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
@@ -441,25 +443,38 @@ 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]
 
         if not exists(net):

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.18
+Version: 0.0.20
 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.20.dist-info/RECORD

@@ -0,0 +1,7 @@
+evolutionary_policy_optimization/__init__.py,sha256=Qavcia0n13jjaWIS_LPW7QrxSLT_BBeKujCjF9kQjbA,133
+evolutionary_policy_optimization/epo.py,sha256=BTBqkgDq-x4dUMlKdSojvV2Yjzf9pDUZGMik32WjdHQ,18361
+evolutionary_policy_optimization/experimental.py,sha256=ktBKxRF27Qsj7WIgBpYlWXqMVxO9zOx2oD1JuDYRAwM,548
+evolutionary_policy_optimization-0.0.20.dist-info/METADATA,sha256=0QNTGATtchVuxVplbrfXAtupcrMKEQD-uisM7CFm7qE,4931
+evolutionary_policy_optimization-0.0.20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.20.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.20.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,,