evolutionary-policy-optimization 0.3.5__tar.gz → 0.4.1__tar.gz

{evolutionary_policy_optimization-0.3.5 → evolutionary_policy_optimization-0.4.1}/.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.3.5 → evolutionary_policy_optimization-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: evolutionary-policy-optimization
-Version: 0.3.5
+Version: 0.4.1
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -40,7 +40,7 @@ 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
@@ -367,6 +367,16 @@ That's it
 }
 ```
 
+```bibtex
+@inproceedings{Sharma2020DynamicsAwareUD,
+    title   = {Dynamics-Aware Unsupervised Discovery of Skills},
+    author  = {Archit Sharma and Shixiang Gu and Sergey Levine and Vikash Kumar and Karol Hausman},
+    booktitle = {International Conference on Learning Representations},
+    year    = {2020},
+    url     = {https://arxiv.org/abs/1907.01657}
+}
+```
+
 ```bibtex
 @article{Zhu2025Hybrid,
     author  = {Zhu, Zimo and Yu, Chuanqiang and Wang, Junti},

{evolutionary_policy_optimization-0.3.5 → evolutionary_policy_optimization-0.4.1}/README.md

@@ -305,6 +305,16 @@ That's it
 }
 ```
 
+```bibtex
+@inproceedings{Sharma2020DynamicsAwareUD,
+    title   = {Dynamics-Aware Unsupervised Discovery of Skills},
+    author  = {Archit Sharma and Shixiang Gu and Sergey Levine and Vikash Kumar and Karol Hausman},
+    booktitle = {International Conference on Learning Representations},
+    year    = {2020},
+    url     = {https://arxiv.org/abs/1907.01657}
+}
+```
+
 ```bibtex
 @article{Zhu2025Hybrid,
     author  = {Zhu, Zimo and Yu, Chuanqiang and Wang, Junti},

{evolutionary_policy_optimization-0.3.5 → evolutionary_policy_optimization-0.4.1}/evolutionary_policy_optimization/__init__.py

@@ -1,3 +1,3 @@
 from evolutionary_policy_optimization.env_wrappers import GymnasiumEnvWrapper, rescale_from_to
-from evolutionary_policy_optimization.epo import EPO, MLP, Actor, Agent, Beta, CategoricalActionDistr, Critic, LatentGenePool, create_agent
+from evolutionary_policy_optimization.epo import EPO, MLP, Actor, Agent, Beta, CategoricalActionDistr, Critic, DiversityDiscr, LatentGenePool, create_agent
 from evolutionary_policy_optimization.mock_env import Env, VectorEnv

{evolutionary_policy_optimization-0.3.5 → evolutionary_policy_optimization-0.4.1}/evolutionary_policy_optimization/epo.py

@@ -52,6 +52,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)
@@ -513,7 +527,7 @@ class MLP(Module):
 
         return self.final_lime(x, prev_layer_inputs)
 
-# discriminator for predicting latent code from states and actions (DIAYN - Eysenbach et al. 2018)
+# discriminator for predicting latent code from states or state transitions (DIAYN / DADS)
 
 class DiversityDiscr(Module):
     def __init__(
@@ -521,15 +535,19 @@ class DiversityDiscr(Module):
         dim_state,
         num_latents,
         dim = 64,
-        depth = 2
+        depth = 2,
+        time_gap = 1
     ):
         super().__init__()
+        self.time_gap = time_gap
+        self.state_transition = time_gap > 0
+
         self.state_proj = nn.Linear(dim_state, dim)
 
         self.net = AttnResidualNormedMLP(
             dim = dim,
             depth = depth,
-            dim_in = dim * 2,
+            dim_in = dim * 2 if self.state_transition else dim,
             dim_out = num_latents
         )
 
@@ -538,8 +556,12 @@ class DiversityDiscr(Module):
             if isinstance(module, (nn.Linear, nn.LayerNorm, nn.RMSNorm)):
                 module.reset_parameters()
 
-    def forward(self, state, next_state):
+    def forward(self, state, next_state = None):
         state_embed = self.state_proj(state)
+
+        if not self.state_transition or not exists(next_state):
+            return self.net(state_embed)
+
         next_state_embed = self.state_proj(next_state)
         return self.net((state_embed, next_state_embed))
 
@@ -1215,6 +1237,7 @@ class Agent(Module):
         critic_optim_kwargs: dict = dict(),
         latent_optim_kwargs: dict = dict(),
         use_diversity_discr = False,
+        diversity_time_gap = 1, # 0 for DIAYN (state only), >= 1 for DADS (state transition with given time gap)
         diversity_discr_kwargs: dict = dict(dim = 64, depth = 2),
         diversity_discr_lr = 3e-4,
         diversity_discr_optim_kwargs: dict = dict(),
@@ -1225,6 +1248,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,
@@ -1236,6 +1260,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)
@@ -1305,14 +1330,16 @@ class Agent(Module):
         # diversity discriminator
 
         self.use_diversity_discr = use_diversity_discr
+        self.diversity_time_gap = diversity_time_gap
 
         if use_diversity_discr:
-            assert exists(latent_gene_pool), 'latent_gene_pool must be present to use DIAYN'
+            assert exists(latent_gene_pool), 'latent_gene_pool must be present to use DIAYN / DADS'
             dim_state = actor.init_layer[0].in_features
 
             self.diversity_discr = DiversityDiscr(
                 dim_state = dim_state,
                 num_latents = latent_gene_pool.num_latents,
+                time_gap = diversity_time_gap,
                 **diversity_discr_kwargs
             )
 
@@ -1443,6 +1470,7 @@ class Agent(Module):
             critic_ema = self.critic_ema.state_dict() if self.use_critic_ema else None,
             latents = unwrap(self.latent_gene_pool).state_dict() if self.has_latent_genes else None,
             diversity_discr = unwrap(self.diversity_discr).state_dict() if self.use_diversity_discr else None,
+            diversity_time_gap = self.diversity_time_gap,
             has_diversity_discr_warmed_up = self.has_diversity_discr_warmed_up.item(),
             actor_optim = unwrap_optim(self.actor_optim).state_dict(),
             critic_optim = unwrap_optim(self.critic_optim).state_dict(),
@@ -1475,6 +1503,9 @@ class Agent(Module):
         if self.use_diversity_discr and 'diversity_discr' in pkg and exists(pkg['diversity_discr']):
             unwrap(self.diversity_discr).load_state_dict(pkg['diversity_discr'])
 
+        if 'diversity_time_gap' in pkg:
+            self.diversity_time_gap = pkg['diversity_time_gap']
+
         if 'has_diversity_discr_warmed_up' in pkg:
             self.has_diversity_discr_warmed_up.copy_(tensor(pkg['has_diversity_discr_warmed_up']))
 
@@ -1577,13 +1608,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()
 
@@ -1625,12 +1655,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
 
@@ -1722,7 +1755,8 @@ class Agent(Module):
                 diversity_discr_loss = self.zero
 
                 if self.use_diversity_discr:
-                    diversity_discr_logits = self.diversity_discr(states, next_states)
+                    discr_inp = (states,) if self.diversity_time_gap == 0 else (states, next_states)
+                    diversity_discr_logits = self.diversity_discr(*discr_inp)
                     diversity_discr_loss = F.cross_entropy(diversity_discr_logits, latent_gene_ids)
 
                     diversity_discr_loss.backward()
@@ -1878,6 +1912,7 @@ def create_agent(
     use_critic_ema = True,
     use_state_norm = False,
     action_is_continuous = False, # continuous control - beta policy
+    diversity_time_gap = 1,
     latent_gene_pool_kwargs: dict = dict(),
     actor_kwargs: dict = dict(),
     critic_kwargs: dict = dict(),
@@ -1923,6 +1958,7 @@ def create_agent(
         state_norm = state_norm,
         latent_gene_pool = latent_gene_pool,
         use_critic_ema = use_critic_ema,
+        diversity_time_gap = diversity_time_gap,
         **kwargs
     )
 
@@ -2142,6 +2178,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)
@@ -2151,7 +2188,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]
 
@@ -2159,7 +2196,16 @@ class EPO(Module):
 
                 if self.agent.use_diversity_discr and self.agent.has_diversity_discr_warmed_up.item():
                     with torch.no_grad():
-                        logits = self.agent.unwrap_model(self.agent.diversity_discr)(rearrange(state, '... -> 1 ...'), rearrange(next_state, '... -> 1 ...'))
+                        discr = self.agent.unwrap_model(self.agent.diversity_discr)
+                        time_gap = self.agent.diversity_time_gap
+
+                        if time_gap == 0:
+                            logits = discr(rearrange(state, '... -> 1 ...'))
+                        elif time_gap == 1 or len(slot.memories) < (time_gap - 1):
+                            logits = discr(rearrange(state, '... -> 1 ...'), rearrange(next_state, '... -> 1 ...'))
+                        else:
+                            past_state = slot.memories[-(time_gap - 1)].state
+                            logits = discr(rearrange(past_state, '... -> 1 ...'), rearrange(next_state, '... -> 1 ...'))
 
                         latent_log_probs = logits.log_softmax(dim = -1)
                         diversity_reward = latent_log_probs[0, latent_id] + math.log(self.agent.num_latents)
@@ -2184,6 +2230,11 @@ class EPO(Module):
 
                 slot.memories.append(memory)
 
+                if self.agent.use_diversity_discr and self.agent.diversity_time_gap > 1:
+                    time_gap = self.agent.diversity_time_gap
+                    if len(slot.memories) >= time_gap:
+                        slot.memories[-time_gap] = slot.memories[-time_gap]._replace(next_state = next_state)
+
                 rewards_per_latent_episode[latent_id, episode_id] += reward
 
                 # episode is done - bootstrap value if truncated, then refill

{evolutionary_policy_optimization-0.3.5 → evolutionary_policy_optimization-0.4.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.3.5"
+version = "0.4.1"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -30,7 +30,7 @@ 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",