I am trying to train a DQN to do optimal energy scheduling. Each state comes as a vector of 4 variables (represented by floats) saved in the replay memory as a state tensor, each action is an integer saved in the memory as a tensor too. I extract the batch of experiences as:
def extract_tensors(experiences):
# Convert batch of Experiences to Experience of batches
batch = Experience(*zip(*experiences))
t1 = torch.cat(batch.state)
t2 = torch.cat(batch.action)
t3 = torch.stack(batch.reward)
t4 = torch.cat(batch.next_state)
return (t1,t2,t3,t4)
I then unpacked them for purposes of updating the prediction and the target networks as:
experiences = memory.sample(batch_size)
states, actions, rewards, next_states = extract_tensors(experiences)
My Qvalues class for update looks like this:
class QValues():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@staticmethod
def get_current(policy_net, states, actions):
return policy_net(states).gather(dim=1, index=actions)
@staticmethod
def get_next(target_net, next_states):
final_state_locations = next_states.flatten(start_dim=1) \
.max(dim=1)[0].eq(0).type(torch.bool)
non_final_state_locations = (final_state_locations == False)
non_final_states = next_states[non_final_state_locations]
batch_size = next_states.shape[0]
values = torch.zeros(batch_size).to(QValues.device)
values[non_final_state_locations] = target_net(non_final_states).max(dim=1)[0].detach()
return values
When I try running the training loop, I get the error below:
<ipython-input-8-4a79494b54ca> in <module>
214
215
--> 216 current_q_values = QValues.get_current(policy_net, states, actions)
217 next_q_values = QValues.get_next(target_net, next_states)
218 target_q_values = (next_q_values * gamma) + rewards
<ipython-input-8-4a79494b54ca> in get_current(policy_net, states, actions)
160 @staticmethod
161 def get_current(policy_net, states, actions):
--> 162 return policy_net(states).gather(dim=1, index=actions)
163
164 @staticmethod
RuntimeError: invalid argument 4: Index tensor must have same dimensions as input tensor at c:
