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

feat(obs_processor): rewrote split method, todo: testing + code cleanup

Browse source
This commit is contained in:
Robin Meersman 2026-05-08 23:46:45 +02:00
parent 07ee32fd3b
commit 34f0b1b80e
10 changed files with 87 additions and 34 deletions
Download patch file Download diff file Expand all files Collapse all files

6
configs/morphology/5_arms_damaged.yaml Normal file
View file

@ -0,0 +1,6 @@
# 5 Arms Full Morphology Configuration
# Baseline 5-arm brittle star.
segments_per_arm: [4, 4, 0, 4, 4]
use_p_control: true
use_torque_control: false

2
scripts/train.py
View file

@ -42,7 +42,7 @@ def main(dict_cfg: DictConfig):
base_dir=os.path.dirname(run_dir),
)
logger = get_logger()
logger.set_level(logging.INFO)
logger.set_level(logging.DEBUG)
logger.info(f"Hydra-initialized run: {run_name}")
logger.info(f"Output directory: {run_dir}")

9
simulate.sh Executable file
View file

@ -0,0 +1,9 @@
#!/usr/bin/env bash
path=$1
uv run scripts/simulate.py \
simulation.model_path="$path"/final_model.flax \
simulation.record_video=True \
simulation.video_output_path=./vids/simulation.mp4 \
simulation.max_steps=10000

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

@ -102,61 +102,47 @@ def create_obs_processor(
return normalized
def _split_to_agents(obs: dict, morph_mode) -> dict:
# TODO: cleanup + MORE testing (works for centralized 5 arms + damaged arms)
output = {}
num_agents = needed_copies # IMPORTANT: number of MLPs
for key, arr in obs.items():
if key not in ordered_keys or arr.size == 0:
if arr.size == 0:
continue
logger.debug(f"[INPUT] {key}: {arr.shape}")
if arr.ndim == 0:
arr = arr.reshape(1)
# -------- CENTRALIZED --------
if morph_mode == MorphMode.CENTRALIZED:
output[key] = arr.reshape(1, -1)
# TODO: padding for centralized
continue
# -------- SEGMENTS --------
if key in _SEGMENT_SCALED_KEYS:
per_agent = []
for i, agent_id in enumerate(agent_indices):
for i, _ 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 (segments per arm?)
taken = jnp.take(arr, idx, axis=0)
pad_len = 4 - taken.shape[0]
padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1))
padded = jnp.pad(taken, [(9, pad_len)] + [(0, 0)] * (taken.ndim - 1))
per_agent.append(padded.reshape(-1))
out = jnp.stack(per_agent)
# -------- JOINTS --------
arr = jnp.stack(per_agent)
elif key in _JOINT_SCALED_KEYS:
per_agent = []
for i, _ in enumerate(agent_indices):
idx = joint_indices[i]
taken = jnp.take(arr, idx, axis=0) # (joint_n, ...)
# pad to 8
taken = jnp.take(arr, idx, axis=0)
pad_len = 8 - taken.shape[0]
padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1))
per_agent.append(padded.reshape(-1))
out = jnp.stack(per_agent)
# -------- GLOBAL --------
arr = jnp.stack(per_agent)
else:
out = jnp.repeat(arr[None, :], num_agents, axis=0)
arr = jnp.repeat(arr[None, :], num_agents, axis=0)
logger.debug(f"[OUTPUT] {key}: {out.shape}")
output[key] = out
if morph_mode == MorphMode.CENTRALIZED:
output[key] = arr.reshape(1, -1)
elif key in _JOINT_SCALED_KEYS:
output[key] = arr.reshape(num_agents, -1)
elif key in _SEGMENT_SCALED_KEYS:
output[key] = arr[:, None]
else:
output[key] = arr
return output

17
src/brittle_star_project/environment/obs_processing_tmp.py Normal file
View file

@ -0,0 +1,17 @@
obs = {}
segs = set()
joints = set()
for key, arr in obs.items():
if key in segs:
# padding using mask 1x
pass
elif key in joints:
# padding using mask 2x
pass
else:
# no padding
pass
# reshape according to segs, joints, or global
pass

1
src/brittle_star_project/trainers/PPOTrainer.py
View file

@ -277,7 +277,6 @@ def apply_shared(net, params, x):
return jax.vmap(lambda xi: net.apply(params, xi))(x_flattened)
# TODO: update to work with extra dimension + message passing
def _rollout_jit(
agent_state,
episode_stats,

0
tests/test_obs_processor.py Normal file
View file

27
tmp.txt Normal file
View file

@ -0,0 +1,27 @@
-- 5 arms, 1 disabled
2026-05-08 14:21:57,035 - DEBUG - [INPUT] joint_actuator_force: (32,)
2026-05-08 14:21:57,035 - DEBUG - [OUTPUT] joint_actuator_force: (32,)
2026-05-08 14:21:57,035 - DEBUG - [INPUT] joint_position: (32,)
2026-05-08 14:21:57,035 - DEBUG - [OUTPUT] joint_position: (32,)
2026-05-08 14:21:57,036 - DEBUG - [INPUT] joint_velocity: (32,)
2026-05-08 14:21:57,036 - DEBUG - [OUTPUT] joint_velocity: (32,)
2026-05-08 14:21:57,036 - DEBUG - [INPUT] segment_contact: (16,)
2026-05-08 14:21:57,036 - DEBUG - [OUTPUT] segment_contact: (16,)
2026-05-08 14:21:57,036 - DEBUG - [INPUT] disk_z_tilt: ()
2026-05-08 14:21:57,036 - DEBUG - [OUTPUT] disk_z_tilt: (1,)
2026-05-08 14:21:57,036 - DEBUG - [INPUT] robot_direction_to_target: (2,)
2026-05-08 14:21:57,036 - DEBUG - [OUTPUT] robot_direction_to_target: (2,)
-- 5 arms, none disabled
2026-05-08 14:27:04,472 - DEBUG - [INPUT] joint_actuator_force: (40,)
2026-05-08 14:27:04,472 - DEBUG - [OUTPUT] joint_actuator_force: (40,)
2026-05-08 14:27:04,472 - DEBUG - [INPUT] joint_position: (40,)
2026-05-08 14:27:04,472 - DEBUG - [OUTPUT] joint_position: (40,)
2026-05-08 14:27:04,473 - DEBUG - [INPUT] joint_velocity: (40,)
2026-05-08 14:27:04,473 - DEBUG - [OUTPUT] joint_velocity: (40,)
2026-05-08 14:27:04,473 - DEBUG - [INPUT] segment_contact: (20,)
2026-05-08 14:27:04,473 - DEBUG - [OUTPUT] segment_contact: (20,)
2026-05-08 14:27:04,473 - DEBUG - [INPUT] disk_z_tilt: ()
2026-05-08 14:27:04,473 - DEBUG - [OUTPUT] disk_z_tilt: (1,)
2026-05-08 14:27:04,474 - DEBUG - [INPUT] robot_direction_to_target: (2,)
2026-05-08 14:27:04,474 - DEBUG - [OUTPUT] robot_direction_to_target: (2,)

9
vids/evaluation_metadata.yaml Normal file
View file

@ -0,0 +1,9 @@
timestamp: '2026-05-07T20:43:29.612567'
morphology_override: null
seed: 1
max_steps: 10000
result:
return: 0.40030128204863225
length: 6705
reached_target: true
final_xy_dist: 0.19969871795136748

BIN
vids/simulation.mp4 Normal file
View file

Binary file not shown.