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

fix(simulate): actor action dimension updated to new way of calculating (as in training)

Browse source
This commit is contained in:
Robin Meersman 2026-05-07 19:53:49 +02:00
parent 6956c5e853
commit 07ee32fd3b
5 changed files with 71 additions and 160 deletions
Download patch file Download diff file Expand all files Collapse all files

107
experiments/debug-experiment-10042026/used_variables.md
View file

@ -1,107 +0,0 @@
## Default envconfig
task: Task = Task.DIRECTED_LOCOMOTION
simulation_time: float = 500.0
num_physics_steps_per_control_step: int = 10
time_scale: int = 2
camera_ids: list[int] = field(default_factory=lambda: [0, 1])
render_size: tuple[int, int] = (480, 640)
joint_randomization_noise_scale: float = 0.0
target_distance: float = 3.0
light_perlin_noise_scale: int = 0
## Default ppoargs
seed: int = 1
torch_deterministic: bool = True
cuda: bool = True
track: bool = False
checkpoint_frequency: int = 100
learning_rate: float = 2.5e-4
anneal_lr: bool = True
gamma: float = 0.99
gae_lambda: float = 0.95
update_epochs: int = 4
norm_adv: bool = True
clip_vloss: bool = True
max_grad_norm: float = 0.5
target_kl: float | None = None
batch_size: int = 0
minibatch_size: int = 0
num_iterations: int = 0
## Used config file: (hpc/debug.yaml)
exp_name: "debug-experiment"
seed: 42
track: true
wandb_project_name: "Let's-find-that-bug"
wandb_entity: "SEL3-2026-Groep-4"
run_dir: "/data/gent/465/vsc46589"
num_envs: 32
num_steps: 32
num_minibatches: 32
total_timesteps: 409600
num_arms: 2
cuda: true
ent_coef: 0.005
vf_coef: 1.0
clip_coef: 0.2
anneal_lr: true
learning_rate: 0.0003
## Arena config:
size: tuple[float, float] = (10.0, 5.0)
sand_ground_color: bool = True
attach_target: bool = True
wall_height: float = 1.5
wall_thickness: float = 0.1
## Morphology:
num_segments_per_arm: int = 4
use_p_control: bool = True
use_torque_control: bool = False
## MLPs:
### Sensor & Feature_extractor:
Both with 3 layers of 300 neurons per layer.
class GenericDenseLayersWithActivation(nn.Module):
layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64])
activation: Callable = nn.tanh
@nn.compact
def __call__(self, x):
for size in self.layer_sizes:
x = nn.Dense(size, kernel_init=orthogonal(jnp.sqrt(2)))(x)
x = self.activation(x)
return x
### Actor:
class Actor(nn.Module):
action_dim: int
@nn.compact
def __call__(self, x):
mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x)
log_std = self.param("log_std", nn.initializers.zeros, (self.action_dim,))
return mean, log_std
### Critic:
class OneDenseLayerMLP(nn.Module):
@nn.compact
def __call__(self, x):
return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
### Observations:
_ALLOWED_OBS_KEYS = {
"joint_position",
"joint_velocity",
"joint_actuator_force",
"actuator_force",
"disk_position",
"disk_rotation",
"disk_linear_velocity",
"disk_angular_velocity",
"unit_xy_direction_to_target",
"xy_distance_to_target",
}

30
scripts/simulate.py
View file

@ -17,12 +17,14 @@ import numpy as np
from omegaconf import DictConfig, OmegaConf
import yaml
import jax.numpy as jnp
from brittle_star_project import Backend, BrittleStarEnv, BrittleStarEnvFactory
from brittle_star_project.configs.main_config import BrittleStarConfig
from brittle_star_project.configs.register_configs import register_configs
from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks
from brittle_star_project.environment.obs_processing import create_obs_processor
from brittle_star_project.environment.env_config import MorphologyConfig
from brittle_star_project.environment.env_config import MorphMode, MorphologyConfig
from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs
from brittle_star_project.evaluation.policy import PolicyAgent
@ -78,10 +80,34 @@ def main(dict_cfg: DictConfig) -> None:
segments_per_arm=env_morphology.segments_per_arm,
reference_segments_per_arm=training.morphology.segments_per_arm,
)
segs_per_arm = jnp.array(env_morphology.segments_per_arm)
needed_copies = 0
agent_indices = [0, 1, 2, 3, 4]
match env_morphology.morph_mode:
case MorphMode.CENTRALIZED:
needed_copies = 1
case MorphMode.FULLY_CONNECTED | MorphMode.RING:
agent_mask = segs_per_arm > 0
agent_indices = jnp.where(agent_mask)[0]
needed_copies = jnp.where(segs_per_arm > 0, 1, 0).sum().item()
case MorphMode.SEGMENT:
agent_mask = segs_per_arm > 0
agent_indices = jnp.where(agent_mask)[0]
needed_copies = jnp.where(segs_per_arm > 0, 1, 0).sum().item()
needed_copies = (segs_per_arm.sum() + jnp.where(segs_per_arm > 0, 1, 0).sum()).item()
num_arms = jnp.where(segs_per_arm > 0, 1, 0).sum().item()
obs_processor = create_obs_processor(
bounds_dict=training.obs_bounds.to_bounds_dict(),
padding_masks=padding_masks,
needed_copies=needed_copies,
num_arms=num_arms,
morph_mode=env_morphology.morph_mode,
segments_per_arm=env_morphology.segments_per_arm,
agent_indices=agent_indices,
)
# 6. Build environment
@ -105,7 +131,7 @@ def main(dict_cfg: DictConfig) -> None:
state0 = env.reset(seed=seed)
# Calculate the action dimension the model was trained with
trained_action_dim = sum(training.morphology.segments_per_arm) * 2
trained_action_dim = raw_env.action_space.shape[0] // needed_copies
# 7. Load policy
policy = PolicyAgent.from_checkpoint(

63
src/brittle_star_project/environment/obs_processing.py
View file

@ -6,6 +6,8 @@ from brittle_star_project.environment.env_config import MorphMode
from experiment_logger import get_logger
logger = get_logger()
_JOINT_SCALED_KEYS = frozenset(
{
"joint_position",
@ -55,8 +57,6 @@ def create_obs_processor(
segments_per_arm=[4, 4, 4, 4, 4],
agent_indices=[0, 1, 2, 3, 4],
):
logger = get_logger()
# made a set to allow O(1) search
ordered_keys = frozenset(
[
@ -87,13 +87,6 @@ def create_obs_processor(
return new_obs
def _prune_features(obs: dict) -> dict:
pruned = {}
for key, arr in obs.items():
if key in ordered_keys:
pruned[key] = arr
return pruned
def _normalize_features(obs: dict) -> dict:
normalized = {}
for key, arr in obs.items():
@ -108,35 +101,11 @@ def create_obs_processor(
normalized[key] = arr
return normalized
def _pad_features(obs: dict, agent_count: int) -> dict:
assert padding_masks is not None
padded = {}
for key, arr in obs.items():
if key in _JOINT_SCALED_KEYS:
target_size = padding_masks["target_size_2x"]
out = jnp.zeros((agent_count, target_size), dtype=arr.dtype)
# place structured values at front, rest stays 0
out = out.at[:, : arr.shape[1]].set(arr)
padded[key] = out
elif key in _SEGMENT_SCALED_KEYS:
target_size = padding_masks["target_size_1x"]
out = jnp.zeros((agent_count, target_size), dtype=arr.dtype)
out = out.at[:, : arr.shape[1]].set(arr)
padded[key] = out
else:
padded[key] = arr
return padded
def _split_to_agents(obs: dict, morph_mode) -> dict:
key_to_agents = {}
output = {}
num_agents = needed_copies # IMPORTANT: number of MLPs
for key, arr in obs.items():
# TODO Should this still be here?
if arr.size == 0:
if key not in ordered_keys or arr.size == 0:
continue
logger.debug(f"[INPUT] {key}: {arr.shape}")
@ -144,17 +113,22 @@ def create_obs_processor(
if arr.ndim == 0:
arr = arr.reshape(1)
# -------- CENTRALIZED --------
if morph_mode == MorphMode.CENTRALIZED:
key_to_agents[key] = arr.reshape(1, -1)
output[key] = arr.reshape(1, -1)
# TODO: padding for centralized
continue
# -------- SEGMENTS --------
if key in _SEGMENT_SCALED_KEYS:
per_agent = []
for agent_id in agent_indices:
taken = jnp.take(arr, agent_id, axis=0) # (segs, ...)
for i, agent_id in enumerate(agent_indices):
idx = segment_indices[i]
taken = jnp.take(arr, idx, axis=0) # (segs, ...)
logger.debug(f"WHY {taken.shape}")
# pad to 4
# pad to 4 (segments per arm?)
pad_len = 4 - taken.shape[0]
padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1))
@ -162,11 +136,13 @@ def create_obs_processor(
out = jnp.stack(per_agent)
# -------- JOINTS --------
elif key in _JOINT_SCALED_KEYS:
per_agent = []
for agent_id in agent_indices:
taken = jnp.take(arr, agent_id, axis=0) # (joint_n, ...)
for i, _ in enumerate(agent_indices):
idx = joint_indices[i]
taken = jnp.take(arr, idx, axis=0) # (joint_n, ...)
# pad to 8
pad_len = 8 - taken.shape[0]
@ -180,9 +156,9 @@ def create_obs_processor(
out = jnp.repeat(arr[None, :], num_agents, axis=0)
logger.debug(f"[OUTPUT] {key}: {out.shape}")
key_to_agents[key] = out
output[key] = out
return key_to_agents
return output
def _flatten_features(obs: dict) -> jnp.ndarray:
"""
@ -214,7 +190,6 @@ def create_obs_processor(
def _process_single(obs_dict: dict) -> jnp.ndarray:
processed = _add_derived_features(obs_dict)
processed = _prune_features(processed)
processed = _normalize_features(processed)
processed = _split_to_agents(processed, morph_mode)
flat = _flatten_features(processed) # (num_arms, total_feat)

27
src/brittle_star_project/evaluation/policy.py
View file

@ -45,7 +45,8 @@ class PolicyAgent:
key = f"Dense_{idx}"
if key not in dense_params:
break
layer_sizes.append(int(np.asarray(dense_params[key]["kernel"]).shape[1]))
layer_sizes.append(int(np.asarray(dense_params[key]["kernel"]).shape[-1]))
idx += 1
if not layer_sizes:
@ -53,8 +54,8 @@ class PolicyAgent:
self._sensor = GenericDenseLayersWithActivation(layer_sizes=layer_sizes)
self._actor = Actor(action_dim=action_dim)
self._sensor_apply = jax.jit(self._sensor.apply)
self._actor_apply = jax.jit(self._actor.apply)
self._sensor.apply = jax.jit(self._sensor.apply)
self._actor.apply = jax.jit(self._actor.apply)
self._params = {
"sensor_params": sensor_params,
"actor_params": actor_params,
@ -79,11 +80,25 @@ class PolicyAgent:
obs_processor=obs_processor,
)
def _apply_per_node(self, net, params, x):
# params: (nodes, ...)
# x: (batch, nodes, feat)
def apply_single_node(p, x_node):
# x_node: (batch, feat)
return jax.vmap(lambda xi: net.apply(p, xi))(x_node)
return jax.vmap(apply_single_node, in_axes=(0, 1), out_axes=1)(params, x)
def act(self, *, observations: dict[str, Any]) -> np.ndarray:
"""Return deterministic action (actor mean, no exploration noise)."""
batched_obs = jax.tree.map(lambda x: jnp.asarray(x)[None, ...], observations)
obs = self._obs_processor(batched_obs)[0]
hidden = self._sensor_apply(self._params["sensor_params"], obs)
mean, _log_std = self._actor_apply(self._params["actor_params"], hidden)
obs = self._obs_processor(batched_obs)
# TODO: message passing
hidden = self._apply_per_node(self._sensor, self._params["sensor_params"], obs)
# hidden = jax.vmap(...)
mean, _log_std = self._apply_per_node(self._actor, self._params["actor_params"], hidden)
return np.asarray(mean, dtype=np.float32).ravel()

4
src/brittle_star_project/evaluation/rollout.py
View file

@ -114,18 +114,20 @@ def rollout_viewer(
episode_return = 0.0
observations = _get_observations(state)
prev_dist = _get_xy_distance_to_target(observations) if observations else None
reached_target = _target_reached(state=state)
steps = 0
with mujoco.viewer.launch_passive(model, data) as viewer:
step_iter = range(int(max_steps)) if max_steps is not None else itertools.count()
for _step_idx in step_iter:
for _ in step_iter:
if not viewer.is_running():
break
step_start = time.time()
obs_dict = observations or {}
action = policy.act(observations=obs_dict)
if action_mask is not None:
action = action[action_mask]