evolutionary-policy-optimization 0.2.3__py3-none-any.whl → 0.2.6__py3-none-any.whl

evolutionary_policy_optimization/experimental.py

@@ -1,11 +1,20 @@
+from random import uniform
+from copy import deepcopy
+
 import torch
 import torch.nn.functional as F
-from einops import rearrange
+from torch.func import vmap, functional_call
+from torch.nn import Module, ParameterList
+
+from einops import rearrange, reduce, repeat
+
+def exists(v):
+    return v is not None
 
 def l2norm(t, dim = -1):
     return F.normalize(t, dim = dim)
 
-def crossover_weights(w1, w2, transpose = False):
+def crossover_weights(w1, w2):
     assert w2.shape == w2.shape
 
     no_batch = w1.ndim == 2
@@ -15,6 +24,9 @@ def crossover_weights(w1, w2, transpose = False):
 
     assert w1.ndim == 3
 
+    i, j = w1.shape[-2:]
+    transpose = i < j
+
     if transpose:
         w1, w2 = tuple(rearrange(t, 'b i j -> b j i') for t in (w1, w2))
 
@@ -45,14 +57,16 @@ def crossover_weights(w1, w2, transpose = False):
 
 def mutate_weight(
     w,
-    transpose = False,
     mutation_strength = 1.
 ):
 
+    i, j = w.shape[-2:]
+    transpose = i < j
+
     if transpose:
         w = w.transpose(-1, -2)
 
-    rank = min(w2.shape[1:])
+    rank = min(w.shape[1:])
     assert rank >= 2
 
     u, s, v = torch.svd(w)
@@ -70,9 +84,132 @@ def mutate_weight(
 
     return out
 
+# wrapper that manages network to population
+# able to receive fitness and employ selection + crossover
+
+class PopulationWrapper(Module):
+    def __init__(
+        self,
+        net: Module,
+        pop_size,
+        num_selected,
+        tournament_size,
+        learning_rate = 1e-3,
+        init_std_dev = 1e-1
+    ):
+        super().__init__()
+        assert num_selected < pop_size
+        assert tournament_size < num_selected
+
+        self.pop_size = pop_size
+        self.num_selected = num_selected
+        self.tournament_size = tournament_size
+        self.num_offsprings = pop_size - num_selected
+
+        self.net = net
+
+        params = dict(net.named_parameters())
+        device = next(iter(params.values())).device
+
+        pop_params = {name: (torch.randn((pop_size, *param.shape), device = device) * init_std_dev).requires_grad_() for name, param in params.items()}
+
+        self.param_names = pop_params.keys()
+        self.param_values = ParameterList(list(pop_params.values()))
+
+        def _forward(params, data):
+            return functional_call(net, params, data)
+
+        self.forward_pop_nets = vmap(_forward, in_dims = (0, None))
+
+    @property
+    def pop_params(self):
+        return dict(zip(self.param_names, self.param_values))
+
+    def individual(self, id) -> Module:
+        assert 0 <= id < self.pop_size
+        state_dict = {key: param[id] for key, param in self.pop_params.items()}
+
+        net = deepcopy(self.net)
+        net.load_state_dict(state_dict)
+        return net
+
+    def parameters(self):
+        return self.pop_params.values()
+
+    def genetic_algorithm_step_(
+        self,
+        fitnesses
+    ):
+        fitnesses = reduce(fitnesses, 'b p -> p', 'mean') # average across samples
+
+        num_selected = self.num_selected
+
+        # selection
+
+        sel_fitnesses, sel_indices = fitnesses.topk(num_selected, dim = -1)
+
+        # tournaments
+
+        tourn_ids = torch.randn((self.num_offsprings, self.tournament_size)).argsort(dim = -1)
+        tourn_scores = sel_fitnesses[tourn_ids]
+
+        winner_ids = tourn_scores.topk(2, dim = -1).indices
+        winner_ids = rearrange(winner_ids, 'offsprings couple -> couple offsprings')
+        parent_ids = sel_indices[winner_ids]
+
+        # crossover
+
+        for param in self.param_values:
+            parents = param[sel_indices]
+            parent1, parent2 = param[parent_ids]
+
+            children = parent1.lerp_(parent2, uniform(0.25, 0.75))
+
+            pop = torch.cat((parents, children))
+
+            param.data.copy_(pop)
+
+    def forward(
+        self,
+        data,
+        *,
+        individual_id = None,
+        labels = None,
+        return_logits_with_loss = False
+    ):
+        # if `individual_id` passed in, will forward for only that one network
+
+        if exists(individual_id):
+            assert 0 <= individual_id < self.pop_size
+            params = {key: param[individual_id] for key, param in self.pop_params.items()}
+            return functional_call(self.net, params, data)
+
+        out = self.forward_pop_nets(dict(self.pop_params), data)
+
+        if not exists(labels):
+            return out
+
+        logits = out
+        pop_size = logits.shape[0]
+
+        losses = F.cross_entropy(
+            rearrange(logits, 'p b ... l -> (p b) l ...'),
+            repeat(labels, 'b ... -> (p b) ...', p = pop_size),
+            reduction = 'none'
+        )
+
+        losses = rearrange(losses, '(p b) ... -> p b ...', p = pop_size)
+
+        if not return_logits_with_loss:
+            return losses
+
+        return losses, logits
+
+# test
+
 if __name__ == '__main__':
-    w1 = torch.randn(32, 16)
-    w2 = torch.randn(32, 16)
+    w1 = torch.randn(2, 32, 16)
+    w2 = torch.randn(2, 32, 16)
 
     child = crossover_weights(w1, w2)
     mutated_w1 = mutate_weight(w1)

{evolutionary_policy_optimization-0.2.3.dist-info → evolutionary_policy_optimization-0.2.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.2.3
+Version: 0.2.6
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -47,6 +47,9 @@ Provides-Extra: examples
 Requires-Dist: numpy; extra == 'examples'
 Requires-Dist: pufferlib>=2.0.6; extra == 'examples'
 Requires-Dist: tqdm; extra == 'examples'
+Provides-Extra: experimental
+Requires-Dist: tensordict; extra == 'experimental'
+Requires-Dist: torchvision; extra == 'experimental'
 Provides-Extra: test
 Requires-Dist: pytest; extra == 'test'
 Requires-Dist: ruff>=0.4.2; extra == 'test'
@@ -56,7 +59,7 @@ Description-Content-Type: text/markdown
 
 ## Evolutionary Policy Optimization
 
-Pytorch implementation of [Evolutionary Policy Optimization](https://web3.arxiv.org/abs/2503.19037), from Wang et al. of the Robotics Institute at Carnegie Mellon University
+Pytorch implementation of [Evolutionary Policy Optimization](https://web3.arxiv.org/abs/2503.19037), from [Wang](https://www.jianrenw.com/) et al. of the Robotics Institute at Carnegie Mellon University
 
 This paper stands out, as I have witnessed the positive effects first hand in an [exploratory project](https://github.com/lucidrains/firefly-torch) (mixing evolution with gradient based methods). Perhaps the Alexnet moment for genetic algorithms has not come to pass yet.
 

{evolutionary_policy_optimization-0.2.3.dist-info → evolutionary_policy_optimization-0.2.6.dist-info}/RECORD

@@ -2,9 +2,9 @@ evolutionary_policy_optimization/__init__.py,sha256=NyiYDYU7DlpmOTM7xiBQET3r1WwX
 evolutionary_policy_optimization/distributed.py,sha256=MxyxqxANAuOm8GYb0Yu09EHd_aVLhK2uwgrfuVWciPU,2342
 evolutionary_policy_optimization/env_wrappers.py,sha256=bDL06o9_b1iW6k3fw2xifnOnYlzs643tdW6Yv2gsIdw,803
 evolutionary_policy_optimization/epo.py,sha256=81rf249ykaPrAEMGk9KsF98qDkCUhW8xL3-2UXIvI2E,51838
-evolutionary_policy_optimization/experimental.py,sha256=ZyOGHbE4dXmt4zCljSzcUklua4vlOwQtslhFEm0JN94,1716
+evolutionary_policy_optimization/experimental.py,sha256=7LOrMIaU4fr2Vme1ZpHNIvlvFEIdWj0-uemhQoNJcPQ,5549
 evolutionary_policy_optimization/mock_env.py,sha256=TLyyRm6tOD0Kdn9QqJJQriaSnsR-YmNQHo4OohmZFG4,1410
-evolutionary_policy_optimization-0.2.3.dist-info/METADATA,sha256=syjis1-9dDCEwfGt7CeMfhC5k7OegIl4BBsKaTnYssQ,8697
-evolutionary_policy_optimization-0.2.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.2.3.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.2.3.dist-info/RECORD,,
+evolutionary_policy_optimization-0.2.6.dist-info/METADATA,sha256=fD68WuKAl76bc0O8w9ThQupw9OpEep0brkPgBlhGBhk,8858
+evolutionary_policy_optimization-0.2.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.2.6.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.2.6.dist-info/RECORD,,