tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
Archived
1
Fork 0
Code Issues Releases Activity

Merge pull request #6 from SELab-3-2026/chore/MLP_Designs

Browse source 
Semi Generic MLP library.
This commit is contained in:
Cedric Mekeirle 2026-04-03 12:55:37 +02:00 • committed by GitHub
commit 4956e1593d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 62 additions and 298 deletions
Download patch file Download diff file Expand all files Collapse all files

33
src/brittle_star_project/rl/DummyAgent.py → src/MLPs/mlps.py
View file

@ -1,36 +1,27 @@
from dataclasses import dataclass, fields
from dataclasses import dataclass, fields, field
import flax
import flax.linen as nn
import jax.numpy as jnp
import jax.tree_util
import numpy as np
from typing import Sequence, Callable
from flax.linen.initializers import constant, orthogonal
class Network(nn.Module):
"""
Dummy model only used for testing purposes
inspired by: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py
"""
hidden_dim: int = 195
# semi generic so we can easily make a config for it in experiments
class GenericDenseLayersWithActivation(nn.Module):
layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64])
activation: Callable = nn.tanh
@nn.compact
def __call__(self, x):
x = nn.Dense(self.hidden_dim, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(
x
)
x = nn.relu(x)
x = nn.Dense(self.hidden_dim, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(
x
)
x = nn.relu(x)
for size in self.layer_sizes:
x = nn.Dense(size, kernel_init=orthogonal(jnp.sqrt(2)))(x)
x = self.activation(x)
return x
class Critic(nn.Module):
class OneDenseLayerMLP(nn.Module):
@nn.compact
def __call__(self, x):
return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
@ -49,10 +40,10 @@ class Actor(nn.Module):
@jax.tree_util.register_dataclass
@dataclass
class AgentParams:
network_params: flax.core.FrozenDict
sensor_params: flax.core.FrozenDict
actor_params: flax.core.FrozenDict
critic_params: flax.core.FrozenDict
critic_network_params: flax.core.FrozenDict
feature_extractor_params: flax.core.FrozenDict
@jax.tree_util.register_dataclass

25
src/brittle_star_project/rl/__init__.py
View file

@ -1,25 +0,0 @@
from .DummyAgent import Network, Critic, Actor, AgentParams, Storage
from .base import (
RLAlgorithm,
RLModel,
Transition,
create_model,
register_rl_model,
registered_model_types,
)
from .random_policy_model import RandomPolicyModel
__all__ = [
"RLAlgorithm",
"RLModel",
"RandomPolicyModel",
"Transition",
"create_model",
"register_rl_model",
"registered_model_types",
"Network",
"Critic",
"Actor",
"AgentParams",
"Storage",
]

162
src/brittle_star_project/rl/base.py
View file

@ -1,162 +0,0 @@
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
import json
from pathlib import Path
from typing import Any
@dataclass(frozen=True, slots=True)
class Transition:
"""A minimal transition container for RL.
This is intentionally generic because the underlying env state type may be a
JAX pytree, a numpy struct, or something library-specific.
"""
obs: Any
action: Any
reward: float
next_obs: Any
terminated: bool
truncated: bool
info: dict[str, Any] | None = None
class RLAlgorithm(ABC):
"""Insertable RL algorithm interface."""
@abstractmethod
def select_action(self, *, obs: Any, rng: Any | None = None) -> Any:
raise NotImplementedError
def observe(self, transition: Transition) -> None:
"""Optional hook to store transitions."""
def update(self, *, rng: Any | None = None) -> dict[str, float]:
"""Optional hook to run one training update."""
return {}
def save(self, path: str) -> None:
raise NotImplementedError("Save not implemented")
def load(self, path: str) -> None:
raise NotImplementedError("Load not implemented")
_RL_MODEL_REGISTRY: dict[str, type["RLModel"]] = {}
def registered_model_types() -> list[str]:
return sorted(_RL_MODEL_REGISTRY)
def create_model(type_name: str, *, payload: dict[str, Any]) -> "RLModel":
model_cls = _RL_MODEL_REGISTRY.get(type_name)
if model_cls is None:
known = ", ".join(sorted(_RL_MODEL_REGISTRY)) or "<none>"
raise ValueError(f"Unknown RLModel type '{type_name}'. Known: {known}")
return model_cls.from_payload(payload)
def get_rl_model_registry() -> dict[str, type["RLModel"]]:
"""Return a copy of the current RLModel registry.
The registry is populated by importing concrete model modules that use the
`@register_rl_model(...)` decorator.
"""
return dict(_RL_MODEL_REGISTRY)
def register_rl_model(*type_names: str):
"""Decorator to register an `RLModel` for generic loading.
Concrete model modules should apply this decorator, so `base.py` never needs
to import concrete models (avoids circular imports).
"""
if not type_names:
raise TypeError("register_rl_model() requires at least one type name")
primary = type_names[0]
def _decorator(cls: type[RLModel]):
for name in type_names:
_RL_MODEL_REGISTRY[name] = cls
cls.type_name = primary
return cls
return _decorator
class RLModel(ABC):
"""Serializable policy/model interface.
This is the artifact that `train.py` writes and `simulate.py` loads.
"""
# Overwritten by the `@register_rl_model(...)` decorator.
type_name: str = "RLModel"
def reset(self, seed: int | None = None) -> None:
"""Optional hook for RNG/stateful models."""
@abstractmethod
def act(self, *, obs: Any | None = None, t: float = 0.0) -> Any:
raise NotImplementedError
def train(self, *, env: Any, num_epochs: int = 1) -> None:
"""Optional training hook.
Many models won't learn; for those this can be a no-op.
"""
_ = (env, num_epochs)
def to_payload(self) -> dict[str, Any]:
"""Return JSON-serializable model parameters."""
return {}
@classmethod
def from_payload(cls, payload: dict[str, Any]) -> "RLModel":
"""Reconstruct a model from `to_payload()` output."""
return cls(**payload) # type: ignore[arg-type]
def save(self, path: str | Path) -> Path:
out = Path(path)
out.parent.mkdir(parents=True, exist_ok=True)
doc = {
"type": self.type_name,
"version": 1,
"payload": self.to_payload(),
}
out.write_text(json.dumps(doc, indent=2, sort_keys=True) + "\n")
return out
@classmethod
def load(cls, path: str | Path) -> "RLModel":
p = Path(path)
doc = json.loads(p.read_text())
type_name = doc.get("type")
if not isinstance(type_name, str):
raise ValueError("Model artifact missing string field 'type'")
model_cls = _RL_MODEL_REGISTRY.get(type_name)
if model_cls is None:
known = ", ".join(sorted(_RL_MODEL_REGISTRY)) or "<none>"
raise ValueError(f"Unknown RLModel type '{type_name}'. Known: {known}")
payload = doc.get("payload")
# Backward compatibility: older artifacts stored fields at top-level.
if payload is None:
payload = {k: v for k, v in doc.items() if k not in ("type", "version")}
if not isinstance(payload, dict):
raise ValueError("Model artifact field 'payload' must be an object")
return model_cls.from_payload(payload)

52
src/brittle_star_project/rl/random_policy_model.py
View file

@ -1,52 +0,0 @@
from __future__ import annotations
from dataclasses import dataclass, field
import numpy as np
from .base import RLModel, register_rl_model
@register_rl_model("random")
@dataclass(slots=True)
class RandomPolicyModel(RLModel):
"""A minimal, serializable policy model that outputs random controls.
This is intentionally *not* a learning algorithm yet. It exists so we can:
- produce a stable model artifact from `train.py`
- load that artifact in `simulate.py`
- drive the MuJoCo viewer with the model's actions
"""
nu: int = 0
seed: int = 0
ctrl_noise_scale: float = 0.5
_rng: np.random.RandomState = field(init=False, repr=False)
def __post_init__(self) -> None:
self.reset(self.seed)
def reset(self, seed: int | None = None) -> None:
if seed is not None:
self.seed = int(seed)
self._rng = np.random.RandomState(self.seed)
def act(self, *, obs: np.ndarray | None = None, t: float = 0.0) -> np.ndarray:
if self.nu <= 0:
return np.zeros((0,), dtype=np.float32)
ctrl = self.ctrl_noise_scale * self._rng.randn(self.nu)
return ctrl.astype(np.float32)
def to_payload(self) -> dict[str, object]:
return {
"seed": int(self.seed),
"ctrl_noise_scale": float(self.ctrl_noise_scale),
}
@classmethod
def from_payload(cls, payload: dict[str, object]) -> RandomPolicyModel:
return cls(
seed=int(payload.get("seed", 0)),
ctrl_noise_scale=float(payload.get("ctrl_noise_scale", 0.5)),
)

40
src/ppo.py
View file

@ -8,9 +8,7 @@ import jax.numpy as jnp
# Chose to use a class as it seemed the easiest way to integrate the CleanRL code style
# with our need to seperate concerns
class PPO:
def __init__(
self, args, input_network, action_network, critic, critic_network, message_passer=None
):
def __init__(self, args, sensor, actor, critic, feature_extractor, message_passer=None):
self.args = args
if not message_passer:
@ -20,10 +18,10 @@ class PPO:
partial(
ppo_loss,
args=args,
input_network_apply=input_network.apply,
action_network_apply=action_network.apply,
sensor_apply=sensor.apply,
actor_apply=actor.apply,
critic_apply=critic.apply,
critic_network_apply=critic_network.apply,
feature_extractor_apply=feature_extractor.apply,
message_passer=message_passer,
),
has_aux=True,
@ -87,20 +85,20 @@ that are now not in the same scope
@partial(jax.jit, static_argnums=(0, 1, 2, 3, 4))
def get_action_and_value2(
input_apply,
action_apply,
def get_action_and_value(
sensor_apply,
actor_apply,
message_passer,
critic_apply,
critic_network_apply,
feature_extractor_apply,
params: flax.core.FrozenDict,
x: jnp.ndarray,
action: jnp.ndarray,
):
hidden_network = input_apply(params["network_params"], x)
hidden_critic = critic_network_apply(params["critic_network_params"], x)
hidden_network = message_passer(hidden_network)
mean, log_std = action_apply(params["actor_params"], hidden_network)
hidden_sensor = sensor_apply(params["sensor_params"], x)
hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x)
hidden_sensor = message_passer(hidden_sensor)
mean, log_std = actor_apply(params["actor_params"], hidden_sensor)
std = jnp.exp(log_std)
logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
@ -118,18 +116,18 @@ def ppo_loss(
mb_advantages,
mb_returns,
args,
input_network_apply,
action_network_apply,
sensor_apply,
actor_apply,
message_passer,
critic_apply,
critic_network_apply,
feature_extractor_apply,
):
newlogprob, entropy, newvalue = get_action_and_value2(
input_network_apply,
action_network_apply,
newlogprob, entropy, newvalue = get_action_and_value(
sensor_apply,
actor_apply,
message_passer,
critic_apply,
critic_network_apply,
feature_extractor_apply,
params,
x,
a,

48
src/train.py
View file

@ -19,7 +19,13 @@ 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 Actor, AgentParams, Critic, Network, Storage
from MLPs.mlps import (
GenericDenseLayersWithActivation,
Actor,
OneDenseLayerMLP,
AgentParams,
Storage,
)
from ppo import PPO
@ -66,7 +72,7 @@ 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, critic_network_key = jax.random.split(key, 5)
key, sensor_key, actor_key, critic_key, feature_extractor_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")
@ -116,10 +122,11 @@ def train(args: PPOArgs):
return args.learning_rate * frac
print("Initializing the models...")
network = Network()
critic_network = Network()
sensor = GenericDenseLayersWithActivation()
feature_extractor = GenericDenseLayersWithActivation()
actor = Actor(action_dim=env.single_action_space.shape[0]) # continuous actions for MJX
critic = Critic()
critic = OneDenseLayerMLP()
# messager = OneDenseLayerMLP()
sample_obs = jnp.concatenate(
[
@ -128,15 +135,17 @@ def train(args: PPOArgs):
if v.size > 0
]
)
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, critic_network.apply(critic_network_params, sample_obs))
sensor_params = sensor.init(sensor_key, sample_obs)
feature_extractor_params = feature_extractor.init(feature_extractor_key, sample_obs)
actor_params = actor.init(actor_key, sensor.apply(sensor_params, sample_obs))
critic_params = critic.init(
critic_key, feature_extractor.apply(feature_extractor_params, sample_obs)
)
agent_state = TrainState.create(
apply_fn=None,
params=asdict(
AgentParams(network_params, actor_params, critic_params, critic_network_params)
AgentParams(sensor_params, actor_params, critic_params, feature_extractor_params)
),
tx=optax.chain(
optax.clip_by_global_norm(args.max_grad_norm),
@ -146,11 +155,11 @@ def train(args: PPOArgs):
),
)
network.apply = jax.jit(network.apply)
critic_network.apply = jax.jit(critic_network.apply)
sensor.apply = jax.jit(sensor.apply)
feature_extractor.apply = jax.jit(feature_extractor.apply)
actor.apply = jax.jit(actor.apply)
critic.apply = jax.jit(critic.apply)
ppo_instance = PPO(args, network, actor, critic, critic_network)
ppo_instance = PPO(args, sensor, actor, critic, feature_extractor)
@jax.jit
def get_action_and_value_noise(
@ -158,7 +167,11 @@ def train(args: PPOArgs):
next_obs: jnp.ndarray,
key: jax.random.PRNGKey,
):
hidden = network.apply(agent_state.params["network_params"], next_obs)
hidden = sensor.apply(agent_state.params["sensor_params"], next_obs)
hidden_critic = feature_extractor.apply(
agent_state.params["feature_extractor_params"], next_obs
)
# Continuous actions: sample from a Gaussian parameterized by the actor
mean, log_std = actor.apply(agent_state.params["actor_params"], hidden)
key, subkey = jax.random.split(key)
@ -166,7 +179,7 @@ def train(args: PPOArgs):
std = jnp.exp(log_std)
action = mean + noise * std
logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
value = critic.apply(agent_state.params["critic_params"], hidden)
value = critic.apply(agent_state.params["critic_params"], hidden_critic)
return action, logprob, value.squeeze(-1), key
@jax.jit
@ -182,7 +195,7 @@ def train(args: PPOArgs):
def compute_gae(agent_state, next_obs, next_done, storage):
next_value = critic.apply(
agent_state.params["critic_params"],
network.apply(agent_state.params["network_params"], next_obs),
sensor.apply(agent_state.params["sensor_params"], next_obs),
).squeeze(-1)
advantages = jnp.zeros((args.num_envs,))
@ -300,9 +313,10 @@ def train(args: PPOArgs):
[
vars(args),
[
agent_state.params["network_params"],
agent_state.params["sensor_params"],
agent_state.params["actor_params"],
agent_state.params["critic_params"],
agent_state.params["feature_extractor_params"],
],
]
)