Notes

num_episodes = 10
max_question_length,max_correct_response_length,max_incorrect_response_length
# Training loop
for episode in range(num_episodes):
episode_states = []
episode_actions = []
episode_rewards = []

state = env.reset()

while True:

# # Extract the question from the state (assuming state contains the question)
question_sequence = state["question"]
text = "".join([tokenizer.token_to_character(token) for token in question_sequence])

# response_text = generate_response(question_text) # Call your generate_response function here
question_sequence = torch.tensor(state["question"], dtype=torch.float32)
question_sequence = question_sequence.unsqueeze(0).unsqueeze(0)
correct_response_sequence = torch.tensor(state["correct_response"], dtype=torch.float32)
correct_response_sequence = correct_response_sequence.unsqueeze(0).unsqueeze(0)

incorrect_response_sequence = torch.tensor(state["incorrect_response"], dtype=torch.float32)
incorrect_response_sequence = incorrect_response_sequence.unsqueeze(0).unsqueeze(0)

question_length = len(question_sequence)
correct_response_length = len(correct_response_sequence)
incorrect_response_length = len(incorrect_response_sequence)
# input_size = max_question_length + max_correct_response_length + max_incorrect_response_length

question = torch.nn.functional.pad(question_sequence, (0, input_size - question_sequence.size(-1)))
correct_response = torch.nn.functional.pad(correct_response_sequence, (0, input_size - correct_response_sequence.size(-1)))
incorrect_response = torch.nn.functional.pad(incorrect_response_sequence, (0, input_size - incorrect_response_sequence.size(-1)))

action_probs = policy.forward(question, correct_response, incorrect_response)
print(text)
action = torch.multinomial(action_probs.view(-1) ,1).item()
# Perform actions, get rewards, and transition to the next state
action = policy.select_action({'question' : question ,
'correct_response' : correct_response,
'incorrect_response' : incorrect_response})
# Assuming you have a reward function that calculates reward given the action and state
reward = calculate_reward(action, state)
next_state, _, done = env.step(action)
episode_states.append(state)
episode_actions.append(action)
episode_rewards.append(reward)
state = next_state
if done:
break

# Compute discounted returns
discounted_returns = []
running_add = 0
for r in reversed(episode_rewards):
running_add = r + 0.99 * running_add # Discount factor: 0.99
discounted_returns.insert(0, running_add)

# Calculate loss and perform policy gradient update
# Create separate tensors for 'question', 'correct_response', and 'incorrect_response'
# Convert lists to tensors
question_tensors = [torch.tensor(state['question'], dtype=torch.float32) for state in episode_states]
correct_response_tensors = [torch.tensor(state['correct_response'], dtype=torch.float32) for state in episode_states]
incorrect_response_tensors = [torch.tensor(state['incorrect_response'], dtype=torch.float32) for state in episode_states]
# Concatenate the tensors
question_tensors = torch.cat(question_tensors)
correct_response_tensors = torch.cat(correct_response_tensors)
incorrect_response_tensors = torch.cat(incorrect_response_tensors)
# Reshape the tensors
question_tensors = question_tensors.view(1, len(episode_states), input_size)
correct_response_tensors = correct_response_tensors.view(1, len(episode_states), input_size)
incorrect_response_tensors = incorrect_response_tensors.view(1, len(episode_states), input_size)

# Forward pass through the policy network
action_probs = policy.forward(question_tensors, correct_response_tensors, incorrect_response_tensors)

# Forward pass through the policy network
selected_action_probs = torch.gather(action_probs, 1, torch.tensor(episode_actions).unsqueeze(1))
loss = -torch.sum(torch.log(selected_action_probs) * torch.FloatTensor(discounted_returns))
optimizer.zero_grad()
loss.backward()
optimizer.step()

# Print episode information
if episode % 10 == 0:
print(f"Episode {episode}, Total Reward: {np.sum(episode_rewards)}")

# Use the trained policy for inference
with torch.no_grad():
test_state = env.reset()
while True:
# Extract the question from the state (assuming state contains the question)
question_text = test_state["question"]
question_text = "".join([tokenizer.token_to_character(token) for token in question_text])

# Generate a response based on the input question_text


generate_response(question_text) # Call your generate_response function here

test_state = torch.tensor([test_state], dtype=torch.float32)
action_probs = policy(test_state)
action = np.argmax(action_probs.detach().numpy()[0])
next_state, _, done = env.step(action)
test_state = next_state
if done:
break