evolutionary-policy-optimization 0.0.45__tar.gz → 0.0.47__tar.gz

{evolutionary_policy_optimization-0.0.45 → evolutionary_policy_optimization-0.0.47}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.45
+Version: 0.0.47
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization

{evolutionary_policy_optimization-0.0.45 → evolutionary_policy_optimization-0.0.47}/evolutionary_policy_optimization/epo.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from functools import partial, wraps
 from pathlib import Path
 from collections import namedtuple
+from random import randrange
 
 import torch
 from torch import nn, cat, stack, is_tensor, tensor
@@ -87,6 +88,11 @@ def temp_batch_dim(fn):
 
     return inner
 
+# fitness related
+
+def get_fitness_scores(cum_rewards, memories):
+    return cum_rewards
+
 # generalized advantage estimate
 
 def calc_generalized_advantage_estimate(
@@ -643,6 +649,7 @@ class Agent(Module):
         actor_optim_kwargs: dict = dict(),
         critic_optim_kwargs: dict = dict(),
         latent_optim_kwargs: dict = dict(),
+        get_fitness_scores: Callable[..., Tensor] = get_fitness_scores
     ):
         super().__init__()
 
@@ -663,6 +670,10 @@ class Agent(Module):
         self.actor_loss = partial(actor_loss, **actor_loss_kwargs)
         self.calc_gae = partial(calc_generalized_advantage_estimate, **calc_gae_kwargs)
 
+        # fitness score related
+
+        self.get_fitness_scores = get_fitness_scores
+
         # learning hparams
 
         self.batch_size = batch_size
@@ -766,10 +777,12 @@ class Agent(Module):
 
     def forward(
         self,
-        memories_and_fitness_scores: MemoriesAndFitnessScores,
+        memories_and_cumulative_rewards: MemoriesAndCumulativeRewards,
         epochs = 2
     ):
-        memories, fitness_scores = memories_and_fitness_scores
+        memories, cumulative_rewards = memories_and_cumulative_rewards
+
+        fitness_scores = self.get_fitness_scores(cumulative_rewards, memories)
 
         (
             episode_ids,
@@ -952,9 +965,9 @@ Memory = namedtuple('Memory', [
     'done'
 ])
 
-MemoriesAndFitnessScores = namedtuple('MemoriesAndFitnessScores', [
+MemoriesAndCumulativeRewards = namedtuple('MemoriesAndCumulativeRewards', [
     'memories',
-    'fitness_scores'
+    'cumulative_rewards'
 ])
 
 class EPO(Module):
@@ -977,8 +990,9 @@ class EPO(Module):
     @torch.no_grad()
     def forward(
         self,
-        env
-    ) -> MemoriesAndFitnessScores:
+        env,
+        fix_seed_across_latents = True
+    ) -> MemoriesAndCumulativeRewards:
 
         self.agent.eval()
 
@@ -986,16 +1000,26 @@ class EPO(Module):
 
         memories: list[Memory] = []
 
-        fitness_scores = torch.zeros((self.num_latents))
+        cumulative_rewards = torch.zeros((self.num_latents))
 
         for episode_id in tqdm(range(self.episodes_per_latent), desc = 'episode'):
 
+            # maybe fix seed for environment across all latents
+
+            env_reset_kwargs = dict()
+
+            if fix_seed_across_latents:
+                seed = randrange(int(1e6))
+                env_reset_kwargs = dict(seed = seed)
+
+            # for each latent (on a single machine for now)
+
             for latent_id in tqdm(range(self.num_latents), desc = 'latent'):
                 time = 0
 
                 # initial state
 
-                state = env.reset()
+                state = env.reset(**env_reset_kwargs)
 
                 # get latent from pool
 
@@ -1019,9 +1043,9 @@ class EPO(Module):
 
                     state, reward, done = env(action)
 
-                    # update fitness for each gene as cumulative reward received, but make this customizable at some point
+                    # update cumulative rewards per latent, to be used as default fitness score
 
-                    fitness_scores[latent_id] += reward
+                    cumulative_rewards[latent_id] += reward
                     
                     # store memories
 
@@ -1051,7 +1075,7 @@ class EPO(Module):
 
                 memories.append(memory_for_gae)
 
-        return MemoriesAndFitnessScores(
+        return MemoriesAndCumulativeRewards(
             memories = memories,
-            fitness_scores = fitness_scores
+            cumulative_rewards = cumulative_rewards
         )

{evolutionary_policy_optimization-0.0.45 → evolutionary_policy_optimization-0.0.47}/evolutionary_policy_optimization/mock_env.py

@@ -25,7 +25,8 @@ class Env(Module):
         return self.dummy.device
 
     def reset(
-        self
+        self,
+        seed
     ):
         state = randn(self.state_shape, device = self.device)
         return state

{evolutionary_policy_optimization-0.0.45 → evolutionary_policy_optimization-0.0.47}/pyproject.toml

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