<jupyter_start><jupyter_text>Generate training data<jupyter_code>num_rollouts = 10
max_timesteps = 2000

# Use the first 10000 (obs, action) pairs as training data
env = gym.make('Hopper-v1')
X_all, y_all = [], []

for i in range(num_rollouts):
    print(i, end=',')
    obs = env.reset()
    action = policy_fn(obs[None, :])
    X_all.append(obs)
    y_all.append(action)   
    for j in range(max_timesteps):
        obs, r, done, _ = env.step(action)
        action = policy_fn(obs[None, :])
        X_all.append(obs)
        y_all.append(action)
        
X_all = np.array(X_all)
# concat instead of array is necessary to keep array in the right shape
y_all = np.concatenate(y_all)
# split index
idx = int(X_all.shape[0] * 4/5)

X_tv = X_all[:idx, :]
X_test = X_all[idx:, :]
y_tv = y_all[:idx, :]
y_test = y_all[idx:, :]
print(X_tv.shape, X_test.shape, y_tv.shape, y_test.shape)
with open('train_test.pkl', 'wb') as opf:
    pickle.dump([X_tv, y_tv, X_test, y_test], opf)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14840637450199204, 'fraction_numerical': 0.014940239043824702, 'mean_word_length': 3.8086124401913874, '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': '652423', 'license': 'MIT', 'n_tokens_mistral': 428, 'n_tokens_neox': 413, 'n_words': 88, 'path': 'hw1/2-warmup/behavior-cloning-generate-train-test-data.ipynb', 'repo_name': 'zyxue/berkeley-cs294-hw', 'size': '1004'}