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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.38
+Version: 0.0.40
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -101,6 +101,45 @@ fitness = torch.randn(128)
 latent_pool.genetic_algorithm_step(fitness) # update latent genes with genetic algorithm
 ```
 
+End to end learning
+
+```python
+import torch
+
+from evolutionary_policy_optimization import (
+    create_agent,
+    EPO,
+    Env
+)
+
+agent = create_agent(
+    dim_state = 512,
+    num_latents = 8,
+    dim_latent = 32,
+    actor_num_actions = 5,
+    actor_dim_hiddens = (256, 128),
+    critic_dim_hiddens = (256, 128, 64)
+)
+
+epo = EPO(
+    agent,
+    episodes_per_latent = 1,
+    max_episode_length = 10,
+    action_sample_temperature = 1.
+)
+
+env = Env((512,))
+
+memories = epo(env)
+
+agent(memories)
+
+# saving and loading
+
+agent.save('./agent.pt', overwrite = True)
+agent.load('./agent.pt')
+```
+
 ## Citations
 
 ```bibtex

{evolutionary_policy_optimization-0.0.38 → evolutionary_policy_optimization-0.0.40}/README.md

@@ -49,6 +49,45 @@ fitness = torch.randn(128)
 latent_pool.genetic_algorithm_step(fitness) # update latent genes with genetic algorithm
 ```
 
+End to end learning
+
+```python
+import torch
+
+from evolutionary_policy_optimization import (
+    create_agent,
+    EPO,
+    Env
+)
+
+agent = create_agent(
+    dim_state = 512,
+    num_latents = 8,
+    dim_latent = 32,
+    actor_num_actions = 5,
+    actor_dim_hiddens = (256, 128),
+    critic_dim_hiddens = (256, 128, 64)
+)
+
+epo = EPO(
+    agent,
+    episodes_per_latent = 1,
+    max_episode_length = 10,
+    action_sample_temperature = 1.
+)
+
+env = Env((512,))
+
+memories = epo(env)
+
+agent(memories)
+
+# saving and loading
+
+agent.save('./agent.pt', overwrite = True)
+agent.load('./agent.pt')
+```
+
 ## Citations
 
 ```bibtex

{evolutionary_policy_optimization-0.0.38 → evolutionary_policy_optimization-0.0.40}/evolutionary_policy_optimization/epo.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from functools import partial
+from functools import partial, wraps
 from pathlib import Path
 from collections import namedtuple
 
@@ -9,6 +9,7 @@ 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
+from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten
 
 import einx
 from einops import rearrange, repeat, einsum, pack
@@ -22,6 +23,8 @@ from hl_gauss_pytorch import HLGaussLayer
 
 from ema_pytorch import EMA
 
+from tqdm import tqdm
+
 # helpers
 
 def exists(v):
@@ -47,9 +50,20 @@ def l2norm(t):
 def log(t, eps = 1e-20):
     return t.clamp(min = eps).log()
 
+def gumbel_noise(t):
+    return -log(-log(torch.rand_like(t)))
+
+def gumbel_sample(t, temperature = 1.):
+    is_greedy = temperature <= 0.
+
+    if not is_greedy:
+        t = (t / temperature) + gumbel_noise(t)
+
+    return t.argmax(dim = -1)
+
 def calc_entropy(logits):
     prob = logits.softmax(dim = -1)
-    return -prob * log(prob)
+    return -(prob * log(prob)).sum(dim = -1)
 
 def gather_log_prob(
     logits, # Float[b l]
@@ -60,11 +74,24 @@ def gather_log_prob(
     log_prob = log_probs.gather(-1, indices)
     return rearrange(log_prob, '... 1 -> ...')
 
+def temp_batch_dim(fn):
+
+    @wraps(fn)
+    def inner(*args, **kwargs):
+        args, kwargs = tree_map(lambda t: rearrange(t, '... -> 1 ...') if is_tensor(t) else t, (args, kwargs))
+
+        out = fn(*args, **kwargs)
+
+        out = tree_map(lambda t: rearrange(t, '1 ... -> ...') if is_tensor(t) else t, out)
+        return out
+
+    return inner
+
 # generalized advantage estimate
 
 def calc_generalized_advantage_estimate(
-    rewards, # Float[g n]
-    values,  # Float[g n+1]
+    rewards, # Float[n]
+    values,  # Float[n+1]
     masks,   # Bool[n]
     gamma = 0.99,
     lam = 0.95,
@@ -75,9 +102,7 @@ def calc_generalized_advantage_estimate(
     use_accelerated = default(use_accelerated, rewards.is_cuda)
     device = rewards.device
 
-    masks = repeat(masks, 'n -> g n', g = rewards.shape[0])
-
-    values, values_next = values[:, :-1], values[:, 1:]
+    values, values_next = values[:-1], values[1:]
 
     delta = rewards + gamma * values_next * masks - values
     gates = gamma * lam * masks
@@ -565,6 +590,8 @@ class LatentGenePool(Module):
         if not exists(latent_id) and self.num_latents == 1:
             latent_id = 0
 
+        assert exists(latent_id)
+
         if not is_tensor(latent_id):
             latent_id = tensor(latent_id, device = device)
 
@@ -681,17 +708,38 @@ class Agent(Module):
     def get_actor_actions(
         self,
         state,
-        latent_id
+        latent_id = None,
+        latent = None,
+        sample = False,
+        temperature = 1.
     ):
-        latent = self.latent_gene_pool(latent_id = latent_id, state = state)
-        return self.actor(state, latent)
+        assert exists(latent_id) or exists(latent)
+
+        if not exists(latent):
+            latent = self.latent_gene_pool(latent_id = latent_id)
+
+        logits = self.actor(state, latent)
+
+        if not sample:
+            return logits
+
+        actions = gumbel_sample(logits, temperature = temperature)
+
+        log_probs = gather_log_prob(logits, actions)
+
+        return actions, log_probs
 
     def get_critic_values(
         self,
         state,
-        latent_id
+        latent_id = None,
+        latent = None
     ):
-        latent = self.latent_gene_pool(latent_id = latent_id, state = state)
+        assert exists(latent_id) or exists(latent)
+
+        if not exists(latent):
+            latent = self.latent_gene_pool(latent_id = latent_id)
+
         return self.critic(state, latent)
 
     def update_latent_gene_pool_(
@@ -702,13 +750,13 @@ class Agent(Module):
 
     def forward(
         self,
-        memories_and_next_value: MemoriesAndNextValue,
+        memories_and_fitness_scores: MemoriesAndFitnessScores,
         epochs = 2
     ):
-        memories, next_value = memories_and_next_value
+        memories, fitness_scores = memories_and_fitness_scores
 
         (
-            _,
+            episode_ids,
             states,
             latent_gene_ids,
             actions,
@@ -718,35 +766,46 @@ class Agent(Module):
             dones
         ) = map(stack, zip(*memories))
 
-        values_with_next, ps = pack((values, next_value), '*')
+        advantages = self.calc_gae(
+            rewards[:-1],
+            values,
+            dones[:-1],
+        )
 
-        advantages = self.calc_gae(rewards, values_with_next, dones)
+        valid_episode = episode_ids >= 0
 
-        dataset = TensorDataset(states, latent_gene_ids, actions, log_probs, advantages, values)
+        dataset = TensorDataset(
+            *[   
+                advantages[valid_episode[:-1]],
+                *[t[valid_episode] for t in (states, latent_gene_ids, actions, log_probs, values)]
+            ]
+        )
 
         dataloader = DataLoader(dataset, batch_size = self.batch_size, shuffle = True)
 
         self.actor.train()
         self.critic.train()
 
-        for _ in range(epochs):
+        for _ in tqdm(range(epochs), desc = 'learning actor/critic epoch'):
             for (
+                advantages,
                 states,
                 latent_gene_ids,
                 actions,
                 log_probs,
-                advantages,
                 old_values
             ) in dataloader:
 
-                latents = self.latent_gene_pool(latent_gene_ids)
+                latents = self.latent_gene_pool(latent_id = latent_gene_ids)
 
                 # learn actor
 
                 logits = self.actor(states, latents)
+
                 actor_loss = self.actor_loss(logits, log_probs, actions, advantages)
 
                 actor_loss.backward()
+
                 self.actor_optim.step()
                 self.actor_optim.zero_grad()
 
@@ -755,7 +814,7 @@ class Agent(Module):
                 critic_loss = self.critic(
                     states,
                     latents,
-                    targets = advantages + old_values
+                    target = advantages + old_values
                 )
 
                 critic_loss.backward()
@@ -763,6 +822,10 @@ class Agent(Module):
                 self.critic_optim.step()
                 self.critic_optim.zero_grad()
 
+        # apply evolution
+
+        self.latent_gene_pool.genetic_algorithm_step(fitness_scores)
+
 # reinforcement learning related - ppo
 
 def actor_loss(
@@ -789,15 +852,7 @@ def actor_loss(
 
     entropy_aux_loss = -entropy_weight * entropy
 
-    return actor_loss + entropy_aux_loss
-
-def critic_loss(
-    pred_values,  # Float[b]
-    advantages,   # Float[b]
-    old_values    # Float[b]
-):
-    discounted_values = advantages + old_values
-    return F.mse_loss(pred_values, discounted_values)
+    return (actor_loss + entropy_aux_loss).mean()
 
 # agent contains the actor, critic, and the latent genetic pool
 
@@ -810,6 +865,11 @@ def create_agent(
     critic_dim_hiddens: int | tuple[int, ...],
 ) -> Agent:
 
+    latent_gene_pool = LatentGenePool(
+        num_latents = num_latents,
+        dim_latent = dim_latent
+    )
+
     actor = Actor(
         num_actions = actor_num_actions,
         dim_state = dim_state,
@@ -821,13 +881,7 @@ def create_agent(
         dim_state = dim_state,
         dim_latent = dim_latent,
         dim_hiddens = critic_dim_hiddens
-    )
-
-    latent_gene_pool = LatentGenePool(
-        num_latents = num_latents,
-        dim_latent = dim_latent,
-    )
-
+    )                
     return Agent(actor = actor, critic = critic, latent_gene_pool = latent_gene_pool)
 
 # EPO - which is just PPO with natural selection of a population of latent variables conditioning the agent
@@ -840,14 +894,13 @@ Memory = namedtuple('Memory', [
     'action',
     'log_prob',
     'reward',
-    'values',
+    'value',
     'done'
 ])
 
-MemoriesAndNextValue = namedtuple('MemoriesAndNextValue', [
+MemoriesAndFitnessScores = namedtuple('MemoriesAndFitnessScores', [
     'memories',
-    'next_value',
-    'cumulative_rewards'
+    'fitness_scores'
 ])
 
 class EPO(Module):
@@ -856,10 +909,12 @@ class EPO(Module):
         self,
         agent: Agent,
         episodes_per_latent,
-        max_episode_length
+        max_episode_length,
+        action_sample_temperature = 1.
     ):
         super().__init__()
         self.agent = agent
+        self.action_sample_temperature = action_sample_temperature
 
         self.num_latents = agent.latent_gene_pool.num_latents
         self.episodes_per_latent = episodes_per_latent
@@ -869,10 +924,80 @@ class EPO(Module):
     def forward(
         self,
         env
-    ) -> MemoriesAndNextValue:
+    ) -> MemoriesAndFitnessScores:
 
         self.agent.eval()
 
+        invalid_episode = tensor(-1) # will use `episode_id` value of `-1` for the `next_value`, needed for not discarding last reward for generalized advantage estimate
+
         memories: list[Memory] = []
 
-        raise NotImplementedError
+        fitness_scores = torch.zeros((self.num_latents))
+
+        for episode_id in tqdm(range(self.episodes_per_latent), desc = 'episode'):
+
+            for latent_id in tqdm(range(self.num_latents), desc = 'latent'):
+                time = 0
+
+                # initial state
+
+                state = env.reset()
+
+                # get latent from pool
+
+                latent = self.agent.latent_gene_pool(latent_id = latent_id)
+
+                # until maximum episode length
+
+                done = tensor(False)
+
+                while time < self.max_episode_length:
+
+                    # sample action
+
+                    action, log_prob = temp_batch_dim(self.agent.get_actor_actions)(state, latent = latent, sample = True, temperature = self.action_sample_temperature)
+
+                    # values
+
+                    value = temp_batch_dim(self.agent.get_critic_values)(state, latent = latent)
+
+                    # get the next state, action, and reward
+
+                    state, reward, done = env(action)
+
+                    # update fitness for each gene as cumulative reward received, but make this customizable at some point
+
+                    fitness_scores[latent_id] += reward
+                    
+                    # store memories
+
+                    memory = Memory(
+                        tensor(episode_id),
+                        state,
+                        tensor(latent_id),
+                        action,
+                        log_prob,
+                        reward,
+                        value,
+                        done
+                    )
+
+                    memories.append(memory)
+
+                    time += 1
+
+                # need the final next value for GAE, iiuc
+
+                next_value = temp_batch_dim(self.agent.get_critic_values)(state, latent = latent)
+
+                memory_for_gae = memory._replace(
+                    episode_id = invalid_episode,
+                    value = next_value
+                )
+
+                memories.append(memory_for_gae)
+
+        return MemoriesAndFitnessScores(
+            memories = memories,
+            fitness_scores = fitness_scores
+        )

{evolutionary_policy_optimization-0.0.38 → evolutionary_policy_optimization-0.0.40}/pyproject.toml

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

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

@@ -71,4 +71,34 @@ def test_create_agent(
     # saving and loading
 
     agent.save('./agent.pt', overwrite = True)
-    agent.load('./agent.pt')
+    agent.load('./agent.pt')
+
+def test_e2e_with_mock_env():
+    from evolutionary_policy_optimization import create_agent, EPO, Env
+
+    agent = create_agent(
+        dim_state = 512,
+        num_latents = 8,
+        dim_latent = 32,
+        actor_num_actions = 5,
+        actor_dim_hiddens = (256, 128),
+        critic_dim_hiddens = (256, 128, 64)
+    )
+
+    epo = EPO(
+        agent,
+        episodes_per_latent = 1,
+        max_episode_length = 10,
+        action_sample_temperature = 1.
+    )
+
+    env = Env((512,))
+
+    memories = epo(env)
+
+    agent(memories)
+
+    # saving and loading
+
+    agent.save('./agent.pt', overwrite = True)
+    agent.load('./agent.pt')