diff --git a/scripts/poster_visualisations/render_poster_videos.py b/scripts/poster_visualisations/render_poster_videos.py index 8a73638..eca60fc 100644 --- a/scripts/poster_visualisations/render_poster_videos.py +++ b/scripts/poster_visualisations/render_poster_videos.py @@ -14,6 +14,12 @@ _ARCH_DIR_MAP = { "SEGMENT": "segment", } +ROBOT_COLOR_MAP = { + "CENTRALIZED": "#0D567C", # Blue + "FULLY_CONNECTED": "#8C0E0F", # Reddish + "RING": "#FCB304", # Pale Yellow +} + def _arch_dir(name: str) -> str: 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( "--robot-color", - default="#2B4162", # ring = #888888, centralized = #2B4162, fully connected = FA9F42 + default="#2B4162", help="Hex color for the brittle star robot", ) args = parser.parse_args() @@ -109,6 +115,9 @@ def main() -> None: 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( env=bundle.env, policy=bundle.policy, @@ -122,7 +131,7 @@ def main() -> None: camera_fovy=camera_fovy, camera_xyz=camera_xyz, target_xy=target_xy, - robot_color=args.robot_color, + robot_color=color, width=args.width, height=args.height, fps=args.fps, diff --git a/scripts/poster_visualisations/render_static_path_image.py b/scripts/poster_visualisations/render_static_path_image.py index 2057133..17fb4b5 100644 --- a/scripts/poster_visualisations/render_static_path_image.py +++ b/scripts/poster_visualisations/render_static_path_image.py @@ -18,6 +18,12 @@ from brittle_star_project.evaluation.video import ( 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: for name in names: @@ -119,9 +125,10 @@ def main() -> None: 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. 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. nbody = int(model.nbody) 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. 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. nbody = int(model.nbody) body_parent = model.body_parentid @@ -187,7 +193,7 @@ def main() -> None: path_step = max(1, int(args.frame_stride)) * 3 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.make_current()