def group_list(actions, agent_states, env_lens):
    """
    Unflat the list of items by lens
    :param items: list of items
    :param lens: list of integers
    :return: list of list of items grouped by lengths
    """
    grouped_actions = []
    grouped_agent_states = []
    cur_ofs = 0
    for g_len in env_lens:
        grouped_actions.append(actions[cur_ofs: cur_ofs + g_len])
        grouped_agent_states.append(agent_states[cur_ofs: cur_ofs + g_len])
        cur_ofs += g_len
    return grouped_actions, grouped_agent_states