<jupyter_start><jupyter_text>Convert one replay at a time<jupyter_code>def convert_replay_to_s_a_r_d_s(replay, reward_type, obs_type, normalize_reward):
    kaggle_env = kaggle_environments.make(
        'mab',
        configuration=replay['configuration'],
        steps=replay['steps'],
        info=replay['info']
    )
    sim_env = ve.KaggleMABEnvTorchVectorized(
        reward_type=reward_type,
        obs_type=obs_type,
        normalize_reward=normalize_reward,
        env_device=torch.device('cpu'),
        out_device=torch.device('cpu')
    )

    actions = []
    thresholds = []
    for step_info in kaggle_env.steps[1:]:
        actions.append(torch.tensor([s['action'] for s in step_info]))
        thresholds.append(torch.tensor(step_info[0]['observation']['thresholds']).float())

    s_batch = []
    a_batch = []
    r_batch = []
    d_batch = []
    next_s_batch = []
    sim_env.reset()
    sim_env.orig_thresholds = torch.tensor(kaggle_env.steps[0][0]['observation']['thresholds'], dtype=sim_env.orig_thresholds.dtype).view(sim_env.orig_thresholds.shape)
    s = sim_env.obs
    for i, a in enumerate(actions):
        a = a.unsqueeze(0)
        next_s, r, done, _ = sim_env.step(a)
        s_batch.append(s.clone())
        a_batch.append(a.clone())
        r_batch.append(r.clone())
        d_batch.append(torch.ones(r.shape) if done else torch.zeros(r.shape))
        next_s_batch.append(next_s.clone())
        s = next_s
        assert torch.allclose(sim_env.thresholds.view(-1), thresholds[i].view(-1)), f'ERROR: {i}'
    
    return (torch.cat(s_batch).view(-1, *s.shape[-2:]),
            torch.cat(a_batch).view(-1),
            torch.cat(r_batch).view(-1),
            torch.cat(d_batch).view(-1),
            torch.cat(next_s_batch).view(-1, *s.shape[-2:]))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15304443225452552, 'fraction_numerical': 0.007679648930334613, 'mean_word_length': 3.0714285714285716, 'pattern_counts': {'":': 0, '<': 5, '<?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': '612779', 'license': 'MIT', 'n_tokens_mistral': 670, 'n_tokens_neox': 665, 'n_words': 81, 'path': 'Convert_replays_to_database.ipynb', 'repo_name': 'IsaiahPressman/Kaggle_Santa_2020', 'size': '1823'}