def walls_to_mujoco()

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

• 43 lines of code
• 5 McCabe index (conditional complexity)

def walls_to_mujoco(floor, floor_size, grid_size, walls, friction=None):
    '''
        Take a list of walls in grid frame and add them to the floor in the worldgen frame.
        Args:
            floor (worldgen.Floor): floor
            floor_size (float): size of floor
            grid_size (int): size of placement grid
            walls (Wall list): list of walls
            friction (float): wall friction
    '''
    wall_width = floor_size / grid_size / 2
    grid_cell_length = floor_size / grid_size
    for i, wall in enumerate(walls):
        if wall.is_vertical:
            wall_length_grid = (wall.pt2[1] - wall.pt1[1] + 1)
            offset = np.array([-1, 1])
        else:
            wall_length_grid = (wall.pt2[0] - wall.pt1[0] + 1)
            offset = np.array([1, -1])

        # Convert to mujoco frame
        wall_length = wall_length_grid * grid_cell_length
        # Subtract 1 grid_cell_length such that walls originate and end in the center of a grid cell
        # Subtract 1 wall_width such that perpendicular walls do not intersect at the center of a
        # grid cell
        wall_length -= grid_cell_length + wall_width

        if wall.is_vertical:
            size = (wall_width, wall_length, wall.height)
        else:
            size = (wall_length, wall_width, wall.height)

        # Position of object should be in the middle of a grid cell (add 0.5) shifted by
        #     the wall width such that corners don't overlap
        pos = np.array([wall.pt1[0] + 0.5, wall.pt1[1] + 0.5]) / grid_size
        pos += offset * wall_width / floor_size / 2

        # Convert from mujoco to worldgen scale
        scale_x = (floor_size - size[0]) / floor_size
        scale_y = (floor_size - size[1]) / floor_size
        pos = pos / np.array([scale_x, scale_y])
        geom = Geom('box', size, name=f"wall{i}")
        geom.mark_static()
        geom.add_transform(set_geom_attr_transform('rgba', wall.rgba))
        geom.add_transform(set_geom_attr_transform('group', 1))
        if friction is not None:
            geom.add_transform(set_geom_attr_transform('friction', friction))
        floor.append(geom, placement_xy=pos)