evolutionary-policy-optimization 0.4.0__tar.gz → 0.4.2__tar.gz

{evolutionary_policy_optimization-0.4.0 → evolutionary_policy_optimization-0.4.2}/.gitignore

@@ -174,3 +174,15 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+
+# microduck runs / artifacts
+checkpoints*/
+experiments/
+videos*/
+recordings/
+wandb/
+scratch/
+data/
+*.pt
+*.mp4
+*.log

{evolutionary_policy_optimization-0.4.0 → evolutionary_policy_optimization-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: evolutionary-policy-optimization
-Version: 0.4.0
+Version: 0.4.2
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -40,10 +40,11 @@ Requires-Dist: einops>=0.8.1
 Requires-Dist: einx>=0.3.0
 Requires-Dist: ema-pytorch>=0.7.7
 Requires-Dist: hl-gauss-pytorch>=0.2.8
-Requires-Dist: mean-conc-beta>=0.0.7
+Requires-Dist: mean-conc-beta>=0.2.7
 Requires-Dist: torch>=2.2
 Requires-Dist: tqdm
 Requires-Dist: x-mlps-pytorch>=0.3.0
+Requires-Dist: x-ppo>=0.0.16
 Provides-Extra: examples
 Requires-Dist: fire; extra == 'examples'
 Requires-Dist: gymnasium[box2d]>=1.0.0; extra == 'examples'

{evolutionary_policy_optimization-0.4.0 → evolutionary_policy_optimization-0.4.2}/evolutionary_policy_optimization/epo.py

@@ -29,6 +29,7 @@ from torch.utils._pytree import tree_map
 from torch.utils.data import DataLoader, TensorDataset
 from tqdm import tqdm
 from x_mlps_pytorch import AttnResidualNormedMLP
+from x_ppo import hpo_actor_loss, ppo_actor_loss, spo_actor_loss
 
 from evolutionary_policy_optimization.distributed import all_gather, get_world_and_rank, is_distributed, maybe_barrier, maybe_sync_seed
 
@@ -52,6 +53,20 @@ def divisible_by(num, den):
 def to_device(inp, device):
     return tree_map(lambda t: t.to(device) if is_tensor(t) else t, inp)
 
+class ShuffledBatches:
+    # fresh permutation on every pass
+
+    def __init__(self, fields, batch_size):
+        self.fields = fields
+        self.batch_size = batch_size
+
+    def __iter__(self):
+        perm = torch.randperm(self.fields[0].shape[0])
+
+        for i in range(0, len(perm), self.batch_size):
+            b = perm[i:i + self.batch_size]
+            yield [field[b] for field in self.fields]
+
 def maybe(fn):
 
     @wraps(fn)
@@ -1234,6 +1249,7 @@ class Agent(Module):
         hidden_spr_ema_update = 0.99,
         get_fitness_scores: Callable[..., Tensor] = get_fitness_scores,
         wrap_with_accelerate: bool = True,
+        use_dataloader = False,
         accelerate_kwargs: dict = dict(),
         accelerator = None,
         quiet: bool = False,
@@ -1245,6 +1261,7 @@ class Agent(Module):
         # hf accelerate
 
         self.wrap_with_accelerate = wrap_with_accelerate
+        self.use_dataloader = use_dataloader
 
         if wrap_with_accelerate:
             self.accelerate = accelerator if exists(accelerator) else Accelerator(**accelerate_kwargs)
@@ -1592,13 +1609,12 @@ class Agent(Module):
         epochs = 2
 
     ):
-        memories_and_cumulative_rewards = to_device(memories_and_cumulative_rewards, self.device)
-
         memories_list, rewards_per_latent_episode = memories_and_cumulative_rewards
 
-        # stack memories
+        # stack, then move once per field
 
-        memories = map(stack, zip(*memories_list))
+        memories = [t.to(self.device) for t in map(stack, zip(*memories_list))]
+        rewards_per_latent_episode = rewards_per_latent_episode.to(self.device)
 
         memories_list.clear()
 
@@ -1640,12 +1656,15 @@ class Agent(Module):
 
         valid_episode = episode_ids >= 0
 
-        dataset = TensorDataset(*[t[valid_episode] for t in (advantages, states, next_states, latent_gene_ids, actions, log_probs, values, dones)])
+        fields = tuple(t[valid_episode] for t in (advantages, states, next_states, latent_gene_ids, actions, log_probs, values, dones))
 
-        dataloader = DataLoader(dataset, batch_size = self.batch_size, shuffle = True)
+        if self.wrap_with_accelerate or self.use_dataloader:
+            dataloader = DataLoader(TensorDataset(*fields), batch_size = self.batch_size, shuffle = True)
 
-        if self.wrap_with_accelerate:
-            dataloader = self.accelerate.prepare(dataloader)
+            if self.wrap_with_accelerate:
+                dataloader = self.accelerate.prepare(dataloader)
+        else:
+            dataloader = ShuffledBatches(fields, self.batch_size)
 
         # updating actor and critic
 
@@ -1845,39 +1864,31 @@ def actor_loss(
     advantages,     # Float[b]
     eps_clip = 0.2,
     entropy_weight = .01,
-    eps = 1e-5,
     norm_advantages = True,
-    use_spo = False
+    use_spo = False,
+    use_hpo = False
 ):
-    batch = advantages.shape[0]
-
     log_probs = sum_to_batch(distr.log_prob(actions))
     entropy = sum_to_batch(distr.entropy())
 
-    ratio = (log_probs - old_log_probs).exp()
-
-    if norm_advantages:
-        advantages = F.layer_norm(advantages, (batch,), eps = eps)
-
-    if use_spo:
-        # simple policy optimization - line 14 Algorithm 1 https://arxiv.org/abs/2401.16025v9
-
-        actor_loss = - (
-            ratio * advantages -
-            advantages.abs() / (2 * eps_clip) * (ratio - 1.).square()
-        )
+    if use_hpo:
+        loss_fn = hpo_actor_loss
+    elif use_spo:
+        loss_fn = spo_actor_loss
     else:
-        # classic clipped surrogate loss from ppo
-
-        clipped_ratio = ratio.clamp(min = 1. - eps_clip, max = 1. + eps_clip)
-
-        actor_loss = -torch.min(clipped_ratio * advantages, ratio * advantages)
-
-    # add entropy loss for exploration
+        loss_fn = ppo_actor_loss
+
+    loss = loss_fn(
+        log_probs,
+        old_log_probs,
+        advantages,
+        eps_clip = eps_clip,
+        normalize_advantages = norm_advantages,
+    )
 
-    entropy_aux_loss = -entropy_weight * entropy
+    # entropy bonus
 
-    return (actor_loss + entropy_aux_loss).mean()
+    return (loss - entropy_weight * entropy).mean()
 
 # agent contains the actor, critic, and the latent genetic pool
 
@@ -2160,6 +2171,7 @@ class EPO(Module):
 
             next_obs, rewards, terminated, truncated = env.step_batch(actions.cpu().numpy(), worker_ids = active)[:4]
 
+            next_states = from_numpy(np.ascontiguousarray(next_obs)).float().to(self.device)
             rewards = from_numpy(np.asarray(rewards)).float().to(self.device)
             terminated = from_numpy(np.asarray(terminated)).to(self.device)
             truncated = from_numpy(np.asarray(truncated)).to(self.device)
@@ -2169,7 +2181,7 @@ class EPO(Module):
 
                 latent_id, episode_id, latent, state = slot.latent_id, slot.episode_id, slot.latent, slot.state
 
-                next_state = from_numpy(np.array(next_obs[k])).float().to(self.device)
+                next_state = next_states[k]
 
                 reward, terminated_at, truncated_at = rewards[k], terminated[k], truncated[k]
 

{evolutionary_policy_optimization-0.4.0 → evolutionary_policy_optimization-0.4.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.4.0"
+version = "0.4.2"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -30,10 +30,11 @@ dependencies = [
     'einops>=0.8.1',
     'ema-pytorch>=0.7.7',
     'hl-gauss-pytorch>=0.2.8',
-    'mean-conc-beta>=0.0.7',
+    'mean-conc-beta>=0.2.7',
     'torch>=2.2',
     'tqdm',
     "x-mlps-pytorch>=0.3.0",
+    "x-ppo>=0.0.16",
 ]
 
 [project.urls]