evolutionary-policy-optimization 0.2.22__tar.gz → 0.3.1__tar.gz

{evolutionary_policy_optimization-0.2.22 → evolutionary_policy_optimization-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: evolutionary-policy-optimization
-Version: 0.2.22
+Version: 0.3.1
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -35,12 +35,11 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: accelerate>=1.6.0
-Requires-Dist: adam-atan2-pytorch
 Requires-Dist: assoc-scan>=0.0.2
 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.1.19
+Requires-Dist: hl-gauss-pytorch>=0.2.8
 Requires-Dist: torch>=2.2
 Requires-Dist: tqdm
 Requires-Dist: x-mlps-pytorch>=0.3.0
@@ -99,7 +98,8 @@ critic = Critic(dim_state = 32, dim = 256, mlp_depth = 3, dim_latent = 32)
 
 latent = latent_pool(latent_id = 2)
 
-actions = actor(state, latent)
+action_distr = actor(state, latent)
+actions = action_distr.sample()
 value = critic(state, latent)
 
 # interact with environment and receive rewards, termination etc
@@ -145,7 +145,7 @@ epo = EPO(
 
 env = Env((512,))
 
-epo(agent, env, num_learning_cycles = 5)
+epo(env, num_learning_cycles = 5)
 
 # saving and loading
 

{evolutionary_policy_optimization-0.2.22 → evolutionary_policy_optimization-0.3.1}/README.md

@@ -37,7 +37,8 @@ critic = Critic(dim_state = 32, dim = 256, mlp_depth = 3, dim_latent = 32)
 
 latent = latent_pool(latent_id = 2)
 
-actions = actor(state, latent)
+action_distr = actor(state, latent)
+actions = action_distr.sample()
 value = critic(state, latent)
 
 # interact with environment and receive rewards, termination etc
@@ -83,7 +84,7 @@ epo = EPO(
 
 env = Env((512,))
 
-epo(agent, env, num_learning_cycles = 5)
+epo(env, num_learning_cycles = 5)
 
 # saving and loading
 

evolutionary_policy_optimization-0.3.1/evolutionary_policy_optimization/__init__.py

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

{evolutionary_policy_optimization-0.2.22 → evolutionary_policy_optimization-0.3.1}/evolutionary_policy_optimization/distributed.py

@@ -1,12 +1,11 @@
+import einx
 import torch
-from torch import nn
+import torch.distributed as dist
 import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
 from torch.autograd import Function
 
-import torch.distributed as dist
-
-import einx
-from einops import rearrange
 
 def exists(val):
     return val is not None

evolutionary_policy_optimization-0.3.1/evolutionary_policy_optimization/env_wrappers.py

@@ -0,0 +1,71 @@
+from math import prod
+
+import torch
+from torch.nn import Module
+
+from evolutionary_policy_optimization.epo import Agent, create_agent, exists
+
+
+def rescale_from_to(x, from_range = (0., 1.), to_range = (-1., 1.)):
+    # e.g. beta actions on (0, 1) -> the env's action bounds
+
+    from_low, from_high = from_range
+    to_low, to_high = to_range
+
+    if torch.is_tensor(x):
+        dd_kwargs = dict(device = x.device, dtype = x.dtype)
+        from_low, from_high, to_low, to_high = [torch.as_tensor(t, **dd_kwargs) for t in (from_low, from_high, to_low, to_high)]
+
+    return to_low + (to_high - to_low) * (x - from_low) / (from_high - from_low)
+
+class GymnasiumEnvWrapper(Module):
+    def __init__(
+        self,
+        env,
+        rescale_to = None # e.g. (-2., 2.) for pendulum's torque bounds
+    ):
+        super().__init__()
+        self.env = env
+
+        if not exists(rescale_to) and not hasattr(env.action_space, 'n'):
+            rescale_to = (env.action_space.low, env.action_space.high)
+
+        self.rescale_to = rescale_to
+
+    def reset(self, *args, **kwargs):
+        return self.env.reset(*args, **kwargs)
+
+    def step(self, actions, *args, **kwargs):
+        # beta lives on (0, 1) - rescale to the env's bounds at the interface
+
+        if exists(self.rescale_to):
+            actions = rescale_from_to(actions, to_range = self.rescale_to)
+
+        return self.env.step(actions, *args, **kwargs)
+
+    def close(self, *args, **kwargs):
+        return self.env.close(*args, **kwargs)
+
+    def to_agent_hparams(self):
+        action_space = self.env.action_space
+        is_continuous = not hasattr(action_space, 'n')
+
+        num_actions = action_space.n if not is_continuous else prod(action_space.shape)
+
+        return dict(
+            dim_state = self.env.observation_space.shape[0],
+            actor_num_actions = num_actions,
+            action_is_continuous = is_continuous
+        )
+
+    def to_epo_agent(
+        self,
+        *args,
+        **kwargs
+    ) -> Agent:
+
+        return create_agent(
+            *args,
+            **self.to_agent_hparams(),
+            **kwargs
+        )