def build_world_step()

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

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

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

        self.curr_n_boxes = env._random_state.randint(self.n_boxes[0], self.n_boxes[1] + 1)

        env.metadata['curr_n_boxes'] = np.zeros((self.n_boxes[1]))
        env.metadata['curr_n_boxes'][:self.curr_n_boxes] = 1
        env.metadata['curr_n_boxes'] = env.metadata['curr_n_boxes'].astype(np.bool)

        self.curr_n_elongated_boxes = env._random_state.randint(
            self.n_elongated_boxes[0], min(self.n_elongated_boxes[1], self.curr_n_boxes) + 1)

        self.box_size_array = self.box_size * np.ones((self.curr_n_boxes, 3))
        if self.curr_n_elongated_boxes > 0:
            # sample number of x-aligned boxes
            n_xaligned = env._random_state.randint(self.curr_n_elongated_boxes + 1)
            self.box_size_array[:n_xaligned, :] = self.box_size * np.array([3.3, 0.3, 1.0])
            self.box_size_array[n_xaligned:self.curr_n_elongated_boxes, :] = (self.box_size * np.array([0.3, 3.3, 1.0]))
        env.metadata['box_size_array'] = self.box_size_array

        successful_placement = True
        for i in range(self.curr_n_boxes):
            char = chr(ord('A') + i % 26)
            geom = Geom("box", self.box_size_array[i, :], name=f'moveable_box{i}')
            geom.set_material(Material(texture="chars/" + char + ".png"))
            geom.add_transform(set_geom_attr_transform('mass', self.box_mass))
            if self.mark_box_corners:
                for j, (x, y) in enumerate([[0, 0], [0, 1], [1, 0], [1, 1]]):
                    geom.mark(f'moveable_box{i}_corner{j}', relative_xyz=(x, y, 0.5),
                              rgba=[1., 1., 1., 0.])
            if self.friction is not None:
                geom.add_transform(set_geom_attr_transform('friction', self.friction))
            if self.box_only_z_rot:
                geom.add_transform(remove_hinge_axis_transform(np.array([1.0, 0.0, 0.0])))
                geom.add_transform(remove_hinge_axis_transform(np.array([0.0, 1.0, 0.0])))

            if self.placement_fn is not None:
                _placement_fn = (self.placement_fn[i]
                                 if isinstance(self.placement_fn, list)
                                 else self.placement_fn)
                pos, _ = rejection_placement(env, _placement_fn, floor_size,
                                             self.box_size_array[i, :2])
                if pos is not None:
                    floor.append(geom, placement_xy=pos)
                else:
                    successful_placement = False
            else:
                floor.append(geom)
        return successful_placement