dataset_utils.py

import collections
import csv
import json
import random
import string
import sys
from datetime import datetime
from pathlib import Path
from typing import Optional

import d4rl

# TODO use ultra only for antmaze-ultra
# import d4rlultra.d4rl as d4rl

import gym
import numpy as np
from tqdm import tqdm

Batch = collections.namedtuple(
    "Batch", ["observations", "initial_observations", "actions", "rewards", "masks", "next_observations"]
)


def split_into_trajectories(
    observations, actions, rewards, masks, dones_float, next_observations
):
    trajs = [[]]

    for i in tqdm(range(len(observations))):
        trajs[-1].append(
            (
                observations[i],
                actions[i],
                rewards[i],
                masks[i],
                dones_float[i],
                next_observations[i],
            )
        )
        if dones_float[i] == 1.0 and i + 1 < len(observations):
            trajs.append([])

    return trajs


def merge_trajectories(trajs):
    observations = []
    actions = []
    rewards = []
    masks = []
    dones_float = []
    next_observations = []

    for traj in trajs:
        for obs, act, rew, mask, done, next_obs in traj:
            observations.append(obs)
            actions.append(act)
            rewards.append(rew)
            masks.append(mask)
            dones_float.append(done)
            next_observations.append(next_obs)

    return (
        np.stack(observations),
        np.stack(actions),
        np.stack(rewards),
        np.stack(masks),
        np.stack(dones_float),
        np.stack(next_observations),
    )


class Dataset(object):
    def __init__(
        self,
        observations: np.ndarray,
        initial_observations: np.ndarray,
        actions: np.ndarray,
        rewards: np.ndarray,
        masks: np.ndarray,
        dones_float: np.ndarray,
        next_observations: np.ndarray,
        size: int,
    ):
        self.observations = observations
        self.initial_observations = initial_observations
        self.actions = actions
        self.rewards = rewards
        self.masks = masks
        self.dones_float = dones_float
        self.next_observations = next_observations
        self.size = size
        self.initial_size = len(initial_observations)

    # initial observation batch size?
    def sample(self, batch_size: int) -> Batch:
        initial_indx = np.random.randint(self.initial_size, size=batch_size)
        indx = np.random.randint(self.size, size=batch_size)
        return Batch(
            initial_observations=self.initial_observations[initial_indx],
            observations=self.observations[indx],
            actions=self.actions[indx],
            rewards=self.rewards[indx],
            masks=self.masks[indx],
            next_observations=self.next_observations[indx],
        )


class D4RLDataset(Dataset):
    def __init__(
        self,
        env: gym.Env,
        add_env: gym.Env = "None",
        expert_ratio: float = 1.0,
        clip_to_eps: bool = True,
        heavy_tail: bool = False,
        heavy_tail_higher: float = 0.0,
        eps: float = 1e-5,
    ):
        dataset = d4rl.qlearning_dataset(env)
        if add_env != "None":
            add_data = d4rl.qlearning_dataset(add_env)
            if expert_ratio >= 1:
                raise ValueError("in the mix setting, the expert_ratio must < 1")
            length_add_data = int(add_data["rewards"].shape[0] * (1 - expert_ratio))
            length_expert_data = int(length_add_data * expert_ratio)
            for k, _ in dataset.items():
                dataset[k] = np.concatenate(
                    [
                        add_data[k][:-length_expert_data],
                        dataset[k][:length_expert_data],
                    ],
                    axis=0,
                )
            print("-------------------------------")
            print(
                f"we are in the mix data regimes, len(expert):{length_expert_data} | len(add_data): {length_add_data} | expert ratio: {expert_ratio}"
            )
            print("-------------------------------")

        if heavy_tail:
            dataset = d4rl.qlearning_dataset(
                env, heavy_tail=True, heavy_tail_higher=heavy_tail_higher
            )
        if clip_to_eps:
            lim = 1 - eps
            dataset["actions"] = np.clip(dataset["actions"], -lim, lim)

        dones_float = np.zeros_like(dataset["rewards"])

        # To get initial observations
        initial_observations = [dataset['observations'][0].astype(np.float32)]
        
        for i in range(len(dones_float) - 1):
            if (
                np.linalg.norm(
                    dataset["observations"][i + 1] - dataset["next_observations"][i]
                )
                > 1e-6
                or dataset["terminals"][i] == 1.0
            ):
                dones_float[i] = 1

                # add initial observation at next step of done signal
                if i < len(dones_float) - 1:
                    initial_observations.append(dataset['observations'][i+1].astype(np.float32))
            else:
                dones_float[i] = 0

        dones_float[-1] = 1

        super().__init__(
            dataset["observations"].astype(np.float32),
            initial_observations = np.array(initial_observations, dtype=np.float32),
            actions=dataset["actions"].astype(np.float32),
            rewards=dataset["rewards"].astype(np.float32),
            masks=1.0 - dataset["terminals"].astype(np.float32),
            dones_float=dones_float.astype(np.float32),
            next_observations=dataset["next_observations"].astype(np.float32),
            size=len(dataset["observations"]),
        )



class ReplayBuffer(Dataset):
    def __init__(
        self, observation_space: gym.spaces.Box, action_dim: int, capacity: int
    ):
        observations = np.empty(
            (capacity, *observation_space.shape), dtype=observation_space.dtype
        )
        actions = np.empty((capacity, action_dim), dtype=np.float32)
        rewards = np.empty((capacity,), dtype=np.float32)
        masks = np.empty((capacity,), dtype=np.float32)
        dones_float = np.empty((capacity,), dtype=np.float32)
        next_observations = np.empty(
            (capacity, *observation_space.shape), dtype=observation_space.dtype
        )
        super().__init__(
            observations=observations,
            actions=actions,
            rewards=rewards,
            masks=masks,
            dones_float=dones_float,
            next_observations=next_observations,
            size=0,
        )

        self.size = 0

        self.insert_index = 0
        self.capacity = capacity

    def initialize_with_dataset(self, dataset: Dataset, num_samples: Optional[int]):
        assert (
            self.insert_index == 0
        ), "Can insert a batch online in an empty replay buffer."

        dataset_size = len(dataset.observations)

        if num_samples is None:
            num_samples = dataset_size
        else:
            num_samples = min(dataset_size, num_samples)
        assert (
            self.capacity >= num_samples
        ), "Dataset cannot be larger than the replay buffer capacity."

        if num_samples < dataset_size:
            perm = np.random.permutation(dataset_size)
            indices = perm[:num_samples]
        else:
            indices = np.arange(num_samples)

        self.observations[:num_samples] = dataset.observations[indices]
        self.actions[:num_samples] = dataset.actions[indices]
        self.rewards[:num_samples] = dataset.rewards[indices]
        self.masks[:num_samples] = dataset.masks[indices]
        self.dones_float[:num_samples] = dataset.dones_float[indices]
        self.next_observations[:num_samples] = dataset.next_observations[indices]

        self.insert_index = num_samples
        self.size = num_samples

    def insert(
        self,
        observation: np.ndarray,
        action: np.ndarray,
        reward: float,
        mask: float,
        done_float: float,
        next_observation: np.ndarray,
    ):
        self.observations[self.insert_index] = observation
        self.actions[self.insert_index] = action
        self.rewards[self.insert_index] = reward
        self.masks[self.insert_index] = mask
        self.dones_float[self.insert_index] = done_float
        self.next_observations[self.insert_index] = next_observation

        self.insert_index = (self.insert_index + 1) % self.capacity
        self.size = min(self.size + 1, self.capacity)


def _gen_dir_name():
    now_str = datetime.now().strftime("%m-%d-%y_%H.%M.%S")
    rand_str = "".join(random.choices(string.ascii_lowercase, k=4))
    return f"{now_str}_{rand_str}"


class Log:
    def __init__(
        self,
        root_log_dir,
        cfg_dict,
        txt_filename="log.txt",
        csv_filename="progress.csv",
        cfg_filename="config.json",
        flush=True,
    ):
        self.dir = Path(root_log_dir) / _gen_dir_name()
        self.dir.mkdir(parents=True)
        self.txt_file = open(self.dir / txt_filename, "w")
        self.csv_file = None
        (self.dir / cfg_filename).write_text(json.dumps(cfg_dict))
        self.txt_filename = txt_filename
        self.csv_filename = csv_filename
        self.cfg_filename = cfg_filename
        self.flush = flush

    def write(self, message, end="\n"):
        now_str = datetime.now().strftime("%H:%M:%S")
        message = f"[{now_str}] " + message
        for f in [sys.stdout, self.txt_file]:
            print(message, end=end, file=f, flush=self.flush)

    def __call__(self, *args, **kwargs):
        self.write(*args, **kwargs)

    def row(self, dict):
        if self.csv_file is None:
            self.csv_file = open(self.dir / self.csv_filename, "w", newline="")
            self.csv_writer = csv.DictWriter(self.csv_file, list(dict.keys()))
            self.csv_writer.writeheader()

        self(str(dict))
        self.csv_writer.writerow(dict)
        if self.flush:
            self.csv_file.flush()

    def close(self):
        self.txt_file.close()
        if self.csv_file is not None:
            self.csv_file.close()