Merge pull request #35 from accuracy-maker/master

add sumtree,PER and PESR

tests/rl/test_prioritizedreplybuffer.py

@@ -0,0 +1,61 @@
+import pytest
+import random
+import torch
+from zeta.rl.PrioritizedReplayBuffer import PrioritizedReplayBuffer, SumTree  # Replace 'your_module' with the actual module where classes are defined
+
+@pytest.fixture
+def replay_buffer():
+    state_size = 4
+    action_size = 2
+    buffer_size = 100
+    device = torch.device("cpu")
+    return PrioritizedReplayBuffer(state_size, action_size, buffer_size, device)
+
+def test_initialization(replay_buffer):
+    assert replay_buffer.eps == 1e-2
+    assert replay_buffer.alpha == 0.1
+    assert replay_buffer.beta == 0.1
+    assert replay_buffer.max_priority == 1.0
+    assert replay_buffer.count == 0
+    assert replay_buffer.real_size == 0
+    assert replay_buffer.size == 100
+    assert replay_buffer.device == torch.device("cpu")
+
+def test_add(replay_buffer):
+    transition = (torch.rand(4), torch.rand(2), 1.0, torch.rand(4), False)
+    replay_buffer.add(transition)
+    assert replay_buffer.count == 1
+    assert replay_buffer.real_size == 1
+
+def test_sample(replay_buffer):
+    for i in range(10):
+        transition = (torch.rand(4), torch.rand(2), 1.0, torch.rand(4), False)
+        replay_buffer.add(transition)
+
+    batch, weights, tree_idxs = replay_buffer.sample(5)
+    assert len(batch) == 5
+    assert len(weights) == 5
+    assert len(tree_idxs) == 5
+
+def test_update_priorities(replay_buffer):
+    for i in range(10):
+        transition = (torch.rand(4), torch.rand(2), 1.0, torch.rand(4), False)
+        replay_buffer.add(transition)
+
+    batch, weights, tree_idxs = replay_buffer.sample(5)
+    new_priorities = torch.rand(5)
+    replay_buffer.update_priorities(tree_idxs, new_priorities)
+
+def test_sample_with_invalid_batch_size(replay_buffer):
+    with pytest.raises(AssertionError):
+        replay_buffer.sample(101)
+
+def test_add_with_max_size(replay_buffer):
+    for i in range(100):
+        transition = (torch.rand(4), torch.rand(2), 1.0, torch.rand(4), False)
+        replay_buffer.add(transition)
+
+    assert replay_buffer.count == 0
+    assert replay_buffer.real_size == 100
+
+# Additional tests for edge cases, exceptions, and more scenarios can be added as needed.

tests/rl/test_prioritizedsequencereplybuffer.py

@@ -0,0 +1,64 @@
+import pytest
+import random
+import torch
+from zeta.rl.PrioritizedSequenceReplayBuffer import PrioritizedSequenceReplayBuffer, SumTree  # Replace 'your_module' with the actual module where classes are defined
+
+@pytest.fixture
+def replay_buffer():
+    state_size = 4
+    action_size = 2
+    buffer_size = 100
+    device = torch.device("cpu")
+    return PrioritizedSequenceReplayBuffer(state_size, action_size, buffer_size, device)
+
+def test_initialization(replay_buffer):
+    assert replay_buffer.eps == 1e-5
+    assert replay_buffer.alpha == 0.1
+    assert replay_buffer.beta == 0.1
+    assert replay_buffer.max_priority == 1.0
+    assert replay_buffer.decay_window == 5
+    assert replay_buffer.decay_coff == 0.4
+    assert replay_buffer.pre_priority == 0.7
+    assert replay_buffer.count == 0
+    assert replay_buffer.real_size == 0
+    assert replay_buffer.size == 100
+    assert replay_buffer.device == torch.device("cpu")
+
+def test_add(replay_buffer):
+    transition = (torch.rand(4), torch.rand(2), 1.0, torch.rand(4), False)
+    replay_buffer.add(transition)
+    assert replay_buffer.count == 1
+    assert replay_buffer.real_size == 1
+
+def test_sample(replay_buffer):
+    for i in range(10):
+        transition = (torch.rand(4), torch.rand(2), 1.0, torch.rand(4), False)
+        replay_buffer.add(transition)
+
+    batch, weights, tree_idxs = replay_buffer.sample(5)
+    assert len(batch) == 5
+    assert len(weights) == 5
+    assert len(tree_idxs) == 5
+
+def test_update_priorities(replay_buffer):
+    for i in range(10):
+        transition = (torch.rand(4), torch.rand(2), 1.0, torch.rand(4), False)
+        replay_buffer.add(transition)
+
+    batch, weights, tree_idxs = replay_buffer.sample(5)
+    new_priorities = torch.rand(5)
+    replay_buffer.update_priorities(tree_idxs, new_priorities)
+
+def test_sample_with_invalid_batch_size(replay_buffer):
+    with pytest.raises(AssertionError):
+        replay_buffer.sample(101)
+
+def test_add_with_max_size(replay_buffer):
+    for i in range(100):
+        transition = (torch.rand(4), torch.rand(2), 1.0, torch.rand(4), False)
+        replay_buffer.add(transition)
+
+    assert replay_buffer.count == 0
+    assert replay_buffer.real_size == 100
+
+# Additional tests for edge cases, exceptions, and more scenarios can be added as needed.

tests/rl/test_sumtree.py

@@ -0,0 +1,56 @@
+import pytest
+from zeta.rl.sumtree import SumTree  # Replace 'your_module' with the actual module where SumTree is defined
+
+# Fixture for initializing SumTree instances with a given size
+@pytest.fixture
+def sum_tree():
+    size = 10  # You can change the size as needed
+    return SumTree(size)
+
+# Basic tests
+def test_initialization(sum_tree):
+    assert sum_tree.size == 10
+    assert sum_tree.count == 0
+    assert sum_tree.real_size == 0
+    assert sum_tree.total == 0
+
+def test_update_and_get(sum_tree):
+    sum_tree.add(5, "data1")
+    assert sum_tree.total == 5
+    data_idx, priority, data = sum_tree.get(5)
+    assert data_idx == 0
+    assert priority == 5
+    assert data == "data1"
+
+def test_add_overflow(sum_tree):
+    for i in range(15):
+        sum_tree.add(i, f"data{i}")
+    assert sum_tree.count == 5
+    assert sum_tree.real_size == 10
+
+# Parameterized testing for various scenarios
+@pytest.mark.parametrize("values, expected_total", [
+    ([1, 2, 3, 4, 5], 15),
+    ([10, 20, 30, 40, 50], 150),
+])
+def test_multiple_updates(sum_tree, values, expected_total):
+    for value in values:
+        sum_tree.add(value, None)
+    assert sum_tree.total == expected_total
+
+# Exception testing
+def test_get_with_invalid_cumsum(sum_tree):
+    with pytest.raises(AssertionError):
+        sum_tree.get(20)
+
+# More tests for specific methods
+def test_get_priority(sum_tree):
+    sum_tree.add(10, "data1")
+    priority = sum_tree.get_priority(0)
+    assert priority == 10
+
+def test_repr(sum_tree):
+    expected_repr = f"SumTree(nodes={sum_tree.nodes}, data={sum_tree.data})"
+    assert repr(sum_tree) == expected_repr
+
+# More test cases can be added as needed

zeta/rl/PrioritizedReplayBuffer.py

@@ -0,0 +1,85 @@
+from sumtree import SumTree
+import torch
+import random
+
+class PrioritizedReplayBuffer:
+    def __init__(self, state_size, action_size, buffer_size, device, eps=1e-2, alpha=0.1, beta=0.1):
+        self.tree = SumTree(size=buffer_size)
+
+
+        self.eps = eps  
+        self.alpha = alpha  
+        self.beta = beta  
+        self.max_priority = 1.  
+
+
+        self.state = torch.empty(buffer_size, state_size, dtype=torch.float)
+        self.action = torch.empty(buffer_size, action_size, dtype=torch.float)
+        self.reward = torch.empty(buffer_size, dtype=torch.float)
+        self.next_state = torch.empty(buffer_size, state_size, dtype=torch.float)
+        self.done = torch.empty(buffer_size, dtype=torch.uint8)
+
+        self.count = 0
+        self.real_size = 0
+        self.size = buffer_size
+
+        # device
+        self.device = device
+
+    def add(self, transition):
+        state, action, reward, next_state, done = transition
+
+
+        self.tree.add(self.max_priority, self.count)
+
+        self.state[self.count] = torch.as_tensor(state)
+        self.action[self.count] = torch.as_tensor(action)
+        self.reward[self.count] = torch.as_tensor(reward)
+        self.next_state[self.count] = torch.as_tensor(next_state)
+        self.done[self.count] = torch.as_tensor(done)
+
+
+        self.count = (self.count + 1) % self.size
+        self.real_size = min(self.size, self.real_size + 1)
+
+    def sample(self, batch_size):
+        assert self.real_size >= batch_size, "buffer contains less samples than batch size"
+
+        sample_idxs, tree_idxs = [], []
+        priorities = torch.empty(batch_size, 1, dtype=torch.float)
+
+
+        segment = self.tree.total / batch_size
+        for i in range(batch_size):
+            a, b = segment * i, segment * (i + 1)
+
+            cumsum = random.uniform(a, b)
+
+            tree_idx, priority, sample_idx = self.tree.get(cumsum)
+
+            priorities[i] = priority
+            tree_idxs.append(tree_idx)
+            sample_idxs.append(sample_idx)
+
+        probs = priorities / self.tree.total
+
+        weights = (self.real_size * probs) ** -self.beta
+
+        weights = weights / weights.max()
+        batch = (
+            self.state[sample_idxs].to(self.device),
+            self.action[sample_idxs].to(self.device),
+            self.reward[sample_idxs].to(self.device),
+            self.next_state[sample_idxs].to(self.device),
+            self.done[sample_idxs].to(self.device)
+        )
+        return batch, weights, tree_idxs
+
+    def update_priorities(self, data_idxs, priorities):
+        if isinstance(priorities, torch.Tensor):
+            priorities = priorities.detach().cpu().numpy()
+
+        for data_idx, priority in zip(data_idxs, priorities):
+            priority = (priority + self.eps) ** self.alpha
+            self.tree.update(data_idx, priority)
+            self.max_priority = max(self.max_priority, priority)

zeta/rl/PrioritizedSequenceReplayBuffer.py

@@ -0,0 +1,112 @@
+from sumtree import SumTree
+import torch
+import random
+
+class PrioritizedSequenceReplayBuffer:
+    def __init__(self,state_size,action_size,buffer_size,device,eps=1e-5,alpha=0.1,beta=0.1,
+                 decay_window=5,
+                 decay_coff=0.4,
+                 pre_priority=0.7):
+        self.tree = SumTree(data_size=buffer_size)
+
+        # PESR params
+        self.eps = eps
+        self.alpha = alpha
+        self.beta = beta
+        self.max_priority = 1.
+        self.decay_window = decay_window
+        self.decay_coff = decay_coff
+        self.pre_priority = pre_priority
+
+        # buffer params
+        self.state = torch.empty(buffer_size, state_size, dtype=torch.float)
+        self.action = torch.empty(buffer_size, action_size, dtype=torch.float)
+        self.reward = torch.empty(buffer_size, dtype=torch.float)
+        self.next_state = torch.empty(buffer_size, state_size, dtype=torch.float)
+        self.done = torch.empty(buffer_size, dtype=torch.uint8)
+
+        self.count = 0
+        self.real_size = 0
+        self.size = buffer_size
+
+        # device
+        self.device = device
+
+    def add(self, transition):
+        state, action, reward, next_state, done = transition
+
+        # store transition index with maximum priority in sum tree
+        self.tree.add(self.max_priority, self.count)
+
+        # store transition in the buffer
+        self.state[self.count] = torch.as_tensor(state)
+        self.action[self.count] = torch.as_tensor(action)
+        self.reward[self.count] = torch.as_tensor(reward)
+        self.next_state[self.count] = torch.as_tensor(next_state)
+        self.done[self.count] = torch.as_tensor(done)
+
+        # update counters
+        self.count = (self.count + 1) % self.size
+        self.real_size = min(self.size, self.real_size + 1)
+
+    def sample(self,batch_size):
+        assert self.real_size >= batch_size, "buffer contains less samples than batch size"
+
+        sample_idxs, tree_idxs = [], []
+        priorities = torch.empty(batch_size, 1, dtype=torch.float)
+
+        segment = self.tree.total_priority / batch_size
+        for i in range(batch_size):
+            a, b = segment * i, segment * (i + 1)
+
+            cumsum = random.uniform(a, b)
+            # sample_idx is a sample index in buffer, needed further to sample actual transitions
+            # tree_idx is a index of a sample in the tree, needed further to update priorities
+            tree_idx, priority, sample_idx = self.tree.get(cumsum)
+
+            priorities[i] = priority
+            tree_idxs.append(tree_idx)
+            sample_idxs.append(sample_idx)
+        """
+        Note:
+        The priorities stored in sumtree are all times alpha
+        """
+        probs = priorities / self.tree.total_priority
+        weights = (self.real_size * probs) ** -self.beta
+        weights = weights / weights.max()
+        batch = (
+            self.state[sample_idxs].to(self.device),
+            self.action[sample_idxs].to(self.device),
+            self.reward[sample_idxs].to(self.device),
+            self.next_state[sample_idxs].to(self.device),
+            self.done[sample_idxs].to(self.device)
+        )
+        return batch, weights, tree_idxs
+
+    def update_priorities(self,data_idxs,abs_td_errors):
+        """
+        when we get the TD-error, we should update the transition priority p_j
+        And update decay_window's transition priorities
+        """
+        if isinstance(abs_td_errors,torch.Tensor):
+            abs_td_errors = abs_td_errors.detach().cpu().numpy()
+
+        for data_idx, td_error in zip(data_idxs,abs_td_errors):
+            # first update the batch: p_j
+            # p_j <- max{|delta_j| + eps, pre_priority * p_j}
+            old_priority = self.pre_priority * self.tree.nodes[data_idx + self.tree.size - 1]
+            priority = (td_error + self.eps) ** self.alpha
+            priority = max(priority,old_priority)
+            self.tree.update(data_idx,priority)
+            self.max_priority = max(self.max_priority,priority)
+
+        # And then apply decay
+        if self.count >= self.decay_window:
+            # count points to the next position
+            # count means the idx in the buffer and number of transition
+            for i in reversed(range(self.decay_window)):
+                idx = (self.count - i - 1) % self.size
+                decayed_priority = priority * (self.decay_coff ** (i + 1))
+                tree_idx = idx + self.tree.size - 1
+                existing_priority = self.tree.nodes[tree_idx]
+                self.tree.update(idx,max(decayed_priority,existing_priority))