Buggy color visualization

Summary

When using timor.utilities.visualization.color_visualization, the base_placement for a robot cannot be set. Robots with a set base placement will be moved mistakenly.

Steps to reproduce

a = any_assembly
viz = a.to_pin_robot(base_placement=Transformation.random()).visualize()
color_visualization(viz, a, cmap)

will result in the robot with wrong base placement.

What is the current bug behavior?

Base placement is ignored. Furthermore, visual geometry data of the modules is overridden.

What is the expected correct behavior?

Ideally, if a robot is pre-visualized in a visualizer, the colors should change in meshcat instead of having to change the mesh colors and re-visualizing the robot.

Relevant logs and/or screenshot

None

Possible fixes

Here's a quick fix for now, I will look into it in more details in the next days:

def color_visualization(viz: MeshcatVisualizer,
                        assembly: ModuleAssembly,
                        color_map: Dict[Union[str, int, module_classification.ModuleType], np.ndarray] = None):
    """
    Adds colors to a visualized robot in viz, depending on module types.

    :param viz: The visualizer to color
    :param assembly: The assembly (that can already be visualized in viz) to color
    :param color_map: A dictionary mapping module integer enumerated module types OR module IDs to colors. If color_map
      maps module IDs to colors, this method first tries to interpret them as custom module IDs matching the internal
      representation in the assembly. If this fails, it will interpret them as original module IDs.
    """
    update_in_viz = True
    if tuple(viz.model.names) != tuple(assembly.robot.model.names):  # Minimum sanity check for pre-initialization
        logging.warning("The assembly should be pre-visualized in the visualizer. Mesh colors will be overridden.")
        update_in_viz = False
    if color_map is None:
        color_map = DEFAULT_COLOR_MAP
    if all(isinstance(k, str) for k in color_map.keys()):
        cmap = color_map
        by_type = False
    else:
        cmap = {int(k): v for k, v in color_map.items()}  # Make sure that both, enumerated types and integers, work
        by_type = True
    viz_cmap = dict()
    for id_custom, id_original, m in zip(assembly.internal_module_ids,
                                         assembly.original_module_ids,
                                         assembly.module_instances):
        if by_type:
            key = int(module_classification.get_module_type(m))
        elif id_custom in cmap:
            key = id_custom
        else:
            key = id_original
        if key in cmap:
            color = cmap[key]
        else:
            color = None

        for body in m.bodies:
            viz_name = '.'.join(body.id)
            viz_cmap[viz_name] = color

    for go in assembly.robot.visual.geometryObjects:
        stem = '.'.join(go.name.split('.')[:-1])
        if stem in viz_cmap:
            if update_in_viz:
                viz.viewer[viz.viewerVisualGroupName + '/' + go.name].set_property(u'color', viz_cmap[stem].tolist())
            else:
                go.meshColor = viz_cmap[stem]
        else:
            logging.warning(f"Could not find color for robot geometry {go.name}.")

/cc @00000000014ADE32