evolutionary-policy-optimization 0.0.41__tar.gz → 0.0.43__tar.gz

{evolutionary_policy_optimization-0.0.41 → evolutionary_policy_optimization-0.0.43}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evolutionary-policy-optimization
-Version: 0.0.41
+Version: 0.0.43
 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.41 → evolutionary_policy_optimization-0.0.43}/evolutionary_policy_optimization/epo.py

@@ -875,25 +875,32 @@ def create_agent(
     actor_num_actions,
     actor_dim_hiddens: int | tuple[int, ...],
     critic_dim_hiddens: int | tuple[int, ...],
+    latent_gene_pool_kwargs: dict = dict(),
+    actor_kwargs: dict = dict(),
+    critic_kwargs: dict = dict(),
 ) -> Agent:
 
     latent_gene_pool = LatentGenePool(
         num_latents = num_latents,
-        dim_latent = dim_latent
+        dim_latent = dim_latent,
+        **latent_gene_pool_kwargs
     )
 
     actor = Actor(
         num_actions = actor_num_actions,
         dim_state = dim_state,
         dim_latent = dim_latent,
-        dim_hiddens = actor_dim_hiddens
+        dim_hiddens = actor_dim_hiddens,
+        **actor_kwargs
     )
 
     critic = Critic(
         dim_state = dim_state,
         dim_latent = dim_latent,
-        dim_hiddens = critic_dim_hiddens
-    )                
+        dim_hiddens = critic_dim_hiddens,
+        **critic_kwargs
+    )
+
     return Agent(actor = actor, critic = critic, latent_gene_pool = latent_gene_pool)
 
 # EPO - which is just PPO with natural selection of a population of latent variables conditioning the agent

{evolutionary_policy_optimization-0.0.41 → evolutionary_policy_optimization-0.0.43}/evolutionary_policy_optimization/mock_env.py

@@ -7,7 +7,7 @@ from torch.nn import Module
 # functions
 
 def cast_tuple(v):
-    return v if isinstance(v, tuple) else v\
+    return v if isinstance(v, tuple) else (v,)
 
 # mock env
 

{evolutionary_policy_optimization-0.0.41 → evolutionary_policy_optimization-0.0.43}/pyproject.toml

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

{evolutionary_policy_optimization-0.0.41 → evolutionary_policy_optimization-0.0.43}/tests/test_epo.py

@@ -73,7 +73,10 @@ def test_create_agent(
     agent.save('./agent.pt', overwrite = True)
     agent.load('./agent.pt')
 
-def test_e2e_with_mock_env():
+@pytest.mark.parametrize('frozen_latents', (False, True))
+def test_e2e_with_mock_env(
+    frozen_latents
+):
     from evolutionary_policy_optimization import create_agent, EPO, Env
 
     agent = create_agent(
@@ -82,7 +85,10 @@ def test_e2e_with_mock_env():
         dim_latent = 32,
         actor_num_actions = 5,
         actor_dim_hiddens = (256, 128),
-        critic_dim_hiddens = (256, 128, 64)
+        critic_dim_hiddens = (256, 128, 64),
+        latent_gene_pool_kwargs = dict(
+            frozen_latents = frozen_latents
+        )
     )
 
     epo = EPO(