Buffer.py

from typing import Optional, Union, Tuple
import threading
import multiprocessing as mp

import numpy as np
import torch as th
from typeguard import typechecked

from utils import auto_device


class DoneHandler:
    @typechecked
    def handle_done(self, buffer: 'Buffer'):
        raise NotImplementedError()

# another example done handler
class RewToGoDoneHandler(DoneHandler):
    def __init__(self, gamma):
        self.gamma = gamma

    @typechecked
    def handle_done(self, buffer: 'Buffer'):
        """turn last episode rewards into reward-to-go"""
        if buffer.ep_start == buffer.ep_end:
            return
        # calculate reward-to-go
        for i in range(buffer.ep_end-1, buffer.ep_start-1, -1):
            buffer.rewards[i] += self.gamma * buffer.rewards[i+1]
        buffer.ep_start = buffer.ep_end


type_to_th_type = {
    int: th.int64,
    np.int64: th.int64,
}
type_to_np_type = {
    int: np.int64,
}


class Buffer:
    """Buffer for storing potentially huge amount of images"""
    @typechecked
    def __init__(
            self,
            buffer_size: int = 20000,
            state:Union[th.Tensor,np.ndarray]=None,
            action:Union[th.Tensor,np.ndarray,int,np.int64]=None,
            done_handlers:Tuple[DoneHandler, ...] = (),
            device:str='auto',
            preload_sample:bool=True,
    ):
        """
        n_samples: int - number of samples to store
        state: th.Tensor or np.ndarray - example state to be stored
        action: th.Tensor or np.ndarray - example action to be stored
        max_in_memory: int - maximum number of samples to store in memory
        img_history_len: int - number of images to store for each sample
        done_handlers: List[Callable] - list of functions to call when an episode is done
        """
        assert isinstance(state, th.Tensor) or isinstance(state, np.ndarray)
        self.state_shape = state.shape
        self.state_dtype = state.dtype if isinstance(state, np.ndarray) else state.dtype
        self.action_shape = (1,) if isinstance(action, (int, np.integer))  else action.shape
        self.action_dtype = action.dtype if isinstance(action, np.ndarray) else type(action)
        self.buffer_size = buffer_size
        self.n_stored = 0
        self.device = auto_device(device)

        self.transforms = {}
        self.clear()
        self.done_handlers = done_handlers
        self.preload_sample = preload_sample
        self.preloaded_sample = None
        self.preload_done = False
        self.proc = None
        self.preload_thread = None
        self.n_missed_preloads = 0

    def _preload(self, batch_size, queue):
        preloaded_sample = self._sample(batch_size)
        queue.put(preloaded_sample)

    def preload(self, batch_size):
        def preload_worker():
            queue = mp.Queue()
            self.proc = mp.Process(target=self._preload, args=(batch_size,queue))
            self.proc.start()
            # self._preload(batch_size, queue)
            self.preloaded_sample = queue.get()
            self.preload_done = True
            # print("Preloading done")
            self.proc.join()
            self.proc = None
            # print("joined")
        if self.proc is not None:
            # print(f"previous preloading still running {self.n_missed_preloads}")
            self.n_missed_preloads += 1
        else:
            self.n_missed_preloads = 0
            self.preload_thread = threading.Thread(target=preload_worker, daemon=True)
            self.preload_thread.start()

    def clear(self):
        self.states =  np.empty((self.buffer_size, *self.state_shape),  dtype=self.state_dtype)
        self.actions = np.empty((self.buffer_size, *self.action_shape), dtype=self.action_dtype)
        self.rewards = np.empty((self.buffer_size, 1), dtype=np.float32)
        self.terminated =   np.empty((self.buffer_size, 1), dtype=bool)
        self.ep_start = 0
        self.ep_end = 0
        self.n_stored = 0

    def add(self, state, action, reward, terminated):
        self.states [self.ep_end] = state
        self.actions[self.ep_end] = action
        self.rewards[self.ep_end] = reward
        self.terminated[self.ep_end] = terminated
        self.n_stored = min(self.buffer_size, self.n_stored + 1)
        self.ep_end += 1
        if self.ep_end == self.buffer_size:
            self.ep_end = 0
        if terminated:
            self._handle_done()

    def calculate_statistics(self):
        # get the histogram for rewards:
        rewards = self.rewards[:self.n_stored]
        rewards = rewards.flatten()
        reward_histogram = {}
        if self.n_stored > 0:
            unique = np.unique(rewards)
            if unique.size <= 50:
                unique, counts = np.unique(rewards, return_counts=True)
                reward_histogram = dict(zip(unique.tolist(), counts.tolist()))
            else:
                # Bin continuous rewards to avoid one-count-per-value
                counts, edges = np.histogram(rewards, bins=50)
                centers = (edges[:-1] + edges[1:]) / 2.0
                reward_histogram = {
                    float(center): int(count)
                    for center, count in zip(centers, counts)
                    if count > 0
                }
        return {'reward_histogram': reward_histogram,
                'terminated_fraction': np.sum(self.terminated[:self.n_stored]) / self.n_stored,
                'n_stored': self.n_stored,
                'buffer_size': self.buffer_size,
                }

    @staticmethod
    def to_device(batch, device):
        return [th.from_numpy(x).to(device) for x in batch]

    def _sample(self, batch_size):
        # Avoid sampling the most recent transition when the buffer is not full,
        # since its next_state has not been written yet.
        if self.n_stored < 2: # need state, next_state
            raise ValueError("Not enough samples in buffer to sample.")
        if self.n_stored < self.buffer_size:
            idxs_valid = np.arange(self.n_stored - 1)
        else:
            idxs_valid = np.arange(self.buffer_size)

        idx = np.random.choice(idxs_valid, batch_size)
        next_idx = (idx + 1) % self.buffer_size
        return (self.states[idx],
                self.actions[idx],
                self.rewards[idx],
                self.states[next_idx],
                self.terminated[idx],)

    def sample(self, batch_size):
        if self.preload_done:
            # get the preloaded sample and start preloading the next one
            self.preload_done = False
            self.preload(batch_size)
            return Buffer.to_device(self.preloaded_sample, self.device)
        else:
            # print('preloading not ready, preparing for the next time')
            self.preload(batch_size)
        return Buffer.to_device(self._sample(batch_size), self.device)

    def _handle_done(self):
        for handler, h_kwargs in self.done_handlers:
            handler(self, **h_kwargs)

    def cleanup(self) -> None:
        """Cleanup method to terminate any active multiprocessing processes and wait for threads."""
        # Wait for preload thread to complete with timeout
        if self.preload_thread is not None and self.preload_thread.is_alive():
            self.preload_thread.join(timeout=5.0)
        
        # Terminate any active processes
        if self.proc is not None:
            try:
                if self.proc.is_alive():
                    self.proc.terminate()
                    self.proc.join(timeout=5.0)
                    if self.proc.is_alive():
                        self.proc.kill()
            except Exception:
                pass
            self.proc = None


class TDBuffer(Buffer):
    def __init__(self, *args,  td_steps:int=1 ,**kwargs):
        super().__init__(*args, **kwargs)
        """
        :param td_steps: int - number of steps to look ahead for TD learning
        """
        self.td_steps = td_steps

    def sample(self, batch_size, preloading=False):
        if not preloading and self.preloaded_sample is not None:
            return self.preloaded_sample
        if self.preload_sample and not preloading:
            print('preloading not ready, preparing for the next time')
            self.preloaded_sample = None
            self._preload(batch_size)
        idx = np.random.randint(low=0,high=self.n_stored, size=(batch_size,))