mirror of
https://github.com/wassname/kair_algorithms_draft.git
synced 2026-09-09 11:25:10 +08:00
Add per (#8)
* Add per and modify etc * Replace pre-commit-config.yaml and add pre-commit hook in .git * Modify .gitignore * Modify .gitignore * Modify buffer and code * Modify replay buffer and per * Modify .gitignore
This commit is contained in:
Kyunghwan Kim
@@ -0,0 +1,151 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Prioritized Replay buffer for baselines.
|
||||
|
||||
- Author: Kh Kim
|
||||
- Contact: kh.kim@medipixel.io
|
||||
- Paper: https://arxiv.org/pdf/1511.05952.pdf
|
||||
https://arxiv.org/pdf/1707.08817.pdf
|
||||
"""
|
||||
|
||||
import random
|
||||
from typing import Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from algorithms.common.buffer.replay_buffer import ReplayBuffer
|
||||
from algorithms.common.buffer.segment_tree import MinSegmentTree, SumSegmentTree
|
||||
|
||||
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
class PrioritizedReplayBuffer(ReplayBuffer):
|
||||
"""Create Prioritized Replay buffer.
|
||||
|
||||
Taken from OpenAI baselines github repository:
|
||||
https://github.com/openai/baselines/blob/master/baselines/deepq/replay_buffer.py
|
||||
|
||||
Attributes:
|
||||
buffer_size (int): size of replay buffer for experience
|
||||
alpha (float): alpha parameter for prioritized replay buffer
|
||||
tree_idx (int): next index of tree
|
||||
sum_tree (SumSegmentTree): sum tree for prior
|
||||
min_tree (MinSegmentTree): min tree for min prior to get max weight
|
||||
init_priority (float): lower bound of priority
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, buffer_size: int, batch_size: int, demo: list = None, alpha: float = 0.6
|
||||
):
|
||||
"""Initialization.
|
||||
|
||||
Args:
|
||||
buffer_size (int): size of replay buffer for experience
|
||||
batch_size (int): size of a batched sampled from replay buffer for training
|
||||
demo (list): demonstration
|
||||
alpha (float): alpha parameter for prioritized replay buffer
|
||||
|
||||
"""
|
||||
super(PrioritizedReplayBuffer, self).__init__(buffer_size, batch_size, demo)
|
||||
assert alpha >= 0
|
||||
self.buffer_size = buffer_size
|
||||
self.alpha = alpha
|
||||
self.tree_idx = 0
|
||||
|
||||
# capacity must be positive and a power of 2.
|
||||
tree_capacity = 1
|
||||
while tree_capacity < self.buffer_size:
|
||||
tree_capacity *= 2
|
||||
|
||||
self.sum_tree = SumSegmentTree(tree_capacity)
|
||||
self.min_tree = MinSegmentTree(tree_capacity)
|
||||
self.init_priority = 1.0
|
||||
|
||||
# for init priority of demo
|
||||
if demo:
|
||||
for _ in range(len(demo)):
|
||||
self.sum_tree[self.tree_idx] = self.init_priority ** self.alpha
|
||||
self.min_tree[self.tree_idx] = self.init_priority ** self.alpha
|
||||
self.tree_idx += 1
|
||||
|
||||
def add(
|
||||
self,
|
||||
state: np.ndarray,
|
||||
action: np.ndarray,
|
||||
reward: np.float64,
|
||||
next_state: np.ndarray,
|
||||
done: bool,
|
||||
):
|
||||
"""Add experience and priority."""
|
||||
idx = self.tree_idx
|
||||
self.tree_idx = (self.tree_idx + 1) % self.buffer_size
|
||||
super().add(state, action, reward, next_state, done)
|
||||
|
||||
self.sum_tree[idx] = self.init_priority ** self.alpha
|
||||
self.min_tree[idx] = self.init_priority ** self.alpha
|
||||
|
||||
def extend(self, transitions: list):
|
||||
"""Add experiences to memory."""
|
||||
raise NotImplementedError
|
||||
|
||||
def _sample_proportional(self, batch_size: int) -> list:
|
||||
"""Sample indices based on proportional."""
|
||||
indices = []
|
||||
p_total = self.sum_tree.sum(0, len(self.buffer) - 1)
|
||||
segment = p_total / batch_size
|
||||
for i in range(batch_size):
|
||||
a = segment * i
|
||||
b = segment * (i + 1)
|
||||
upperbound = random.uniform(a, b)
|
||||
idx = self.sum_tree.retrieve(upperbound)
|
||||
indices.append(idx)
|
||||
return indices
|
||||
|
||||
def sample(self, beta: float = 0.4) -> Tuple[torch.Tensor, ...]:
|
||||
"""Sample a batch of experiences."""
|
||||
assert beta > 0
|
||||
|
||||
indices = self._sample_proportional(self.batch_size)
|
||||
states, actions, rewards, next_states, dones, weights = [], [], [], [], [], []
|
||||
|
||||
# get max weight
|
||||
p_min = self.min_tree.min() / self.sum_tree.sum()
|
||||
max_weight = (p_min * len(self.buffer)) ** (-beta)
|
||||
|
||||
for i in indices:
|
||||
s, a, r, n_s, d = self.buffer[i]
|
||||
states.append(np.array(s, copy=False))
|
||||
actions.append(np.array(a, copy=False))
|
||||
rewards.append(np.array(r, copy=False))
|
||||
next_states.append(np.array(n_s, copy=False))
|
||||
dones.append(np.array(float(d), copy=False))
|
||||
|
||||
# calculate weights
|
||||
p_sample = self.sum_tree[i] / self.sum_tree.sum()
|
||||
weight = (p_sample * len(self.buffer)) ** (-beta)
|
||||
weights.append(weight / max_weight)
|
||||
|
||||
states = torch.FloatTensor(np.array(states)).to(device)
|
||||
actions = torch.FloatTensor(np.array(actions)).to(device)
|
||||
rewards = torch.FloatTensor(np.array(rewards).reshape(-1, 1)).to(device)
|
||||
next_states = torch.FloatTensor(np.array(next_states)).to(device)
|
||||
dones = torch.FloatTensor(np.array(dones).reshape(-1, 1)).to(device)
|
||||
weights = torch.FloatTensor(np.array(weights).reshape(-1, 1)).to(device)
|
||||
|
||||
experiences = (states, actions, rewards, next_states, dones, weights, indices)
|
||||
|
||||
return experiences
|
||||
|
||||
def update_priorities(self, indices: list, priorities: np.ndarray):
|
||||
"""Update priorities of sampled transitions."""
|
||||
assert len(indices) == len(priorities)
|
||||
|
||||
for idx, priority in zip(indices, priorities):
|
||||
assert priority > 0
|
||||
assert 0 <= idx < len(self.buffer)
|
||||
|
||||
self.sum_tree[idx] = priority ** self.alpha
|
||||
self.min_tree[idx] = priority ** self.alpha
|
||||
|
||||
self.init_priority = max(self.init_priority, priority)
|
||||
@@ -1,8 +1,7 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Replay buffer for baselines."""
|
||||
|
||||
import random
|
||||
from collections import deque
|
||||
from typing import Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@@ -18,55 +17,70 @@ class ReplayBuffer:
|
||||
ddpg-pendulum/ddpg_agent.py
|
||||
|
||||
Attributes:
|
||||
buffer (deque): deque of replay buffer
|
||||
buffer (list): list of replay buffer
|
||||
batch_size (int): size of a batched sampled from replay buffer for training
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, buffer_size, batch_size, seed, demo=None):
|
||||
def __init__(self, buffer_size: int, batch_size: int, demo: list = None):
|
||||
"""Initialize a ReplayBuffer object.
|
||||
|
||||
Args:
|
||||
buffer_size (int): size of replay buffer for experience
|
||||
batch_size (int): size of a batched sampled from replay buffer for training
|
||||
seed (int): random seed
|
||||
demo (deque) : demonstration deque
|
||||
demo (list) : demonstration list
|
||||
|
||||
"""
|
||||
self.buffer = deque(maxlen=buffer_size) if not demo else demo
|
||||
|
||||
self.buffer = list() if not demo else demo
|
||||
self.buffer_size = buffer_size
|
||||
self.batch_size = batch_size
|
||||
random.seed(seed)
|
||||
self.idx = 0
|
||||
|
||||
def add(self, state, action, reward, next_state, done):
|
||||
def add(
|
||||
self,
|
||||
state: np.ndarray,
|
||||
action: np.ndarray,
|
||||
reward: np.float64,
|
||||
next_state: np.ndarray,
|
||||
done: bool,
|
||||
):
|
||||
"""Add a new experience to memory."""
|
||||
self.buffer.append((state, action, reward, next_state, done))
|
||||
data = (state, action, reward, next_state, done)
|
||||
|
||||
def extend(self, transitions):
|
||||
if len(self.buffer) == self.buffer_size:
|
||||
self.buffer[self.idx] = data
|
||||
self.idx = (self.idx + 1) % self.buffer_size
|
||||
else:
|
||||
self.buffer.append(data)
|
||||
|
||||
def extend(self, transitions: list):
|
||||
"""Add experiences to memory."""
|
||||
self.buffer.extend(transitions)
|
||||
|
||||
def sample(self):
|
||||
def sample(
|
||||
self
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Randomly sample a batch of experiences from memory."""
|
||||
experiences = random.sample(self.buffer, k=self.batch_size)
|
||||
idxs = np.random.randint(0, len(self.buffer), size=self.batch_size)
|
||||
|
||||
states, actions, rewards, next_states, dones = [], [], [], [], []
|
||||
|
||||
for e in experiences:
|
||||
states.append(np.expand_dims(e[0], axis=0))
|
||||
actions.append(e[1])
|
||||
rewards.append(e[2])
|
||||
next_states.append(np.expand_dims(e[3], axis=0))
|
||||
dones.append(e[4])
|
||||
for i in idxs:
|
||||
s, a, r, n_s, d = self.buffer[i]
|
||||
states.append(np.array(s, copy=False))
|
||||
actions.append(np.array(a, copy=False))
|
||||
rewards.append(np.array(r, copy=False))
|
||||
next_states.append(np.array(n_s, copy=False))
|
||||
dones.append(np.array(float(d), copy=False))
|
||||
|
||||
states = torch.from_numpy(np.vstack(states)).float().to(device)
|
||||
actions = torch.from_numpy(np.vstack(actions)).float().to(device)
|
||||
rewards = torch.from_numpy(np.vstack(rewards)).float().to(device)
|
||||
next_states = torch.from_numpy(np.vstack(next_states)).float().to(device)
|
||||
dones = torch.from_numpy(np.vstack(dones).astype(np.uint8)).float().to(device)
|
||||
states = torch.FloatTensor(np.array(states)).to(device)
|
||||
actions = torch.FloatTensor(np.array(actions)).to(device)
|
||||
rewards = torch.FloatTensor(np.array(rewards).reshape(-1, 1)).to(device)
|
||||
next_states = torch.FloatTensor(np.array(next_states)).to(device)
|
||||
dones = torch.FloatTensor(np.array(dones).reshape(-1, 1)).to(device)
|
||||
|
||||
return (states, actions, rewards, next_states, dones)
|
||||
return states, actions, rewards, next_states, dones
|
||||
|
||||
def __len__(self):
|
||||
def __len__(self) -> int:
|
||||
"""Return the current size of internal memory."""
|
||||
return len(self.buffer)
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Segment tree for Proirtized Replay Buffer."""
|
||||
|
||||
import operator
|
||||
from typing import Callable
|
||||
|
||||
|
||||
class SegmentTree:
|
||||
""" Create SegmentTree.
|
||||
|
||||
Taken from OpenAI baselines github repository:
|
||||
https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py
|
||||
|
||||
Attributes:
|
||||
capacity (int)
|
||||
tree (list)
|
||||
operation (function)
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, capacity: int, operation: Callable, init_value: float):
|
||||
"""Initialization.
|
||||
|
||||
Args:
|
||||
capacity (int)
|
||||
operation (function)
|
||||
init_value (float)
|
||||
|
||||
"""
|
||||
assert (
|
||||
capacity > 0 and capacity & (capacity - 1) == 0
|
||||
), "capacity must be positive and a power of 2."
|
||||
self.capacity = capacity
|
||||
self.tree = [init_value for _ in range(2 * capacity)]
|
||||
self.operation = operation
|
||||
|
||||
def _operate_helper(
|
||||
self, start: int, end: int, node: int, node_start: int, node_end: int
|
||||
) -> float:
|
||||
"""Returns result of operation in segment."""
|
||||
if start == node_start and end == node_end:
|
||||
return self.tree[node]
|
||||
mid = (node_start + node_end) // 2
|
||||
if end <= mid:
|
||||
return self._operate_helper(start, end, 2 * node, node_start, mid)
|
||||
else:
|
||||
if mid + 1 <= start:
|
||||
return self._operate_helper(start, end, 2 * node + 1, mid + 1, node_end)
|
||||
else:
|
||||
return self.operation(
|
||||
self._operate_helper(start, mid, 2 * node, node_start, mid),
|
||||
self._operate_helper(mid + 1, end, 2 * node + 1, mid + 1, node_end),
|
||||
)
|
||||
|
||||
def operate(self, start: int = 0, end: int = 0) -> float:
|
||||
"""Returns result of applying `self.operation`."""
|
||||
if end <= 0:
|
||||
end += self.capacity
|
||||
end -= 1
|
||||
|
||||
return self._operate_helper(start, end, 1, 0, self.capacity - 1)
|
||||
|
||||
def __setitem__(self, idx: int, val: float):
|
||||
"""Set value in tree."""
|
||||
idx += self.capacity
|
||||
self.tree[idx] = val
|
||||
|
||||
idx //= 2
|
||||
while idx >= 1:
|
||||
self.tree[idx] = self.operation(self.tree[2 * idx], self.tree[2 * idx + 1])
|
||||
idx //= 2
|
||||
|
||||
def __getitem__(self, idx: int) -> float:
|
||||
"""Get real value in leaf node of tree."""
|
||||
assert 0 <= idx < self.capacity
|
||||
|
||||
return self.tree[self.capacity + idx]
|
||||
|
||||
|
||||
class SumSegmentTree(SegmentTree):
|
||||
""" Create SumSegmentTree.
|
||||
|
||||
Taken from OpenAI baselines github repository:
|
||||
https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, capacity: int):
|
||||
"""Initialization.
|
||||
|
||||
Args:
|
||||
capacity (int)
|
||||
|
||||
"""
|
||||
super(SumSegmentTree, self).__init__(
|
||||
capacity=capacity, operation=operator.add, init_value=0.0
|
||||
)
|
||||
|
||||
def sum(self, start: int = 0, end: int = 0) -> float:
|
||||
"""Returns arr[start] + ... + arr[end]."""
|
||||
return super(SumSegmentTree, self).operate(start, end)
|
||||
|
||||
def retrieve(self, upperbound: float) -> int:
|
||||
"""Find the highest index `i` about upper bound in the tree"""
|
||||
assert 0 <= upperbound <= self.sum() + 1e-5
|
||||
|
||||
idx = 1
|
||||
|
||||
while idx < self.capacity: # while non-leaf
|
||||
left = 2 * idx
|
||||
right = left + 1
|
||||
if self.tree[left] > upperbound:
|
||||
idx = 2 * idx
|
||||
else:
|
||||
upperbound -= self.tree[left]
|
||||
idx = right
|
||||
return idx - self.capacity
|
||||
|
||||
|
||||
class MinSegmentTree(SegmentTree):
|
||||
""" Create SegmentTree.
|
||||
|
||||
Taken from OpenAI baselines github repository:
|
||||
https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, capacity: int):
|
||||
"""Initialization.
|
||||
|
||||
Args:
|
||||
capacity (int)
|
||||
|
||||
"""
|
||||
super(MinSegmentTree, self).__init__(
|
||||
capacity=capacity, operation=min, init_value=float("inf")
|
||||
)
|
||||
|
||||
def min(self, start: int = 0, end: int = 0) -> float:
|
||||
"""Returns min(arr[start], ..., arr[end])."""
|
||||
return super(MinSegmentTree, self).operate(start, end)
|
||||
Reference in New Issue
Block a user