def collect_traj(env, agent, T, visualize=False):
    """
        agent must have methods: sample_action and intended_action
        Run trajectory on sampled actions
        record states, sampled actions, intended actions and reward
    """
    states = []
    intended_actions = []
    taken_actions = []

    s = env.reset()

    reward = 0.0

    for t in range(T):

        a_intended = agent.intended_action(s)
        a = agent.sample_action(s)
        next_s, r, done, _ = env.step(a[None, :])
        reward += r

        states.append(s.ravel())
        intended_actions.append(a_intended)
        taken_actions.append(a)

        s = next_s

        if visualize:
            env.render()


        if done == True:
            break
            
    return states, intended_actions, taken_actions, reward