rebrac_pytorch.py

# source: https://github.com/tinkoff-ai/ReBRAC
# https://arxiv.org/abs/2305.09836

import math
import os
import uuid
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Dict, Sequence, Tuple, Union

import gymnasium as gym
import h5py
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import tyro
from tqdm.auto import trange

import wandb


@dataclass
class Config:
    # wandb params
    project: str = "ReBRAC"
    group: str = "rebrac-halfcheetah-medium-v2"
    name: str = "rebrac-new"
    # model params
    actor_learning_rate: float = 1e-3
    critic_learning_rate: float = 1e-3
    hidden_dim: int = 256
    actor_n_hiddens: int = 3
    critic_n_hiddens: int = 3
    gamma: float = 0.99
    tau: float = 0.005
    actor_bc_coef: float = 0.001
    critic_bc_coef: float = 0.01
    actor_ln: bool = False
    critic_ln: bool = True
    policy_noise: float = 0.2
    noise_clip: float = 0.5
    policy_freq: int = 2
    normalize_q: bool = True
    # training params
    env_name: str = "HalfCheetah-v5"
    dataset_path: str = "~/.d4rl/datasets/halfcheetah_medium-v2.hdf5"
    batch_size: int = 1024
    num_epochs: int = 1000
    num_updates_on_epoch: int = 1000
    normalize_reward: bool = False
    normalize_states: bool = False
    # evaluation params
    eval_episodes: int = 10
    eval_every: int = 5
    # general params
    train_seed: int = 0
    eval_seed: int = 42

    def __post_init__(self):
        self.name = (
            f"{self.name}-{Path(self.dataset_path).stem}-{str(uuid.uuid4())[:8]}"
        )


class DetActor(nn.Module):
    def __init__(
        self,
        state_dim: int,
        action_dim: int,
        hidden_dim: int = 256,
        layernorm: bool = True,
        n_hiddens: int = 3,
    ):
        super().__init__()
        self.layernorm = layernorm
        layers = []
        bound = math.sqrt(1 / state_dim)
        fc = nn.Linear(state_dim, hidden_dim)
        nn.init.uniform_(fc.weight, -bound, bound)
        nn.init.constant_(fc.bias, 0.1)
        layers.append(fc)
        layers.append(nn.ReLU())
        if layernorm:
            layers.append(nn.LayerNorm(hidden_dim))
        for _ in range(n_hiddens - 1):
            bound = math.sqrt(1 / hidden_dim)
            fc = nn.Linear(hidden_dim, hidden_dim)
            nn.init.uniform_(fc.weight, -bound, bound)
            nn.init.constant_(fc.bias, 0.1)
            layers.append(fc)
            layers.append(nn.ReLU())
            if layernorm:
                layers.append(nn.LayerNorm(hidden_dim))
        last_fc = nn.Linear(hidden_dim, action_dim)
        nn.init.uniform_(last_fc.weight, -1e-3, 1e-3)
        nn.init.uniform_(last_fc.bias, -1e-3, 1e-3)
        layers.append(last_fc)
        layers.append(nn.Tanh())
        self.net = nn.Sequential(*layers)

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        return self.net(state)


class Critic(nn.Module):
    def __init__(
        self,
        state_dim: int,
        action_dim: int,
        hidden_dim: int = 256,
        layernorm: bool = True,
        n_hiddens: int = 3,
    ):
        super().__init__()
        self.layernorm = layernorm
        input_dim = state_dim + action_dim
        layers = []
        bound = math.sqrt(1 / input_dim)
        fc = nn.Linear(input_dim, hidden_dim)
        nn.init.uniform_(fc.weight, -bound, bound)
        nn.init.constant_(fc.bias, 0.1)
        layers.append(fc)
        layers.append(nn.ReLU())
        if layernorm:
            layers.append(nn.LayerNorm(hidden_dim))
        for _ in range(n_hiddens - 1):
            bound = math.sqrt(1 / hidden_dim)
            fc = nn.Linear(hidden_dim, hidden_dim)
            nn.init.uniform_(fc.weight, -bound, bound)
            nn.init.constant_(fc.bias, 0.1)
            layers.append(fc)
            layers.append(nn.ReLU())
            if layernorm:
                layers.append(nn.LayerNorm(hidden_dim))
        last_fc = nn.Linear(hidden_dim, 1)
        nn.init.uniform_(last_fc.weight, -3e-3, 3e-3)
        nn.init.uniform_(last_fc.bias, -3e-3, 3e-3)
        layers.append(last_fc)
        self.net = nn.Sequential(*layers)

    def forward(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
        state_action = torch.cat([state, action], dim=-1)
        return self.net(state_action).squeeze(-1)


class EnsembleCritic(nn.Module):
    def __init__(
        self,
        state_dim: int,
        action_dim: int,
        hidden_dim: int = 256,
        num_critics: int = 10,
        layernorm: bool = True,
        n_hiddens: int = 3,
    ):
        super().__init__()
        self.num_critics = num_critics
        self.critics = nn.ModuleList(
            [
                Critic(state_dim, action_dim, hidden_dim, layernorm, n_hiddens)
                for _ in range(num_critics)
            ]
        )

    def forward(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
        return torch.stack([critic(state, action) for critic in self.critics], dim=0)


def qlearning_dataset(p: str) -> Dict:
    obs_data = []
    next_obs_data = []
    action_data = []
    next_action_data = []
    reward_data = []
    terminal_data = []
    timeout_data = []

    # The newer version of the dataset adds an explicit
    # timeouts field. Keep old method for backwards compatability.
    with h5py.File(Path(p).expanduser()) as f:
        dataset = {k: f[k][:] for k in f.keys() if isinstance(f[k], h5py.Dataset)}
        N = dataset["rewards"].shape[0]
        for i in trange(N - 1):
            obs = dataset["observations"][i].astype(np.float32)
            new_obs = dataset["next_observations"][i].astype(np.float32)
            action = dataset["actions"][i].astype(np.float32)
            new_action = dataset["actions"][i + 1].astype(np.float32)
            reward = dataset["rewards"][i].astype(np.float32)
            t1 = bool(dataset["terminals"][i])
            t2 = bool(dataset["timeouts"][i])

            if t1 or t2:
                continue

            obs_data.append(obs)
            next_obs_data.append(new_obs)
            action_data.append(action)
            next_action_data.append(new_action)
            reward_data.append(reward)
            terminal_data.append(t1)
            timeout_data.append(t2)

    return {
        "observations": np.array(obs_data),
        "actions": np.array(action_data),
        "next_observations": np.array(next_obs_data),
        "next_actions": np.array(next_action_data),
        "rewards": np.array(reward_data),
        "terminals": np.array(terminal_data),
        "timeouts": np.array(timeout_data),
    }


class ReplayBuffer:
    data: Dict[str, np.ndarray]
    mean: np.ndarray
    std: np.ndarray
    device: torch.device

    def __init__(self, device: torch.device):
        self.mean = np.array(0.0)
        self.std = np.array(1.0)
        self.device = device

    def create_from_trans(
        self,
        dataset_path: str,
        normalize_reward: bool = False,
        is_normalize: bool = False,
    ):
        d4rl_data = qlearning_dataset(dataset_path)
        buffer = {
            "states": torch.tensor(
                d4rl_data["observations"], dtype=torch.float32, device=self.device
            ),
            "actions": torch.tensor(
                d4rl_data["actions"], dtype=torch.float32, device=self.device
            ),
            "rewards": torch.tensor(
                d4rl_data["rewards"], dtype=torch.float32, device=self.device
            ),
            "next_states": torch.tensor(
                d4rl_data["next_observations"], dtype=torch.float32, device=self.device
            ),
            "next_actions": torch.tensor(
                d4rl_data["next_actions"], dtype=torch.float32, device=self.device
            ),
            "dones": torch.tensor(
                d4rl_data["terminals"], dtype=torch.float32, device=self.device
            ),
        }

        if is_normalize:
            self.mean, self.std = self.compute_mean_std(
                d4rl_data["observations"], eps=1e-3
            )
            buffer["states"] = torch.tensor(
                self.normalize_states(d4rl_data["observations"], self.mean, self.std),
                dtype=torch.float32,
                device=self.device,
            )
            buffer["next_states"] = torch.tensor(
                self.normalize_states(
                    d4rl_data["next_observations"], self.mean, self.std
                ),
                dtype=torch.float32,
                device=self.device,
            )

        if normalize_reward:
            buffer["rewards"] = self.normalize_reward(dataset_path, buffer["rewards"])

        self.data = buffer

    @staticmethod
    def compute_mean_std(
        states: np.ndarray, eps: float = 1e-3
    ) -> Tuple[np.ndarray, np.ndarray]:
        mean = states.mean(0)
        std = states.std(0) + eps
        return mean, std

    @property
    def size(self) -> int:
        return self.data["states"].shape[0]

    def sample_batch(self, batch_size: int) -> Dict[str, torch.Tensor]:
        indices = torch.randint(0, self.size, (batch_size,), device=self.device)
        batch = {key: value[indices] for key, value in self.data.items()}
        return batch

    @staticmethod
    def normalize_reward(dataset_name: str, rewards: torch.Tensor) -> torch.Tensor:
        if "antmaze" in dataset_name:
            return rewards * 100.0
        else:
            raise NotImplementedError(
                "Reward normalization is implemented only for AntMaze yet!"
            )

    @staticmethod
    def normalize_states(
        states: np.ndarray, mean: np.ndarray, std: np.ndarray
    ) -> np.ndarray:
        return (states - mean) / std


class Metrics:
    def __init__(self, metrics: Sequence[str]):
        self.accumulators: Dict[str, list] = {key: [0.0, 0] for key in metrics}

    def update(self, updates: Dict[str, float]):
        for key, value in updates.items():
            acc, steps = self.accumulators[key]
            self.accumulators[key] = [acc + value, steps + 1]

    def compute(self) -> Dict[str, float]:
        return {k: v[0] / v[1] for k, v in self.accumulators.items()}


def normalize(
    arr: torch.Tensor, mean: torch.Tensor, std: torch.Tensor, eps: float = 1e-8
) -> torch.Tensor:
    return (arr - mean) / (std + eps)


def make_eval_env(env_name: str, seed: int) -> gym.Env:
    env = gym.make(env_name)
    env.action_space.seed(seed)
    env.observation_space.seed(seed)
    return env


def wrap_env(
    env: gym.Env,
    state_mean: Union[np.ndarray, float] = 0.0,
    state_std: Union[np.ndarray, float] = 1.0,
    reward_scale: float = 1.0,
) -> gym.Env:
    # PEP 8: E731 do not assign a lambda expression, use a def
    def normalize_state(state: np.ndarray) -> np.ndarray:
        return (
            state - state_mean
        ) / state_std  # epsilon should be already added in std.

    def scale_reward(reward: float) -> float:
        # Please be careful, here reward is multiplied by scale!
        return reward_scale * reward

    env = gym.wrappers.TransformObservation(env, normalize_state, None)
    env = gym.wrappers.RescaleAction(env, -1.0, 1.0)
    if reward_scale != 1.0:
        env = gym.wrappers.TransformReward(env, scale_reward)
    return env


def evaluate(
    env: gym.Env,
    actor: nn.Module,
    num_episodes: int,
    seed: int,
) -> np.ndarray:
    env.action_space.seed(seed)
    env.observation_space.seed(seed)

    returns = []
    device = next(actor.parameters()).device
    for _ in trange(num_episodes, desc="Eval", leave=False):
        (obs, _), t1, t2 = env.reset(), False, False
        total_reward = 0.0
        while not (t1 or t2):
            obs_tensor = torch.from_numpy(obs).float().unsqueeze(0).to(device)
            with torch.no_grad():
                action = actor(obs_tensor).squeeze(0).cpu().numpy()
            obs, reward, t1, t2,_ = env.step(action)
            total_reward += reward
        returns.append(total_reward)

    return np.array(returns)


def soft_update(target: nn.Module, source: nn.Module, tau: float):
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau)


def update_actor(
    actor: nn.Module,
    target_actor: nn.Module,
    critic: nn.Module,
    target_critic: nn.Module,
    batch: Dict[str, torch.Tensor],
    beta: float,
    tau: float,
    normalize_q: bool,
    metrics: Metrics,
    actor_optimizer: optim.Optimizer,
):
    actor_optimizer.zero_grad()
    actions = actor(batch["states"])
    bc_penalty = ((actions - batch["actions"]) ** 2).sum(-1)
    q_values = torch.min(critic(batch["states"], actions), dim=0)[0]
    lmbda = 1.0
    if normalize_q:
        lmbda = (1 / torch.abs(q_values).mean()).detach()
    loss = (beta * bc_penalty - lmbda * q_values).mean()
    loss.backward()
    actor_optimizer.step()

    random_actions = torch.rand_like(batch["actions"]) * 2 - 1
    metrics.update(
        {
            "actor_loss": loss.item(),
            "bc_mse_policy": bc_penalty.mean().item(),
            "bc_mse_random": ((random_actions - batch["actions"]) ** 2)
            .sum(-1)
            .mean()
            .item(),
            "action_mse": ((actions - batch["actions"]) ** 2).mean().item(),
        }
    )

    soft_update(target_actor, actor, tau)
    soft_update(target_critic, critic, tau)


def update_critic(
    actor: nn.Module,
    target_actor: nn.Module,
    critic: nn.Module,
    target_critic: nn.Module,
    batch: Dict[str, torch.Tensor],
    gamma: float,
    beta: float,
    policy_noise: float,
    noise_clip: float,
    metrics: Metrics,
    critic_optimizer: optim.Optimizer,
):
    critic_optimizer.zero_grad()
    next_actions = target_actor(batch["next_states"])
    noise = torch.clamp(
        torch.randn_like(next_actions) * policy_noise, -noise_clip, noise_clip
    )
    next_actions = torch.clamp(next_actions + noise, -1, 1)
    bc_penalty = ((next_actions - batch["next_actions"]) ** 2).sum(-1)
    next_q = torch.min(target_critic(batch["next_states"], next_actions), dim=0)[0]
    next_q = next_q - beta * bc_penalty
    target_q = batch["rewards"] + (1 - batch["dones"]) * gamma * next_q
    q = critic(batch["states"], batch["actions"])
    loss = ((q - target_q.unsqueeze(0)) ** 2).mean(dim=1).sum()
    q_min = torch.min(q, dim=0)[0].mean()
    loss.backward()
    critic_optimizer.step()
    metrics.update(
        {
            "critic_loss": loss.item(),
            "q_min": q_min.item(),
        }
    )


def train(config: Config):
    dict_config = asdict(config)
    dict_config["mlc_job_name"] = os.environ.get("PLATFORM_JOB_NAME")

    wandb.init(
        config=dict_config,
        project=config.project,
        group=config.group,
        name=config.name,
        id=str(uuid.uuid4()),
    )
    wandb.mark_preempting()

    # use cuda if available
    if torch.cuda.is_available() is True:
        device = torch.device("cuda")
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    else:
        device = torch.device("cpu")

    buffer = ReplayBuffer(device)
    buffer.create_from_trans(
        config.dataset_path, config.normalize_reward, config.normalize_states
    )

    torch.manual_seed(config.train_seed)
    np.random.seed(config.train_seed)

    eval_env = make_eval_env(config.env_name, seed=config.eval_seed)
    eval_env = wrap_env(eval_env, buffer.mean, buffer.std)

    state_dim = buffer.data["states"].shape[-1]
    action_dim = buffer.data["actions"].shape[-1]

    actor = DetActor(
        state_dim,
        action_dim,
        hidden_dim=config.hidden_dim,
        layernorm=config.actor_ln,
        n_hiddens=config.actor_n_hiddens,
    ).to(device)
    target_actor = DetActor(
        state_dim,
        action_dim,
        hidden_dim=config.hidden_dim,
        layernorm=config.actor_ln,
        n_hiddens=config.actor_n_hiddens,
    ).to(device)
    target_actor.load_state_dict(actor.state_dict())

    critic = EnsembleCritic(
        state_dim,
        action_dim,
        hidden_dim=config.hidden_dim,
        num_critics=2,
        layernorm=config.critic_ln,
        n_hiddens=config.critic_n_hiddens,
    ).to(device)
    target_critic = EnsembleCritic(
        state_dim,
        action_dim,
        hidden_dim=config.hidden_dim,
        num_critics=2,
        layernorm=config.critic_ln,
        n_hiddens=config.critic_n_hiddens,
    ).to(device)
    target_critic.load_state_dict(critic.state_dict())

    actor_optimizer = optim.Adam(actor.parameters(), lr=config.actor_learning_rate)
    critic_optimizer = optim.Adam(critic.parameters(), lr=config.critic_learning_rate)

    # metrics
    bc_metrics_to_log = [
        "critic_loss",
        "q_min",
        "actor_loss",
        "bc_mse_policy",
        "bc_mse_random",
        "action_mse",
    ]

    for epoch in trange(config.num_epochs, desc="ReBRAC Epochs"):
        # metrics for accumulation during epoch and logging to wandb
        # we need to reset them every epoch
        metrics_obj = Metrics(bc_metrics_to_log)

        for step in range(config.num_updates_on_epoch):
            batch = buffer.sample_batch(config.batch_size)

            update_critic(
                actor,
                target_actor,
                critic,
                target_critic,
                batch,
                config.gamma,
                config.critic_bc_coef,
                config.policy_noise,
                config.noise_clip,
                metrics_obj,
                critic_optimizer,
            )

            if step % config.policy_freq == 0:
                update_actor(
                    actor,
                    target_actor,
                    critic,
                    target_critic,
                    batch,
                    config.actor_bc_coef,
                    config.tau,
                    config.normalize_q,
                    metrics_obj,
                    actor_optimizer,
                )

        # log mean over epoch for each metric
        mean_metrics = metrics_obj.compute()
        wandb.log(
            {"epoch": epoch, **{f"ReBRAC/{k}": v for k, v in mean_metrics.items()}}
        )

        if epoch % config.eval_every == 0 or epoch == config.num_epochs - 1:
            eval_returns = evaluate(
                eval_env,
                actor,
                config.eval_episodes,
                seed=config.eval_seed,
            )
            wandb.log(
                {
                    "epoch": epoch,
                    "eval/return_mean": np.mean(eval_returns),
                    "eval/return_std": np.std(eval_returns),
                }
            )


if __name__ == "__main__":
    cfg = tyro.cli(Config)
    train(cfg)