Old version looks faster, verify
@@ -39,19 +39,13 @@ EPSILON_FINAL = 0.02
def unpack_batch(batch):
- states, actions, rewards, dones, last_states = [], [], [], [], []
- for exp in batch:
- state = np.array(exp.state, copy=False)
- states.append(state)
- actions.append(exp.action)
- rewards.append(exp.reward)
- dones.append(exp.last_state is None)
- if exp.last_state is None:
- last_states.append(state) # the result will be masked anyway
- else:
- last_states.append(np.array(exp.last_state, copy=False))
+ states = [np.array(exp[0].state, copy=False) for exp in batch]
+ next_states = [np.array(exp[-1].state, copy=False) for exp in batch]
+ actions = [exp[0].action for exp in batch]
+ rewards = [exp[0].reward for exp in batch]
+ dones = [exp[0].done for exp in batch]
return np.array(states, copy=False), np.array(actions), np.array(rewards, dtype=np.float32), \
- np.array(dones, dtype=np.uint8), np.array(last_states, copy=False)
+ np.array(dones, dtype=np.uint8), np.array(next_states, copy=False)
def calc_loss(batch, net, tgt_net, cuda=False):
@@ -95,7 +89,8 @@ if __name__ == "__main__":
epsilon_greedy_selector = ptan.actions.EpsilonGreedyActionSelector(epsilon=1.0)
agent = ptan.agent.DQNAgent(net, epsilon_greedy_selector, cuda=args.cuda)
- exp_source = ptan.experience.ExperienceSourceFirstLast(env, agent, gamma=GAMMA, steps_count=1)
+# exp_source = ptan.experience.ExperienceSourceFirstLast(env, agent, gamma=GAMMA, steps_count=1)
+ exp_source = ptan.experience.ExperienceSource(env, agent, steps_count=2)
buffer = ptan.experience.ExperienceReplayBuffer(exp_source, buffer_size=REPLAY_SIZE)
optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE)
