data_container.py

""""
The implementation is based on SegmentTree, and Priority sampling from https://github.com/labmlai/annotated_deep_learning_paper_implementations
"""
import numpy as np
import torch
import random
import time



class SegmenTreeSum:
    def __init__(self, capacity, alpha):
        self.alpha = alpha
        self.capacity = capacity
        self.priority_sum = [0 for _ in range(2 * self.capacity)]
        self.priorities = np.zeros(self.capacity, dtype=np.float32)
        self.size = 0
        self.next_idx = 0
        self.max_priority = 1.0
        
    def add(self, priority=None):
        idx = self.next_idx
        self.next_idx = (idx + 1) % self.capacity
        self.size = min(self.capacity, self.size + 1)
        if priority is not None:
            priority_alpha = priority ** self.alpha
        else:
            priority_alpha = self.max_priority ** self.alpha
        self._set_priority_sum(idx, priority_alpha)
        self.priorities[idx] = priority_alpha
        
    def sample(self, batch_size, beta=1.0):
        """
        ### Sample from buffer
        """
        # Initialize samples
        samples = {
            'indexes': np.zeros(shape=batch_size, dtype=np.int32)
        }

        for i in range(batch_size):
            p = random.random() * self._sum()
            idx = self.find_prefix_sum_idx(p)
            samples['indexes'][i] = idx
        
        return samples['indexes']

    
    def _set_priority_min(self, idx, priority_alpha):
        idx += self.capacity
        self.priority_min[idx] = priority_alpha

        while idx >= 2:
            idx //= 2
            self.priority_min[idx] = min(self.priority_min[2 * idx], self.priority_min[2 * idx + 1])
            
    def _set_priority_sum(self, idx, priority):
        idx += self.capacity
        self.priority_sum[idx] = priority

        while idx >= 2:
            idx //= 2
            self.priority_sum[idx] = self.priority_sum[2 * idx] + self.priority_sum[2 * idx + 1]

    def _sum(self):
        return self.priority_sum[1]
    
    def _min(self):
        return self.priority_min[1]
    
    def find_prefix_sum_idx(self, prefix_sum):
        idx = 1
        while idx < self.capacity:
            if self.priority_sum[idx * 2] > prefix_sum:
                idx = 2 * idx
            else:
                prefix_sum -= self.priority_sum[idx * 2]
                idx = 2 * idx + 1
            
        return idx - self.capacity
    
    def update_priorities(self, indexes, priorities):
        for idx, priority in zip(indexes, priorities):
            priority = abs(priority)
            self.max_priority = max(self.max_priority, priority)
            priority_alpha = priority ** self.alpha
            self._set_priority_sum(idx, priority_alpha)
            self.priorities[idx] = priority_alpha
            
    def is_full(self):
        return self.capacity == self.size


class DATA(object):
    def __init__(self, state_dim, action_dim, device, max_size=int(1e6)):
        self.max_size = max_size
        self.ptr = 0
        self.size = 0

        self.state = np.zeros((max_size, state_dim))
        self.action = np.zeros((max_size, action_dim))
        self.vae_action = np.zeros((max_size, action_dim))
        self.next_state = np.zeros((max_size, state_dim))
        self.reward = np.zeros((max_size, 1))
        self.not_done = np.zeros((max_size, 1))
        
        self.state_dim = state_dim
        self.action_dim = action_dim

        self.device = device


    def add(self, state, action, next_state, reward, done):
        self.state[self.ptr] = np.array(state).copy()
        self.action[self.ptr] = np.array(action).copy()
        self.vae_action[self.ptr] = np.array(action).copy()
        self.next_state[self.ptr] = np.array(next_state).copy()
        self.reward[self.ptr] = np.array(reward).copy()
        self.not_done[self.ptr] = np.array(1. - done).copy()

        self.ptr = (self.ptr + 1) % self.max_size
        self.size = min(self.size + 1, self.max_size)


    def sample(self, batch_size, use_bootstrap=False, use_vae_training=False):
        ind = np.random.randint(0, self.size, size=batch_size)

        if use_bootstrap:
            return (
            torch.FloatTensor(self.state[ind]).to(self.device),
            torch.FloatTensor(self.action[ind]).to(self.device),
            torch.FloatTensor(self.next_state[ind]).to(self.device),
            torch.FloatTensor(self.reward[ind]).to(self.device),
            torch.FloatTensor(self.not_done[ind]).to(self.device),
            torch.FloatTensor(self.bootstrap_mask[ind]).to(self.device),
        )
        
        if use_vae_training:
            return (
            torch.FloatTensor(self.state[ind]).to(self.device),
            torch.FloatTensor(self.action[ind]).to(self.device),
            torch.FloatTensor(self.vae_action[ind]).to(self.device),
            torch.FloatTensor(self.next_state[ind]).to(self.device),
            torch.FloatTensor(self.reward[ind]).to(self.device),
            torch.FloatTensor(self.not_done[ind]).to(self.device)
        )

        return (
            torch.FloatTensor(self.state[ind]).to(self.device),
            torch.FloatTensor(self.action[ind]).to(self.device),
            torch.FloatTensor(self.next_state[ind]).to(self.device),
            torch.FloatTensor(self.reward[ind]).to(self.device),
            torch.FloatTensor(self.not_done[ind]).to(self.device)
        )
        
    def sample_batch(self, batch_size):
        ind = np.random.randint(0, self.size, size=batch_size)
        return (
            self.state[ind],
            self.action[ind],
            self.next_state[ind],
            self.reward[ind],
            self.not_done[ind]
        )
    
    def sample_all(self):
        # Just return the whole, not transfer them to the gpu device
        return (
        self.state[:self.size].copy(),
        self.action[:self.size].copy(),
        self.next_state[:self.size].copy(),
        self.reward[:self.size].copy(),
        self.not_done[:self.size].copy()
        )


    def save(self, save_folder):
        np.save(f"{save_folder}_state.npy", self.state[:self.size])
        np.save(f"{save_folder}_action.npy", self.action[:self.size])
        np.save(f"{save_folder}_next_state.npy", self.next_state[:self.size])
        np.save(f"{save_folder}_reward.npy", self.reward[:self.size])
        np.save(f"{save_folder}_not_done.npy", self.not_done[:self.size])
        np.save(f"{save_folder}_ptr.npy", self.ptr)

    
    def load_append(self, save_folder, size):
        current_ptr = self.ptr
        current_size = self.size
        print(f"size before loading: {current_size}")
        print(f"loading data from  {save_folder} with the size {size}")
        reward_buffer = np.load(f"{save_folder}_reward.npy")
        load_size = min(reward_buffer.shape[0], size)
        free_space = self.max_size - current_size
        load_size = min(load_size, free_space)

        self.state[current_ptr:current_ptr+load_size] = np.load(f"{save_folder}_state.npy")[:load_size]
        self.action[current_ptr:current_ptr+load_size] = np.load(f"{save_folder}_action.npy")[:load_size]
        self.vae_action[current_ptr:current_ptr+load_size] = np.load(f"{save_folder}_action.npy")[:load_size]
        self.next_state[current_ptr:current_ptr+load_size] = np.load(f"{save_folder}_next_state.npy")[:load_size]
        self.reward[current_ptr:current_ptr+load_size] = reward_buffer[:load_size]
        self.not_done[current_ptr:current_ptr+load_size] = np.load(f"{save_folder}_not_done.npy")[:load_size]

        self.ptr = (current_ptr + load_size) % self.max_size
        self.size = min(current_size + load_size, self.max_size)
        print(f"size after loading: {self.size}")

    
    def load_a_part(self, save_folder, size=-1, bootstrap_dim=None, ratio=1.0):
        print(f"loading data from  {save_folder}")
        reward_buffer = np.load(f"{save_folder}_reward.npy")
        # Adjust crt_size if we're using a custom size
        size = min(int(size), self.max_size) if size > 0 else self.max_size
        self.size = min(reward_buffer.shape[0], size)
        self.ptr = (self.ptr + self.size) % self.max_size
        
        
        
        print("-----------------")
        print(f"Data size: {self.size}")
        print("-----------------")

        # then load!
        self.state[:self.size] = np.load(f"{save_folder}_state.npy")[:self.size]
        self.action[:self.size] = np.load(f"{save_folder}_action.npy")[:self.size]
        self.vae_action[:self.size] = np.load(f"{save_folder}_action.npy")[:self.size]
        self.next_state[:self.size] = np.load(f"{save_folder}_next_state.npy")[:self.size]
        self.reward[:self.size] = reward_buffer[:self.size]
        self.not_done[:self.size] = np.load(f"{save_folder}_not_done.npy")[:self.size]
        
        if bootstrap_dim is not None:
            self.bootstrap_dim = bootstrap_dim
            bootstrap_mask = np.random.binomial(n=1, size=(1, self.size, bootstrap_dim,), p=0.8)
            bootstrap_mask = np.squeeze(bootstrap_mask, axis=0)
            self.bootstrap_mask = bootstrap_mask[:self.size]

        print(f"ptr: {self.ptr}")
        print(f"size: {self.size}")
        print(f"max_size: {self.max_size}")
        print(f"state: {self.state[self.ptr]}, {self.state[self.ptr-1]}")


    def load(self, save_folder, size=-1, bootstrap_dim=None):
        print(f"loading data from  {save_folder}")
        reward_buffer = np.load(f"{save_folder}_reward.npy")
        # Adjust crt_size if we're using a custom size
        size = min(int(size), self.max_size) if size > 0 else self.max_size
        self.size = min(reward_buffer.shape[0], size)
        
        
        print("-----------------")
        print(f"Data size: {self.size}")
        print("-----------------")

        # resize data first!
        self.state.resize((self.size, self.state_dim))
        self.action.resize((self.size, self.action_dim))
        self.vae_action.resize((self.size, self.action_dim))
        self.next_state.resize((self.size, self.state_dim))
        self.reward.resize((self.size, 1))
        self.not_done.resize((self.size, 1))
        
        # then load!
        self.state[:self.size] = np.load(f"{save_folder}_state.npy")[:self.size]
        self.action[:self.size] = np.load(f"{save_folder}_action.npy")[:self.size]
        self.vae_action[:self.size] = np.load(f"{save_folder}_action.npy")[:self.size]
        self.next_state[:self.size] = np.load(f"{save_folder}_next_state.npy")[:self.size]
        self.reward[:self.size] = reward_buffer[:self.size]
        self.not_done[:self.size] = np.load(f"{save_folder}_not_done.npy")[:self.size]
        
        if bootstrap_dim is not None:
            self.bootstrap_dim = bootstrap_dim
            bootstrap_mask = np.random.binomial(n=1, size=(1, self.size, bootstrap_dim,), p=0.8)
            bootstrap_mask = np.squeeze(bootstrap_mask, axis=0)
            self.bootstrap_mask = bootstrap_mask[:self.size]



class PRIORITIZED_SUM_DATA(DATA):
    def __init__(self, state_dim, action_dim, device, max_size=int(1e6), alpha=0.6):
        super(PRIORITIZED_SUM_DATA, self).__init__(state_dim, action_dim, device, max_size)
        self.alpha = alpha
        self.tree = SegmenTreeSum(max_size, alpha)
    
    def add(self, state, action, next_state, reward, done, priority=None):
        super(PRIORITIZED_SUM_DATA, self).add(state, action, next_state, reward, done)
        if priority is not None:
            self.tree.add(priority)
        else:
            self.tree.add()
    
    def sample(self, batch_size, beta=1.0, use_bootstrap=False, use_vae_training=False, get_priorities=False):
        indexes = self.tree.sample(batch_size)
        if get_priorities:
            return (
            torch.FloatTensor(self.state[indexes]).to(self.device),
            torch.FloatTensor(self.action[indexes]).to(self.device),
            torch.FloatTensor(self.next_state[indexes]).to(self.device),
            torch.FloatTensor(self.reward[indexes]).to(self.device),
            torch.FloatTensor(self.not_done[indexes]).to(self.device),
            indexes,
            self.tree.priorities[indexes],
        )
        else:
            return (
                torch.FloatTensor(self.state[indexes]).to(self.device),
                torch.FloatTensor(self.action[indexes]).to(self.device),
                torch.FloatTensor(self.next_state[indexes]).to(self.device),
                torch.FloatTensor(self.reward[indexes]).to(self.device),
                torch.FloatTensor(self.not_done[indexes]).to(self.device),
                indexes,
            )

    def sample_randomly(self, batch_size, use_bootstrap=False, use_vae_training=False, get_priorities=False):
        ind = np.random.randint(0, self.size, size=batch_size)

        if use_bootstrap:
            return (
            torch.FloatTensor(self.state[ind]).to(self.device),
            torch.FloatTensor(self.action[ind]).to(self.device),
            torch.FloatTensor(self.next_state[ind]).to(self.device),
            torch.FloatTensor(self.reward[ind]).to(self.device),
            torch.FloatTensor(self.not_done[ind]).to(self.device),
            torch.FloatTensor(self.bootstrap_mask[ind]).to(self.device),
        )
        
        if use_vae_training:
            return (
            torch.FloatTensor(self.state[ind]).to(self.device),
            torch.FloatTensor(self.action[ind]).to(self.device),
            torch.FloatTensor(self.vae_action[ind]).to(self.device),
            torch.FloatTensor(self.next_state[ind]).to(self.device),
            torch.FloatTensor(self.reward[ind]).to(self.device),
            torch.FloatTensor(self.not_done[ind]).to(self.device)
        )
        
        if get_priorities:
            return (
                torch.FloatTensor(self.state[ind]).to(self.device), 
                torch.FloatTensor(self.action[ind]).to(self.device),
                torch.FloatTensor(self.next_state[ind]).to(self.device),
                torch.FloatTensor(self.reward[ind]).to(self.device),
                torch.FloatTensor(self.not_done[ind]).to(self.device),
                ind,
                self.tree.priorities[ind]
            )

        return (
            torch.FloatTensor(self.state[ind]).to(self.device),
            torch.FloatTensor(self.action[ind]).to(self.device),
            torch.FloatTensor(self.next_state[ind]).to(self.device),
            torch.FloatTensor(self.reward[ind]).to(self.device),
            torch.FloatTensor(self.not_done[ind]).to(self.device),
            ind
        )
    
    def update_priorities(self, indexes, priorities):
        self.tree.update_priorities(indexes, priorities)

    def is_full(self):
        return self.tree.is_full()