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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.72
+Version: 0.1.0
 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.72 → evolutionary_policy_optimization-0.1.0}/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
@@ -1176,7 +1176,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 +1210,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 +1232,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 +1250,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.0}/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.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "evolutionary-policy-optimization"
-version = "0.0.72"
+version = "0.1.0"
 description = "EPO - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{evolutionary_policy_optimization-0.0.72 → evolutionary_policy_optimization-0.1.0}/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.0/train_gym.py

@@ -0,0 +1,44 @@
+# gymnasium
+
+import gymnasium as gym
+
+env = gym.make(
+    'LunarLander-v3',
+    render_mode = 'rgb_array'
+)
+
+state_dim = env.observation_space.shape[0]
+num_actions = env.action_space.n
+
+# epo
+
+import torch
+
+from evolutionary_policy_optimization import (
+    create_agent,
+    EPO,
+    Env
+)
+
+agent = create_agent(
+    dim_state = state_dim,
+    num_latents = 1,
+    dim_latent = 32,
+    actor_num_actions = num_actions,
+    actor_dim_hiddens = (256, 128),
+    critic_dim_hiddens = (256, 128, 64),
+    latent_gene_pool_kwargs = dict(
+        frac_natural_selected = 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 = 5)