71 lines
2.4 KiB
Python
71 lines
2.4 KiB
Python
import torch
|
|
import pickle
|
|
|
|
class ExperienceCollector:
|
|
def __init__(self):
|
|
self.states = []
|
|
self.actions = []
|
|
self.rewards = []
|
|
self.advantages = []
|
|
self._current_episode_states = []
|
|
self._current_episode_actions = []
|
|
self._current_episode_estimated_values = []
|
|
|
|
def begin_episode(self):
|
|
self._current_episode_states = []
|
|
self._current_episode_actions = []
|
|
self._current_episode_estimated_values = []
|
|
|
|
def record_decision(self, state, action, estimated_value=0):
|
|
self._current_episode_states.append(state)
|
|
self._current_episode_actions.append(action)
|
|
self._current_episode_estimated_values.append(estimated_value)
|
|
|
|
def complete_episode(self, reward):
|
|
num_states = len(self._current_episode_states)
|
|
self.states += self._current_episode_states
|
|
self.actions += self._current_episode_actions
|
|
self.rewards += [reward for _ in range(num_states)]
|
|
|
|
for i in range(num_states):
|
|
advantage = reward - self._current_episode_estimated_values[i]
|
|
self.advantages.append(advantage)
|
|
|
|
self._current_episode_states = []
|
|
self._current_episode_actions = []
|
|
self._current_episode_estimated_values = []
|
|
|
|
|
|
class ExperienceBuffer:
|
|
def __init__(self, states, actions, rewards, advantages):
|
|
self.states = states
|
|
self.actions = actions
|
|
self.rewards = rewards
|
|
self.advantages = advantages
|
|
|
|
def save(self, filename):
|
|
with open(filename, 'wb') as f:
|
|
pickle.dump((self.states, self.actions, self.rewards, self.advantages), f)
|
|
|
|
@classmethod
|
|
def load(cls, filename):
|
|
with open(filename, 'rb') as f:
|
|
states, actions, rewards, advantages = pickle.load(f)
|
|
return cls(
|
|
states=states,
|
|
actions=actions,
|
|
rewards=rewards,
|
|
advantages=advantages)
|
|
|
|
def combine_experience(collectors):
|
|
combined_states = torch.cat([torch.tensor(c.states) for c in collectors])
|
|
combined_actions = torch.cat([torch.tensor(c.actions) for c in collectors])
|
|
combined_rewards = torch.cat([torch.tensor(c.rewards) for c in collectors])
|
|
combined_advantages = torch.cat([torch.tensor(c.advantages) for c in collectors])
|
|
|
|
return ExperienceBuffer(
|
|
combined_states,
|
|
combined_actions,
|
|
combined_rewards,
|
|
combined_advantages)
|