evolutionary-policy-optimization 0.0.39__py3-none-any.whl → 0.0.40__py3-none-any.whl

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
@@ -73,6 +74,19 @@ 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(
@@ -939,20 +953,13 @@ class EPO(Module):
 
                 while time < self.max_episode_length:
 
-                    batched_state = rearrange(state, '... -> 1 ...')
-
                     # sample action
 
-                    action, log_prob = self.agent.get_actor_actions(batched_state, latent = latent, sample = True, temperature = self.action_sample_temperature)
-
-                    action = rearrange(action, '1 ... -> ...')
-                    log_prob = rearrange(log_prob, '1 ... -> ...')
+                    action, log_prob = temp_batch_dim(self.agent.get_actor_actions)(state, latent = latent, sample = True, temperature = self.action_sample_temperature)
 
                     # values
 
-                    value = self.agent.get_critic_values(batched_state, latent = latent)
-
-                    value = rearrange(value, '1 ... -> ...')
+                    value = temp_batch_dim(self.agent.get_critic_values)(state, latent = latent)
 
                     # get the next state, action, and reward
 
@@ -981,10 +988,7 @@ class EPO(Module):
 
                 # need the final next value for GAE, iiuc
 
-                batched_state = rearrange(state, '... -> 1 ...')
-
-                next_value = self.agent.get_critic_values(batched_state, latent = latent)
-                next_value = rearrange(next_value, '1 ... -> ...')
+                next_value = temp_batch_dim(self.agent.get_critic_values)(state, latent = latent)
 
                 memory_for_gae = memory._replace(
                     episode_id = invalid_episode,

{evolutionary_policy_optimization-0.0.39.dist-info → evolutionary_policy_optimization-0.0.40.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.39
+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

{evolutionary_policy_optimization-0.0.39.dist-info → evolutionary_policy_optimization-0.0.40.dist-info}/RECORD

@@ -1,8 +1,8 @@
 evolutionary_policy_optimization/__init__.py,sha256=0q0aBuFgWi06MLMD8FiHzBYQ3_W4LYWrwmCtF3u5H2A,201
-evolutionary_policy_optimization/epo.py,sha256=lzxPamJahE5KqBwzyYlGOwNeUoB2vONLwtRcWqCI_Jw,29800
+evolutionary_policy_optimization/epo.py,sha256=VZOT1-jdZBE39awP7nhE-I1lGKMTfhUv4Dls9ptNsWk,29854
 evolutionary_policy_optimization/experimental.py,sha256=9FrJGviLESlYysHI3i83efT9g2ZB9ha4u3K9HXN98_w,1100
 evolutionary_policy_optimization/mock_env.py,sha256=QqVPZVJtrvQmSDcnYDTob_A5sDwiUzGj6_tmo6BII5c,918
-evolutionary_policy_optimization-0.0.39.dist-info/METADATA,sha256=TTNQD7sTWIgpVwnrQrFFBD-cyySkvwJr_J3ABxTpor8,5409
-evolutionary_policy_optimization-0.0.39.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-evolutionary_policy_optimization-0.0.39.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-evolutionary_policy_optimization-0.0.39.dist-info/RECORD,,
+evolutionary_policy_optimization-0.0.40.dist-info/METADATA,sha256=5ruqqTCmYto8tqkRlc_peBgRhWkhmRdzUef2ot67ky0,5409
+evolutionary_policy_optimization-0.0.40.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+evolutionary_policy_optimization-0.0.40.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+evolutionary_policy_optimization-0.0.40.dist-info/RECORD,,