Spaces:
Sleeping
Sleeping
File size: 5,048 Bytes
6fa23b0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
# NOT FINISHED
from .data_structures.Action_Balanced_Replay_Buffer import Action_Balanced_Replay_Buffer
from .data_structures.Replay_Buffer import Replay_Buffer
import numpy as np
import random
class Memory_Shaper(object):
"""Takes in the experience of full episodes and reshapes it according to macro-actions you define. Then it provides
a replay buffer with this reshaped data to learn from"""
def __init__(self, buffer_size, batch_size, seed, new_reward_fn, action_balanced_replay_buffer=True):
self.reset()
self.buffer_size = buffer_size
self.batch_size = batch_size
self.seed = seed
self.new_reward_fn = new_reward_fn
self.action_balanced_replay_buffer = action_balanced_replay_buffer
def put_adapted_experiences_in_a_replay_buffer(self, action_id_to_actions):
"""Adds experiences to the replay buffer after re-imagining that the actions taken were macro-actions according to
action_rules as well as primitive actions.
NOTE that we want to put both primitive actions and macro-actions into replay buffer so that it can learn that
its better to do a macro-action rather than the same primitive actions (which we will enforce with reward penalty)
"""
actions_to_action_id = {v: k for k, v in action_id_to_actions.items()}
self.num_actions = len(action_id_to_actions)
print(actions_to_action_id)
for key in actions_to_action_id.keys():
assert isinstance(key, tuple)
assert isinstance(actions_to_action_id[key], int)
episodes = len(self.states)
for data_type in [self.states, self.next_states, self.rewards, self.actions, self.dones]:
assert len(data_type) == episodes
max_action_length = self.calculate_max_action_length(actions_to_action_id)
if self.action_balanced_replay_buffer:
print("Using action balanced replay buffer")
replay_buffer = Action_Balanced_Replay_Buffer(self.buffer_size, self.batch_size, self.seed, num_actions=self.num_actions)
else:
print("Using ordinary replay buffer")
replay_buffer = Replay_Buffer(self.buffer_size, self.batch_size, self.seed)
for episode_ix in range(episodes):
self.add_adapted_experience_for_an_episode(episode_ix, actions_to_action_id, max_action_length, replay_buffer)
return replay_buffer
def calculate_max_action_length(self, actions_to_action_id):
"""Calculates the max length of the provided macro-actions"""
max_length = 0
for key in actions_to_action_id.keys():
action_length = len(key)
if action_length > max_length:
max_length = action_length
return max_length
def add_adapted_experience_for_an_episode(self, episode_ix, action_rules, max_action_length, replay_buffer):
"""Adds all the experiences we have been given to a replay buffer after adapting experiences that involved doing a
macro action"""
states = self.states[episode_ix]
next_states = self.next_states[episode_ix]
rewards = self.rewards[episode_ix]
actions = self.actions[episode_ix]
dones = self.dones[episode_ix]
assert len(states) == len(next_states) == len(rewards) == len(dones) == len(actions), "{} {} {} {} {} = {}".format(len(states), len(next_states), len(rewards), len(dones), len(actions), actions)
steps = len(states)
for step in range(steps):
replay_buffer.add_experience(states[step], actions[step], rewards[step], next_states[step], dones[step])
for action_length in range(2, max_action_length + 1):
if step < action_length - 1: continue
action_sequence = tuple(actions[step - action_length + 1 : step + 1])
assert all([action in range(self.num_actions) for action in action_sequence]), "All actions should be primitive here"
if action_sequence in action_rules.keys():
new_action = action_rules[action_sequence]
new_state = states[step - action_length + 1]
new_reward = np.sum(rewards[step - action_length + 1:step + 1])
new_reward = self.new_reward_fn(new_reward, len(action_sequence))
new_next_state = next_states[step]
new_dones = dones[step]
replay_buffer.add_experience(new_state, new_action, new_reward, new_next_state, new_dones)
def add_episode_experience(self, states, next_states, rewards, actions, dones):
"""Adds in an episode of experience"""
self.states.append(states)
self.next_states.append(next_states)
self.rewards.append(rewards)
self.actions.append(actions)
self.dones.append(dones)
def reset(self):
self.states = []
self.next_states = []
self.rewards = []
self.actions = []
self.dones = []
|