8c17debc9dc6552da8c317c0c852b55f03c80c61,slm_lab/spec/spec_util.py,,resolve_aeb,#Any#,141

Before Change

            range(agent_num), range(env_num), range(body_num)))
    elif ae_product == "inner":
        ae_coor_itr = zip(range(agent_num), range(env_num))
        aeb_coor_list = list(itertools.product(
            ae_coor_itr, range(body_num)))
        aeb_coor_list = [(a, e, b) for ((a, e), b) in aeb_coor_list]
    else:  // custom AEB, body_num is a coor_list
        aeb_coor_list = [tuple(aeb) for aeb in sorted(body_num)]
    assert is_aeb_compact(
        aeb_coor_list), "Failed check: for a, e, uniq count == len (shape), and for each a,e hash, b uniq count == b len (shape)"

After Change

    body_num = _.get(spec, "body.num")
    body_num_list = body_num if _.is_list(body_num) else [body_num] * env_num

    aeb_list = []
    if ae_product == "outer":
        for e in range(env_num):
            sub_aeb_list = list(itertools.product(
                range(agent_num), [e], range(body_num_list[e])))
            aeb_list.extend(sub_aeb_list)
    elif ae_product == "inner":
        for a, e in zip(range(agent_num), range(env_num)):
            sub_aeb_list = list(
                itertools.product([a], [e], range(body_num_list[e])))
            aeb_list.extend(sub_aeb_list)
    else:  // custom AEB, body_num is a aeb_list
        aeb_list = [tuple(aeb) for aeb in body_num]
    aeb_list.sort()
    assert is_aeb_compact(
        aeb_list), "Failed check: for a, e, uniq count == len (shape), and for each a,e hash, b uniq count == b len (shape)"
    return aeb_list

Instances