def iterate_batches(env, net, batch_size, file):
    batch = []
    episode_reward = 0.0
    episode_steps = []
    obs = env.reset()
    sm = nn.Softmax(dim=1)
    flag_injected_bug_spotted = [False, False]
    while True:
        obs_v = torch.FloatTensor([obs])
        act_probs_v = sm(net(obs_v))
        act_probs = act_probs_v.data.numpy()[0]
        action = np.random.choice(len(act_probs), p=act_probs)
        (next_obs, reward, is_done, _) = env.step(action)
        if (((- 0.5) < next_obs[0] < (- 0.45)) and (not flag_injected_bug_spotted[0])):
            file.write('BUG1 ')
            flag_injected_bug_spotted[0] = True
        if ((0.45 < next_obs[0] < 0.5) and (not flag_injected_bug_spotted[1])):
            file.write('BUG2 ')
            flag_injected_bug_spotted[1] = True
        episode_reward += reward
        episode_steps.append(EpisodeStep(observation=obs, action=action))
        if is_done:
            file.write('\n')
            batch.append(Episode(reward=episode_reward, steps=episode_steps))
            episode_reward = 0.0
            episode_steps = []
            next_obs = env.reset()
            flag_injected_bug_spotted = [False, False]
            if (len(batch) == batch_size):
                (yield batch)
                batch = []
        obs = next_obs
def iterate_batches(env, net, batch_size, file):
    batch = []
    episode_reward = 0.0
    episode_steps = []
    obs = env.reset()
    sm = nn.Softmax(dim=1)
    flag_injected_bug_spotted = [False, False]
    while True:
        obs_v = torch.FloatTensor([obs])
        act_probs_v = sm(net(obs_v))
        act_probs = act_probs_v.data.numpy()[0]
        action = np.random.choice(len(act_probs), p=act_probs)
        (next_obs, reward, is_done, _) = env.step(action)
        if (((- 0.5) < next_obs[0] < (- 0.45)) and (not flag_injected_bug_spotted[0])):
            file.write('BUG1 ')
            flag_injected_bug_spotted[0] = True
        if ((0.45 < next_obs[0] < 0.5) and (not flag_injected_bug_spotted[1])):
            file.write('BUG2 ')
            flag_injected_bug_spotted[1] = True
        episode_reward += reward
        episode_steps.append(EpisodeStep(observation=obs, action=action))
        if is_done:
            file.write('\n')
            batch.append(Episode(reward=episode_reward, steps=episode_steps))
            episode_reward = 0.0
            episode_steps = []
            next_obs = env.reset()
            flag_injected_bug_spotted = [False, False]
            if (len(batch) == batch_size):
                (yield batch)
                batch = []
        obs = next_obs