From f594cafeadf2647ab4c66f5d0fb510d4f8e9259e Mon Sep 17 00:00:00 2001 From: RobinMeersman <77965843+RobinMeersman@users.noreply.github.com> Date: Wed, 1 Apr 2026 19:17:42 +0200 Subject: [PATCH] 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 --- .pre-commit-config.yaml | 2 +- configs/example.json | 8 ++ experiments/simulate.py | 23 ++-- ruff.toml | 1 - src/__init__.py | 25 +++++ src/brittle_star_project/__init__.py | 4 + .../dataclasses/PPOArgs.py | 7 +- .../environment/BrittleStarJaxEnvWrapper.py | 19 ++++ .../environment/__init__.py | 3 +- .../environment/env_config.py | 12 +- src/brittle_star_project/render/__init__.py | 3 + src/brittle_star_project/rl/DummyAgent.py | 1 + src/train.py | 105 +++++------------- 13 files changed, 113 insertions(+), 100 deletions(-) create mode 100644 configs/example.json diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b751166..c289bb6 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -18,4 +18,4 @@ repos: rev: v5.0.0 hooks: - id: no-commit-to-branch - args: ['--branch', 'main', '--branch', 'dev'] + args: ['--branch', 'main', '--branch', 'dev'] \ No newline at end of file diff --git a/configs/example.json b/configs/example.json new file mode 100644 index 0000000..8b646db --- /dev/null +++ b/configs/example.json @@ -0,0 +1,8 @@ +{ + "morphology": { + "num_arms": 2, + "num_segments_per_arm": 4, + "use_p_control": true, + "use_torque_control": false + } +} \ No newline at end of file diff --git a/experiments/simulate.py b/experiments/simulate.py index 346bcda..28062fa 100644 --- a/experiments/simulate.py +++ b/experiments/simulate.py @@ -4,17 +4,15 @@ import argparse from pathlib import Path from brittle_star_project import ( - ArenaConfig, Backend, BrittleStarEnv, BrittleStarEnvFactory, - EnvConfig, - MorphologyConfig, - Task, + SimulationConfig, + simulate_policy, ) +from brittle_star_project.environment import from_json from brittle_star_project.rl import RLModel # imports concrete models via rl.__init__ from brittle_star_project.rl.base import get_rl_model_registry -from brittle_star_project.renderer import SimulationConfig, simulate_policy MODEL_BY_NAME = get_rl_model_registry() MODEL_OPTIONS = sorted(MODEL_BY_NAME) @@ -35,9 +33,9 @@ def parse_args() -> argparse.Namespace: help="Which model class to instantiate when --model is omitted.", ) p.add_argument( - "--task", - choices=[t.value for t in Task], - default=Task.DIRECTED_LOCOMOTION.value, + "--backend", + choices=[b for b in Backend], + default=Backend.MJX, ) p.add_argument("--seed", type=int, default=None) return p.parse_args() @@ -46,14 +44,11 @@ def parse_args() -> argparse.Namespace: def main() -> None: args = parse_args() + morphology_cfg, arena_cfg, env_cfg = from_json("../configs/test.json") + # ======= ENVIRONMENT SETUP ======= - backend = Backend.MJC - task = Task(args.task) - - morphology_cfg = MorphologyConfig() - arena_cfg = ArenaConfig(attach_target=(task == Task.DIRECTED_LOCOMOTION)) - env_cfg = EnvConfig(task=task) + backend = args.backend factory = BrittleStarEnvFactory() raw_env = factory.create_environment(backend, morphology_cfg, arena_cfg, env_cfg) diff --git a/ruff.toml b/ruff.toml index cf728e1..3149b75 100644 --- a/ruff.toml +++ b/ruff.toml @@ -393,4 +393,3 @@ extend-ignore = [ "PLW0603", # global-statement # "PLW1404", # implicit-str-concat ] - diff --git a/src/__init__.py b/src/__init__.py index e69de29..35ecd10 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -0,0 +1,25 @@ +from .brittle_star_project import ( + ArenaConfig, + Backend, + BrittleStarEnv, + BrittleStarEnvFactory, + EnvConfig, + MorphologyConfig, + Task, + simulate_policy, + SimulationConfig, + ControlPolicy, +) + +__all__ = [ + "ArenaConfig", + "Backend", + "BrittleStarEnv", + "BrittleStarEnvFactory", + "EnvConfig", + "MorphologyConfig", + "Task", + "simulate_policy", + "SimulationConfig", + "ControlPolicy", +] diff --git a/src/brittle_star_project/__init__.py b/src/brittle_star_project/__init__.py index b6b123a..3902cbb 100644 --- a/src/brittle_star_project/__init__.py +++ b/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", ] diff --git a/src/brittle_star_project/dataclasses/PPOArgs.py b/src/brittle_star_project/dataclasses/PPOArgs.py index 4fc3f2e..03c220e 100644 --- a/src/brittle_star_project/dataclasses/PPOArgs.py +++ b/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 diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py index 1ea07c7..259f1a5 100644 --- a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py +++ b/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})" + ) diff --git a/src/brittle_star_project/environment/__init__.py b/src/brittle_star_project/environment/__init__.py index 78896cb..f87eae5 100644 --- a/src/brittle_star_project/environment/__init__.py +++ b/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", ] diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index f50f4d8..6f55f35 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/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 diff --git a/src/brittle_star_project/render/__init__.py b/src/brittle_star_project/render/__init__.py index e69de29..51e0aa1 100644 --- a/src/brittle_star_project/render/__init__.py +++ b/src/brittle_star_project/render/__init__.py @@ -0,0 +1,3 @@ +from .renderer import simulate_policy, SimulationConfig, ControlPolicy + +__all__ = ["simulate_policy", "SimulationConfig", "ControlPolicy"] diff --git a/src/brittle_star_project/rl/DummyAgent.py b/src/brittle_star_project/rl/DummyAgent.py index fff4068..bc027dc 100644 --- a/src/brittle_star_project/rl/DummyAgent.py +++ b/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 diff --git a/src/train.py b/src/train.py index eb64bbe..ac6849e 100644 --- a/src/train.py +++ b/src/train.py @@ -7,6 +7,7 @@ from typing import Callable import flax import jax import jax.numpy as jnp +import matplotlib.pyplot as plt import numpy as np import optax import torch @@ -18,7 +19,8 @@ from torch.utils.tensorboard import SummaryWriter from brittle_star_project.dataclasses import PPOArgs from brittle_star_project.dataclasses.EpisodeStatistics import EpisodeStatistics from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper -from brittle_star_project.rl import Network, Actor, Critic, AgentParams, Storage +from brittle_star_project.rl import Actor, AgentParams, Critic, Network, Storage +from ppo import PPO def convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray: @@ -27,9 +29,11 @@ def convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray: ) -def make_env(num_envs: int) -> Callable: +def make_env(config_path: str | None, num_envs: int) -> Callable: def thunk(): - return BrittleStarJaxEnvWrapper.default(num_envs=num_envs) + if config_path is None: + return BrittleStarJaxEnvWrapper.default(num_envs=num_envs) + return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs) return thunk @@ -62,14 +66,15 @@ def train(args: PPOArgs): random.seed(args.seed) np.random.seed(args.seed) key = jax.random.PRNGKey(args.seed) - key, network_key, actor_key, critic_key = jax.random.split(key, 4) + key, network_key, actor_key, critic_key, critic_network_key = jax.random.split(key, 5) torch.backends.cudnn.deterministic = args.torch_deterministic device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu") print(f"Running on device: {device}") print("Creating the environment...") - env = make_env(num_envs=args.num_envs)() + env = make_env(config_path=args.config_path, num_envs=args.num_envs)() + print(f"Environment: {env}") episode_stats = EpisodeStatistics( episode_returns=jnp.zeros(args.num_envs, dtype=jnp.float32), @@ -112,6 +117,7 @@ def train(args: PPOArgs): print("Initializing the models...") network = Network() + critic_network = Network() actor = Actor(action_dim=env.single_action_space.shape[0]) # continuous actions for MJX critic = Critic() @@ -123,12 +129,15 @@ def train(args: PPOArgs): ] ) network_params = network.init(network_key, sample_obs) + critic_network_params = critic_network.init(critic_network_key, sample_obs) actor_params = actor.init(actor_key, network.apply(network_params, sample_obs)) - critic_params = critic.init(critic_key, network.apply(network_params, sample_obs)) + critic_params = critic.init(critic_key, critic_network.apply(critic_network_params, sample_obs)) agent_state = TrainState.create( apply_fn=None, - params=asdict(AgentParams(network_params, actor_params, critic_params)), + params=asdict( + AgentParams(network_params, actor_params, critic_params, critic_network_params) + ), tx=optax.chain( optax.clip_by_global_norm(args.max_grad_norm), optax.inject_hyperparams(optax.adam)( @@ -138,8 +147,10 @@ def train(args: PPOArgs): ) network.apply = jax.jit(network.apply) + critic_network.apply = jax.jit(critic_network.apply) actor.apply = jax.jit(actor.apply) critic.apply = jax.jit(critic.apply) + ppo_instance = PPO(args, network, actor, critic, critic_network) @jax.jit def get_action_and_value_noise( @@ -158,20 +169,6 @@ def train(args: PPOArgs): value = critic.apply(agent_state.params["critic_params"], hidden) return action, logprob, value.squeeze(-1), key - @jax.jit - def get_action_and_value( - params: flax.core.FrozenDict, - x: jnp.ndarray, - action: np.ndarray, - ): - hidden = network.apply(params["network_params"], x) - mean, log_std = actor.apply(params["actor_params"], hidden) - std = jnp.exp(log_std) - logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1) - entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(-1) - value = critic.apply(params["critic_params"], hidden).squeeze(-1) - return logprob, entropy, value - @jax.jit def compute_gae_once(carry, inp, gamma, gae_lambda): advantages = carry @@ -199,61 +196,6 @@ def train(args: PPOArgs): ) return storage.replace(advantages=advantages, returns=advantages + storage.values) - def ppo_loss(params, x, a, logp, mb_advantages, mb_returns): - newlogprob, entropy, newvalue = get_action_and_value(params, x, a) - logratio = newlogprob - logp - ratio = jnp.exp(logratio) - approx_kl = ((ratio - 1) - logratio).mean() - - if args.norm_adv: - mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8) - - pg_loss1 = -mb_advantages * ratio - pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef) - pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean() - v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean() - entropy_loss = entropy.mean() - loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef - return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl)) - - ppo_loss_grad_fn = jax.value_and_grad(ppo_loss, has_aux=True) - - @jax.jit - def update_ppo(agent_state, storage, key): - def update_epoch(carry, _): - agent_state, key = carry - key, subkey = jax.random.split(key) - - def flatten(x): - return x.reshape((-1,) + x.shape[2:]) - - def convert_data(x): - x = jax.random.permutation(subkey, x) - return jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:]) - - flatten_storage = jax.tree.map(flatten, storage) - shuffled_storage = jax.tree.map(convert_data, flatten_storage) - - def update_minibatch(agent_state, minibatch): - (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn( - agent_state.params, - minibatch.obs, - minibatch.actions, - minibatch.logprobs, - minibatch.advantages, - minibatch.returns, - ) - agent_state = agent_state.apply_gradients(grads=grads) - return agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) - - agent_state, metrics = jax.lax.scan(update_minibatch, agent_state, shuffled_storage) - return (agent_state, key), metrics - - (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan( - update_epoch, (agent_state, key), (), length=args.update_epochs - ) - return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key - # --- Main training loop --- global_step = 0 start_time = time.time() @@ -303,6 +245,7 @@ def train(args: PPOArgs): print("Starting training...") iters_bar = tqdm.tqdm(range(1, args.num_iterations + 1)) + losses = [] for _ in iters_bar: iteration_time_start = time.time() @@ -312,10 +255,12 @@ def train(args: PPOArgs): global_step += args.num_steps * args.num_envs storage = compute_gae(agent_state, next_obs, next_done, storage) - agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key = update_ppo( + agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key = ppo_instance.update_ppo( agent_state, storage, key ) + losses.append(jnp.mean(loss)) + avg_episodic_return = np.mean(jax.device_get(episode_stats.returned_episode_returns)) iters_bar.set_postfix_str( f"global_step={global_step}, avg_episodic_return={avg_episodic_return}" @@ -367,6 +312,12 @@ def train(args: PPOArgs): env.close() writer.close() + print("Saving loss plot...") + plt.plot(losses) + plt.title("PPO Loss, mean over minibatches") + plt.savefig(f"runs/{run_name}/{args.exp_name}_losses.png") + plt.close() + def main() -> None: args = tyro.cli(PPOArgs)