def _hide_geoms()

in robogym/envs/rearrange/simulation/base.py [0:0]

• 34 lines of code
• 13 McCabe index (conditional complexity)

    def _hide_geoms(self, hide_targets=False, hide_objects=False, hide_robot=False):
        sim = self.mj_sim

        geom_ids_to_hide = []

        # Hide sites.
        site_rgba = sim.model.site_rgba.copy()
        sim.model.site_rgba[:] = np.zeros_like(site_rgba)

        if hide_targets:
            # Hide targets
            target_ids = [
                target_id
                for i in range(self.num_objects)
                for target_id in geom_ids_of_body(sim, f"target:object{i}")
            ]

            assert len(target_ids) > 0

            geom_ids_to_hide += target_ids

        if hide_objects:
            # Hide targets
            object_ids = [
                target_id
                for i in range(self.num_objects)
                for target_id in geom_ids_of_body(sim, f"object{i}")
            ]

            assert len(object_ids) > 0

            geom_ids_to_hide += object_ids

        if hide_robot:
            robot_geom_ids = [
                sim.model.geom_name2id(name)
                for name in sim.model.geom_names
                if name.startswith("robot0:")
            ]
            geom_ids_to_hide += robot_geom_ids

        geom_rgba = sim.model.geom_rgba.copy()
        sim.model.geom_rgba[geom_ids_to_hide, -1] = 0.0

        yield

        # If sim becomes stale (e.g. because of reset while being yield), we don't have to reset
        # it to the original state
        if not sim.is_stale():
            # Restore sites and targets
            sim.model.geom_rgba[:] = geom_rgba
            sim.model.site_rgba[:] = site_rgba