Morphology/brittle star 2 arms (#17)

Provide the features to read in Environment (Morphology, arena, ...) config files in the JSON format. The example JSON file contains the config for a brittle star with 2 arms

src/brittle_star_project/__init__.py

@@ -2,6 +2,7 @@ from .environment.env_types import Backend, Task
 from .environment.env_config import ArenaConfig, EnvConfig, MorphologyConfig
 from .environment.factory import BrittleStarEnvFactory
 from .environment.env_wrapper import BrittleStarEnv
+from .render import simulate_policy, SimulationConfig, ControlPolicy
 
 __all__ = [
     "ArenaConfig",
@@ -11,4 +12,7 @@ __all__ = [
     "EnvConfig",
     "MorphologyConfig",
     "Task",
+    "simulate_policy",
+    "SimulationConfig",
+    "ControlPolicy",
 ]

src/brittle_star_project/dataclasses/PPOArgs.py

@@ -7,6 +7,9 @@ class PPOArgs:
     source: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py
     """
 
+    # path to environment config file, if None, use default config
+    config_path: str | None = None
+
     # the name of this experiment
     exp_name: str = "brittle_star_ppo"
 
@@ -41,8 +44,6 @@ class PPOArgs:
     hf_entity: str = ""
 
     # ==== Algorithm specific dataclasses ====
-    # the id of the environment
-    env_id: str = ""  # todo
 
     # total timesteps of the experiments
     total_timesteps: int = 10000000
@@ -51,7 +52,7 @@ class PPOArgs:
     learning_rate: float = 2.5e-4
 
     # the number of parallel game environments
-    num_envs: int = 16
+    num_envs: int = 100
 
     # the number of steps to run in each environment per policy rollout
     num_steps: int = 128

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -8,6 +8,7 @@ from brittle_star_project import (
     ArenaConfig,
     Backend,
 )
+from brittle_star_project.environment import from_json
 
 
 class BrittleStarJaxEnvWrapper:
@@ -76,3 +77,21 @@ class BrittleStarJaxEnvWrapper:
         return BrittleStarJaxEnvWrapper(
             morphology, arena, env_config, num_envs=num_envs, backend=backend
         )
+
+    @staticmethod
+    def from_config(
+        config_path: str, num_envs: int, backend: Backend = Backend.MJX
+    ) -> "BrittleStarJaxEnvWrapper":
+        morphology_cfg, arena_cfg, env_cfg = from_json(config_path)
+        return BrittleStarJaxEnvWrapper(
+            morphology_cfg, arena_cfg, env_cfg, num_envs=num_envs, backend=backend
+        )
+
+    def __str__(self):
+        morphology_str = str(self._morphology)
+        arena_str = str(self._arena)
+        env_config_str = str(self._env_config)
+        return (
+            f"BrittleStarJaxEnvWrapper(backend={self._backend}, num_envs={self._num_envs}, "
+            + f"morphology={morphology_str}, arena={arena_str}, env_config={env_config_str})"
+        )

src/brittle_star_project/environment/__init__.py

@@ -1,4 +1,4 @@
-from .env_config import ArenaConfig, EnvConfig, MorphologyConfig
+from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, from_json
 from .env_types import Backend, Task
 from .env_wrapper import BrittleStarEnv, StepResult
 from .factory import BrittleStarEnvFactory
@@ -12,4 +12,5 @@ __all__ = [
     "BrittleStarEnv",
     "StepResult",
     "BrittleStarEnvFactory",
+    "from_json",
 ]

src/brittle_star_project/environment/env_config.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+import json
 
 from .env_types import Task
 
@@ -48,6 +49,11 @@ class EnvConfig:
     # Per docs in upstream env config: integer factors of 200.
     light_perlin_noise_scale: int = 0
 
-    @staticmethod
-    def from_json(path: str) -> EnvConfig:
-        pass
+
+def from_json(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]:
+    with open(path, "r") as f:
+        config_json = json.load(f)
+        morphology = MorphologyConfig(**config_json.get("morphology", {}))
+        arena = ArenaConfig(**config_json.get("arena", {}))
+        env = EnvConfig(**config_json.get("env", {}))
+        return morphology, arena, env

src/brittle_star_project/render/__init__.py

@@ -0,0 +1,3 @@
+from .renderer import simulate_policy, SimulationConfig, ControlPolicy
+
+__all__ = ["simulate_policy", "SimulationConfig", "ControlPolicy"]

src/brittle_star_project/rl/DummyAgent.py

@@ -52,6 +52,7 @@ class AgentParams:
     network_params: flax.core.FrozenDict
     actor_params: flax.core.FrozenDict
     critic_params: flax.core.FrozenDict
+    critic_network_params: flax.core.FrozenDict
 
 
 @jax.tree_util.register_dataclass