<jupyter_start><jupyter_text>REINFORCE with (adaptive) Value Function Baseline and Aggregated Gradients<jupyter_code>LEARNING_RATE = .0000005
LEARNING_RATE_V = .00005

# Run the algorithm multiple times to see performance across random seeds
performance = np.zeros((NUM_ALGO_TRIALS,NUM_ITERATIONS))

for algo_trial in range(NUM_ALGO_TRIALS):
    # Initialise policy as a feedforward NN
    pi = Policy(STATE_DIM, ACTION_DIM)
    v = ValueFunction(STATE_DIM)

    # 1 iteration == deploying a fixed policy in the environment for ROLLOUT_LENGTH timesteps and then 
    # using the collected data from that rollout to perform ROLLOUT_LENGTH policy gradient updates to 
    # find a new policy for the next iteration
    for iteration in range(NUM_ITERATIONS):
        # Always reset the state and history for each new policy
        state = env.reset()
        history = []

        # Collect a rollout using the current policy
        for timestep in range(ROLLOUT_LENGTH): 
            action = pi.act(np.asarray(state).flatten())
            next_state, reward, done, info = env.step(action)
            history.append([state, action, reward, next_state, done])
            state = next_state
            if done:
                state = env.reset()
                
        # Sweep backwards through the collected data and use sample_return to accumulate the rewards 
        # to form the Monte Carlo estimates of the returns
        sample_return = 0
        
        # Reset the gradient aggregators
        param_grads = []
        torch.log(pi(state,'mean') + pi(state,'std')).backward() # need to backprop something before zero_grad()
        pi.zero_grad()
        with torch.no_grad():
            for index, param in enumerate(pi.parameters()):
                param_grads.append(param.grad)
        
        # Sweep back through collected trajectory and aggregate gradients
        for entry in reversed(history):
            state, action, reward, next_state, done = entry
            sample_return += reward

            # Update the Value function baseline
            v_loss = (v(state) - sample_return)**2
            v.zero_grad()
            v_loss.backward()
            
            with torch.no_grad():            
                for index, param in enumerate(v.parameters()):
                    param -= LEARNING_RATE_V * param.grad                                                    

            # Clear the gradients in PyTorch computational graph
            pi.zero_grad()

            # Calculate the gradient of the log-density under the current policy at state, action
            p1 = (1/ torch.sqrt((torch.det(torch.diag(pi(state,'std'))))))
            M = torch.inverse(torch.diag(pi(state,'std')))
            p2 = (torch.from_numpy(action).float() - pi(state,'mean'))
            p3 = p2.T @ M @ p2
            l = torch.log(p1) + (-0.5*p3)
            l.backward()
            
            with torch.no_grad():
                for index, param in enumerate(pi.parameters()):
                    # Aggregate gradients
                    param_grads[index] += LEARNING_RATE * param.grad * (reward + v(next_state) - v(state))
            
            # Clear sample return if done
            if done:
                sample_return = 0

        # After aggregating all the gradients through the log, take a policy gradient step
        with torch.no_grad():
            for index, param in enumerate(pi.parameters()):
                param += param_grads[index]                

        # Reset the state so we can perform an independent evaluation rollout of the updated policy
        state = env.reset()
        evaluation_return = 0
        evaluation_returns = []

        for test_index in range(NUM_TESTS):
            for timestep in range(ROLLOUT_LENGTH):
                action = pi.act(np.asarray(state).flatten())            
                next_state, reward, done, info = env.step(action)
                state = next_state
                evaluation_return += reward
                if done:
                    evaluation_returns.append(evaluation_return)
                    evaluation_return = 0
                    state = env.reset()
                    break
        performance[algo_trial][iteration] = np.mean(evaluation_returns)
        print('REINFORCE Iteration: ', iteration,'\t', 'Evaluation: ', "{:.2f}".format(performance[algo_trial][iteration]) , '\tNum Evals: ', len(evaluation_returns))

pi(state,'mean')
LEARNING_RATE = .00001
LEARNING_RATE_V = .0005

# Run the algorithm multiple times to see performance across random seeds
performance = np.zeros((NUM_ALGO_TRIALS,NUM_ITERATIONS))

for algo_trial in range(NUM_ALGO_TRIALS):
    # Initialise policy as a feedforward NN
    pi = Policy(STATE_DIM, ACTION_DIM)
    v = ValueFunction(STATE_DIM)
    now = datetime.datetime.now()
    writer = SummaryWriter("runs/scalar-{}".format(now.strftime("%Y%m%d-%H%M%S")))

    # 1 iteration == deploying a fixed policy in the environment for ROLLOUT_LENGTH timesteps and then 
    # using the collected data from that rollout to perform ROLLOUT_LENGTH policy gradient updates to 
    # find a new policy for the next iteration
    for iteration in range(NUM_ITERATIONS):
        # Always reset the state and history for each new policy
        state = env.reset()
        history = []

        for timestep in range(ROLLOUT_LENGTH): 
            action = pi.act(np.asarray(state).flatten())
            next_state, reward, done, info = env.step(action)
            reward = reward + 0.5*abs(action.item())
            history.append([state, action, reward, next_state, done])
            state = next_state
            if done:
                state = env.reset()

        # Sweep backwards through the collected data and use sample_return to accumulate the rewards 
        # to form the Monte Carlo estimates of the returns
        sample_return = 0
        likelihoods = []
        for entry in reversed(history):
            state, action, reward, next_state, done = entry
            sample_return += reward

            # Update the Value function baseline
            v_loss = (v(state) - sample_return)**2
            v.zero_grad()
            v_loss.backward()
            
            with torch.no_grad():            
                for index, param in enumerate(v.parameters()):
                    param -= LEARNING_RATE_V * param.grad                                                    

            # Initialise the gradients stored in PyTorch auto-grad computation graph to 0
            pi.zero_grad()

            # Calculate the gradient of the log-density under the current policy at state, action
            p1 = (1/ torch.sqrt((torch.det(torch.diag(pi(state,'std'))))))
            M = torch.inverse(torch.diag(pi(state,'std')))
            p2 = (torch.from_numpy(action).float() - pi(state,'mean'))
            p3 = p2.T @ M @ p2
            l = torch.log(p1) + (-0.5*p3)
            if l.isnan().any():
                raise Exception('About to Blow!')
            l.backward()
            
            with torch.no_grad():
                likelihoods.append(l)
                grads = []
                for index, param in enumerate(pi.parameters()):
                    # Perform a gradient ascent step on the policy parameters
                    param += LEARNING_RATE * param.grad * (sample_return - v(state))
                    grads.extend(param.grad.flatten().tolist())
                writer.add_histogram('policy parameters', torch.nn.utils.parameters_to_vector(pi.parameters()), iteration) 
                writer.add_histogram('policy gradient', np.asarray(grads), iteration)
                writer.add_histogram('log-likelihoods', np.asarray(likelihoods), iteration)                
                writer.flush()
                    
            if done:
                sample_return = 0

                
                
    param_grads = []
    torch.log(pi(state)[action]).backward()
    with torch.no_grad():
        for index, param in enumerate(pi.parameters()):
            param_grads.append(param.grad * 0)

#     for entry in reversed(history):
#         state, action, reward, next_state, done = entry
#         sample_return += reward

#         pi.zero_grad()
#         torch.log(pi(state)[action]).backward()
        
#         with torch.no_grad():
#             for index, param in enumerate(pi.parameters()):
                param_grads[index] += LEARNING_RATE * param.grad * sample_return

#         if done:
#             sample_return = 0
    
    with torch.no_grad():
        for index, param in enumerate(pi.parameters()):
            param += param_grads[index]                
                
                
                
                
                
                
                
                
                
                
                
                
                
        # Reset the state so we can perform an independent evaluation rollout of the updated policy
        state = env.reset()
        evaluation_return = 0
        evaluation_returns = []

        for test_index in range(NUM_TESTS):
            for timestep in range(ROLLOUT_LENGTH):
                action = pi.act(np.asarray(state).flatten())            
                next_state, reward, done, info = env.step(action)
                state = next_state
                evaluation_return += reward
                if done:
                    evaluation_returns.append(evaluation_return)
                    evaluation_return = 0
                    state = env.reset()
                    break
        performance[algo_trial][iteration] = np.mean(evaluation_returns)
        print('REINFORCE Iteration: ', iteration,'\t', 'Evaluation: ', "{:.2f}".format(performance[algo_trial][iteration]) , '\tNum Evals: ', len(evaluation_returns))
    writer.close()        
l
vals = []
with torch.no_grad():
    for i in range(1000):
        state = env.reset()
        vals.append(v(state))
pi(state,'mean')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09103557940767243, 'fraction_numerical': 0.006360564500099384, 'mean_word_length': 1.7960544595721033, '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': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '280559', 'license': 'MIT', 'n_tokens_mistral': 2748, 'n_tokens_neox': 2610, 'n_words': 726, 'path': 'examples/Continuous Cartpole.ipynb', 'repo_name': 'cvigoe/rlkit', 'size': '10062'}