diff --git a/configs/morphology/5_arms_damaged.yaml b/configs/morphology/5_arms_damaged.yaml new file mode 100644 index 0000000..2afd643 --- /dev/null +++ b/configs/morphology/5_arms_damaged.yaml @@ -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 diff --git a/scripts/train.py b/scripts/train.py index f82d6bc..8e61bd0 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -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}") diff --git a/simulate.sh b/simulate.sh new file mode 100755 index 0000000..86a96c1 --- /dev/null +++ b/simulate.sh @@ -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 \ No newline at end of file diff --git a/src/brittle_star_project/environment/obs_processing.py b/src/brittle_star_project/environment/obs_processing.py index 550624d..74c60cd 100644 --- a/src/brittle_star_project/environment/obs_processing.py +++ b/src/brittle_star_project/environment/obs_processing.py @@ -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 diff --git a/src/brittle_star_project/environment/obs_processing_tmp.py b/src/brittle_star_project/environment/obs_processing_tmp.py new file mode 100644 index 0000000..e2c6c82 --- /dev/null +++ b/src/brittle_star_project/environment/obs_processing_tmp.py @@ -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 diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index f82ba83..960a176 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -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, diff --git a/tests/test_obs_processor.py b/tests/test_obs_processor.py new file mode 100644 index 0000000..e69de29 diff --git a/tmp.txt b/tmp.txt new file mode 100644 index 0000000..899a63b --- /dev/null +++ b/tmp.txt @@ -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,) \ No newline at end of file diff --git a/vids/evaluation_metadata.yaml b/vids/evaluation_metadata.yaml new file mode 100644 index 0000000..0f59928 --- /dev/null +++ b/vids/evaluation_metadata.yaml @@ -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 diff --git a/vids/simulation.mp4 b/vids/simulation.mp4 new file mode 100644 index 0000000..9c6c807 Binary files /dev/null and b/vids/simulation.mp4 differ