evolutionary-policy-optimization 0.0.72__tar.gz → 0.1.1__tar.gz

{evolutionary_policy_optimization-0.0.72 → evolutionary_policy_optimization-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.72
+Version: 0.1.1
 Summary: EPO - Pytorch
 Project-URL: Homepage, https://pypi.org/project/evolutionary-policy-optimization/
 Project-URL: Repository, https://github.com/lucidrains/evolutionary-policy-optimization
@@ -34,6 +34,7 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
+Requires-Dist: accelerate
 Requires-Dist: adam-atan2-pytorch
 Requires-Dist: assoc-scan>=0.0.2
 Requires-Dist: einops>=0.8.1

{evolutionary_policy_optimization-0.0.72 → evolutionary_policy_optimization-0.1.1}/evolutionary_policy_optimization/__init__.py

@@ -9,3 +9,7 @@ from evolutionary_policy_optimization.epo import (
 )
 
 from evolutionary_policy_optimization.mock_env import Env
+
+from evolutionary_policy_optimization.env_wrappers import (
+    GymnasiumEnvWrapper
+)

evolutionary_policy_optimization-0.1.1/evolutionary_policy_optimization/env_wrappers.py

@@ -0,0 +1,36 @@
+import torch
+from torch.nn import Module
+
+from evolutionary_policy_optimization.epo import create_agent, Agent
+
+class GymnasiumEnvWrapper(Module):
+    def __init__(
+        self,
+        env
+    ):
+        super().__init__()
+        self.env = env
+
+    def reset(self, *args, **kwargs):
+        return self.env.reset(*args, **kwargs)
+
+    def step(self, *args, **kwargs):
+        return self.env.step(*args, **kwargs)
+
+    def to_agent_hparams(self):
+        return dict(
+            dim_state = self.env.observation_space.shape[0],
+            actor_num_actions = self.env.action_space.n
+        )
+
+    def to_epo_agent(
+        self,
+        *args,
+        **kwargs
+    ) -> Agent:
+
+        return create_agent(
+            *args,
+            **self.to_agent_hparams(),
+            **kwargs
+        )

{evolutionary_policy_optimization-0.0.72 → evolutionary_policy_optimization-0.1.1}/evolutionary_policy_optimization/epo.py

@@ -70,7 +70,7 @@ def interface_torch_numpy(fn, device):
 
         out = fn(*args, **kwargs)
 
-        out = tree_map(lambda t: from_numpy(t).to(device) if isinstance(t, np.ndarray) else t, out)
+        out = tree_map(lambda t: from_numpy(np.array(t)).to(device) if isinstance(t, (np.ndarray, np.float64)) else t, out)
         return out
 
     return decorated_fn
@@ -1040,6 +1040,7 @@ def actor_loss(
 # agent contains the actor, critic, and the latent genetic pool
 
 def create_agent(
+    *,
     dim_state,
     num_latents,
     dim_latent,
@@ -1176,7 +1177,7 @@ class EPO(Module):
             if fix_environ_across_latents:
                 maybe_seed = environment_seeds[episode_id]
 
-            yield latent_id, episode_id, maybe_seed
+            yield latent_id, episode_id, maybe_seed.item()
 
     @torch.no_grad()
     def gather_experience_from(
@@ -1210,7 +1211,7 @@ class EPO(Module):
             if fix_environ_across_latents:
                 reset_kwargs.update(seed = maybe_seed)
 
-            state = interface_torch_numpy(env.reset, device = self.device)(**reset_kwargs)
+            state, _ = interface_torch_numpy(env.reset, device = self.device)(**reset_kwargs)
 
             # get latent from pool
 
@@ -1232,7 +1233,7 @@ class EPO(Module):
 
                 # get the next state, action, and reward
 
-                state, reward, truncated, terminated = interface_torch_numpy(env.forward, device = self.device)(action)
+                state, reward, truncated, terminated, _ = interface_torch_numpy(env.step, device = self.device)(action)
 
                 done = truncated or terminated
 
@@ -1250,7 +1251,7 @@ class EPO(Module):
                     log_prob,
                     reward,
                     value,
-                    terminated
+                    tensor(terminated)
                 )
 
                 memory = Memory(*tuple(t.cpu() for t in memory))

{evolutionary_policy_optimization-0.0.72 → evolutionary_policy_optimization-0.1.1}/evolutionary_policy_optimization/mock_env.py

@@ -22,34 +22,34 @@ class Env(Module):
         self.state_shape = cast_tuple(state_shape)
 
         self.can_terminate_after = can_terminate_after
-        self.register_buffer('step', tensor(0))
+        self.register_buffer('_step', tensor(0))
 
     @property
     def device(self):
-        return self.step.device
+        return self._step.device
 
     def reset(
         self,
         seed = None
     ):
         state = randn(self.state_shape, device = self.device)
-        self.step.zero_()
-        return state.numpy()
+        self._step.zero_()
+        return state.numpy(), None
 
-    def forward(
+    def step(
         self,
         actions,
     ):
         state = randn(self.state_shape, device = self.device)
         reward = randint(0, 5, (), device = self.device).float()
 
-        if self.step > self.can_terminate_after:
+        if self._step > self.can_terminate_after:
             truncated = tensor(choice((True, False)), device =self.device)
             terminated = tensor(choice((True, False)), device =self.device)
         else:
             truncated = terminated = tensor(False, device = self.device)
 
-        self.step.add_(1)
+        self._step.add_(1)
 
-        out = state, reward, truncated, terminated
-        return tuple(t.numpy() for t in out)
+        out = (state, reward, truncated, terminated)
+        return (*tuple(t.numpy() for t in out), None)

{evolutionary_policy_optimization-0.0.72 → evolutionary_policy_optimization-0.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.0.72"
+version = "0.1.1"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -24,6 +24,7 @@ classifiers = [
   ]
 
 dependencies = [
+    "accelerate",
     "adam-atan2-pytorch",
     'assoc-scan>=0.0.2',
     'einx>=0.3.0',

{evolutionary_policy_optimization-0.0.72 → evolutionary_policy_optimization-0.1.1}/requirements.txt

@@ -1,3 +1,5 @@
 box2d-py
 gymnasium[box2d]>=1.0.0
+moviepy>=1.0.3
+numpy>=2.2.5
 tqdm

evolutionary_policy_optimization-0.1.1/train_gym.py

@@ -0,0 +1,43 @@
+import torch
+
+from evolutionary_policy_optimization import (
+    EPO,
+    GymnasiumEnvWrapper
+)
+
+# gymnasium
+
+import gymnasium as gym
+
+env = gym.make(
+    'LunarLander-v3',
+    render_mode = 'rgb_array'
+)
+
+env = GymnasiumEnvWrapper(env)
+
+# epo
+
+agent = env.to_epo_agent(
+    num_latents = 8,
+    dim_latent = 32,
+    actor_dim_hiddens = (256, 128),
+    critic_dim_hiddens = (256, 128, 64),
+    latent_gene_pool_kwargs = dict(
+        frac_natural_selected = 0.5,
+        frac_tournaments = 0.5
+    )
+)
+
+epo = EPO(
+    agent,
+    episodes_per_latent = 1,
+    max_episode_length = 10,
+    action_sample_temperature = 1.
+)
+
+epo.to('cpu' if not torch.cuda.is_available() else 'cuda')
+
+epo(agent, env, num_learning_cycles = 1)
+
+agent.save('./agent.pt', overwrite = True)