common/pref_replay_buffer.py

from typing import Union, Dict, Optional
import os
import math
import numbers
import zarr
import numcodecs
import numpy as np
from functools import cached_property

def check_chunks_compatible(chunks: tuple, shape: tuple):
    assert len(shape) == len(chunks)
    for c in chunks:
        assert isinstance(c, numbers.Integral)
        assert c > 0

def rechunk_recompress_array(group, name, 
        chunks=None, chunk_length=None,
        compressor=None, tmp_key='_temp'):
    old_arr = group[name]
    if chunks is None:
        if chunk_length is not None:
            chunks = (chunk_length,) + old_arr.chunks[1:]
        else:
            chunks = old_arr.chunks
    check_chunks_compatible(chunks, old_arr.shape)
    
    if compressor is None:
        compressor = old_arr.compressor
    
    if (chunks == old_arr.chunks) and (compressor == old_arr.compressor):
        # no change
        return old_arr

    # rechunk recompress
    group.move(name, tmp_key)
    old_arr = group[tmp_key]
    n_copied, n_skipped, n_bytes_copied = zarr.copy(
        source=old_arr,
        dest=group,
        name=name,
        chunks=chunks,
        compressor=compressor,
    )
    del group[tmp_key]
    arr = group[name]
    return arr

def get_optimal_chunks(shape, dtype, 
        target_chunk_bytes=2e6, 
        max_chunk_length=None):
    """
    Common shapes
    T,D
    T,N,D
    T,H,W,C
    T,N,H,W,C
    """
    itemsize = np.dtype(dtype).itemsize
    # reversed
    rshape = list(shape[::-1])
    if max_chunk_length is not None:
        rshape[-1] = int(max_chunk_length)
    split_idx = len(shape)-1
    for i in range(len(shape)-1):
        this_chunk_bytes = itemsize * np.prod(rshape[:i])
        next_chunk_bytes = itemsize * np.prod(rshape[:i+1])
        if this_chunk_bytes <= target_chunk_bytes \
            and next_chunk_bytes > target_chunk_bytes:
            split_idx = i

    rchunks = rshape[:split_idx]
    item_chunk_bytes = itemsize * np.prod(rshape[:split_idx])
    this_max_chunk_length = rshape[split_idx]
    next_chunk_length = min(this_max_chunk_length, math.ceil(
            target_chunk_bytes / item_chunk_bytes))
    rchunks.append(next_chunk_length)
    len_diff = len(shape) - len(rchunks)
    rchunks.extend([1] * len_diff)
    chunks = tuple(rchunks[::-1])
    # print(np.prod(chunks) * itemsize / target_chunk_bytes)
    return chunks


class PrefReplayBuffer:
    """
    Zarr-based temporal data structure specifically for preference dataset.
    Stores pairs of trajectories (observations, actions) along with votes.
    """
    
    def __init__(self, root: Union[zarr.Group, Dict[str, dict]]):
        """
        Initialize the preference replay buffer. Use class methods to create or load buffers.
        """
        assert 'data' in root
        assert 'meta' in root
        for key, value in root['data'].items():
            assert value.shape[0] == root['meta']['votes'].shape[0]
        self.root = root

    # ============= create constructors ===============
    @classmethod
    def create_empty_zarr(cls, storage=None, root=None):
        if root is None:
            if storage is None:
                storage = zarr.MemoryStore()
            root = zarr.group(store=storage)
        data = root.require_group('data', overwrite=False)
        meta = root.require_group('meta', overwrite=False)
        
        if 'votes' not in meta:
            votes = meta.zeros('votes', shape=(0,), dtype=np.float32, compressor=None, overwrite=False)
        if 'votes_2' not in meta:
            votes_2 = meta.zeros('votes_2', shape=(0,), dtype=np.float32, compressor=None, overwrite=False)
        
        return cls(root=root)

    @classmethod
    def create_empty_numpy(cls):
        root = {
            'data': dict(),
            'meta': {
                'episode_ends': np.zeros((0,), dtype=np.int64),
                'votes': np.zeros((0,), dtype=np.float32), 
                'votes_2': np.zeros((0,), dtype=np.float32),
                'length': np.zeros((0,), dtype=np.int64),
                'length_2': np.zeros((0,), dtype=np.int64),
                'beta_priori': np.zeros((0,), dtype=np.float32),
                'beta_priori_2': np.zeros((0,), dtype=np.float32),
            }
        }
        return cls(root=root)

    @classmethod
    def create_from_group(cls, group, **kwargs):
        if 'data' not in group:
            # create from scratch
            buffer = cls.create_empty_zarr(root=group, **kwargs)
        else:
            # already exists
            buffer = cls(root=group, **kwargs)
        return buffer

    @classmethod
    def create_from_path(cls, zarr_path, mode='r', **kwargs):
        """
        Open a Zarr file from disk for large datasets that cannot fit in memory.
        """
        group = zarr.open(zarr_path, mode=mode)
        return cls.create_from_group(group=group, **kwargs)

    # ============= Add episodes ===============
    def add_pref_episode(self, data: Dict[str, np.ndarray], 
                         meta_data: Optional[Dict[str, Union[np.ndarray, int]]] = None,
                         chunks: Optional[Dict[str, tuple]] = dict(),
                         compressors: Union[str, numcodecs.abc.Codec, dict] = dict()):
        """
        Add a pair of episodes (obs/action for each trajectory) along with metadata (votes).
        """
        assert 'obs' in data and 'obs_2' in data, "obs and obs_2 keys are required"
        assert 'action' in data and 'action_2' in data, "action and action_2 keys are required"
        
        is_zarr = isinstance(self.root, zarr.Group)
        curr_len = len(self.root['meta']['votes'])
        episode_length = len(data['obs'])
        new_len = curr_len + 1

        # Add trajectory 1
        for key in ['obs', 'action']:
            # Create the new shape to accommodate all time steps
            value = data[key]
            new_shape = (new_len,) + (episode_length,) + data[key].shape[1:]  # This will set (new_len, T, dim)

            if key not in self.root['data']:
                # Create a new array if it doesn't exist
                if is_zarr:
                    cks = self._resolve_array_chunks(chunks, key, data[key])
                    cpr = self._resolve_array_compressor(compressors, key, data[key])
                    arr = self.root['data'].zeros(name=key, shape=new_shape, chunks=cks, dtype=data[key].dtype, compressor=cpr)
                else:
                    arr = np.zeros(new_shape, dtype=data[key].dtype)
                    self.root['data'][key] = arr
            else:
                arr = self.root['data'][key]
                if is_zarr:
                    arr.resize(new_shape)
                else:
                    arr.resize(new_shape, refcheck=False)

            # Store the full sequence, adjusting the shape to match the time steps in data[key]
            arr[new_len-1, -value.shape[0]:, :] = value # Now this assumes data[key] has shape (T, dim)


        # Add trajectory 2 (obs_2, action_2)
        for key in ['obs_2', 'action_2']:
            value = data[key]
            # Create the new shape to accommodate all time steps
            new_shape = (new_len,) + (episode_length,) + data[key].shape[1:]  # This will set (new_len, T, dim)

            if key not in self.root['data']:
                # Create a new array if it doesn't exist
                if is_zarr:
                    cks = self._resolve_array_chunks(chunks, key, data[key])
                    cpr = self._resolve_array_compressor(compressors, key, data[key])
                    arr = self.root['data'].zeros(name = key, shape=new_shape, chunks=cks, dtype=data[key].dtype, compressor=cpr)
                else:
                    arr = np.zeros(new_shape, dtype=data[key].dtype)
                    self.root['data'][key] = arr
            else:
                arr = self.root['data'][key]
                if is_zarr:
                    arr.resize(new_shape)
                else:
                    arr.resize(new_shape, refcheck=False)

            # Store the full sequence, adjusting the shape to match the time steps in data[key]
            arr[new_len-1, -value.shape[0]:, :] = value  # Now this assumes data[key] has shape (T, dim)

        # Add votes to meta
        if meta_data:
            for key in ['votes', 'votes_2']:
                new_shape = (new_len,) + (1,)
                if key not in self.root['meta']:
                    if is_zarr:
                        self.root['meta'].zeros(name=key, shape=new_shape, chunks=new_shape, dtype=np.float32)
                    else:
                        self.root['meta'][key] = np.zeros(new_shape, dtype=np.float32)
                arr = self.root['meta'][key]
                if is_zarr:
                    arr.resize(new_shape)
                else:
                    arr.resize(new_shape, refcheck=False)
                arr[new_len-1] = meta_data[key]

            for key in ['length', 'length_2']:
                new_shape = (new_len,) + (1,)
                if key not in self.root['meta']:
                    if is_zarr:
                        self.root['meta'].zeros(name=key, shape=new_shape, chunks=new_shape, dtype=np.float32)
                    else:
                        self.root['meta'][key] = np.zeros(new_shape, dtype=np.float32)
                arr = self.root['meta'][key]
                if is_zarr:
                    arr.resize(new_shape)
                else:
                    arr.resize(new_shape, refcheck=False)
                arr[new_len-1] = meta_data[key]

            for key in ['beta_priori', 'beta_priori_2']:
                new_shape = (new_len,) + (2,)
                if key not in self.root['meta']:
                    if is_zarr:
                        self.root['meta'].zeros(name=key, shape=new_shape, chunks=new_shape, dtype=np.float32)
                    else: 
                        self.root['meta'][key] = np.zeros(new_shape, dtype=np.float32)
                arr = self.root['meta'][key]
                if is_zarr:
                    arr.resize(new_shape)
                else:
                    arr.resize(new_shape, refcheck=False)
                arr[new_len-1] = meta_data[key]

    # ============= Get episodes ===============
    def get_pref_episode(self, idx: int, copy: bool = False):
        """
        Get a pair of episodes by index, including observation and action sequences for both trajectories.
        """
        if copy:
            return {
                'obs': self.root['data']['obs'][idx].copy(),
                'action': self.root['data']['action'][idx].copy(),
                'obs_2': self.root['data']['obs_2'][idx].copy(),
                'action_2': self.root['data']['action_2'][idx].copy(),
                'votes': self.root['meta']['votes'][idx].copy(),
                'votes_2': self.root['meta']['votes_2'][idx].copy(),
                'length': self.root['meta']['length'][idx].copy(),
                'length_2': self.root['meta']['length_2'][idx].copy(),
                'beta_priori': self.root['meta']['beta_priori'][idx].copy(),
                'beta_priori_2': self.root['meta']['beta_priori_2'][idx].copy(),
            }
        else:
            return {
                'obs': self.root['data']['obs'][idx],
                'action': self.root['data']['action'][idx],
                'obs_2': self.root['data']['obs_2'][idx],
                'action_2': self.root['data']['action_2'][idx],
                'votes': self.root['meta']['votes'][idx],
                'votes_2': self.root['meta']['votes_2'][idx],
                'length': self.root['meta']['length'][idx],
                'length_2': self.root['meta']['length_2'][idx],
                'beta_priori': self.root['meta']['beta_priori'][idx],
                'beta_priori_2': self.root['meta']['beta_priori_2'][idx],
            }

    def get_episode_slice(self, idx):
        """
        Get the slice range for an episode based on the index for slicing observation and action arrays.
        """
        return slice(idx, idx + 1)

    # ============= Save methods ===============
    def save_to_store(self, store, chunks: Optional[Dict[str, tuple]] = dict(),
                      compressors: Union[str, numcodecs.abc.Codec, dict] = dict(),
                      if_exists='replace', **kwargs):
        """
        Save the replay buffer to the given store with optional chunking and compression.
        """
        root = zarr.group(store=store)
        # Save data and meta with chunking and compression if provided
        for key, value in self.root['data'].items():
            cks = self._resolve_array_chunks(chunks, key, value)
            cpr = self._resolve_array_compressor(compressors, key, value)
            arr = self.root['data'][key]
            zarr.copy(arr, root['data'], name=key, chunks=cks, compressor=cpr)
        
        # Save meta
        for key, value in self.root['meta'].items():
            zarr.copy(value, root['meta'], name=key)
    
    def save_to_path(self, zarr_path, chunks: Optional[Dict[str, tuple]] = dict(),
                     compressors: Union[str, numcodecs.abc.Codec, dict] = dict(), 
                     if_exists='replace', **kwargs):
        """
        Save the replay buffer to a path.
        """
        store = zarr.DirectoryStore(zarr_path)
        return self.save_to_store(store, chunks=chunks, compressors=compressors, if_exists=if_exists, **kwargs)

    # ============= Helper methods ===============
    @classmethod
    def _resolve_array_chunks(cls,
            chunks: Union[dict, tuple], key, array):
        cks = None
        if isinstance(chunks, dict):
            if key in chunks:
                cks = chunks[key]
            elif isinstance(array, zarr.Array):
                cks = array.chunks
        elif isinstance(chunks, tuple):
            cks = chunks
        else:
            raise TypeError(f"Unsupported chunks type {type(chunks)}")
        # backup default
        if cks is None:
            cks = get_optimal_chunks(shape=array.shape, dtype=array.dtype)
        # check
        check_chunks_compatible(chunks=cks, shape=array.shape)
        return cks

    @classmethod
    def _resolve_array_compressor(cls, compressors: Union[dict, str, numcodecs.abc.Codec], key, array):
        # Resolves the compressor for the array
        return compressors.get(key, array.compressor if isinstance(array, zarr.Array) else numcodecs.Blosc())

    @property
    def n_steps(self):
        return len(self.root['meta']['votes'])

    @property
    def data(self):
        return self.root['data']

    @property
    def meta(self):
        return self.root['meta']