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

feat: color map for archs

Browse source
This commit is contained in:
Jona Reynaert 2026-05-16 01:50:35 +02:00
parent 3eb6e0fb45
commit 3489448a3d
2 changed files with 20 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

13
scripts/poster_visualisations/render_poster_videos.py
View file

@ -14,6 +14,12 @@ _ARCH_DIR_MAP = {
"SEGMENT": "segment", "SEGMENT": "segment",
} }
ROBOT_COLOR_MAP = {
"CENTRALIZED": "#0D567C", # Blue
"FULLY_CONNECTED": "#8C0E0F", # Reddish
"RING": "#FCB304", # Pale Yellow
}
def _arch_dir(name: str) -> str: def _arch_dir(name: str) -> str:
return _ARCH_DIR_MAP.get(name, name.lower()) return _ARCH_DIR_MAP.get(name, name.lower())
@ -54,7 +60,7 @@ def main() -> None:
parser.add_argument("--fps", type=int, default=60) parser.add_argument("--fps", type=int, default=60)
parser.add_argument( parser.add_argument(
"--robot-color", "--robot-color",
default="#2B4162", # ring = #888888, centralized = #2B4162, fully connected = FA9F42 default="#2B4162",
help="Hex color for the brittle star robot", help="Hex color for the brittle star robot",
) )
args = parser.parse_args() args = parser.parse_args()
@ -109,6 +115,9 @@ def main() -> None:
args.follow_camera: out_dir / "follow.mp4", args.follow_camera: out_dir / "follow.mp4",
} }
print(bundle.architecture)
color = ROBOT_COLOR_MAP.get(bundle.architecture, args.robot_color)
result = record_episode_multi_camera( result = record_episode_multi_camera(
env=bundle.env, env=bundle.env,
policy=bundle.policy, policy=bundle.policy,
@ -122,7 +131,7 @@ def main() -> None:
camera_fovy=camera_fovy, camera_fovy=camera_fovy,
camera_xyz=camera_xyz, camera_xyz=camera_xyz,
target_xy=target_xy, target_xy=target_xy,
robot_color=args.robot_color, robot_color=color,
width=args.width, width=args.width,
height=args.height, height=args.height,
fps=args.fps, fps=args.fps,

12
scripts/poster_visualisations/render_static_path_image.py
View file

@ -18,6 +18,12 @@ from brittle_star_project.evaluation.video import (
hex_to_rgba, hex_to_rgba,
) )
ROBOT_COLOR_MAP = {
"CENTRALIZED": "#0D567C", # Blue
"FULLY_CONNECTED": "#8C0E0F", # Reddish
"RING": "#FCB304", # Pale Yellow
}
def _enum_value(enum_obj, *names: str) -> int: def _enum_value(enum_obj, *names: str) -> int:
for name in names: for name in names:
@ -119,9 +125,10 @@ def main() -> None:
model, mujoco.mjtObj.mjOBJ_BODY, "BrittleStarMorphology/central_disk" model, mujoco.mjtObj.mjOBJ_BODY, "BrittleStarMorphology/central_disk"
) )
robot_rgba = hex_to_rgba(ROBOT_COLOR_MAP.get(bundle.architecture, args.robot_color), 1.0)
# Optionally override robot color by recoloring geoms belonging to the robot's body subtree. # Optionally override robot color by recoloring geoms belonging to the robot's body subtree.
if args.robot_color is not None: if args.robot_color is not None:
robot_rgba = hex_to_rgba(args.robot_color, 1.0)
# Collect body IDs in the subtree rooted at `body_id` by walking parent links. # Collect body IDs in the subtree rooted at `body_id` by walking parent links.
nbody = int(model.nbody) nbody = int(model.nbody)
body_parent = model.body_parentid body_parent = model.body_parentid
@ -141,7 +148,6 @@ def main() -> None:
# Optionally override robot color by recoloring geoms belonging to the robot's body subtree. # Optionally override robot color by recoloring geoms belonging to the robot's body subtree.
if args.robot_color is not None: if args.robot_color is not None:
robot_rgba = hex_to_rgba(args.robot_color, 1.0)
# Collect body IDs in the subtree rooted at `body_id` by walking parent links. # Collect body IDs in the subtree rooted at `body_id` by walking parent links.
nbody = int(model.nbody) nbody = int(model.nbody)
body_parent = model.body_parentid body_parent = model.body_parentid
@ -187,7 +193,7 @@ def main() -> None:
path_step = max(1, int(args.frame_stride)) * 3 path_step = max(1, int(args.frame_stride)) * 3
path_points_visible = path_points[::path_step] path_points_visible = path_points[::path_step]
path_rgba = hex_to_rgba(args.path_color, 0.92) path_rgba = hex_to_rgba(ROBOT_COLOR_MAP.get(bundle.architecture, args.path_color), 0.92)
ctx = mujoco.GLContext(args.width, args.height) ctx = mujoco.GLContext(args.width, args.height)
ctx.make_current() ctx.make_current()