diff --git a/configs/example.yaml b/configs/example.yaml index 14bf3d7..baeb1b0 100644 --- a/configs/example.yaml +++ b/configs/example.yaml @@ -1,5 +1,4 @@ morphology: - num_arms: 2 - num_segments_per_arm: 4 + segments_per_arm: [4, 4] use_p_control: true use_torque_control: false diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index 78083e9..206e902 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -7,11 +7,24 @@ from .env_types import Task @dataclass(frozen=True, slots=True) class MorphologyConfig: - num_arms: int = 5 - num_segments_per_arm: int = 4 + """Brittle star morphology configuration. + + segments_per_arm defines the number of segments for each arm. The length of + this list implicitly sets the number of arms. Use 0 segments to represent + a fully amputated arm (e.g., [4, 0, 4, 2, 4] for a 5-arm morphology with + arm 1 removed and arm 3 shortened). + + The upstream biorobot library natively supports per-arm segment counts. + """ + + segments_per_arm: tuple[int, ...] = (4, 4, 4, 4, 4) use_p_control: bool = True use_torque_control: bool = False + @property + def num_arms(self) -> int: + return len(self.segments_per_arm) + @dataclass(frozen=True, slots=True) class ArenaConfig: diff --git a/src/brittle_star_project/environment/factory.py b/src/brittle_star_project/environment/factory.py index 1a891ea..523feb2 100644 --- a/src/brittle_star_project/environment/factory.py +++ b/src/brittle_star_project/environment/factory.py @@ -22,7 +22,7 @@ class BrittleStarEnvFactory: spec = default_brittle_star_morphology_specification( num_arms=config.num_arms, - num_segments_per_arm=config.num_segments_per_arm, + num_segments_per_arm=list(config.segments_per_arm), use_p_control=config.use_p_control, use_torque_control=config.use_torque_control, )