def build_world_step()

in mae_envs/modules/agents.py [0:0]

• 30 lines of code
• 10 McCabe index (conditional complexity)

    def build_world_step(self, env, floor, floor_size):
        env.metadata['n_agents'] = self.n_agents
        successful_placement = True

        for i in range(self.n_agents):
            env.metadata.pop(f"agent{i}_initpos", None)

        for i in range(self.n_agents):
            obj = ObjFromXML("particle_hinge", name=f"agent{i}")
            if self.friction is not None:
                obj.add_transform(set_geom_attr_transform('friction', self.friction))
            if self.color is not None:
                _color = (self.color[i]
                          if isinstance(self.color[0], (list, tuple, np.ndarray))
                          else self.color)
                obj.add_transform(set_geom_attr_transform('rgba', _color))
            if not self.damp_z:
                obj.add_transform(set_joint_damping_transform(1, 'tz'))

            if self.placement_fn is not None:
                _placement_fn = (self.placement_fn[i]
                                 if isinstance(self.placement_fn, list)
                                 else self.placement_fn)
                obj_size = get_size_from_xml(obj)
                pos, pos_grid = rejection_placement(env, _placement_fn, floor_size, obj_size)
                if pos is not None:
                    floor.append(obj, placement_xy=pos)
                    # store spawn position in metadata. This allows sampling subsequent agents
                    # close to previous agents
                    env.metadata[f"agent{i}_initpos"] = pos_grid
                else:
                    successful_placement = False
            else:
                floor.append(obj)
        return successful_placement