def recolor_occlusion_geoms()

in robogym/utils/sensor_utils.py [0:0]

• 14 lines of code
• 7 McCabe index (conditional complexity)

def recolor_occlusion_geoms(sim, robot_occlusion_data):
    """
    Color the occlusion geoms differently according to whether the simulator and the
    phasespace tracker matches.
    """
    colormap = [
        [0, 0, 0, 0.1],  # transparent grey for both off
        [1, 0, 0, 0.7],  # red for robot not but sim occluded
        [0, 0, 1, 0.7],  # blue for robot occluded but sim not
        [1, 1, 0, 1.0],  # solid yellow for both occluded
    ]
    sim_occlusion_data = check_occlusion(sim)
    geom_ids = [sim.model.geom_name2id(m) for m in OCCLUSION_MARKERS]

    for g_id, robot_occluded, sim_occluded in zip(
        geom_ids, robot_occlusion_data, sim_occlusion_data
    ):
        category = 2 * int(robot_occluded) + int(sim_occluded)
        sim.model.geom_rgba[g_id] = colormap[category]