<jupyter_start><jupyter_text>for i in range(arg_num_trajs):    print("trajectory ", i)    all_ob, all_agent_ob, all_ac, all_rew, ep_ret = [], [], [], [], [0 for k in range(n_agents)]    for k in range(n_agents):        all_ob.append([])        all_ac.append([])        all_rew.append([])    obs = env.reset()    obs = [ob[None, :] for ob in obs]    action = [np.zeros([1]) for _ in range(n_agents)]    step = 0    done = False    episode_r_n = [0.] * n_agents    while (not done) or (step<=2):        if arg_algo != "random":            action, _, _ = model.step(obs, action)            actions_list = [onehot(action[k][0], n_actions[k]) for k in range(n_agents)]        else:            actions_list = [onehot(np.random.randint(n_actions[k]), n_actions[k]) for k in range(n_agents)]        for k in range(n_agents):            all_ob[k].append(obs[k])            all_ac[k].append(actions_list[k])        all_agent_ob.append(np.concatenate(obs, axis=1))        obs, rew, done, _ = env.step(actions_list)        print(obs)        episode_r_n = [episode_r_n[i]+rew[i] for i in range(n_agents)]        for k in range(n_agents):            all_rew[k].append(rew[k])            ep_ret[k] += rew[k]            all_pos[k].append(np.array(env.world.agents[k].state.p_pos))        obs = [ob[None, :] for ob in obs]        step += 1        if arg_image:            img = env.render(mode='rgb_array')            images.append(img[0])            time.sleep(0.02)        if step == max_steps or True in done:            done = True             step = 0        else:            done = False     print("\n--- episode-{} | [sum-reward]: {}".format(i, np.sum(episode_r_n)))    episode_r_n_all.append(episode_r_n)    episode_r_n_sum.append(np.sum(episode_r_n))    for k in range(n_agents):        all_ob[k] = np.squeeze(all_ob[k])    all_agent_ob = np.squeeze(all_agent_ob)    traj_data = {        "ob": all_ob, "ac": all_ac, "rew": all_rew,        "ep_ret": ep_ret, "all_ob": all_agent_ob    }    sample_trajs.append(traj_data)     print('episode', i, 'expected_return', ep_ret)    for k in range(n_agents):        avg_ret[k].append(ep_ret[k])<jupyter_code>data_path = './results/'+arg_algo+'/'+env_name+'/l-0.1-b-1000/seed-'+str(1)+ '/checkpoint' + '%05d'%arg_epoch + "-%d"%arg_num_trajs + "tra-1.pkl"
with open(data_path, 'rb') as f:
    data = pkl.load(f)
data[0]["ob"][3][0][2:4]
env.world.entities
def update_env_uppos_tag(state_data, timestep):
    pray_perspective = state_data[3][timestep]
    
    pray_pos = pray_perspective[2:4]
    landmark_1_pos = pray_perspective[4:6] + pray_pos
    landmark_2_pos = pray_perspective[6:8] + pray_pos
    predator_1_pos = pray_perspective[8:10] + pray_pos
    predator_2_pos = pray_perspective[10:12] + pray_pos
    predator_3_pos = pray_perspective[12:14] + pray_pos
    
    return [predator_1_pos, predator_2_pos, predator_3_pos, pray_pos, landmark_1_pos,landmark_2_pos]
    
    
obs = env.reset()
images = []
for i in range(50):
    if(arg_env == "simple_tag"):
        update_list = update_env_uppos_tag(data[0]["ob"],i)
    for j, pos in enumerate(update_list):
        env.world.entities[j].state.p_pos = pos
        
        
    img = env.render(mode='rgb_array')
    images.append(img[0])
    time.sleep(0.05)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15365853658536585, 'fraction_numerical': 0.01920731707317073, 'mean_word_length': 2.9963459196102313, 'pattern_counts': {'":': 6, '<': 6, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '648190', 'license': 'Apache-2.0', 'n_tokens_mistral': 1275, 'n_tokens_neox': 1244, 'n_words': 224, 'path': 'multi-agent-irl/Rendering.ipynb', 'repo_name': 'renos/CoDAIL', 'size': '3280'}