diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/classifier/qgpo_classifier.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/classifier/qgpo_classifier.py new file mode 100644 index 0000000000000000000000000000000000000000..8b581ab7ad84fb5b314a712f862cc5e9fe0e43c1 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/classifier/qgpo_classifier.py @@ -0,0 +1,77 @@ +from typing import Dict + +import torch +import torch.nn.functional as F + +from .base import BaseClassifier + + +class QGPOClassifier(BaseClassifier): + """ + QGPOClassifier defines logp(y | x, t) using negative MSE. + Assuming nn_classifier is a NN used to predict y through x and t, i.e, pred_y = nn_classifier(x, t), + logp is defined as - temperature * MSE(nn_classifier(x, t), y). + """ + + def loss(self, x: torch.Tensor, t: torch.Tensor, y: Dict[str, torch.Tensor]): + """ In-support Contrastive Energy Prediction Loss in https://arxiv.org/pdf/2304.12824 + + Args: + x: torch.Tensor, + Noisy support actions. Shape: (batch_size, K, act_dim) + t: torch.Tensor, + Diffusion time steps. Shape: (batch_size, ) + y: Dict[str, torch.Tensor], + { + "soft_label": torch.Tensor, + Softmax of Q values. Shape: (batch_size, K, 1) + "obs": torch.Tensor, + Observations. Shape: (batch_size, obs_dim)} + + Returns: + loss: torch.Tensor + """ + b, k = x.shape[:2] + + soft_label, obs = y["soft_label"], y["obs"] + + f = self.model(x, t.unsqueeze(1).repeat(1, k), obs.unsqueeze(1).repeat(1, k, 1)) + + loss = - (soft_label * F.log_softmax(f, 1)).sum(1).mean() + + with torch.no_grad(): + f_max = f.max(1)[0].mean().item() + f_mean = f.mean().item() + f_min = f.min(1)[0].mean().item() + + return loss, {"f_max": f_max, "f_mean": f_mean, "f_min": f_min} + + def update(self, x: torch.Tensor, noise: torch.Tensor, y: Dict[str, torch.Tensor], update_ema: bool = True): + loss, log = self.loss(x, noise, y) + self.optim.zero_grad() + loss.backward() + if isinstance(self.grad_clip_norm, float): + grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_norm).item() + else: + grad_norm = None + self.optim.step() + if update_ema: + self.ema_update() + log.update({"loss": loss.item(), "grad_norm": grad_norm}) + return log + + def logp(self, x: torch.Tensor, t: torch.Tensor, c: torch.Tensor): + """ + Args: + x: torch.Tensor, + Noisy actions. Shape (batch_size, act_dim) + t: torch.Tensor, + Timestep. Shape (batch_size, ) + c: torch.Tensor, + Observations. Shape (batch_size, obs_dim) + + Returns: + f: torch.Tensor, + Energy prediction. Shape (batch_size, 1) + """ + return self.model_ema(x, t, c) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/classifier/rew_classifiers.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/classifier/rew_classifiers.py new file mode 100644 index 0000000000000000000000000000000000000000..b57f6f36e4356c9031256d152c17697e52b22a84 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/classifier/rew_classifiers.py @@ -0,0 +1,29 @@ +from typing import Optional + +from cleandiffuser.nn_classifier import BaseNNClassifier +from .base import BaseClassifier + + +class CumRewClassifier(BaseClassifier): + def __init__( + self, + nn_classifier: BaseNNClassifier, + device: str = "cpu", + optim_params: Optional[dict] = None, + ): + super().__init__(nn_classifier, 0.995, None, optim_params, device) + + def loss(self, x, noise, R): + pred_R = self.model(x, noise, None) + return ((pred_R - R) ** 2).mean() + + def update(self, x, noise, R): + self.optim.zero_grad() + loss = self.loss(x, noise, R) + loss.backward() + self.optim.step() + self.ema_update() + return {"loss": loss.item()} + + def logp(self, x, noise, c=None): + return self.model_ema(x, noise) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/base_dataset.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/base_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..ffce5924afa0dce30f7fdd2057a8295446bccf63 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/base_dataset.py @@ -0,0 +1,41 @@ +from typing import Dict + +import torch.nn +from torch.utils.data import Dataset + + +# Adapted from the datasets on: https://github.com/real-stanford/diffusion_policy + +# Observation Horizon: To|n_obs_steps +# Action Horizon: Ta|n_action_steps +# Prediction Horizon: T|horizon +# To = 3 +# Ta = 4 +# T = 6 +# |o|o|o| +# | | |a|a|a|a| +# To = 2 +# Ta = 5 +# T = 6 +# |o|o| +# | |a|a|a|a|a| + + +class BaseDataset(Dataset): + def get_normalizer(self, **kwargs): + raise NotImplementedError() + + def __len__(self) -> int: + return 0 + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + """ + output: + obs: + key1: T, Do1 # default key: state + key2: T, Do2 + action: T, Da + reward: T, 1 + info: + """ + raise NotImplementedError() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/block_push_dataset.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/block_push_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..c8e698e1a3b09cfd9e456c76412cc14d62040b5e --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/block_push_dataset.py @@ -0,0 +1,76 @@ +from typing import Dict +import torch +import numpy as np +import copy +import pathlib +from cleandiffuser.dataset.base_dataset import BaseDataset +from cleandiffuser.dataset.replay_buffer import ReplayBuffer +from cleandiffuser.dataset.dataset_utils import SequenceSampler, MinMaxNormalizer, EmptyNormalizer, dict_apply + +# dev/block_pushing/multimodal_push_seed.zarr +# ├── data +# │ ├── action (114962, 2) float32 +# │ └── obs (114962, 16) float32 +# └── meta +# └── episode_ends (1000,) int64 + +class BlockPushDataset(BaseDataset): + def __init__(self, + zarr_path, + horizon=1, + pad_before=0, + pad_after=0, + obs_keys=['obs', 'action'] + ): + super().__init__() + self.replay_buffer = ReplayBuffer.copy_from_path( + zarr_path, keys=obs_keys) + + self.sampler = SequenceSampler( + replay_buffer=self.replay_buffer, + sequence_length=horizon, + pad_before=pad_before, + pad_after=pad_after) + + self.horizon = horizon + self.pad_before = pad_before + self.pad_after = pad_after + + self.normalizer = self.get_normalizer() + + def get_normalizer(self): + state_normalizer = MinMaxNormalizer(self.replay_buffer['obs'][:]) # (N, obs_dim) + action_normalizer = MinMaxNormalizer(self.replay_buffer['action'][:]) # (N, action_dim) + + return { + "obs": { + "state": state_normalizer + }, + "action": action_normalizer + } + + def __str__(self) -> str: + return f"Keys: {self.replay_buffer.keys()} Steps: {self.replay_buffer.n_steps} Episodes: {self.replay_buffer.n_episodes}" + + def __len__(self) -> int: + return len(self.sampler) + + def _sample_to_data(self, sample): + state = sample['obs'].astype(np.float32) # (T, 5) + state = self.normalizer['obs']['state'].normalize(state) + + action = sample['action'].astype(np.float32) # (T, 2) + action = self.normalizer['action'].normalize(action) + data = { + 'obs': { + 'state': state, # T, D_o + }, + 'action': action, # T, D_a + } + return data + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + sample = self.sampler.sample_sequence(idx) + data = self._sample_to_data(sample) + torch_data = dict_apply(data, torch.tensor) + return torch_data \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/d4rl_antmaze_dataset.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/d4rl_antmaze_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..28422bfc00f015039f241d808a522b355e580cd0 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/d4rl_antmaze_dataset.py @@ -0,0 +1,563 @@ +from typing import Dict + +import numpy as np +import torch + +from cleandiffuser.dataset.base_dataset import BaseDataset +from cleandiffuser.utils import GaussianNormalizer, dict_apply + + +class D4RLAntmazeDataset(BaseDataset): + """ **D4RL-Antmaze Sequential Dataset** + + torch.utils.data.Dataset wrapper for D4RL-Antmaze dataset. + Chunk the dataset into sequences of length `horizon` with obs-repeat/act-zero/reward-zero padding. + Use GaussianNormalizer to normalize the observations as default. + Each batch contains + - batch["obs"]["state"], observations of shape (batch_size, horizon, o_dim) + - batch["act"], actions of shape (batch_size, horizon, a_dim) + - batch["rew"], rewards of shape (batch_size, horizon, 1) + - batch["val"], Monte Carlo return of shape (batch_size, 1) + + Args: + dataset: Dict[str, np.ndarray], + D4RL-Antmaze dataset. Obtained by calling `env.get_dataset()`. + horizon: int, + Length of each sequence. Default is 1. + max_path_length: int, + Maximum length of the episodes. Default is 1001. + noreaching_penalty: float, + Penalty for not reaching the goal. Default is -100. + discount: float, + Discount factor. Default is 0.99. + + Examples: + >>> env = gym.make("antmaze-medium-play-v2") + >>> dataset = D4RLAntmazeDataset(env.get_dataset(), horizon=32) + >>> dataloader = DataLoader(dataset, batch_size=32, shuffle=True) + >>> batch = next(iter(dataloader)) + >>> obs = batch["obs"]["state"] # (32, 32, 29) + >>> act = batch["act"] # (32, 32, 8) + >>> rew = batch["rew"] # (32, 32, 1) + >>> val = batch["val"] # (32, 1) + + >>> normalizer = dataset.get_normalizer() + >>> obs = env.reset()[None, :] + >>> normed_obs = normalizer.normalize(obs) + >>> unnormed_obs = normalizer.unnormalize(normed_obs) + """ + def __init__( + self, + dataset: Dict[str, np.ndarray], + horizon: int = 1, + max_path_length: int = 1001, + noreaching_penalty: float = -100., + discount: float = 0.99, + ): + super().__init__() + + observations, actions, rewards, timeouts, terminals = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["timeouts"], + dataset["terminals"]) + rewards -= 1 # -1 for each step and 0 for reaching the goal + dones = np.logical_or(timeouts, terminals) + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + + self.horizon = horizon + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + + self.indices = [] + self.seq_obs, self.seq_act, self.seq_rew = [], [], [] + self.tml_and_not_timeout = [] + + self.path_lengths, ptr = [], 0 + path_idx = 0 + for i in range(timeouts.shape[0]): + + if i != 0 and ((dones[i - 1] and not dones[i]) or timeouts[i - 1]): + + path_length = i - ptr + self.path_lengths.append(path_length) + + if terminals[i] and not timeouts[i]: + self.tml_and_not_timeout.append([path_idx, i - ptr]) + + # 1. agent walks out of the goal + if path_length < max_path_length: + + _seq_obs = np.zeros((max_path_length, self.o_dim), dtype=np.float32) + _seq_act = np.zeros((max_path_length, self.a_dim), dtype=np.float32) + _seq_rew = np.zeros((max_path_length, 1), dtype=np.float32) + + _seq_obs[:i - ptr] = normed_observations[ptr:i] + _seq_act[:i - ptr] = actions[ptr:i] + _seq_rew[:i - ptr] = rewards[ptr:i][:, None] + + # repeat padding + _seq_obs[i - ptr:] = normed_observations[i] # repeat last state + _seq_act[i - ptr:] = 0 # repeat zero action + _seq_rew[i - ptr:] = 0 # repeat zero reward + + self.seq_obs.append(_seq_obs) + self.seq_act.append(_seq_act) + self.seq_rew.append(_seq_rew) + + # 2. agent never reaches the goal during the episode + elif path_length == max_path_length: + + self.seq_obs.append(normed_observations[ptr:i]) + self.seq_act.append(actions[ptr:i]) + self.seq_rew.append(rewards[ptr:i][:, None]) + + # panelty for not reaching the goal + self.seq_rew[-1][-1] = noreaching_penalty + + else: + raise ValueError(f"path_length: {path_length} > max_path_length: {max_path_length}") + + max_start = min(self.path_lengths[-1] - 1, max_path_length - horizon) + self.indices += [(path_idx, start, start + horizon) for start in range(max_start + 1)] + + ptr = i + path_idx += 1 + + self.seq_obs = np.array(self.seq_obs) + self.seq_act = np.array(self.seq_act) + self.seq_rew = np.array(self.seq_rew) + + self.seq_val = np.copy(self.seq_rew) + for i in range(max_path_length - 1): + self.seq_val[:, - 2 - i] = self.seq_rew[:, -2 - i] + discount * self.seq_val[:, -1 - i] + self.tml_and_not_timeout = np.array(self.tml_and_not_timeout, dtype=np.int64) + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return len(self.indices) + + def __getitem__(self, idx: int): + path_idx, start, end = self.indices[idx] + + data = { + 'obs': { + 'state': self.seq_obs[path_idx, start:end]}, + 'act': self.seq_act[path_idx, start:end], + 'rew': self.seq_rew[path_idx, start:end], + 'val': self.seq_val[path_idx, start], + } + + torch_data = dict_apply(data, torch.tensor) + + return torch_data + + +class D4RLAntmazeTDDataset(BaseDataset): + """ **D4RL-Antmaze Transition Dataset** + + torch.utils.data.Dataset wrapper for D4RL-Antmaze dataset. + Chunk the dataset into transitions. + Use GaussianNormalizer to normalize the observations as default. + Each batch contains + - batch["obs"]["state"], observation of shape (batch_size, o_dim) + - batch["next_obs"]["state"], next observation of shape (batch_size, o_dim) + - batch["act"], action of shape (batch_size, a_dim) + - batch["rew"], reward of shape (batch_size, 1) + - batch["tml"], terminal of shape (batch_size, 1) + + Args: + dataset: Dict[str, np.ndarray], + D4RL-MuJoCo TD dataset. Obtained by calling `d4rl.qlearning_dataset(env)`. + reward_tune: str, + Reward tuning method. Can be "iql", "cql", "antmaze", or "none". Default is "iql". + + Examples: + >>> env = gym.make("antmaze-medium-play-v2") + >>> dataset = D4RLAntmazeTDDataset(d4rl.qlearning_dataset(env)) + >>> dataloader = DataLoader(dataset, batch_size=32, shuffle=True) + >>> batch = next(iter(dataloader)) + >>> obs = batch["obs"]["state"] # (32, 29) + >>> act = batch["act"] # (32, 8) + >>> rew = batch["rew"] # (32, 1) + >>> tml = batch["tml"] # (32, 1) + >>> next_obs = batch["next_obs"]["state"] # (32, 29) + + >>> normalizer = dataset.get_normalizer() + >>> obs = env.reset()[None, :] + >>> normed_obs = normalizer.normalize(obs) + >>> unnormed_obs = normalizer.unnormalize(normed_obs) + """ + def __init__(self, dataset: Dict[str, np.ndarray], reward_tune: str = "iql"): + super().__init__() + + observations, actions, next_observations, rewards, terminals = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["next_observations"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["terminals"].astype(np.float32)) + + if reward_tune == "iql": + rewards = rewards - 1. + elif reward_tune == "cql": + rewards = (rewards - 0.5) * 4. + elif reward_tune == "antmaze": + rewards = (rewards - 0.25) * 2. + elif reward_tune == "none": + rewards = rewards + else: + raise ValueError(f"reward_tune: {reward_tune} is not supported.") + + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + normed_next_observations = self.normalizers["state"].normalize(next_observations) + + self.obs = torch.tensor(normed_observations) + self.act = torch.tensor(actions) + self.rew = torch.tensor(rewards)[:, None] + self.tml = torch.tensor(terminals)[:, None] + self.next_obs = torch.tensor(normed_next_observations) + + self.size = self.obs.shape[0] + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return self.size + + def __getitem__(self, idx: int): + + data = { + 'obs': { + 'state': self.obs[idx], }, + 'next_obs': { + 'state': self.next_obs[idx], }, + 'act': self.act[idx], + 'rew': self.rew[idx], + 'tml': self.tml[idx], } + + return data + + +class MultiHorizonD4RLAntmazeDataset(BaseDataset): + def __init__( + self, + dataset, + horizons=(10, 20), + max_path_length=1001, + noreaching_penalty=-100, + discount=0.99, + ): + super().__init__() + + observations, actions, rewards, timeouts, terminals = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["timeouts"], + dataset["terminals"]) + rewards -= 1 + dones = np.logical_or(timeouts, terminals) + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + + self.horizons = horizons + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + self.discount = discount ** np.arange(max_path_length, dtype=np.float32) + + self.indices = [[] for _ in range(len(horizons))] + self.seq_obs, self.seq_act, self.seq_rew = [], [], [] + + self.path_lengths, ptr = [], 0 + path_idx = 0 + for i in range(timeouts.shape[0]): + + if i != 0 and ((dones[i - 1] and not dones[i]) or timeouts[i - 1]): + + path_length = i - ptr + self.path_lengths.append(path_length) + + # 1. agent walks out of the goal + if path_length < max_path_length: + + _seq_obs = np.zeros((max_path_length, self.o_dim), dtype=np.float32) + _seq_act = np.zeros((max_path_length, self.a_dim), dtype=np.float32) + _seq_rew = np.zeros((max_path_length, 1), dtype=np.float32) + + _seq_obs[:i - ptr] = normed_observations[ptr:i] + _seq_act[:i - ptr] = actions[ptr:i] + _seq_rew[:i - ptr] = rewards[ptr:i][:, None] + + # repeat padding + _seq_obs[i - ptr:] = normed_observations[i] # repeat last state + _seq_act[i - ptr:] = 0 # repeat zero action + _seq_rew[i - ptr:] = 0 # repeat zero reward + + self.seq_obs.append(_seq_obs) + self.seq_act.append(_seq_act) + self.seq_rew.append(_seq_rew) + + # 2. agent never reaches the goal during the episode + elif path_length == max_path_length: + + self.seq_obs.append(normed_observations[ptr:i]) + self.seq_act.append(actions[ptr:i]) + self.seq_rew.append(rewards[ptr:i][:, None]) + + # panelty for not reaching the goal + self.seq_rew[-1][-1] = noreaching_penalty + + else: + raise ValueError(f"path_length: {path_length} > max_path_length: {max_path_length}") + + max_starts = [min(self.path_lengths[-1] - 1, max_path_length - horizon) for horizon in horizons] + for k in range(len(horizons)): + self.indices[k] += [(path_idx, start, start + horizons[k]) for start in range(max_starts[k] + 1)] + + ptr = i + path_idx += 1 + + self.seq_obs = np.array(self.seq_obs) + self.seq_act = np.array(self.seq_act) + self.seq_rew = np.array(self.seq_rew) + + self.len_each_horizon = [len(indices) for indices in self.indices] + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return max(self.len_each_horizon) + + def __getitem__(self, idx: int): + + indices = [ + int(self.len_each_horizon[i] * (idx / self.len_each_horizon[-1])) for i in range(len(self.horizons))] + + torch_datas = [] + + for i, horizon in enumerate(self.horizons): + + path_idx, start, end = self.indices[i][indices[i]] + + rewards = self.seq_rew[path_idx, start:] + values = (rewards * self.discount[:rewards.shape[0], None]).sum(0) + + data = { + 'obs': { + 'state': self.seq_obs[path_idx, start:end]}, + 'act': self.seq_act[path_idx, start:end], + 'rew': self.seq_rew[path_idx, start:end], + 'val': values} + + torch_data = dict_apply(data, torch.tensor) + + torch_datas.append({ + "horizon": horizon, + "data": torch_data, + }) + + return torch_datas + +class DV_D4RLAntmazeSeqDataset(BaseDataset): + """ **D4RL-Antmaze Sequential Dataset** + + torch.utils.data.Dataset wrapper for D4RL-Antmaze dataset. + Chunk the dataset into sequences of length `horizon` with obs-repeat/act-zero/reward-zero padding. + Use GaussianNormalizer to normalize the observations as default. + Each batch contains + - batch["obs"]["state"], observations of shape (batch_size, horizon, o_dim) + - batch["act"], actions of shape (batch_size, horizon, a_dim) + - batch["rew"], rewards of shape (batch_size, horizon, 1) + - batch["val"], Monte Carlo return of shape (batch_size, 1) + + Args: + dataset: Dict[str, np.ndarray], + D4RL-Antmaze dataset. Obtained by calling `env.get_dataset()`. + horizon: int, + Length of each sequence. Default is 1. + max_path_length: int, + Maximum length of the episodes. Default is 1001. + noreaching_penalty: float, + Penalty for not reaching the goal. Default is -100. + discount: float, + Discount factor. Default is 0.99. + + Examples: + >>> env = gym.make("antmaze-medium-play-v2") + >>> dataset = D4RLAntmazeDataset(env.get_dataset(), horizon=32) + >>> dataloader = DataLoader(dataset, batch_size=32, shuffle=True) + >>> batch = next(iter(dataloader)) + >>> obs = batch["obs"]["state"] # (32, 32, 29) + >>> act = batch["act"] # (32, 32, 8) + >>> rew = batch["rew"] # (32, 32, 1) + >>> val = batch["val"] # (32, 1) + + >>> normalizer = dataset.get_normalizer() + >>> obs = env.reset()[None, :] + >>> normed_obs = normalizer.normalize(obs) + >>> unnormed_obs = normalizer.unnormalize(normed_obs) + """ + def __init__( + self, + dataset: Dict[str, np.ndarray], + horizon: int = 1, + max_path_length: int = 1001, + discount: float = 0.99, + continous_reward_at_done: bool = False, + reward_tune: str = "iql", + center_mapping: bool = True, + learn_policy: bool = False, + stride: int = 1, + only_learn_reached_policy: bool = False, + ): + super().__init__() + + observations, actions, rewards, timeouts, terminals = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["timeouts"].astype(np.float32), + dataset["terminals"].astype(np.float32)) + + self.learn_policy = learn_policy + self.stride = stride + + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + + self.horizon = horizon + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + + self.indices = [] + self.seq_obs, self.seq_act, self.seq_rew, self.seq_tml = [], [], [], [] + + ptr = 0 # ptr: t=0 pos + path_idx = 0 + + timeouts_indices = np.where(timeouts == 1)[0] + + for index in timeouts_indices: + terminal_index = -1 + assert index - ptr + 1 == max_path_length + for i in range(ptr, index + 1): + if terminals[i]: + terminal_index = i + break + + if terminal_index == -1 and self.learn_policy and not only_learn_reached_policy: + path_length = index - ptr + 1 + assert path_length == max_path_length + + _seq_obs = np.zeros((max_path_length + (horizon - 1) * stride, self.o_dim), dtype=np.float32) + _seq_act = np.zeros((max_path_length + (horizon - 1) * stride, self.a_dim), dtype=np.float32) + _seq_rew = np.zeros((max_path_length + (horizon - 1) * stride, 1), dtype=np.float32) + _seq_tml = np.zeros((max_path_length + (horizon - 1) * stride, 1), dtype=np.float32) + + _seq_obs[:path_length] = normed_observations[ptr:index+1] + _seq_act[:path_length] = actions[ptr:index+1] + _seq_rew[:path_length] = rewards[ptr:index+1][:, None] + _seq_tml[:path_length] = terminals[ptr:index+1][:, None] + + self.seq_obs.append(_seq_obs) + self.seq_act.append(_seq_act) + self.seq_rew.append(_seq_rew) + self.seq_tml.append(_seq_tml) + + max_start = max_path_length - (horizon - 1) * stride - 1 + self.indices += [(path_idx, start, start + (horizon - 1) * stride + 1) for start in range(max_start + 1)] + path_idx += 1 + + elif terminal_index != -1: + path_length = terminal_index - ptr + 1 + assert path_length <= max_path_length + assert rewards[terminal_index] == 1 + + _seq_obs = np.zeros((max_path_length + (horizon - 1) * stride, self.o_dim), dtype=np.float32) + _seq_act = np.zeros((max_path_length + (horizon - 1) * stride, self.a_dim), dtype=np.float32) + _seq_rew = np.zeros((max_path_length + (horizon - 1) * stride, 1), dtype=np.float32) + _seq_tml = np.zeros((max_path_length + (horizon - 1) * stride, 1), dtype=np.float32) + + _seq_obs[:path_length] = normed_observations[ptr:terminal_index+1] + _seq_act[:path_length] = actions[ptr:terminal_index+1] + _seq_rew[:path_length] = rewards[ptr:terminal_index+1][:, None] + _seq_tml[:path_length] = terminals[ptr:terminal_index+1][:, None] + + _seq_obs[path_length:] = normed_observations[terminal_index] # repeat state + _seq_act[path_length:] = 0 # zero action + _seq_rew[path_length:] = 1 if continous_reward_at_done else 0 + _seq_tml[path_length:] = 1 + + self.seq_obs.append(_seq_obs) + self.seq_act.append(_seq_act) + self.seq_rew.append(_seq_rew) + self.seq_tml.append(_seq_tml) + + max_start = path_length - 1 + self.indices += [(path_idx, start, start + (horizon - 1) * stride + 1) for start in range(max_start + 1)] + path_idx += 1 + + ptr = index + 1 + + self.seq_obs = np.array(self.seq_obs) + self.seq_act = np.array(self.seq_act) + self.seq_rew = np.array(self.seq_rew) + self.seq_tml = np.array(self.seq_tml) + + if reward_tune == "iql": + self.seq_rew += -1 + elif reward_tune == "none": + self.seq_rew = self.seq_rew + else: + raise ValueError(f"reward_tune: {reward_tune} is not supported.") + + self.seq_val = np.copy(self.seq_rew) + for i in reversed(range(max_path_length - 1)): + self.seq_val[:, i] = self.seq_rew[:, i] + discount * self.seq_val[:, i+1] + + print(f"max discounted return: {self.seq_val.max()}") + print(f"min discounted return: {self.seq_val.min()}") + + # val \in [-1, 1] + self.seq_val = (self.seq_val - self.seq_val.min()) / (self.seq_val.max() - self.seq_val.min()) + if center_mapping: + self.seq_val = self.seq_val * 2 - 1 + print(f"max normed discounted return: {self.seq_val.max()}") + print(f"min normed discounted return: {self.seq_val.min()}") + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return len(self.indices) + + def __getitem__(self, idx: int): + path_idx, start, end = self.indices[idx] + + if self.learn_policy: + horizon_state = self.seq_obs[path_idx, start:end:self.stride].copy() + horizon_state[:, :2] -= horizon_state[0, :2] + else: + horizon_state = self.seq_obs[path_idx, start:end:self.stride] + + data = { + 'obs': {'state': horizon_state}, + 'act': self.seq_act[path_idx, start:end:self.stride], + 'rew': self.seq_rew[path_idx, start:end:self.stride], + 'val': self.seq_val[path_idx, start], + 'tml': self.seq_tml[path_idx, start], + } + + torch_data = dict_apply(data, torch.tensor) + + return torch_data \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/d4rl_kitchen_dataset.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/d4rl_kitchen_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..b5dae5b31b92e020020d6f4501301e3798e7ab31 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/d4rl_kitchen_dataset.py @@ -0,0 +1,444 @@ +from typing import Dict + +import numpy as np +import torch + +from cleandiffuser.dataset.base_dataset import BaseDataset +from cleandiffuser.utils import GaussianNormalizer, dict_apply + + +class D4RLKitchenDataset(BaseDataset): + """ **D4RL-Kitchen Sequential Dataset** + + torch.utils.data.Dataset wrapper for D4RL-Kitchen dataset. + Chunk the dataset into sequences of length `horizon` with obs-repeat/act-zero/reward-repeat padding. + Use GaussianNormalizer to normalize the observations as default. + Each batch contains + - batch["obs"]["state"], observations of shape (batch_size, horizon, o_dim) + - batch["act"], actions of shape (batch_size, horizon, a_dim) + - batch["rew"], rewards of shape (batch_size, horizon, 1) + - batch["val"], Monte Carlo return of shape (batch_size, 1) + + Args: + dataset: Dict[str, np.ndarray], + D4RL-Kitchen dataset. Obtained by calling `env.get_dataset()`. + horizon: int, + Length of each sequence. Default is 1. + max_path_length: int, + Maximum length of the episodes. Default is 280. + discount: float, + Discount factor. Default is 0.99. + + Examples: + >>> env = gym.make("kitchen-mixed-v0") + >>> dataset = D4RLKitchenDataset(env.get_dataset(), horizon=32) + >>> dataloader = DataLoader(dataset, batch_size=32, shuffle=True) + >>> batch = next(iter(dataloader)) + >>> obs = batch["obs"]["state"] # (32, 32, 60) + >>> act = batch["act"] # (32, 32, 9) + >>> rew = batch["rew"] # (32, 32, 1) + >>> val = batch["val"] # (32, 1) + + >>> normalizer = dataset.get_normalizer() + >>> obs = env.reset()[None, :] + >>> normed_obs = normalizer.normalize(obs) + >>> unnormed_obs = normalizer.unnormalize(normed_obs) + """ + def __init__( + self, + dataset: Dict[str, np.ndarray], + horizon: int = 1, + max_path_length: int = 280, + discount: float = 0.99, + ): + super().__init__() + + observations, actions, rewards, timeouts, terminals = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["timeouts"], + dataset["terminals"]) + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + + self.horizon = horizon + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + + self.indices = [] + self.seq_obs, self.seq_act, self.seq_rew = [], [], [] + self.tml_and_not_timeout = [] + + self.path_lengths, ptr = [], 0 + path_idx = 0 + for i in range(timeouts.shape[0]): + if timeouts[i] or terminals[i] or i == timeouts.shape[0] - 1: + self.path_lengths.append(i - ptr + 1) + + if terminals[i] and not timeouts[i]: + self.tml_and_not_timeout.append([path_idx, i - ptr]) + + _seq_obs = np.zeros((max_path_length, self.o_dim), dtype=np.float32) + _seq_act = np.zeros((max_path_length, self.a_dim), dtype=np.float32) + _seq_rew = np.zeros((max_path_length, 1), dtype=np.float32) + + _seq_obs[:i - ptr + 1] = normed_observations[ptr:i + 1] + _seq_act[:i - ptr + 1] = actions[ptr:i + 1] + _seq_rew[:i - ptr + 1] = rewards[ptr:i + 1][:, None] + + # repeat padding + _seq_obs[i - ptr + 1:] = normed_observations[i] # repeat last state + _seq_act[i - ptr + 1:] = 0 # repeat zero action + _seq_rew[i - ptr + 1:] = rewards[i] # repeat last reward + + self.seq_obs.append(_seq_obs) + self.seq_act.append(_seq_act) + self.seq_rew.append(_seq_rew) + + max_start = min(self.path_lengths[-1] - 1, max_path_length - horizon) + self.indices += [(path_idx, start, start + horizon) for start in range(max_start + 1)] + + ptr = i + 1 + path_idx += 1 + + self.seq_obs = np.array(self.seq_obs) + self.seq_act = np.array(self.seq_act) + self.seq_rew = np.array(self.seq_rew) + + self.seq_val = np.copy(self.seq_rew) + for i in range(max_path_length - 1): + self.seq_val[:, - 2 - i] = self.seq_rew[:, -2 - i] + discount * self.seq_val[:, -1 - i] + self.tml_and_not_timeout = np.array(self.tml_and_not_timeout, dtype=np.int64) + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return len(self.indices) + + def __getitem__(self, idx: int): + path_idx, start, end = self.indices[idx] + + data = { + 'obs': { + 'state': self.seq_obs[path_idx, start:end]}, + 'act': self.seq_act[path_idx, start:end], + 'rew': self.seq_rew[path_idx, start:end], + 'val': self.seq_val[path_idx, start], + } + + torch_data = dict_apply(data, torch.tensor) + + return torch_data + + +class D4RLKitchenTDDataset(BaseDataset): + """ **D4RL-Kitchen Transition Dataset** + + torch.utils.data.Dataset wrapper for D4RL-Kitchen dataset. + Chunk the dataset into transitions. + Use GaussianNormalizer to normalize the observations as default. + Each batch contains + - batch["obs"]["state"], observation of shape (batch_size, o_dim) + - batch["next_obs"]["state"], next observation of shape (batch_size, o_dim) + - batch["act"], action of shape (batch_size, a_dim) + - batch["rew"], reward of shape (batch_size, 1) + - batch["tml"], terminal of shape (batch_size, 1) + + Args: + dataset: Dict[str, np.ndarray], + D4RL-MuJoCo TD dataset. Obtained by calling `d4rl.qlearning_dataset(env)`. + + Examples: + >>> env = gym.make("kitchen-mixed-v0") + >>> dataset = D4RLKitchenTDDataset(d4rl.qlearning_dataset(env)) + >>> dataloader = DataLoader(dataset, batch_size=32, shuffle=True) + >>> batch = next(iter(dataloader)) + >>> obs = batch["obs"]["state"] # (32, 60) + >>> act = batch["act"] # (32, 9) + >>> rew = batch["rew"] # (32, 1) + >>> tml = batch["tml"] # (32, 1) + >>> next_obs = batch["next_obs"]["state"] # (32, 60) + + >>> normalizer = dataset.get_normalizer() + >>> obs = env.reset()[None, :] + >>> normed_obs = normalizer.normalize(obs) + >>> unnormed_obs = normalizer.unnormalize(normed_obs) + """ + def __init__(self, dataset: Dict[str, np.ndarray]): + super().__init__() + + observations, actions, next_observations, rewards, terminals = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["next_observations"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["terminals"].astype(np.float32)) + + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + normed_next_observations = self.normalizers["state"].normalize(next_observations) + + self.obs = torch.tensor(normed_observations) + self.act = torch.tensor(actions) + self.rew = torch.tensor(rewards)[:, None] + self.tml = torch.tensor(terminals)[:, None] + self.next_obs = torch.tensor(normed_next_observations) + + self.size = self.obs.shape[0] + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return self.size + + def __getitem__(self, idx: int): + data = { + 'obs': { + 'state': self.obs[idx], }, + 'next_obs': { + 'state': self.next_obs[idx], }, + 'act': self.act[idx], + 'rew': self.rew[idx], + 'tml': self.tml[idx], } + + return data + + +class MultiHorizonD4RLKitchenDataset(BaseDataset): + def __init__( + self, + dataset, + horizons=(10, 20), + max_path_length=280, + discount=0.99, + ): + super().__init__() + + observations, actions, rewards, timeouts, terminals = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["timeouts"], + dataset["terminals"]) + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + + self.horizons = horizons + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + self.discount = discount ** np.arange(max_path_length, dtype=np.float32) + + self.indices = [[] for _ in range(len(horizons))] + self.seq_obs, self.seq_act, self.seq_rew = [], [], [] + + self.path_lengths, ptr = [], 0 + path_idx = 0 + for i in range(timeouts.shape[0]): + if timeouts[i] or terminals[i] or i == timeouts.shape[0] - 1: + self.path_lengths.append(i - ptr + 1) + + _seq_obs = np.zeros((max_path_length, self.o_dim), dtype=np.float32) + _seq_act = np.zeros((max_path_length, self.a_dim), dtype=np.float32) + _seq_rew = np.zeros((max_path_length, 1), dtype=np.float32) + + _seq_obs[:i - ptr + 1] = normed_observations[ptr:i + 1] + _seq_act[:i - ptr + 1] = actions[ptr:i + 1] + _seq_rew[:i - ptr + 1] = rewards[ptr:i + 1][:, None] + + # repeat padding + _seq_obs[i - ptr + 1:] = normed_observations[i] # repeat last state + _seq_act[i - ptr + 1:] = 0 # repeat zero action + _seq_rew[i - ptr + 1:] = rewards[i] # repeat last reward + + self.seq_obs.append(_seq_obs) + self.seq_act.append(_seq_act) + self.seq_rew.append(_seq_rew) + + max_starts = [min(self.path_lengths[-1] - 1, max_path_length - horizon) for horizon in horizons] + for k in range(len(horizons)): + self.indices[k] += [(path_idx, start, start + horizons[k]) for start in range(max_starts[k] + 1)] + + ptr = i + 1 + path_idx += 1 + + self.seq_obs = np.array(self.seq_obs) + self.seq_act = np.array(self.seq_act) + self.seq_rew = np.array(self.seq_rew) + + self.len_each_horizon = [len(indices) for indices in self.indices] + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return max(self.len_each_horizon) + + def __getitem__(self, idx: int): + + indices = [ + int(self.len_each_horizon[i] * (idx / self.len_each_horizon[-1])) for i in range(len(self.horizons))] + + torch_datas = [] + + for i, horizon in enumerate(self.horizons): + path_idx, start, end = self.indices[i][indices[i]] + + rewards = self.seq_rew[path_idx, start:] + values = (rewards * self.discount[:rewards.shape[0], None]).sum(0) + + data = { + 'obs': { + 'state': self.seq_obs[path_idx, start:end]}, + 'act': self.seq_act[path_idx, start:end], + 'rew': self.seq_rew[path_idx, start:end], + 'val': values} + + torch_data = dict_apply(data, torch.tensor) + + torch_datas.append({ + "horizon": horizon, + "data": torch_data, + }) + + return torch_datas + +class DV_D4RLKitchenSeqDataset(BaseDataset): + """ **D4RL-Kitchen Sequential Dataset** + + torch.utils.data.Dataset wrapper for D4RL-Kitchen dataset. + Chunk the dataset into sequences of length `horizon` with obs-repeat/act-zero/reward-repeat padding. + Use GaussianNormalizer to normalize the observations as default. + Each batch contains + - batch["obs"]["state"], observations of shape (batch_size, horizon, o_dim) + - batch["act"], actions of shape (batch_size, horizon, a_dim) + - batch["rew"], rewards of shape (batch_size, horizon, 1) + - batch["val"], Monte Carlo return of shape (batch_size, 1) + + Args: + dataset: Dict[str, np.ndarray], + D4RL-Kitchen dataset. Obtained by calling `env.get_dataset()`. + horizon: int, + Length of each sequence. Default is 1. + max_path_length: int, + Maximum length of the episodes. Default is 280. + discount: float, + Discount factor. Default is 0.99. + + Examples: + >>> env = gym.make("kitchen-mixed-v0") + >>> dataset = D4RLKitchenDataset(env.get_dataset(), horizon=32) + >>> dataloader = DataLoader(dataset, batch_size=32, shuffle=True) + >>> batch = next(iter(dataloader)) + >>> obs = batch["obs"]["state"] # (32, 32, 60) + >>> act = batch["act"] # (32, 32, 9) + >>> rew = batch["rew"] # (32, 32, 1) + >>> val = batch["val"] # (32, 1) + + >>> normalizer = dataset.get_normalizer() + >>> obs = env.reset()[None, :] + >>> normed_obs = normalizer.normalize(obs) + >>> unnormed_obs = normalizer.unnormalize(normed_obs) + """ + def __init__( + self, + dataset: Dict[str, np.ndarray], + horizon: int = 1, + max_path_length: int = 280, + discount: float = 0.99, + center_mapping: bool = True, + stride: int = 1, + ): + super().__init__() + + observations, actions, rewards, timeouts, terminals = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["timeouts"].astype(np.float32), + dataset["terminals"].astype(np.float32)) + self.stride = stride + + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + + self.horizon = horizon + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + + self.indices = [] + self.seq_obs, self.seq_act, self.seq_rew = [], [], [] + + ptr = 0 + path_idx = 0 + for i in range(timeouts.shape[0]): + if timeouts[i] or terminals[i] or i == timeouts.shape[0] - 1: + path_length = i - ptr + 1 + assert path_length <= max_path_length + + _seq_obs = np.zeros((max_path_length + (horizon - 1) * stride, self.o_dim), dtype=np.float32) + _seq_act = np.zeros((max_path_length + (horizon - 1) * stride, self.a_dim), dtype=np.float32) + _seq_rew = np.zeros((max_path_length + (horizon - 1) * stride, 1), dtype=np.float32) + + _seq_obs[:path_length] = normed_observations[ptr:i + 1] + _seq_act[:path_length] = actions[ptr:i + 1] + _seq_rew[:path_length] = rewards[ptr:i + 1][:, None] + + # repeat padding + _seq_obs[path_length:] = normed_observations[i] # repeat last state + _seq_act[path_length:] = 0 + _seq_rew[path_length:] = rewards[i] + + self.seq_obs.append(_seq_obs) + self.seq_act.append(_seq_act) + self.seq_rew.append(_seq_rew) + + max_start = path_length - 1 + self.indices += [(path_idx, start, start + (horizon - 1) * stride + 1) for start in range(max_start + 1)] + + ptr = i + 1 + path_idx += 1 + + self.seq_obs = np.array(self.seq_obs) + self.seq_act = np.array(self.seq_act) + self.seq_rew = np.array(self.seq_rew) + + self.seq_val = np.copy(self.seq_rew) + for i in reversed(range(max_path_length - 1)): + self.seq_val[:, i] = self.seq_rew[:, i] + discount * self.seq_val[:, i+1] + + print(f"max discounted return: {self.seq_val.max()}") + print(f"min discounted return: {self.seq_val.min()}") + + # val \in [-1, 1] + self.seq_val = (self.seq_val - self.seq_val.min()) / (self.seq_val.max() - self.seq_val.min()) + if center_mapping: + self.seq_val = self.seq_val * 2 - 1 + print(f"max normed discounted return: {self.seq_val.max()}") + print(f"min normed discounted return: {self.seq_val.min()}") + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return len(self.indices) + + def __getitem__(self, idx: int): + path_idx, start, end = self.indices[idx] + + horizon_state = self.seq_obs[path_idx, start:end:self.stride] + + data = { + 'obs': {'state': horizon_state}, + 'act': self.seq_act[path_idx, start:end:self.stride], + 'rew': self.seq_rew[path_idx, start:end:self.stride], + 'val': self.seq_val[path_idx, start], + } + + torch_data = dict_apply(data, torch.tensor) + + return torch_data \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/d4rl_maze2d_dataset.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/d4rl_maze2d_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..6a29505b87fb43e07a60512ba3784e48f8e64c6d --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/d4rl_maze2d_dataset.py @@ -0,0 +1,290 @@ +from typing import Dict + +import numpy as np +import torch + +from cleandiffuser.dataset.base_dataset import BaseDataset +from cleandiffuser.utils import GaussianNormalizer, dict_apply + +class DV_D4RLMaze2DSeqDataset(BaseDataset): + """ **D4RL-Maze2D Sequential Dataset** + + torch.utils.data.Dataset wrapper for D4RL-Maze2D dataset. + Chunk the dataset into sequences of length `horizon` with obs-repeat/act-zero/reward-zero padding. + Use GaussianNormalizer to normalize the observations as default. + Each batch contains + - batch["obs"]["state"], observations of shape (batch_size, horizon, o_dim) + - batch["act"], actions of shape (batch_size, horizon, a_dim) + - batch["rew"], rewards of shape (batch_size, horizon, 1) + - batch["val"], Monte Carlo return of shape (batch_size, 1) + + Args: + dataset: Dict[str, np.ndarray], + D4RL-Maze2D dataset. Obtained by calling `env.get_dataset()`. + horizon: int, + Length of each sequence. Default is 1. + max_path_length: int, + Maximum length of the episodes. Default is 1001. + noreaching_penalty: float, + Penalty for not reaching the goal. Default is -100. + discount: float, + Discount factor. Default is 0.99. + + Examples: + >>> env = gym.make("Maze2D-medium-play-v2") + >>> dataset = D4RLMaze2DDataset(env.get_dataset(), horizon=32) + >>> dataloader = DataLoader(dataset, batch_size=32, shuffle=True) + >>> batch = next(iter(dataloader)) + >>> obs = batch["obs"]["state"] # (32, 32, 29) + >>> act = batch["act"] # (32, 32, 8) + >>> rew = batch["rew"] # (32, 32, 1) + >>> val = batch["val"] # (32, 1) + + >>> normalizer = dataset.get_normalizer() + >>> obs = env.reset()[None, :] + >>> normed_obs = normalizer.normalize(obs) + >>> unnormed_obs = normalizer.unnormalize(normed_obs) + """ + def __init__( + self, + dataset: Dict[str, np.ndarray], + horizon: int = 1, + max_path_length: int = 800, + discount: float = 0.99, + continous_reward_at_done: bool = False, + reward_tune: str = "iql", + center_mapping: bool = True, + learn_policy: bool = False, + stride: int = 1, + ): + super().__init__() + + self.max_path_length = max_path_length + + observations, actions, rewards, timeouts = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["timeouts"].astype(np.float32)) + self.learn_policy = learn_policy + self.stride = stride + + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + + self.horizon = horizon + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + + self.indices = [] + self.seq_obs, self.seq_act, self.seq_rew = [], [], [] + + path_idx = 0 + + self.paths = [] + next_end = [-1] * (timeouts.shape[0] + 1) + next_start = [-1] * (timeouts.shape[0] + 1) + for index in reversed(range(timeouts.shape[0])): + if rewards[index] == 1.0: + next_end[index] = index + next_start[index] = next_start[index+1] + else: + next_end[index] = next_end[index+1] + next_start[index] = index + + path_start = next_start[0] + path_end = next_end[path_start] + + if self.learn_policy: + for path_start in range(0, timeouts.shape[0], max_path_length): + path_end = min(path_start + max_path_length - 1, timeouts.shape[0] - 1) + path_length = path_end - path_start + 1 + + _seq_obs = np.zeros((max_path_length + (horizon - 1) * stride, self.o_dim), dtype=np.float32) + _seq_act = np.zeros((max_path_length + (horizon - 1) * stride, self.a_dim), dtype=np.float32) + _seq_rew = np.zeros((max_path_length + (horizon - 1) * stride, 1), dtype=np.float32) + + _seq_obs[:path_length] = normed_observations[path_start:path_end+1] + _seq_act[:path_length] = actions[path_start:path_end+1] + _seq_rew[:path_length] = rewards[path_start:path_end+1][:, None] + + _seq_obs[path_length:] = normed_observations[path_end] # repeat state + _seq_act[path_length:] = 0 # zero action + _seq_rew[path_length:] = 1 if continous_reward_at_done else 0 + + self.seq_obs.append(_seq_obs) + self.seq_act.append(_seq_act) + self.seq_rew.append(_seq_rew) + + max_start = path_length - 1 + self.indices += [(path_idx, start, start + (horizon - 1) * stride + 1) for start in range(max_start + 1)] + self.paths.append((path_start, path_end)) + path_idx += 1 + + else: + while path_end != -1: + path_start = max(path_start, path_end - max_path_length + 1) + path_length = path_end - path_start + 1 + assert path_length >= 2 + + _seq_obs = np.zeros((max_path_length + (horizon - 1) * stride, self.o_dim), dtype=np.float32) + _seq_act = np.zeros((max_path_length + (horizon - 1) * stride, self.a_dim), dtype=np.float32) + _seq_rew = np.zeros((max_path_length + (horizon - 1) * stride, 1), dtype=np.float32) + + _seq_obs[:path_length] = normed_observations[path_start:path_end+1] + _seq_act[:path_length] = actions[path_start:path_end+1] + _seq_rew[:path_length] = rewards[path_start:path_end+1][:, None] + + _seq_obs[path_length:] = normed_observations[path_end] # repeat state + _seq_act[path_length:] = 0 # zero action + _seq_rew[path_length:] = 1 if continous_reward_at_done else 0 + + self.seq_obs.append(_seq_obs) + self.seq_act.append(_seq_act) + self.seq_rew.append(_seq_rew) + + max_start = path_length - 1 + self.indices += [(path_idx, start, start + (horizon - 1) * stride + 1) for start in range(max_start + 1)] + self.paths.append((path_start, path_end)) + path_idx += 1 + + path_start = next_start[path_end] + path_end = next_end[path_start] + + self.seq_obs = np.array(self.seq_obs) + self.seq_act = np.array(self.seq_act) + self.seq_rew = np.array(self.seq_rew) + + if reward_tune == "iql": + self.seq_rew += -1 + elif reward_tune == "none": + self.seq_rew = self.seq_rew + else: + raise ValueError(f"reward_tune: {reward_tune} is not supported.") + + self.seq_val = np.copy(self.seq_rew) + print(self.seq_obs.shape) + for i in reversed(range(max_path_length - 1)): + self.seq_val[:, i] = self.seq_rew[:, i] + discount * self.seq_val[:, i+1] + + print(f"max discounted return: {self.seq_val.max()}") + print(f"min discounted return: {self.seq_val.min()}") + + # val \in [-1, 1] + self.seq_val = (self.seq_val - self.seq_val.min()) / (self.seq_val.max() - self.seq_val.min()) + if center_mapping: + self.seq_val = self.seq_val * 2 - 1 + print(f"max normed discounted return: {self.seq_val.max()}") + print(f"min normed discounted return: {self.seq_val.min()}") + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return len(self.indices) + + def __getitem__(self, idx: int): + path_idx, start, end = self.indices[idx] + + if self.learn_policy: + horizon_state = self.seq_obs[path_idx, start:end:self.stride].copy() + horizon_state[:, :2] -= horizon_state[0, :2] + else: + horizon_state = self.seq_obs[path_idx, start:end:self.stride] + + data = { + 'obs': {'state': horizon_state}, + 'act': self.seq_act[path_idx, start:end:self.stride], + 'rew': self.seq_rew[path_idx, start:end:self.stride], + 'val': self.seq_val[path_idx, start], + } + + torch_data = dict_apply(data, torch.tensor) + + return torch_data + +class D4RLMaze2DTDDataset(BaseDataset): + """ **D4RL-Kitchen Transition Dataset** + + torch.utils.data.Dataset wrapper for D4RL-Kitchen dataset. + Chunk the dataset into transitions. + Use GaussianNormalizer to normalize the observations as default. + Each batch contains + - batch["obs"]["state"], observation of shape (batch_size, o_dim) + - batch["next_obs"]["state"], next observation of shape (batch_size, o_dim) + - batch["act"], action of shape (batch_size, a_dim) + - batch["rew"], reward of shape (batch_size, 1) + - batch["tml"], terminal of shape (batch_size, 1) + + Args: + dataset: Dict[str, np.ndarray], + D4RL-MuJoCo TD dataset. Obtained by calling `d4rl.qlearning_dataset(env)`. + + Examples: + >>> env = gym.make("kitchen-mixed-v0") + >>> dataset = D4RLKitchenTDDataset(d4rl.qlearning_dataset(env)) + >>> dataloader = DataLoader(dataset, batch_size=32, shuffle=True) + >>> batch = next(iter(dataloader)) + >>> obs = batch["obs"]["state"] # (32, 60) + >>> act = batch["act"] # (32, 9) + >>> rew = batch["rew"] # (32, 1) + >>> tml = batch["tml"] # (32, 1) + >>> next_obs = batch["next_obs"]["state"] # (32, 60) + + >>> normalizer = dataset.get_normalizer() + >>> obs = env.reset()[None, :] + >>> normed_obs = normalizer.normalize(obs) + >>> unnormed_obs = normalizer.unnormalize(normed_obs) + """ + def __init__(self, dataset: Dict[str, np.ndarray], reward_tune: str = "iql"): + super().__init__() + + observations, actions, next_observations, rewards, terminals = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["next_observations"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["terminals"].astype(np.float32)) + + if reward_tune == "iql": + rewards = rewards - 1. + elif reward_tune == "cql": + rewards = (rewards - 0.5) * 4. + elif reward_tune == "antmaze": + rewards = (rewards - 0.25) * 2. + elif reward_tune == "none": + rewards = rewards + else: + raise ValueError(f"reward_tune: {reward_tune} is not supported.") + + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + normed_next_observations = self.normalizers["state"].normalize(next_observations) + + self.obs = torch.tensor(normed_observations) + self.act = torch.tensor(actions) + self.rew = torch.tensor(rewards)[:, None] + self.tml = torch.tensor(terminals)[:, None] + self.next_obs = torch.tensor(normed_next_observations) + + self.size = self.obs.shape[0] + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return self.size + + def __getitem__(self, idx: int): + data = { + 'obs': { + 'state': self.obs[idx], }, + 'next_obs': { + 'state': self.next_obs[idx], }, + 'act': self.act[idx], + 'rew': self.rew[idx], + 'tml': self.tml[idx], } + + return data \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/d4rl_mujoco_dataset.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/d4rl_mujoco_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..45d3f770af7fc6d964050609b8656230f4ea6c3b --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/d4rl_mujoco_dataset.py @@ -0,0 +1,453 @@ +from typing import Dict + +import numpy as np +import torch + +from cleandiffuser.dataset.base_dataset import BaseDataset +from cleandiffuser.utils import GaussianNormalizer, dict_apply + + +def return_reward_range(dataset, max_episode_steps): + """ Return the range of episodic returns in the D4RL-MuJoCo dataset. """ + returns, lengths = [], [] + ep_ret, ep_len = 0.0, 0 + for r, d in zip(dataset["rewards"], dataset["terminals"]): + ep_ret += float(r) + ep_len += 1 + if d or ep_len == max_episode_steps: + returns.append(ep_ret) + lengths.append(ep_len) + ep_ret, ep_len = 0.0, 0 + lengths.append(ep_len) # but still keep track of number of steps + assert sum(lengths) == len(dataset["rewards"]) + return min(returns), max(returns) + + +def modify_reward(dataset, max_episode_steps=1000): + """ Modify the episodic return scale of the D4RL-MuJoCo dataset to be within [0, max_episode_steps]. """ + min_ret, max_ret = return_reward_range(dataset, max_episode_steps) + dataset["rewards"] /= max_ret - min_ret + dataset["rewards"] *= max_episode_steps + return dataset + + +class D4RLMuJoCoDataset(BaseDataset): + """ **D4RL-MuJoCo Sequential Dataset** + + torch.utils.data.Dataset wrapper for D4RL-MuJoCo dataset. + Chunk the dataset into sequences of length `horizon` without padding. + Use GaussianNormalizer to normalize the observations as default. + Each batch contains + - batch["obs"]["state"], observations of shape (batch_size, horizon, o_dim) + - batch["act"], actions of shape (batch_size, horizon, a_dim) + - batch["rew"], rewards of shape (batch_size, horizon, 1) + - batch["val"], Monte Carlo return of shape (batch_size, 1) + + Args: + dataset: Dict[str, np.ndarray], + D4RL-MuJoCo dataset. Obtained by calling `env.get_dataset()`. + terminal_penalty: float, + Penalty reward for early-terminal states. Default is -100. + horizon: int, + Length of each sequence. Default is 1. + max_path_length: int, + Maximum length of the episodes. Default is 1000. + discount: float, + Discount factor. Default is 0.99. + + Examples: + >>> env = gym.make("halfcheetah-medium-expert-v2") + >>> dataset = D4RLMuJoCoDataset(env.get_dataset(), horizon=32) + >>> dataloader = DataLoader(dataset, batch_size=32, shuffle=True) + >>> batch = next(iter(dataloader)) + >>> obs = batch["obs"]["state"] # (32, 32, 17) + >>> act = batch["act"] # (32, 32, 6) + >>> rew = batch["rew"] # (32, 32, 1) + >>> val = batch["val"] # (32, 1) + + >>> normalizer = dataset.get_normalizer() + >>> obs = env.reset()[None, :] + >>> normed_obs = normalizer.normalize(obs) + >>> unnormed_obs = normalizer.unnormalize(normed_obs) + """ + def __init__( + self, + dataset: Dict[str, np.ndarray], + terminal_penalty: float = -100., + horizon: int = 1, + max_path_length: int = 1000, + discount: float = 0.99, + ): + super().__init__() + + observations, actions, rewards, timeouts, terminals = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["timeouts"], + dataset["terminals"]) + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + + self.horizon = horizon + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + + n_paths = np.sum(np.logical_or(terminals, timeouts)) + self.seq_obs = np.zeros((n_paths, max_path_length, self.o_dim), dtype=np.float32) + self.seq_act = np.zeros((n_paths, max_path_length, self.a_dim), dtype=np.float32) + self.seq_rew = np.zeros((n_paths, max_path_length, 1), dtype=np.float32) + self.seq_val = np.zeros((n_paths, max_path_length, 1), dtype=np.float32) + self.tml_and_not_timeout = [] + self.indices = [] + + path_lengths, ptr = [], 0 + path_idx = 0 + for i in range(timeouts.shape[0]): + if timeouts[i] or terminals[i]: + path_lengths.append(i - ptr + 1) + + if terminals[i] and not timeouts[i]: + rewards[i] = terminal_penalty if terminal_penalty is not None else rewards[i] + self.tml_and_not_timeout.append([path_idx, i - ptr]) + + self.seq_obs[path_idx, :i - ptr + 1] = normed_observations[ptr:i + 1] + self.seq_act[path_idx, :i - ptr + 1] = actions[ptr:i + 1] + self.seq_rew[path_idx, :i - ptr + 1] = rewards[ptr:i + 1][:, None] + + max_start = min(path_lengths[-1] - 1, max_path_length - horizon) + self.indices += [(path_idx, start, start + horizon) for start in range(max_start + 1)] + + ptr = i + 1 + path_idx += 1 + + self.seq_val[:, -1] = self.seq_rew[:, -1] + for i in range(max_path_length - 1): + self.seq_val[:, - 2 - i] = self.seq_rew[:, -2 - i] + discount * self.seq_val[:, -1 - i] + self.path_lengths = np.array(path_lengths) + self.tml_and_not_timeout = np.array(self.tml_and_not_timeout, dtype=np.int64) + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return len(self.indices) + + def __getitem__(self, idx: int): + path_idx, start, end = self.indices[idx] + + data = { + 'obs': { + 'state': self.seq_obs[path_idx, start:end]}, + 'act': self.seq_act[path_idx, start:end], + 'rew': self.seq_rew[path_idx, start:end], + 'val': self.seq_val[path_idx, start], + } + + torch_data = dict_apply(data, torch.tensor) + + return torch_data + + +class D4RLMuJoCoTDDataset(BaseDataset): + """ **D4RL-MuJoCo Transition Dataset** + + torch.utils.data.Dataset wrapper for D4RL-MuJoCo dataset. + Chunk the dataset into transitions. + Use GaussianNormalizer to normalize the observations as default. + Each batch contains + - batch["obs"]["state"], observation of shape (batch_size, o_dim) + - batch["next_obs"]["state"], next observation of shape (batch_size, o_dim) + - batch["act"], action of shape (batch_size, a_dim) + - batch["rew"], reward of shape (batch_size, 1) + - batch["tml"], terminal of shape (batch_size, 1) + + Args: + dataset: Dict[str, np.ndarray], + D4RL-MuJoCo TD dataset. Obtained by calling `d4rl.qlearning_dataset(env)`. + normalize_reward: bool, + Normalize the reward. Default is False. + + Examples: + >>> env = gym.make("halfcheetah-medium-expert-v2") + >>> dataset = D4RLMuJoCoTDDataset(d4rl.qlearning_dataset(env)) + >>> dataloader = DataLoader(dataset, batch_size=32, shuffle=True) + >>> batch = next(iter(dataloader)) + >>> obs = batch["obs"]["state"] # (32, 17) + >>> act = batch["act"] # (32, 6) + >>> rew = batch["rew"] # (32, 1) + >>> tml = batch["tml"] # (32, 1) + >>> next_obs = batch["next_obs"]["state"] # (32, 17) + + >>> normalizer = dataset.get_normalizer() + >>> obs = env.reset()[None, :] + >>> normed_obs = normalizer.normalize(obs) + >>> unnormed_obs = normalizer.unnormalize(normed_obs) + """ + def __init__(self, dataset: Dict[str, np.ndarray], normalize_reward: bool = False): + super().__init__() + if normalize_reward: + dataset = modify_reward(dataset, 1000) + + observations, actions, next_observations, rewards, terminals = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["next_observations"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["terminals"].astype(np.float32)) + + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + normed_next_observations = self.normalizers["state"].normalize(next_observations) + + self.obs = torch.tensor(normed_observations, dtype=torch.float32) + self.act = torch.tensor(actions, dtype=torch.float32) + self.rew = torch.tensor(rewards, dtype=torch.float32)[:, None] + self.tml = torch.tensor(terminals, dtype=torch.float32)[:, None] + self.next_obs = torch.tensor(normed_next_observations, dtype=torch.float32) + + self.size = self.obs.shape[0] + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return self.size + + def __getitem__(self, idx: int): + data = { + 'obs': { + 'state': self.obs[idx], }, + 'next_obs': { + 'state': self.next_obs[idx], }, + 'act': self.act[idx], + 'rew': self.rew[idx], + 'tml': self.tml[idx], } + + return data + + +class MultiHorizonD4RLMuJoCoDataset(BaseDataset): + def __init__( + self, + dataset, + terminal_penalty=-100, + horizons=(10, 20), + max_path_length=1000, + discount=0.99, + ): + super().__init__() + + observations, actions, rewards, timeouts, terminals = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["timeouts"], + dataset["terminals"]) + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + + self.horizons = horizons + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + self.discount = discount ** np.arange(max_path_length, dtype=np.float32) + + n_paths = np.sum(np.logical_or(terminals, timeouts)) + self.seq_obs = np.zeros((n_paths, max_path_length, self.o_dim), dtype=np.float32) + self.seq_act = np.zeros((n_paths, max_path_length, self.a_dim), dtype=np.float32) + self.seq_rew = np.zeros((n_paths, max_path_length, 1), dtype=np.float32) + self.seq_val = np.zeros((n_paths, max_path_length, 1), dtype=np.float32) + self.indices = [[] for _ in range(len(horizons))] + + path_lengths, ptr = [], 0 + path_idx = 0 + for i in range(timeouts.shape[0]): + if timeouts[i] or terminals[i]: + path_lengths.append(i - ptr + 1) + + if terminals[i] and not timeouts[i]: + rewards[i] = terminal_penalty if terminal_penalty is not None else rewards[i] + + self.seq_obs[path_idx, :i - ptr + 1] = normed_observations[ptr:i + 1] + self.seq_act[path_idx, :i - ptr + 1] = actions[ptr:i + 1] + self.seq_rew[path_idx, :i - ptr + 1] = rewards[ptr:i + 1][:, None] + + max_starts = [min(path_lengths[-1] - 1, max_path_length - horizon) for horizon in horizons] + for k in range(len(horizons)): + self.indices[k] += [(path_idx, start, start + horizons[k]) for start in range(max_starts[k] + 1)] + + ptr = i + 1 + path_idx += 1 + + self.seq_val[:, -1] = self.seq_rew[:, -1] + for i in range(max_path_length - 1): + self.seq_val[:, - 2 - i] = self.seq_rew[:, -2 - i] + discount * self.seq_val[:, -1 - i] + self.path_lengths = np.array(path_lengths) + self.len_each_horizon = [len(indices) for indices in self.indices] + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return max(self.len_each_horizon) + + def __getitem__(self, idx: int): + + indices = [ + int(self.len_each_horizon[i] * (idx / self.len_each_horizon[-1])) for i in range(len(self.horizons))] + + torch_datas = [] + + for i, horizon in enumerate(self.horizons): + + path_idx, start, end = self.indices[i][indices[i]] + + data = { + 'obs': { + 'state': self.seq_obs[path_idx, start:end]}, + 'act': self.seq_act[path_idx, start:end], + 'val': self.seq_val[path_idx, start]} + + torch_data = dict_apply(data, torch.tensor) + + torch_datas.append({ + "horizon": horizon, + "data": torch_data, + }) + + return torch_datas + +class DV_D4RLMuJoCoSeqDataset(BaseDataset): + """ **D4RL-MuJoCo Sequential Dataset** + + torch.utils.data.Dataset wrapper for D4RL-MuJoCo dataset. + Chunk the dataset into sequences of length `horizon` without padding. + Use GaussianNormalizer to normalize the observations as default. + Each batch contains + - batch["obs"]["state"], observations of shape (batch_size, horizon, o_dim) + - batch["act"], actions of shape (batch_size, horizon, a_dim) + - batch["rew"], rewards of shape (batch_size, horizon, 1) + - batch["val"], Monte Carlo return of shape (batch_size, 1) + + Args: + dataset: Dict[str, np.ndarray], + D4RL-MuJoCo dataset. Obtained by calling `env.get_dataset()`. + terminal_penalty: float, + Penalty reward for early-terminal states. Default is -100. + horizon: int, + Length of each sequence. Default is 1. + max_path_length: int, + Maximum length of the episodes. Default is 1000. + discount: float, + Discount factor. Default is 0.99. + + Examples: + >>> env = gym.make("halfcheetah-medium-expert-v2") + >>> dataset = D4RLMuJoCoDataset(env.get_dataset(), horizon=32) + >>> dataloader = DataLoader(dataset, batch_size=32, shuffle=True) + >>> batch = next(iter(dataloader)) + >>> obs = batch["obs"]["state"] # (32, 32, 17) + >>> act = batch["act"] # (32, 32, 6) + >>> rew = batch["rew"] # (32, 32, 1) + >>> val = batch["val"] # (32, 1) + + >>> normalizer = dataset.get_normalizer() + >>> obs = env.reset()[None, :] + >>> normed_obs = normalizer.normalize(obs) + >>> unnormed_obs = normalizer.unnormalize(normed_obs) + """ + def __init__( + self, + dataset: Dict[str, np.ndarray], + terminal_penalty: float = -100, + horizon: int = 1, + max_path_length: int = 1000, + discount: float = 0.99, + center_mapping: bool = True, + stride: int = 1, + full_traj_bonus: float = 100, + ): + super().__init__() + + observations, actions, rewards, timeouts, terminals = ( + dataset["observations"].astype(np.float32), + dataset["actions"].astype(np.float32), + dataset["rewards"].astype(np.float32), + dataset["timeouts"].astype(np.float32), + dataset["terminals"].astype(np.float32)) + self.stride = stride + + self.normalizers = { + "state": GaussianNormalizer(observations)} + normed_observations = self.normalizers["state"].normalize(observations) + + self.horizon = horizon + self.o_dim, self.a_dim = observations.shape[-1], actions.shape[-1] + + n_paths = np.sum(np.logical_or(terminals, timeouts)) + self.seq_obs = np.zeros((n_paths+1, max_path_length, self.o_dim), dtype=np.float32) + self.seq_act = np.zeros((n_paths+1, max_path_length, self.a_dim), dtype=np.float32) + self.seq_rew = np.zeros((n_paths+1, max_path_length, 1), dtype=np.float32) + self.seq_val = np.zeros((n_paths+1, max_path_length, 1), dtype=np.float32) + self.indices = [] + + ptr = 0 + path_idx = 0 + for i in range(timeouts.shape[0]): + if timeouts[i] or terminals[i] or i == timeouts.shape[0] - 1: + path_length = i - ptr + 1 + assert path_length <= max_path_length, f"current path length {path_length}" + + if terminals[i]: + rewards[i] = terminal_penalty if terminal_penalty is not None else rewards[i] + + if path_length == max_path_length: + rewards[i] = rewards[i] + full_traj_bonus if full_traj_bonus is not None else rewards[i] + + self.seq_obs[path_idx, :path_length] = normed_observations[ptr:i + 1] + self.seq_act[path_idx, :path_length] = actions[ptr:i + 1] + self.seq_rew[path_idx, :path_length] = rewards[ptr:i + 1][:, None] + + max_start = path_length - (horizon - 1) * stride - 1 + self.indices += [(path_idx, start, start + (horizon - 1) * stride + 1) for start in range(max_start + 1)] + + ptr = i + 1 + path_idx += 1 + + self.seq_val[:, -1] = self.seq_rew[:, -1] + for i in reversed(range(max_path_length-1)): + self.seq_val[:, i] = self.seq_rew[:, i] + discount * self.seq_val[:, i+1] + + print(f"max discounted return: {self.seq_val.max()}") + print(f"min discounted return: {self.seq_val.min()}") + + # val \in [-1, 1] + self.seq_val = (self.seq_val - self.seq_val.min()) / (self.seq_val.max() - self.seq_val.min()) + if center_mapping: + self.seq_val = self.seq_val * 2 - 1 + print(f"max normed discounted return: {self.seq_val.max()}") + print(f"min normed discounted return: {self.seq_val.min()}") + + def get_normalizer(self): + return self.normalizers["state"] + + def __len__(self): + return len(self.indices) + + def __getitem__(self, idx: int): + path_idx, start, end = self.indices[idx] + + horizon_state = self.seq_obs[path_idx, start:end:self.stride] + + data = { + 'obs': {'state': horizon_state}, + 'act': self.seq_act[path_idx, start:end:self.stride], + 'rew': self.seq_rew[path_idx, start:end:self.stride], + 'val': self.seq_val[path_idx, start], + } + + torch_data = dict_apply(data, torch.tensor) + + return torch_data \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/dataset_utils.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/dataset_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..16a8920318a0ec2ff310e4c7876b552f8c32b80e --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/dataset_utils.py @@ -0,0 +1,409 @@ +from typing import Dict, Callable, List, Union +import numpy as np +import torch +import scipy.interpolate as interpolate +import numba +from cleandiffuser.dataset.replay_buffer import ReplayBuffer +import cleandiffuser.dataset.rotation_conversions as rc +import functools + + +# -----------------------------------------------------------------------------# +# ------------------------------ SequenceSampler ------------------------------# +# -----------------------------------------------------------------------------# + +# Original implemetation: https://github.com/real-stanford/diffusion_policy +# Observation Horizon: To|n_obs_steps +# Action Horizon: Ta|n_action_steps +# Prediction Horizon: T|horizon +# To = 3 +# Ta = 4 +# T = 6 +# |o|o|o| +# | | |a|a|a|a| +# pad_before = 2 +# pad_after = 3 + +@numba.jit(nopython=True) +def create_indices( + episode_ends: np.ndarray, + sequence_length: int, + pad_before: int = 0, pad_after: int = 0, + debug: bool = True) -> np.ndarray: + pad_before = min(max(pad_before, 0), sequence_length - 1) + pad_after = min(max(pad_after, 0), sequence_length - 1) + + indices = list() + for i in range(len(episode_ends)): + start_idx = 0 # episode start index + if i > 0: + start_idx = episode_ends[i - 1] + end_idx = episode_ends[i] # episode end index + episode_length = end_idx - start_idx # episode length + + min_start = -pad_before + max_start = episode_length - sequence_length + pad_after + + # range stops one idx before end + for idx in range(min_start, max_start + 1): + buffer_start_idx = max(idx, 0) + start_idx + buffer_end_idx = min(idx + sequence_length, episode_length) + start_idx + start_offset = buffer_start_idx - (idx + start_idx) + end_offset = (idx + sequence_length + start_idx) - buffer_end_idx + sample_start_idx = 0 + start_offset + sample_end_idx = sequence_length - end_offset + if debug: + assert (start_offset >= 0) + assert (end_offset >= 0) + assert (sample_end_idx - sample_start_idx) == (buffer_end_idx - buffer_start_idx) + indices.append([ + buffer_start_idx, buffer_end_idx, + sample_start_idx, sample_end_idx]) + indices = np.array(indices) + return indices + + +class SequenceSampler: + def __init__( + self, + replay_buffer: ReplayBuffer, + sequence_length: int, + pad_before: int = 0, + pad_after: int = 0, + keys=None, + key_first_k=dict(), + zero_padding: bool = False, + ): + """ + key_first_k: dict str: int + Only take first k data from these keys (to improve perf) + """ + super().__init__() + assert (sequence_length >= 1) + + # all keys + if keys is None: + keys = list(replay_buffer.keys()) + + episode_ends = replay_buffer.episode_ends[:] + + # create indices + # indices (buffer_start_idx, buffer_end_idx, sample_start_idx, sample_end_idx) + # buffer_start_idx and buffer_end_idx define the actual start and end positions of the sample sequence within the original dataset. + # sample_start_idx and sample_end_idx define the relative start and end positions within the sample sequence, + # which is particularly useful when dealing with padding as it can affect the actual length of the sequence. + indices = create_indices( + episode_ends=episode_ends, + sequence_length=sequence_length, + pad_before=pad_before, + pad_after=pad_after, + ) + + self.indices = indices + self.keys = list(keys) # prevent OmegaConf list performance problem + self.sequence_length = sequence_length + self.replay_buffer = replay_buffer + self.zero_padding = zero_padding + self.key_first_k = key_first_k + + def __len__(self): + return len(self.indices) + + def sample_sequence(self, idx): + buffer_start_idx, buffer_end_idx, sample_start_idx, sample_end_idx = self.indices[idx] + result = dict() + for key in self.keys: + input_arr = self.replay_buffer[key] + # performance optimization, avoid small allocation if possible + if key not in self.key_first_k: + sample = input_arr[buffer_start_idx:buffer_end_idx] + else: + # performance optimization, only load used obs steps + n_data = buffer_end_idx - buffer_start_idx + k_data = min(self.key_first_k[key], n_data) + # fill value with Nan to catch bugs + # the non-loaded region should never be used + sample = np.full((n_data,) + input_arr.shape[1:], + fill_value=np.nan, dtype=input_arr.dtype) + sample[:k_data] = input_arr[buffer_start_idx:buffer_start_idx + k_data] + data = sample + if (sample_start_idx > 0) or (sample_end_idx < self.sequence_length): + data = np.zeros( + shape=(self.sequence_length,) + input_arr.shape[1:], + dtype=input_arr.dtype) + if not self.zero_padding: + if sample_start_idx > 0: + data[:sample_start_idx] = sample[0] + if sample_end_idx < self.sequence_length: + data[sample_end_idx:] = sample[-1] + data[sample_start_idx:sample_end_idx] = sample + result[key] = data + return result + + +# -----------------------------------------------------------------------------# +# ---------------------------- Rotation Transformer ---------------------------# +# -----------------------------------------------------------------------------# + +class RotationTransformer: + valid_reps = [ + 'axis_angle', + 'euler_angles', + 'quaternion', + 'rotation_6d', + 'matrix' + ] + + def __init__(self, + from_rep='axis_angle', + to_rep='rotation_6d', + from_convention=None, + to_convention=None): + """ + Valid representations + + Always use matrix as intermediate representation. + """ + assert from_rep != to_rep + assert from_rep in self.valid_reps + assert to_rep in self.valid_reps + if from_rep == 'euler_angles': + assert from_convention is not None + if to_rep == 'euler_angles': + assert to_convention is not None + + forward_funcs = list() + inverse_funcs = list() + + if from_rep != 'matrix': + funcs = [ + getattr(rc, f'{from_rep}_to_matrix'), + getattr(rc, f'matrix_to_{from_rep}') + ] + if from_convention is not None: + funcs = [functools.partial(func, convention=from_convention) + for func in funcs] + forward_funcs.append(funcs[0]) + inverse_funcs.append(funcs[1]) + + if to_rep != 'matrix': + funcs = [ + getattr(rc, f'matrix_to_{to_rep}'), + getattr(rc, f'{to_rep}_to_matrix') + ] + if to_convention is not None: + funcs = [functools.partial(func, convention=to_convention) + for func in funcs] + forward_funcs.append(funcs[0]) + inverse_funcs.append(funcs[1]) + + inverse_funcs = inverse_funcs[::-1] + + self.forward_funcs = forward_funcs + self.inverse_funcs = inverse_funcs + + @staticmethod + def _apply_funcs(x: Union[np.ndarray, torch.Tensor], funcs: list) -> Union[np.ndarray, torch.Tensor]: + x_ = x + if isinstance(x, np.ndarray): + x_ = torch.tensor(x) + x_: torch.Tensor + for func in funcs: + x_ = func(x_) + y = x_ + if isinstance(x, np.ndarray): + y = x_.numpy() + return y + + def forward(self, x: Union[np.ndarray, torch.Tensor] + ) -> Union[np.ndarray, torch.Tensor]: + return self._apply_funcs(x, self.forward_funcs) + + def inverse(self, x: Union[np.ndarray, torch.Tensor] + ) -> Union[np.ndarray, torch.Tensor]: + return self._apply_funcs(x, self.inverse_funcs) + + +# -----------------------------------------------------------------------------# +# --------------------------- multi-field normalizer --------------------------# +# -----------------------------------------------------------------------------# + +def empirical_cdf(sample): + """ https://stackoverflow.com/a/33346366 """ + + # find the unique values and their corresponding counts + quantiles, counts = np.unique(sample, return_counts=True) + + # take the cumulative sum of the counts and divide by the sample size to + # get the cumulative probabilities between 0 and 1 + cumprob = np.cumsum(counts).astype(np.double) / sample.size + + return quantiles, cumprob + + +class CDFNormalizer1d: + """ + CDF normalizer for a single dimension + """ + + def __init__(self, X): + assert X.ndim == 1 + self.X = X.astype(np.float32) + quantiles, cumprob = empirical_cdf(self.X) + self.fn = interpolate.interp1d(quantiles, cumprob) + self.inv = interpolate.interp1d(cumprob, quantiles) + self.xmin, self.xmax = quantiles.min(), quantiles.max() + self.ymin, self.ymax = cumprob.min(), cumprob.max() + + def normalize(self, x): + x = np.clip(x, self.xmin, self.xmax) + y = self.fn(x) + y = 2 * y - 1 + return y + + def unnormalize(self, x, eps=1e-4): + x = (x + 1) / 2. + if (x < self.ymin - eps).any() or (x > self.ymax + eps).any(): + print( + f'''[ dataset/normalization ] Warning: out of range in unnormalize: ''' + f'''[{x.min()}, {x.max()}] | ''' + f'''x : [{self.xmin}, {self.xmax}] | ''' + f'''y: [{self.ymin}, {self.ymax}]''') + x = np.clip(x, self.ymin, self.ymax) + y = self.inv(x) + return y + + +class CDFNormalizer: + """ + makes training data uniform (over each dimension) by transforming it with marginal CDFs + """ + + def __init__(self, X): + self.X = X.astype(np.float32) + self.mins, self.maxs = X.min(0), X.max(0) + self.dim = X.shape[-1] + self.cdfs = [ + CDFNormalizer1d(self.X[:, i]) + for i in range(self.dim)] + + def wrap(self, fn_name, x): + shape = x.shape + x = x.reshape(-1, self.dim) + out = np.zeros_like(x) + for i, cdf in enumerate(self.cdfs): + fn = getattr(cdf, fn_name) + out[:, i] = fn(x[:, i]) + return out.reshape(shape) + + def normalize(self, x): + return self.wrap('normalize', x) + + def unnormalize(self, x): + return self.wrap('unnormalize', x) + + +class GaussianNormalizer: + """ + normalizes data to have zero mean and unit variance + """ + + def __init__(self, X): + self.X = X.astype(np.float32) + self.means, self.stds = X.mean(0), X.std(0) + self.stds[self.stds == 0] = 1. + + def normalize(self, x): + return (x - self.means[None,]) / self.stds[None,] + + def unnormalize(self, x): + return x * self.stds[None,] + self.means[None,] + + +class ImageNormalizer: + """ + normalizes image data from range [0, 1] to [-1, 1]. + """ + + def __init__(self): + pass + + def normalize(self, x): + return x * 2.0 - 1.0 + + def unnormalize(self, x): + return (x + 1.0) / 2.0 + + +class MinMaxNormalizer: + """ + normalizes data through maximum and minimum expansion. + """ + + def __init__(self, X): + X = X.reshape(-1, X.shape[-1]).astype(np.float32) + self.min, self.max = np.min(X, axis=0), np.max(X, axis=0) + self.range = self.max - self.min + if np.any(self.range == 0): + self.range = self.max - self.min + print("Warning: Some features have the same min and max value. These will be set to 0.") + self.range[self.range == 0] = 1 + + def normalize(self, x): + x = x.astype(np.float32) + # nomalize to [0,1] + nx = (x - self.min) / self.range + # normalize to [-1, 1] + nx = nx * 2 - 1 + return nx + + def unnormalize(self, x): + x = x.astype(np.float32) + nx = (x + 1) / 2 + x = nx * self.range + self.min + return x + + +class EmptyNormalizer: + """ + do nothing and change nothing + """ + + def __init__(self): + pass + + def normalize(self, x): + return x + + def unnormalize(self, x): + return x + + +# -----------------------------------------------------------------------------# +# ------------------------------- useful tool ---------------------------------# +# -----------------------------------------------------------------------------# + +def dict_apply( + x: Dict[str, torch.Tensor], + func: Callable[[torch.Tensor], torch.Tensor] +) -> Dict[str, torch.Tensor]: + result = dict() + for key, value in x.items(): + if isinstance(value, dict): + result[key] = dict_apply(value, func) + elif value is None: + result[key] = None + else: + result[key] = func(value) + return result + + +def loop_dataloader(dl): + while True: + for b in dl: + yield b + +def loop_two_dataloaders(dl1, dl2): + while True: + for b1, b2 in zip(dl1, dl2): + yield b1, b2 \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/imagecodecs.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/imagecodecs.py new file mode 100644 index 0000000000000000000000000000000000000000..bfe621e78e919298d05acd740ab955feef473e0d --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/imagecodecs.py @@ -0,0 +1,1385 @@ +# imagecodecs/numcodecs.py + +# Copyright (c) 2021-2022, Christoph Gohlke +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. + +"""Additional numcodecs implemented using imagecodecs.""" + +__version__ = '2022.9.26' + +__all__ = ('register_codecs',) + +import numpy +from numcodecs.abc import Codec +from numcodecs.registry import register_codec, get_codec + +import imagecodecs + + +def protective_squeeze(x: numpy.ndarray): + """ + Squeeze dim only if it's not the last dim. + Image dim expected to be *, H, W, C + """ + img_shape = x.shape[-3:] + if len(x.shape) > 3: + n_imgs = numpy.prod(x.shape[:-3]) + if n_imgs > 1: + img_shape = (-1,) + img_shape + return x.reshape(img_shape) + +def get_default_image_compressor(**kwargs): + if imagecodecs.JPEGXL: + # has JPEGXL + this_kwargs = { + 'effort': 3, + 'distance': 0.3, + # bug in libjxl, invalid codestream for non-lossless + # when decoding speed > 1 + 'decodingspeed': 1 + } + this_kwargs.update(kwargs) + return JpegXl(**this_kwargs) + else: + this_kwargs = { + 'level': 50 + } + this_kwargs.update(kwargs) + return Jpeg2k(**this_kwargs) + +class Aec(Codec): + """AEC codec for numcodecs.""" + + codec_id = 'imagecodecs_aec' + + def __init__( + self, bitspersample=None, flags=None, blocksize=None, rsi=None + ): + self.bitspersample = bitspersample + self.flags = flags + self.blocksize = blocksize + self.rsi = rsi + + def encode(self, buf): + return imagecodecs.aec_encode( + buf, + bitspersample=self.bitspersample, + flags=self.flags, + blocksize=self.blocksize, + rsi=self.rsi, + ) + + def decode(self, buf, out=None): + return imagecodecs.aec_decode( + buf, + bitspersample=self.bitspersample, + flags=self.flags, + blocksize=self.blocksize, + rsi=self.rsi, + out=_flat(out), + ) + + +class Apng(Codec): + """APNG codec for numcodecs.""" + + codec_id = 'imagecodecs_apng' + + def __init__(self, level=None, photometric=None, delay=None): + self.level = level + self.photometric = photometric + self.delay = delay + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.apng_encode( + buf, + level=self.level, + photometric=self.photometric, + delay=self.delay, + ) + + def decode(self, buf, out=None): + return imagecodecs.apng_decode(buf, out=out) + + +class Avif(Codec): + """AVIF codec for numcodecs.""" + + codec_id = 'imagecodecs_avif' + + def __init__( + self, + level=None, + speed=None, + tilelog2=None, + bitspersample=None, + pixelformat=None, + numthreads=None, + index=None, + ): + self.level = level + self.speed = speed + self.tilelog2 = tilelog2 + self.bitspersample = bitspersample + self.pixelformat = pixelformat + self.numthreads = numthreads + self.index = index + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.avif_encode( + buf, + level=self.level, + speed=self.speed, + tilelog2=self.tilelog2, + bitspersample=self.bitspersample, + pixelformat=self.pixelformat, + numthreads=self.numthreads, + ) + + def decode(self, buf, out=None): + return imagecodecs.avif_decode( + buf, index=self.index, numthreads=self.numthreads, out=out + ) + + +class Bitorder(Codec): + """Bitorder codec for numcodecs.""" + + codec_id = 'imagecodecs_bitorder' + + def encode(self, buf): + return imagecodecs.bitorder_encode(buf) + + def decode(self, buf, out=None): + return imagecodecs.bitorder_decode(buf, out=_flat(out)) + + +class Bitshuffle(Codec): + """Bitshuffle codec for numcodecs.""" + + codec_id = 'imagecodecs_bitshuffle' + + def __init__(self, itemsize=1, blocksize=0): + self.itemsize = itemsize + self.blocksize = blocksize + + def encode(self, buf): + return imagecodecs.bitshuffle_encode( + buf, itemsize=self.itemsize, blocksize=self.blocksize + ).tobytes() + + def decode(self, buf, out=None): + return imagecodecs.bitshuffle_decode( + buf, + itemsize=self.itemsize, + blocksize=self.blocksize, + out=_flat(out), + ) + + +class Blosc(Codec): + """Blosc codec for numcodecs.""" + + codec_id = 'imagecodecs_blosc' + + def __init__( + self, + level=None, + compressor=None, + typesize=None, + blocksize=None, + shuffle=None, + numthreads=None, + ): + self.level = level + self.compressor = compressor + self.typesize = typesize + self.blocksize = blocksize + self.shuffle = shuffle + self.numthreads = numthreads + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.blosc_encode( + buf, + level=self.level, + compressor=self.compressor, + typesize=self.typesize, + blocksize=self.blocksize, + shuffle=self.shuffle, + numthreads=self.numthreads, + ) + + def decode(self, buf, out=None): + return imagecodecs.blosc_decode( + buf, numthreads=self.numthreads, out=_flat(out) + ) + + +class Blosc2(Codec): + """Blosc2 codec for numcodecs.""" + + codec_id = 'imagecodecs_blosc2' + + def __init__( + self, + level=None, + compressor=None, + typesize=None, + blocksize=None, + shuffle=None, + numthreads=None, + ): + self.level = level + self.compressor = compressor + self.typesize = typesize + self.blocksize = blocksize + self.shuffle = shuffle + self.numthreads = numthreads + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.blosc2_encode( + buf, + level=self.level, + compressor=self.compressor, + typesize=self.typesize, + blocksize=self.blocksize, + shuffle=self.shuffle, + numthreads=self.numthreads, + ) + + def decode(self, buf, out=None): + return imagecodecs.blosc2_decode( + buf, numthreads=self.numthreads, out=_flat(out) + ) + + +class Brotli(Codec): + """Brotli codec for numcodecs.""" + + codec_id = 'imagecodecs_brotli' + + def __init__(self, level=None, mode=None, lgwin=None): + self.level = level + self.mode = mode + self.lgwin = lgwin + + def encode(self, buf): + return imagecodecs.brotli_encode( + buf, level=self.level, mode=self.mode, lgwin=self.lgwin + ) + + def decode(self, buf, out=None): + return imagecodecs.brotli_decode(buf, out=_flat(out)) + + +class ByteShuffle(Codec): + """ByteShuffle codec for numcodecs.""" + + codec_id = 'imagecodecs_byteshuffle' + + def __init__( + self, shape, dtype, axis=-1, dist=1, delta=False, reorder=False + ): + self.shape = tuple(shape) + self.dtype = numpy.dtype(dtype).str + self.axis = axis + self.dist = dist + self.delta = bool(delta) + self.reorder = bool(reorder) + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + assert buf.shape == self.shape + assert buf.dtype == self.dtype + return imagecodecs.byteshuffle_encode( + buf, + axis=self.axis, + dist=self.dist, + delta=self.delta, + reorder=self.reorder, + ).tobytes() + + def decode(self, buf, out=None): + if not isinstance(buf, numpy.ndarray): + buf = numpy.frombuffer(buf, dtype=self.dtype).reshape(*self.shape) + return imagecodecs.byteshuffle_decode( + buf, + axis=self.axis, + dist=self.dist, + delta=self.delta, + reorder=self.reorder, + out=out, + ) + + +class Bz2(Codec): + """Bz2 codec for numcodecs.""" + + codec_id = 'imagecodecs_bz2' + + def __init__(self, level=None): + self.level = level + + def encode(self, buf): + return imagecodecs.bz2_encode(buf, level=self.level) + + def decode(self, buf, out=None): + return imagecodecs.bz2_decode(buf, out=_flat(out)) + + +class Cms(Codec): + """CMS codec for numcodecs.""" + + codec_id = 'imagecodecs_cms' + + def __init__(self, *args, **kwargs): + pass + + def encode(self, buf, out=None): + # return imagecodecs.cms_transform(buf) + raise NotImplementedError + + def decode(self, buf, out=None): + # return imagecodecs.cms_transform(buf) + raise NotImplementedError + + +class Deflate(Codec): + """Deflate codec for numcodecs.""" + + codec_id = 'imagecodecs_deflate' + + def __init__(self, level=None, raw=False): + self.level = level + self.raw = bool(raw) + + def encode(self, buf): + return imagecodecs.deflate_encode(buf, level=self.level, raw=self.raw) + + def decode(self, buf, out=None): + return imagecodecs.deflate_decode(buf, out=_flat(out), raw=self.raw) + + +class Delta(Codec): + """Delta codec for numcodecs.""" + + codec_id = 'imagecodecs_delta' + + def __init__(self, shape=None, dtype=None, axis=-1, dist=1): + self.shape = None if shape is None else tuple(shape) + self.dtype = None if dtype is None else numpy.dtype(dtype).str + self.axis = axis + self.dist = dist + + def encode(self, buf): + if self.shape is not None or self.dtype is not None: + buf = protective_squeeze(numpy.asarray(buf)) + assert buf.shape == self.shape + assert buf.dtype == self.dtype + return imagecodecs.delta_encode( + buf, axis=self.axis, dist=self.dist + ).tobytes() + + def decode(self, buf, out=None): + if self.shape is not None or self.dtype is not None: + buf = numpy.frombuffer(buf, dtype=self.dtype).reshape(*self.shape) + return imagecodecs.delta_decode( + buf, axis=self.axis, dist=self.dist, out=out + ) + + +class Float24(Codec): + """Float24 codec for numcodecs.""" + + codec_id = 'imagecodecs_float24' + + def __init__(self, byteorder=None, rounding=None): + self.byteorder = byteorder + self.rounding = rounding + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.float24_encode( + buf, byteorder=self.byteorder, rounding=self.rounding + ) + + def decode(self, buf, out=None): + return imagecodecs.float24_decode( + buf, byteorder=self.byteorder, out=out + ) + + +class FloatPred(Codec): + """Floating Point Predictor codec for numcodecs.""" + + codec_id = 'imagecodecs_floatpred' + + def __init__(self, shape, dtype, axis=-1, dist=1): + self.shape = tuple(shape) + self.dtype = numpy.dtype(dtype).str + self.axis = axis + self.dist = dist + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + assert buf.shape == self.shape + assert buf.dtype == self.dtype + return imagecodecs.floatpred_encode( + buf, axis=self.axis, dist=self.dist + ).tobytes() + + def decode(self, buf, out=None): + if not isinstance(buf, numpy.ndarray): + buf = numpy.frombuffer(buf, dtype=self.dtype).reshape(*self.shape) + return imagecodecs.floatpred_decode( + buf, axis=self.axis, dist=self.dist, out=out + ) + + +class Gif(Codec): + """GIF codec for numcodecs.""" + + codec_id = 'imagecodecs_gif' + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.gif_encode(buf) + + def decode(self, buf, out=None): + return imagecodecs.gif_decode(buf, asrgb=False, out=out) + + +class Heif(Codec): + """HEIF codec for numcodecs.""" + + codec_id = 'imagecodecs_heif' + + def __init__( + self, + level=None, + bitspersample=None, + photometric=None, + compression=None, + numthreads=None, + index=None, + ): + self.level = level + self.bitspersample = bitspersample + self.photometric = photometric + self.compression = compression + self.numthreads = numthreads + self.index = index + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.heif_encode( + buf, + level=self.level, + bitspersample=self.bitspersample, + photometric=self.photometric, + compression=self.compression, + numthreads=self.numthreads, + ) + + def decode(self, buf, out=None): + return imagecodecs.heif_decode( + buf, + index=self.index, + photometric=self.photometric, + numthreads=self.numthreads, + out=out, + ) + + +class Jetraw(Codec): + """Jetraw codec for numcodecs.""" + + codec_id = 'imagecodecs_jetraw' + + def __init__( + self, + shape, + identifier, + parameters=None, + verbosity=None, + errorbound=None, + ): + self.shape = shape + self.identifier = identifier + self.errorbound = errorbound + imagecodecs.jetraw_init(parameters, verbosity) + + def encode(self, buf): + return imagecodecs.jetraw_encode( + buf, identifier=self.identifier, errorbound=self.errorbound + ) + + def decode(self, buf, out=None): + if out is None: + out = numpy.empty(self.shape, numpy.uint16) + return imagecodecs.jetraw_decode(buf, out=out) + + +class Jpeg(Codec): + """JPEG codec for numcodecs.""" + + codec_id = 'imagecodecs_jpeg' + + def __init__( + self, + bitspersample=None, + tables=None, + header=None, + colorspace_data=None, + colorspace_jpeg=None, + level=None, + subsampling=None, + optimize=None, + smoothing=None, + ): + self.tables = tables + self.header = header + self.bitspersample = bitspersample + self.colorspace_data = colorspace_data + self.colorspace_jpeg = colorspace_jpeg + self.level = level + self.subsampling = subsampling + self.optimize = optimize + self.smoothing = smoothing + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.jpeg_encode( + buf, + level=self.level, + colorspace=self.colorspace_data, + outcolorspace=self.colorspace_jpeg, + subsampling=self.subsampling, + optimize=self.optimize, + smoothing=self.smoothing, + ) + + def decode(self, buf, out=None): + out_shape = None + if out is not None: + out_shape = out.shape + out = protective_squeeze(out) + img = imagecodecs.jpeg_decode( + buf, + bitspersample=self.bitspersample, + tables=self.tables, + header=self.header, + colorspace=self.colorspace_jpeg, + outcolorspace=self.colorspace_data, + out=out, + ) + if out_shape is not None: + img = img.reshape(out_shape) + return img + + def get_config(self): + """Return dictionary holding configuration parameters.""" + config = dict(id=self.codec_id) + for key in self.__dict__: + if not key.startswith('_'): + value = getattr(self, key) + if value is not None and key in ('header', 'tables'): + import base64 + + value = base64.b64encode(value).decode() + config[key] = value + return config + + @classmethod + def from_config(cls, config): + """Instantiate codec from configuration object.""" + for key in ('header', 'tables'): + value = config.get(key, None) + if value is not None and isinstance(value, str): + import base64 + + config[key] = base64.b64decode(value.encode()) + return cls(**config) + + +class Jpeg2k(Codec): + """JPEG 2000 codec for numcodecs.""" + + codec_id = 'imagecodecs_jpeg2k' + + def __init__( + self, + level=None, + codecformat=None, + colorspace=None, + tile=None, + reversible=None, + bitspersample=None, + resolutions=None, + numthreads=None, + verbose=0, + ): + self.level = level + self.codecformat = codecformat + self.colorspace = colorspace + self.tile = None if tile is None else tuple(tile) + self.reversible = reversible + self.bitspersample = bitspersample + self.resolutions = resolutions + self.numthreads = numthreads + self.verbose = verbose + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.jpeg2k_encode( + buf, + level=self.level, + codecformat=self.codecformat, + colorspace=self.colorspace, + tile=self.tile, + reversible=self.reversible, + bitspersample=self.bitspersample, + resolutions=self.resolutions, + numthreads=self.numthreads, + verbose=self.verbose, + ) + + def decode(self, buf, out=None): + return imagecodecs.jpeg2k_decode( + buf, verbose=self.verbose, numthreads=self.numthreads, out=out + ) + + +class JpegLs(Codec): + """JPEG LS codec for numcodecs.""" + + codec_id = 'imagecodecs_jpegls' + + def __init__(self, level=None): + self.level = level + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.jpegls_encode(buf, level=self.level) + + def decode(self, buf, out=None): + return imagecodecs.jpegls_decode(buf, out=out) + + +class JpegXl(Codec): + """JPEG XL codec for numcodecs.""" + + codec_id = 'imagecodecs_jpegxl' + + def __init__( + self, + # encode + level=None, + effort=None, + distance=None, + lossless=None, + decodingspeed=None, + photometric=None, + planar=None, + usecontainer=None, + # decode + index=None, + keeporientation=None, + # both + numthreads=None, + ): + """ + Return JPEG XL image from numpy array. + Float must be in nominal range 0..1. + + Currently L, LA, RGB, RGBA images are supported in contig mode. + Extra channels are only supported for grayscale images in planar mode. + + Parameters + ---------- + level : Default to None, i.e. not overwriting lossess and decodingspeed options. + When < 0: Use lossless compression + When in [0,1,2,3,4]: Sets the decoding speed tier for the provided options. + Minimum is 0 (slowest to decode, best quality/density), and maximum + is 4 (fastest to decode, at the cost of some quality/density). + effort : Default to 3. + Sets encoder effort/speed level without affecting decoding speed. + Valid values are, from faster to slower speed: 1:lightning 2:thunder + 3:falcon 4:cheetah 5:hare 6:wombat 7:squirrel 8:kitten 9:tortoise. + Speed: lightning, thunder, falcon, cheetah, hare, wombat, squirrel, kitten, tortoise + control the encoder effort in ascending order. + This also affects memory usage: using lower effort will typically reduce memory + consumption during encoding. + lightning and thunder are fast modes useful for lossless mode (modular). + falcon disables all of the following tools. + cheetah enables coefficient reordering, context clustering, and heuristics for selecting DCT sizes and quantization steps. + hare enables Gaborish filtering, chroma from luma, and an initial estimate of quantization steps. + wombat enables error diffusion quantization and full DCT size selection heuristics. + squirrel (default) enables dots, patches, and spline detection, and full context clustering. + kitten optimizes the adaptive quantization for a psychovisual metric. + tortoise enables a more thorough adaptive quantization search. + distance : Default to 1.0 + Sets the distance level for lossy compression: target max butteraugli distance, + lower = higher quality. Range: 0 .. 15. 0.0 = mathematically lossless + (however, use JxlEncoderSetFrameLossless instead to use true lossless, + as setting distance to 0 alone is not the only requirement). + 1.0 = visually lossless. Recommended range: 0.5 .. 3.0. + lossess : Default to False. + Use lossess encoding. + decodingspeed : Default to 0. + Duplicate to level. [0,4] + photometric : Return JxlColorSpace value. + Default logic is quite complicated but works most of the time. + Accepted value: + int: [-1,3] + str: ['RGB', + 'WHITEISZERO', 'MINISWHITE', + 'BLACKISZERO', 'MINISBLACK', 'GRAY', + 'XYB', 'KNOWN'] + planar : Enable multi-channel mode. + Default to false. + usecontainer : + Forces the encoder to use the box-based container format (BMFF) + even when not necessary. + When using JxlEncoderUseBoxes, JxlEncoderStoreJPEGMetadata or + JxlEncoderSetCodestreamLevel with level 10, the encoder will + automatically also use the container format, it is not necessary + to use JxlEncoderUseContainer for those use cases. + By default this setting is disabled. + index : Selectively decode frames for animation. + Default to 0, decode all frames. + When set to > 0, decode that frame index only. + keeporientation : + Enables or disables preserving of as-in-bitstream pixeldata orientation. + Some images are encoded with an Orientation tag indicating that the + decoder must perform a rotation and/or mirroring to the encoded image data. + + If skip_reorientation is JXL_FALSE (the default): the decoder will apply + the transformation from the orientation setting, hence rendering the image + according to its specified intent. When producing a JxlBasicInfo, the decoder + will always set the orientation field to JXL_ORIENT_IDENTITY (matching the + returned pixel data) and also align xsize and ysize so that they correspond + to the width and the height of the returned pixel data. + + If skip_reorientation is JXL_TRUE: the decoder will skip applying the + transformation from the orientation setting, returning the image in + the as-in-bitstream pixeldata orientation. This may be faster to decode + since the decoder doesnt have to apply the transformation, but can + cause wrong display of the image if the orientation tag is not correctly + taken into account by the user. + + By default, this option is disabled, and the returned pixel data is + re-oriented according to the images Orientation setting. + threads : Default to 1. + If <= 0, use all cores. + If > 32, clipped to 32. + """ + + self.level = level + self.effort = effort + self.distance = distance + self.lossless = bool(lossless) + self.decodingspeed = decodingspeed + self.photometric = photometric + self.planar = planar + self.usecontainer = usecontainer + self.index = index + self.keeporientation = keeporientation + self.numthreads = numthreads + + def encode(self, buf): + # TODO: only squeeze all but last dim + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.jpegxl_encode( + buf, + level=self.level, + effort=self.effort, + distance=self.distance, + lossless=self.lossless, + decodingspeed=self.decodingspeed, + photometric=self.photometric, + planar=self.planar, + usecontainer=self.usecontainer, + numthreads=self.numthreads, + ) + + def decode(self, buf, out=None): + return imagecodecs.jpegxl_decode( + buf, + index=self.index, + keeporientation=self.keeporientation, + numthreads=self.numthreads, + out=out, + ) + + +class JpegXr(Codec): + """JPEG XR codec for numcodecs.""" + + codec_id = 'imagecodecs_jpegxr' + + def __init__( + self, + level=None, + photometric=None, + hasalpha=None, + resolution=None, + fp2int=None, + ): + self.level = level + self.photometric = photometric + self.hasalpha = hasalpha + self.resolution = resolution + self.fp2int = fp2int + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.jpegxr_encode( + buf, + level=self.level, + photometric=self.photometric, + hasalpha=self.hasalpha, + resolution=self.resolution, + ) + + def decode(self, buf, out=None): + return imagecodecs.jpegxr_decode(buf, fp2int=self.fp2int, out=out) + + +class Lerc(Codec): + """LERC codec for numcodecs.""" + + codec_id = 'imagecodecs_lerc' + + def __init__(self, level=None, version=None, planar=None): + self.level = level + self.version = version + self.planar = bool(planar) + # TODO: support mask? + # self.mask = None + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.lerc_encode( + buf, + level=self.level, + version=self.version, + planar=self.planar, + ) + + def decode(self, buf, out=None): + return imagecodecs.lerc_decode(buf, out=out) + + +class Ljpeg(Codec): + """LJPEG codec for numcodecs.""" + + codec_id = 'imagecodecs_ljpeg' + + def __init__(self, bitspersample=None): + self.bitspersample = bitspersample + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.ljpeg_encode(buf, bitspersample=self.bitspersample) + + def decode(self, buf, out=None): + return imagecodecs.ljpeg_decode(buf, out=out) + + +class Lz4(Codec): + """LZ4 codec for numcodecs.""" + + codec_id = 'imagecodecs_lz4' + + def __init__(self, level=None, hc=False, header=True): + self.level = level + self.hc = hc + self.header = bool(header) + + def encode(self, buf): + return imagecodecs.lz4_encode( + buf, level=self.level, hc=self.hc, header=self.header + ) + + def decode(self, buf, out=None): + return imagecodecs.lz4_decode(buf, header=self.header, out=_flat(out)) + + +class Lz4f(Codec): + """LZ4F codec for numcodecs.""" + + codec_id = 'imagecodecs_lz4f' + + def __init__( + self, + level=None, + blocksizeid=False, + contentchecksum=None, + blockchecksum=None, + ): + self.level = level + self.blocksizeid = blocksizeid + self.contentchecksum = contentchecksum + self.blockchecksum = blockchecksum + + def encode(self, buf): + return imagecodecs.lz4f_encode( + buf, + level=self.level, + blocksizeid=self.blocksizeid, + contentchecksum=self.contentchecksum, + blockchecksum=self.blockchecksum, + ) + + def decode(self, buf, out=None): + return imagecodecs.lz4f_decode(buf, out=_flat(out)) + + +class Lzf(Codec): + """LZF codec for numcodecs.""" + + codec_id = 'imagecodecs_lzf' + + def __init__(self, header=True): + self.header = bool(header) + + def encode(self, buf): + return imagecodecs.lzf_encode(buf, header=self.header) + + def decode(self, buf, out=None): + return imagecodecs.lzf_decode(buf, header=self.header, out=_flat(out)) + + +class Lzma(Codec): + """LZMA codec for numcodecs.""" + + codec_id = 'imagecodecs_lzma' + + def __init__(self, level=None): + self.level = level + + def encode(self, buf): + return imagecodecs.lzma_encode(buf, level=self.level) + + def decode(self, buf, out=None): + return imagecodecs.lzma_decode(buf, out=_flat(out)) + + +class Lzw(Codec): + """LZW codec for numcodecs.""" + + codec_id = 'imagecodecs_lzw' + + def encode(self, buf): + return imagecodecs.lzw_encode(buf) + + def decode(self, buf, out=None): + return imagecodecs.lzw_decode(buf, out=_flat(out)) + + +class PackBits(Codec): + """PackBits codec for numcodecs.""" + + codec_id = 'imagecodecs_packbits' + + def __init__(self, axis=None): + self.axis = axis + + def encode(self, buf): + if not isinstance(buf, (bytes, bytearray)): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.packbits_encode(buf, axis=self.axis) + + def decode(self, buf, out=None): + return imagecodecs.packbits_decode(buf, out=_flat(out)) + + +class Pglz(Codec): + """PGLZ codec for numcodecs.""" + + codec_id = 'imagecodecs_pglz' + + def __init__(self, header=True, strategy=None): + self.header = bool(header) + self.strategy = strategy + + def encode(self, buf): + return imagecodecs.pglz_encode( + buf, strategy=self.strategy, header=self.header + ) + + def decode(self, buf, out=None): + return imagecodecs.pglz_decode(buf, header=self.header, out=_flat(out)) + + +class Png(Codec): + """PNG codec for numcodecs.""" + + codec_id = 'imagecodecs_png' + + def __init__(self, level=None): + self.level = level + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.png_encode(buf, level=self.level) + + def decode(self, buf, out=None): + return imagecodecs.png_decode(buf, out=out) + + +class Qoi(Codec): + """QOI codec for numcodecs.""" + + codec_id = 'imagecodecs_qoi' + + def __init__(self): + pass + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.qoi_encode(buf) + + def decode(self, buf, out=None): + return imagecodecs.qoi_decode(buf, out=out) + + +class Rgbe(Codec): + """RGBE codec for numcodecs.""" + + codec_id = 'imagecodecs_rgbe' + + def __init__(self, header=False, shape=None, rle=None): + if not header and shape is None: + raise ValueError('must specify data shape if no header') + if shape and shape[-1] != 3: + raise ValueError('invalid shape') + self.shape = shape + self.header = bool(header) + self.rle = None if rle is None else bool(rle) + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.rgbe_encode(buf, header=self.header, rle=self.rle) + + def decode(self, buf, out=None): + if out is None and not self.header: + out = numpy.empty(self.shape, numpy.float32) + return imagecodecs.rgbe_decode( + buf, header=self.header, rle=self.rle, out=out + ) + + +class Rcomp(Codec): + """Rcomp codec for numcodecs.""" + + codec_id = 'imagecodecs_rcomp' + + def __init__(self, shape, dtype, nblock=None): + self.shape = tuple(shape) + self.dtype = numpy.dtype(dtype).str + self.nblock = nblock + + def encode(self, buf): + return imagecodecs.rcomp_encode(buf, nblock=self.nblock) + + def decode(self, buf, out=None): + return imagecodecs.rcomp_decode( + buf, + shape=self.shape, + dtype=self.dtype, + nblock=self.nblock, + out=out, + ) + + +class Snappy(Codec): + """Snappy codec for numcodecs.""" + + codec_id = 'imagecodecs_snappy' + + def encode(self, buf): + return imagecodecs.snappy_encode(buf) + + def decode(self, buf, out=None): + return imagecodecs.snappy_decode(buf, out=_flat(out)) + + +class Spng(Codec): + """SPNG codec for numcodecs.""" + + codec_id = 'imagecodecs_spng' + + def __init__(self, level=None): + self.level = level + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.spng_encode(buf, level=self.level) + + def decode(self, buf, out=None): + return imagecodecs.spng_decode(buf, out=out) + + +class Tiff(Codec): + """TIFF codec for numcodecs.""" + + codec_id = 'imagecodecs_tiff' + + def __init__(self, index=None, asrgb=None, verbose=0): + self.index = index + self.asrgb = bool(asrgb) + self.verbose = verbose + + def encode(self, buf): + # TODO: not implemented + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.tiff_encode(buf) + + def decode(self, buf, out=None): + return imagecodecs.tiff_decode( + buf, + index=self.index, + asrgb=self.asrgb, + verbose=self.verbose, + out=out, + ) + + +class Webp(Codec): + """WebP codec for numcodecs.""" + + codec_id = 'imagecodecs_webp' + + def __init__(self, level=None, lossless=None, method=None, hasalpha=None): + self.level = level + self.hasalpha = bool(hasalpha) + self.method = method + self.lossless = lossless + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + return imagecodecs.webp_encode( + buf, level=self.level, lossless=self.lossless, method=self.method + ) + + def decode(self, buf, out=None): + return imagecodecs.webp_decode(buf, hasalpha=self.hasalpha, out=out) + + +class Xor(Codec): + """XOR codec for numcodecs.""" + + codec_id = 'imagecodecs_xor' + + def __init__(self, shape=None, dtype=None, axis=-1): + self.shape = None if shape is None else tuple(shape) + self.dtype = None if dtype is None else numpy.dtype(dtype).str + self.axis = axis + + def encode(self, buf): + if self.shape is not None or self.dtype is not None: + buf = protective_squeeze(numpy.asarray(buf)) + assert buf.shape == self.shape + assert buf.dtype == self.dtype + return imagecodecs.xor_encode(buf, axis=self.axis).tobytes() + + def decode(self, buf, out=None): + if self.shape is not None or self.dtype is not None: + buf = numpy.frombuffer(buf, dtype=self.dtype).reshape(*self.shape) + return imagecodecs.xor_decode(buf, axis=self.axis, out=_flat(out)) + + +class Zfp(Codec): + """ZFP codec for numcodecs.""" + + codec_id = 'imagecodecs_zfp' + + def __init__( + self, + shape=None, + dtype=None, + strides=None, + level=None, + mode=None, + execution=None, + numthreads=None, + chunksize=None, + header=True, + ): + if header: + self.shape = None + self.dtype = None + self.strides = None + elif shape is None or dtype is None: + raise ValueError('invalid shape or dtype') + else: + self.shape = tuple(shape) + self.dtype = numpy.dtype(dtype).str + self.strides = None if strides is None else tuple(strides) + self.level = level + self.mode = mode + self.execution = execution + self.numthreads = numthreads + self.chunksize = chunksize + self.header = bool(header) + + def encode(self, buf): + buf = protective_squeeze(numpy.asarray(buf)) + if not self.header: + assert buf.shape == self.shape + assert buf.dtype == self.dtype + return imagecodecs.zfp_encode( + buf, + level=self.level, + mode=self.mode, + execution=self.execution, + header=self.header, + numthreads=self.numthreads, + chunksize=self.chunksize, + ) + + def decode(self, buf, out=None): + if self.header: + return imagecodecs.zfp_decode(buf, out=out) + return imagecodecs.zfp_decode( + buf, + shape=self.shape, + dtype=numpy.dtype(self.dtype), + strides=self.strides, + numthreads=self.numthreads, + out=out, + ) + + +class Zlib(Codec): + """Zlib codec for numcodecs.""" + + codec_id = 'imagecodecs_zlib' + + def __init__(self, level=None): + self.level = level + + def encode(self, buf): + return imagecodecs.zlib_encode(buf, level=self.level) + + def decode(self, buf, out=None): + return imagecodecs.zlib_decode(buf, out=_flat(out)) + + +class Zlibng(Codec): + """Zlibng codec for numcodecs.""" + + codec_id = 'imagecodecs_zlibng' + + def __init__(self, level=None): + self.level = level + + def encode(self, buf): + return imagecodecs.zlibng_encode(buf, level=self.level) + + def decode(self, buf, out=None): + return imagecodecs.zlibng_decode(buf, out=_flat(out)) + + +class Zopfli(Codec): + """Zopfli codec for numcodecs.""" + + codec_id = 'imagecodecs_zopfli' + + def encode(self, buf): + return imagecodecs.zopfli_encode(buf) + + def decode(self, buf, out=None): + return imagecodecs.zopfli_decode(buf, out=_flat(out)) + + +class Zstd(Codec): + """ZStandard codec for numcodecs.""" + + codec_id = 'imagecodecs_zstd' + + def __init__(self, level=None): + self.level = level + + def encode(self, buf): + return imagecodecs.zstd_encode(buf, level=self.level) + + def decode(self, buf, out=None): + return imagecodecs.zstd_decode(buf, out=_flat(out)) + + +def _flat(out): + """Return numpy array as contiguous view of bytes if possible.""" + if out is None: + return None + view = memoryview(out) + if view.readonly or not view.contiguous: + return None + return view.cast('B') + + +def register_codecs(codecs=None, force=False, verbose=True): + """Register codecs in this module with numcodecs.""" + for name, cls in globals().items(): + if not hasattr(cls, 'codec_id') or name == 'Codec': + continue + if codecs is not None and cls.codec_id not in codecs: + continue + try: + try: + get_codec({'id': cls.codec_id}) + except TypeError: + # registered, but failed + pass + except ValueError: + # not registered yet + pass + else: + if not force: + if verbose: + log_warning( + f'numcodec {cls.codec_id!r} already registered' + ) + continue + if verbose: + log_warning(f'replacing registered numcodec {cls.codec_id!r}') + register_codec(cls) + + +def log_warning(msg, *args, **kwargs): + """Log message with level WARNING.""" + import logging + + logging.getLogger(__name__).warning(msg, *args, **kwargs) \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/kitchen_dataset.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/kitchen_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..b8ce06ad22b6a8a3a727915ab0d5fdeffe9a874a --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/kitchen_dataset.py @@ -0,0 +1,247 @@ +import copy +import pathlib +from typing import Dict + +import numpy as np +import torch +from tqdm import tqdm + +from cleandiffuser.dataset.base_dataset import BaseDataset +from cleandiffuser.dataset.dataset_utils import SequenceSampler, dict_apply +from cleandiffuser.dataset.replay_buffer import ReplayBuffer +from cleandiffuser.env.kitchen.kitchen_util import parse_mjl_logs +from cleandiffuser.utils import EmptyNormalizer, MinMaxNormalizer + + +class KitchenDataset(BaseDataset): + def __init__(self, + dataset_dir, + horizon=1, + pad_before=0, + pad_after=0, + ): + super().__init__() + + data_directory = pathlib.Path(dataset_dir) + observations = np.load(data_directory / "observations_seq.npy") + actions = np.load(data_directory / "actions_seq.npy") + masks = np.load(data_directory / "existence_mask.npy") + + self.replay_buffer = ReplayBuffer.create_empty_numpy() + for i in range(len(masks)): + eps_len = int(masks[i].sum()) + obs = observations[i, :eps_len].astype(np.float32) + action = actions[i, :eps_len].astype(np.float32) + data = { + 'state': obs, + 'action': action + } + self.replay_buffer.add_episode(data) + + self.sampler = SequenceSampler( + replay_buffer=self.replay_buffer, + sequence_length=horizon, + pad_before=pad_before, + pad_after=pad_after) + + self.horizon = horizon + self.pad_before = pad_before + self.pad_after = pad_after + + self.normalizer = self.get_normalizer() + + def get_normalizer(self): + state_normalizer = MinMaxNormalizer( + self.replay_buffer['state'][:]) # (N, obs_dim) + action_normalizer = MinMaxNormalizer( + self.replay_buffer['action'][:]) # (N, action_dim) + return { + "obs": { + "state": state_normalizer + }, + "action": action_normalizer + } + + def sample_to_data(self, sample): + state = sample['state'].astype(np.float32) + state = self.normalizer['obs']['state'].normalize(state) + + action = sample['action'].astype(np.float32) + action = self.normalizer['action'].normalize(action) + data = { + 'obs': { + 'state': state, + }, + 'action': action, + } + return data + + def __str__(self) -> str: + return f"Keys: {self.replay_buffer.keys()} Steps: {self.replay_buffer.n_steps} Episodes: {self.replay_buffer.n_episodes}" + + def __len__(self) -> int: + return len(self.sampler) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + sample = self.sampler.sample_sequence(idx) + data = self.sample_to_data(sample) + torch_data = dict_apply(data, torch.tensor) + return torch_data + + +class KitchenDatasetV2(BaseDataset): + def __init__(self, + dataset_dir, + horizon=1, + pad_before=0, + pad_after=0, + ): + super().__init__() + + data_directory = pathlib.Path(dataset_dir) + observations = np.load( + data_directory / "observations_seq.npy").astype(np.float32) + actions = np.load( + data_directory / "actions_seq.npy").astype(np.float32) + masks = np.load(data_directory / "existence_mask.npy").astype(np.int64) + + self.state_normalizer = MinMaxNormalizer(observations) + self.action_normalizer = MinMaxNormalizer(actions) + + normed_observations = self.state_normalizer.normalize(observations) + normed_actions = self.action_normalizer.normalize(actions) + + self.replay_buffer = ReplayBuffer.create_empty_numpy() + for i in range(len(masks)): + eps_len = int(masks[i].sum()) + obs = normed_observations[i, :eps_len] + action = normed_actions[i, :eps_len] + data = {'state': obs, 'action': action} + self.replay_buffer.add_episode(data) + + self.sampler = SequenceSampler( + replay_buffer=self.replay_buffer, + sequence_length=horizon, + pad_before=pad_before, + pad_after=pad_after) + + self.horizon = horizon + self.pad_before = pad_before + self.pad_after = pad_after + + self.normalizer = self.get_normalizer() + self.obs_dim, self.act_dim = observations.shape[-1], actions.shape[-1] + + self.normed_observations = normed_observations[np.where(masks)].reshape(-1, self.obs_dim) + self.normed_actions = normed_actions[np.where(masks)].reshape(-1, self.act_dim) + self.observations = observations + self.actions = actions + + def get_normalizer(self): + return { + "state": self.state_normalizer, + "action": self.action_normalizer} + + def __str__(self) -> str: + return f"Keys: {self.replay_buffer.keys()} Steps: {self.replay_buffer.n_steps} Episodes: {self.replay_buffer.n_episodes}" + + def __len__(self) -> int: + return len(self.sampler) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + data = self.sampler.sample_sequence(idx) + torch_data = dict_apply(data, torch.tensor) + return torch_data + + +class KitchenMjlDataset(BaseDataset): + def __init__(self, + dataset_dir, + horizon=1, + pad_before=0, + pad_after=0, + abs_action=True, + robot_noise_ratio=0.1, + ): + super().__init__() + + data_directory = pathlib.Path(dataset_dir) + robot_pos_noise_amp = np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, + 0.1, 0.005, 0.005, 0.0005, 0.0005, 0.0005, 0.0005, 0.0005, + 0.0005, 0.005, 0.005, 0.005, 0.1, 0.1, 0.1, 0.005, + 0.005, 0.005, 0.1, 0.1, 0.1, 0.005], dtype=np.float32) + rng = np.random.default_rng(seed=42) + + data_directory = pathlib.Path(dataset_dir) + self.replay_buffer = ReplayBuffer.create_empty_numpy() + for i, mjl_path in enumerate(tqdm(list(data_directory.glob('*/*.mjl')))): + try: + data = parse_mjl_logs(str(mjl_path.absolute()), skipamount=40) + qpos = data['qpos'].astype(np.float32) + obs = np.concatenate([ + qpos[:, :9], + qpos[:, -21:], + np.zeros((len(qpos), 30), dtype=np.float32) + ], axis=-1) + if robot_noise_ratio > 0: + # add observation noise to match real robot + noise = robot_noise_ratio * robot_pos_noise_amp * rng.uniform( + low=-1., high=1., size=(obs.shape[0], 30)) + obs[:, :30] += noise + episode = { + 'state': obs, + 'action': data['ctrl'].astype(np.float32) + } + self.replay_buffer.add_episode(episode) + except Exception as e: + print(i, e) + + self.sampler = SequenceSampler( + replay_buffer=self.replay_buffer, + sequence_length=horizon, + pad_before=pad_before, + pad_after=pad_after) + + self.horizon = horizon + self.pad_before = pad_before + self.pad_after = pad_after + + self.normalizer = self.get_normalizer() + + def get_normalizer(self): + state_normalizer = MinMaxNormalizer( + self.replay_buffer['state'][:]) # (N, obs_dim) + action_normalizer = MinMaxNormalizer( + self.replay_buffer['action'][:]) # (N, action_dim) + return { + "obs": { + "state": state_normalizer + }, + "action": action_normalizer + } + + def sample_to_data(self, sample): + state = sample['state'].astype(np.float32) + state = self.normalizer['obs']['state'].normalize(state) + + action = sample['action'].astype(np.float32) + action = self.normalizer['action'].normalize(action) + data = { + 'obs': { + 'state': state, + }, + 'action': action, + } + return data + + def __str__(self) -> str: + return f"Keys: {self.replay_buffer.keys()} Steps: {self.replay_buffer.n_steps} Episodes: {self.replay_buffer.n_episodes}" + + def __len__(self) -> int: + return len(self.sampler) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + sample = self.sampler.sample_sequence(idx) + data = self.sample_to_data(sample) + torch_data = dict_apply(data, torch.tensor) + return torch_data diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/pusht_dataset.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/pusht_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..7819845d8e69bd6c33b1756f99134db0ac17d5cc --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/pusht_dataset.py @@ -0,0 +1,233 @@ +from typing import Dict +import torch +import numpy as np +import copy +from cleandiffuser.dataset.base_dataset import BaseDataset +from cleandiffuser.dataset.replay_buffer import ReplayBuffer +from cleandiffuser.dataset.dataset_utils import SequenceSampler, MinMaxNormalizer, ImageNormalizer, dict_apply + + +# data/pusht_cchi_v7_replay.zarr +# ├── data +# │ ├── action (25650, 2) float32 +# │ ├── img (25650, 96, 96, 3) float32 +# │ ├── keypoint (25650, 9, 2) float32 +# │ ├── n_contacts (25650, 1) float32 +# │ └── state (25650, 5) float32 +# └── meta +# └── episode_ends (206,) int64 +# assert('data' in root) +# assert('meta' in root) +# assert('episode_ends' in root['meta']) + + +class PushTStateDataset(BaseDataset): + def __init__(self, + dataset_path, + obs_keys=["state", "action"], + horizon=1, + pad_before=0, + pad_after=0, + abs_action=False + ): + + super().__init__() + self.replay_buffer = ReplayBuffer.copy_from_path( + dataset_path, keys=obs_keys) + + self.sampler = SequenceSampler( + replay_buffer=self.replay_buffer, + sequence_length=horizon, + pad_before=pad_before, + pad_after=pad_after) + + self.horizon = horizon + self.pad_before = pad_before + self.pad_after = pad_after + + self.normalizer = self.get_normalizer() + + def get_normalizer(self): + state_normalizer = MinMaxNormalizer(self.replay_buffer['state'][:]) # (N, obs_dim) + action_normalizer = MinMaxNormalizer(self.replay_buffer['action'][:]) # (N, action_dim) + return { + "obs": { + "state": state_normalizer + }, + "action": action_normalizer + } + + def __str__(self) -> str: + return f"Keys: {self.replay_buffer.keys()} Steps: {self.replay_buffer.n_steps} Episodes: {self.replay_buffer.n_episodes}" + + def __len__(self) -> int: + return len(self.sampler) + + def _sample_to_data(self, sample): + state = sample['state'].astype(np.float32) # (T, 5) + state = self.normalizer['obs']['state'].normalize(state) + + action = sample['action'].astype(np.float32) # (T, 2) + action = self.normalizer['action'].normalize(action) + data = { + 'obs': { + 'state': state, # T, 5 + }, + 'action': action, # T, 2 + } + return data + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + sample = self.sampler.sample_sequence(idx) + data = self._sample_to_data(sample) + torch_data = dict_apply(data, torch.tensor) + return torch_data + + +class PushTKeypointDataset(BaseDataset): + def __init__(self, + dataset_path, + obs_keys=['keypoint', 'state', 'action'], + horizon=1, + pad_before=0, + pad_after=0, + abs_action=False + ): + + super().__init__() + self.replay_buffer = ReplayBuffer.copy_from_path( + dataset_path, keys=obs_keys) + + self.sampler = SequenceSampler( + replay_buffer=self.replay_buffer, + sequence_length=horizon, + pad_before=pad_before, + pad_after=pad_after) + + self.horizon = horizon + self.pad_before = pad_before + self.pad_after = pad_after + + self.normalizer = self.get_normalizer() + + def get_normalizer(self): + agent_pos_normalizer = MinMaxNormalizer(self.replay_buffer['state'][:, :2]) # (N, 2) + keypoint_normalizer = MinMaxNormalizer(self.replay_buffer['keypoint'][:]) # (N, 9, 2) + action_normalizer = MinMaxNormalizer(self.replay_buffer['action'][:]) # (N, 2) + return { + "obs": { + "keypoint": keypoint_normalizer, + "agent_pos": agent_pos_normalizer + }, + "action": action_normalizer + } + + def __str__(self) -> str: + return f"Keys: {self.replay_buffer.keys()} Steps: {self.replay_buffer.n_steps} Episodes: {self.replay_buffer.n_episodes}" + + def __len__(self) -> int: + return len(self.sampler) + + def _sample_to_data(self, sample): + # keypoint + data_size = sample['keypoint'].shape[0] # (T, 9, 2) + keypoint = sample['keypoint'].reshape(-1, sample['keypoint'].shape[-1]).astype(np.float32) + keypoint = self.normalizer['obs']['keypoint'].normalize(keypoint) + keypoint = keypoint.reshape(data_size, -1) # (T, 18) + + # agent_pos + agent_pos = sample['state'][:,:2].astype(np.float32) # (T, 2) + agent_pos = self.normalizer['obs']['agent_pos'].normalize(agent_pos) + + # action + action = sample['action'].astype(np.float32) # (T, 2) + action = self.normalizer['action'].normalize(action) + + data = { + 'obs': { + 'keypoint': keypoint, # T, 18 + 'agent_pos': agent_pos, # T, 2 + }, + 'action': action, # T, 2 + } + return data + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + sample = self.sampler.sample_sequence(idx) + data = self._sample_to_data(sample) + torch_data = dict_apply(data, torch.tensor) + return torch_data + + +class PushTImageDataset(BaseDataset): + def __init__(self, + zarr_path, + obs_keys=['img', 'state', 'action'], + horizon=1, + pad_before=0, + pad_after=0, + abs_action=False + ): + + super().__init__() + self.replay_buffer = ReplayBuffer.copy_from_path( + zarr_path, keys=obs_keys) + + self.sampler = SequenceSampler( + replay_buffer=self.replay_buffer, + sequence_length=horizon, + pad_before=pad_before, + pad_after=pad_after) + + self.horizon = horizon + self.pad_before = pad_before + self.pad_after = pad_after + + self.normalizer = self.get_normalizer() + + def get_normalizer(self): + agent_pos_normalizer = MinMaxNormalizer(self.replay_buffer['state'][...,:2]) + image_normalizer = ImageNormalizer() + action_normalizer = MinMaxNormalizer(self.replay_buffer['action'][:]) + + return { + "obs": { + "image": image_normalizer, + "agent_pos": agent_pos_normalizer + }, + "action": action_normalizer + } + + def __str__(self) -> str: + return f"Keys: {self.replay_buffer.keys()} Steps: {self.replay_buffer.n_steps} Episodes: {self.replay_buffer.n_episodes}" + + def __len__(self) -> int: + return len(self.sampler) + + def _sample_to_data(self, sample): + # image + image = np.moveaxis(sample['img'], -1, 1) / 255 + image = self.normalizer['obs']['image'].normalize(image) + + # agent_pos + agent_pos = sample['state'][:,:2].astype(np.float32) # (T, 2) + agent_pos = self.normalizer['obs']['agent_pos'].normalize(agent_pos) + + # action + action = sample['action'].astype(np.float32) # (T, 2) + action = self.normalizer['action'].normalize(action) + + data = { + 'obs': { + 'image': image, # T, 3, 96, 96 + 'agent_pos': agent_pos, # T, 2 + }, + 'action': action, # T, 2 + } + return data + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + sample = self.sampler.sample_sequence(idx) + data = self._sample_to_data(sample) + torch_data = dict_apply(data, torch.tensor) + return torch_data diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/replay_buffer.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/replay_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..1fbb1baaf51be3a0701b454bf7f23058ca30ff6a --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/replay_buffer.py @@ -0,0 +1,592 @@ +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 ReplayBuffer: + """ + Zarr-based temporal datastructure. + Assumes first dimension to be time. Only chunk in time dimension. + """ + + def __init__(self, root: Union[zarr.Group, Dict[str, dict]]): + """ + Dummy constructor. Use copy_from* and create_from* class methods instead. + """ + assert ('data' in root) + assert ('meta' in root) + assert ('episode_ends' in root['meta']) + for key, value in root['data'].items(): + assert (value.shape[0] == root['meta']['episode_ends'][-1]) + 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 'episode_ends' not in meta: + episode_ends = meta.zeros('episode_ends', shape=(0,), dtype=np.int64, + 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) + } + } + return cls(root=root) + + @classmethod + def create_from_group(cls, group, **kwargs): + if 'data' not in group: + # create from stratch + buffer = cls.create_empty_zarr(root=group, **kwargs) + else: + # already exist + buffer = cls(root=group, **kwargs) + return buffer + + @classmethod + def create_from_path(cls, zarr_path, mode='r', **kwargs): + """ + Open an on-disk zarr directly (for dataset larger than memory). + Slower. + """ + group = zarr.open(os.path.expanduser(zarr_path), mode) + return cls.create_from_group(group, **kwargs) + + # ============= copy constructors =============== + @classmethod + def copy_from_store(cls, src_store, store=None, keys=None, + chunks: Dict[str, tuple] = dict(), + compressors: Union[dict, str, numcodecs.abc.Codec] = dict(), + if_exists='replace', + **kwargs): + """ + Load to memory. + """ + src_root = zarr.group(src_store) + root = None + if store is None: + # numpy backend + meta = dict() + for key, value in src_root['meta'].items(): + if len(value.shape) == 0: + meta[key] = np.array(value) + else: + meta[key] = value[:] + + if keys is None: + keys = src_root['data'].keys() + data = dict() + for key in keys: + arr = src_root['data'][key] + data[key] = arr[:] + + root = { + 'meta': meta, + 'data': data + } + else: + root = zarr.group(store=store) + # copy without recompression + n_copied, n_skipped, n_bytes_copied = zarr.copy_store(source=src_store, dest=store, + source_path='/meta', dest_path='/meta', + if_exists=if_exists) + data_group = root.create_group('data', overwrite=True) + if keys is None: + keys = src_root['data'].keys() + for key in keys: + value = src_root['data'][key] + cks = cls._resolve_array_chunks( + chunks=chunks, key=key, array=value) + cpr = cls._resolve_array_compressor( + compressors=compressors, key=key, array=value) + if cks == value.chunks and cpr == value.compressor: + # copy without recompression + this_path = '/data/' + key + n_copied, n_skipped, n_bytes_copied = zarr.copy_store( + source=src_store, dest=store, + source_path=this_path, dest_path=this_path, + if_exists=if_exists + ) + else: + # copy with recompression + n_copied, n_skipped, n_bytes_copied = zarr.copy( + source=value, dest=data_group, name=key, + chunks=cks, compressor=cpr, if_exists=if_exists + ) + buffer = cls(root=root) + return buffer + + @classmethod + def copy_from_path(cls, zarr_path, backend=None, store=None, keys=None, + chunks: Dict[str, tuple] = dict(), + compressors: Union[dict, str, numcodecs.abc.Codec] = dict(), + if_exists='replace', + **kwargs): + """ + Copy a on-disk zarr to in-memory compressed. + Recommended + """ + if backend == 'numpy': + print('backend argument is deprecated!') + store = None + group = zarr.open(os.path.expanduser(zarr_path), 'r') + return cls.copy_from_store(src_store=group.store, store=store, + keys=keys, chunks=chunks, compressors=compressors, + if_exists=if_exists, **kwargs) + + # ============= 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): + + root = zarr.group(store) + if self.backend == 'zarr': + # recompression free copy + n_copied, n_skipped, n_bytes_copied = zarr.copy_store( + source=self.root.store, dest=store, + source_path='/meta', dest_path='/meta', if_exists=if_exists) + else: + meta_group = root.create_group('meta', overwrite=True) + # save meta, no chunking + for key, value in self.root['meta'].items(): + _ = meta_group.array( + name=key, + data=value, + shape=value.shape, + chunks=value.shape) + + # save data, chunk + data_group = root.create_group('data', overwrite=True) + for key, value in self.root['data'].items(): + cks = self._resolve_array_chunks( + chunks=chunks, key=key, array=value) + cpr = self._resolve_array_compressor( + compressors=compressors, key=key, array=value) + if isinstance(value, zarr.Array): + if cks == value.chunks and cpr == value.compressor: + # copy without recompression + this_path = '/data/' + key + n_copied, n_skipped, n_bytes_copied = zarr.copy_store( + source=self.root.store, dest=store, + source_path=this_path, dest_path=this_path, if_exists=if_exists) + else: + # copy with recompression + n_copied, n_skipped, n_bytes_copied = zarr.copy( + source=value, dest=data_group, name=key, + chunks=cks, compressor=cpr, if_exists=if_exists + ) + else: + # numpy + _ = data_group.array( + name=key, + data=value, + chunks=cks, + compressor=cpr + ) + return store + + 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): + store = zarr.DirectoryStore(os.path.expanduser(zarr_path)) + return self.save_to_store(store, chunks=chunks, + compressors=compressors, if_exists=if_exists, **kwargs) + + @staticmethod + def resolve_compressor(compressor='default'): + if compressor == 'default': + compressor = numcodecs.Blosc(cname='lz4', clevel=5, + shuffle=numcodecs.Blosc.NOSHUFFLE) + elif compressor == 'disk': + compressor = numcodecs.Blosc('zstd', clevel=5, + shuffle=numcodecs.Blosc.BITSHUFFLE) + return compressor + + @classmethod + def _resolve_array_compressor(cls, + compressors: Union[dict, str, numcodecs.abc.Codec], key, array): + # allows compressor to be explicitly set to None + cpr = 'nil' + if isinstance(compressors, dict): + if key in compressors: + cpr = cls.resolve_compressor(compressors[key]) + elif isinstance(array, zarr.Array): + cpr = array.compressor + else: + cpr = cls.resolve_compressor(compressors) + # backup default + if cpr == 'nil': + cpr = cls.resolve_compressor('default') + return cpr + + @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 + + # ============= properties ================= + @cached_property + def data(self): + return self.root['data'] + + @cached_property + def meta(self): + return self.root['meta'] + + def update_meta(self, data): + # sanitize data + np_data = dict() + for key, value in data.items(): + if isinstance(value, np.ndarray): + np_data[key] = value + else: + arr = np.array(value) + if arr.dtype == object: + raise TypeError(f"Invalid value type {type(value)}") + np_data[key] = arr + + meta_group = self.meta + if self.backend == 'zarr': + for key, value in np_data.items(): + _ = meta_group.array( + name=key, + data=value, + shape=value.shape, + chunks=value.shape, + overwrite=True) + else: + meta_group.update(np_data) + + return meta_group + + @property + def episode_ends(self): + return self.meta['episode_ends'] + + def get_episode_idxs(self): + import numba + numba.jit(nopython=True) + + def _get_episode_idxs(episode_ends): + result = np.zeros((episode_ends[-1],), dtype=np.int64) + for i in range(len(episode_ends)): + start = 0 + if i > 0: + start = episode_ends[i - 1] + end = episode_ends[i] + for idx in range(start, end): + result[idx] = i + return result + + return _get_episode_idxs(self.episode_ends) + + @property + def backend(self): + backend = 'numpy' + if isinstance(self.root, zarr.Group): + backend = 'zarr' + return backend + + # =========== dict-like API ============== + def __repr__(self) -> str: + if self.backend == 'zarr': + return str(self.root.tree()) + else: + return super().__repr__() + + def keys(self): + return self.data.keys() + + def values(self): + return self.data.values() + + def items(self): + return self.data.items() + + def __getitem__(self, key): + return self.data[key] + + def __contains__(self, key): + return key in self.data + + # =========== our API ============== + @property + def n_steps(self): + if len(self.episode_ends) == 0: + return 0 + return self.episode_ends[-1] + + @property + def n_episodes(self): + return len(self.episode_ends) + + @property + def chunk_size(self): + if self.backend == 'zarr': + return next(iter(self.data.arrays()))[-1].chunks[0] + return None + + @property + def episode_lengths(self): + ends = self.episode_ends[:] + ends = np.insert(ends, 0, 0) + lengths = np.diff(ends) + return lengths + + def add_episode(self, + data: Dict[str, np.ndarray], + chunks: Optional[Dict[str, tuple]] = dict(), + compressors: Union[str, numcodecs.abc.Codec, dict] = dict()): + assert (len(data) > 0) + is_zarr = (self.backend == 'zarr') + + curr_len = self.n_steps + episode_length = None + for key, value in data.items(): + assert (len(value.shape) >= 1) + if episode_length is None: + episode_length = len(value) + else: + assert (episode_length == len(value)) + new_len = curr_len + episode_length + + for key, value in data.items(): + new_shape = (new_len,) + value.shape[1:] + # create array + if key not in self.data: + if is_zarr: + cks = self._resolve_array_chunks( + chunks=chunks, key=key, array=value) + cpr = self._resolve_array_compressor( + compressors=compressors, key=key, array=value) + arr = self.data.zeros(name=key, + shape=new_shape, + chunks=cks, + dtype=value.dtype, + compressor=cpr) + else: + # copy data to prevent modify + arr = np.zeros(shape=new_shape, dtype=value.dtype) + self.data[key] = arr + else: + arr = self.data[key] + assert (value.shape[1:] == arr.shape[1:]) + # same method for both zarr and numpy + if is_zarr: + arr.resize(new_shape) + else: + arr.resize(new_shape, refcheck=False) + # copy data + arr[-value.shape[0]:] = value + + # append to episode ends + episode_ends = self.episode_ends + if is_zarr: + episode_ends.resize(episode_ends.shape[0] + 1) + else: + episode_ends.resize(episode_ends.shape[0] + 1, refcheck=False) + episode_ends[-1] = new_len + + # rechunk + if is_zarr: + if episode_ends.chunks[0] < episode_ends.shape[0]: + rechunk_recompress_array(self.meta, 'episode_ends', + chunk_length=int(episode_ends.shape[0] * 1.5)) + + def drop_episode(self): + is_zarr = (self.backend == 'zarr') + episode_ends = self.episode_ends[:].copy() + assert (len(episode_ends) > 0) + start_idx = 0 + if len(episode_ends) > 1: + start_idx = episode_ends[-2] + for key, value in self.data.items(): + new_shape = (start_idx,) + value.shape[1:] + if is_zarr: + value.resize(new_shape) + else: + value.resize(new_shape, refcheck=False) + if is_zarr: + self.episode_ends.resize(len(episode_ends) - 1) + else: + self.episode_ends.resize(len(episode_ends) - 1, refcheck=False) + + def pop_episode(self): + assert (self.n_episodes > 0) + episode = self.get_episode(self.n_episodes - 1, copy=True) + self.drop_episode() + return episode + + def extend(self, data): + self.add_episode(data) + + def get_episode(self, idx, copy=False): + idx = list(range(len(self.episode_ends)))[idx] + start_idx = 0 + if idx > 0: + start_idx = self.episode_ends[idx - 1] + end_idx = self.episode_ends[idx] + result = self.get_steps_slice(start_idx, end_idx, copy=copy) + return result + + def get_episode_slice(self, idx): + start_idx = 0 + if idx > 0: + start_idx = self.episode_ends[idx - 1] + end_idx = self.episode_ends[idx] + return slice(start_idx, end_idx) + + def get_steps_slice(self, start, stop, step=None, copy=False): + _slice = slice(start, stop, step) + + result = dict() + for key, value in self.data.items(): + x = value[_slice] + if copy and isinstance(value, np.ndarray): + x = x.copy() + result[key] = x + return result + + # =========== chunking ============= + def get_chunks(self) -> dict: + assert self.backend == 'zarr' + chunks = dict() + for key, value in self.data.items(): + chunks[key] = value.chunks + return chunks + + def set_chunks(self, chunks: dict): + assert self.backend == 'zarr' + for key, value in chunks.items(): + if key in self.data: + arr = self.data[key] + if value != arr.chunks: + check_chunks_compatible(chunks=value, shape=arr.shape) + rechunk_recompress_array(self.data, key, chunks=value) + + def get_compressors(self) -> dict: + assert self.backend == 'zarr' + compressors = dict() + for key, value in self.data.items(): + compressors[key] = value.compressor + return compressors + + def set_compressors(self, compressors: dict): + assert self.backend == 'zarr' + for key, value in compressors.items(): + if key in self.data: + arr = self.data[key] + compressor = self.resolve_compressor(value) + if compressor != arr.compressor: + rechunk_recompress_array(self.data, key, compressor=compressor) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/robomimic_dataset.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/robomimic_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..63258612e0f1b06a5335f6e166483106dcadd682 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/robomimic_dataset.py @@ -0,0 +1,506 @@ +from typing import Dict +import torch +import numpy as np +import copy +import pathlib +import h5py +import zarr +from tqdm import tqdm +import concurrent.futures +from collections import defaultdict +from cleandiffuser.dataset.imagecodecs import register_codecs, Jpeg2k +from cleandiffuser.dataset.base_dataset import BaseDataset +from cleandiffuser.dataset.replay_buffer import ReplayBuffer +from cleandiffuser.dataset.dataset_utils import SequenceSampler, EmptyNormalizer, RotationTransformer, dict_apply, \ + MinMaxNormalizer, ImageNormalizer + +register_codecs() + + +class RobomimicDataset(BaseDataset): + def __init__(self, + dataset_dir, + horizon=1, + pad_before=0, + pad_after=0, + obs_keys=('object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos'), + abs_action=False, + rotation_rep='rotation_6d', + ): + super().__init__() + self.rotation_transformer = RotationTransformer( + from_rep='axis_angle', to_rep=rotation_rep) + + self.replay_buffer = ReplayBuffer.create_empty_numpy() + with h5py.File(dataset_dir) as file: + demos = file['data'] + for i in tqdm(range(len(demos)), desc="Loading hdf5 to ReplayBuffer"): + demo = demos[f'demo_{i}'] + episode = _data_to_obs( + raw_obs=demo['obs'], + raw_actions=demo['actions'][:].astype(np.float32), + obs_keys=obs_keys, + abs_action=abs_action, + rotation_transformer=self.rotation_transformer) + self.replay_buffer.add_episode(episode) + + self.sampler = SequenceSampler( + replay_buffer=self.replay_buffer, + sequence_length=horizon, + pad_before=pad_before, + pad_after=pad_after) + + self.horizon = horizon + self.pad_before = pad_before + self.pad_after = pad_after + self.abs_action = abs_action + self.normalizer = self.get_normalizer() + + def undo_transform_action(self, action): + raw_shape = action.shape + if raw_shape[-1] == 20: + # dual arm + action = action.reshape(-1, 2, 10) + + d_rot = action.shape[-1] - 4 + pos = action[..., :3] + rot = action[..., 3:3 + d_rot] + gripper = action[..., [-1]] + rot = self.rotation_transformer.inverse(rot) + uaction = np.concatenate([ + pos, rot, gripper + ], axis=-1) + + if raw_shape[-1] == 20: + # dual arm + uaction = uaction.reshape(*raw_shape[:-1], 14) + + return uaction + + def get_normalizer(self): + if self.abs_action: + state_normalizer = MinMaxNormalizer(self.replay_buffer['obs'][:]) # (N, obs_dim) + action_normalizer = MinMaxNormalizer(self.replay_buffer['action'][:]) # (N, action_dim) + else: + state_normalizer = MinMaxNormalizer(self.replay_buffer['obs'][:]) # (N, obs_dim) + action_normalizer = MinMaxNormalizer(self.replay_buffer['action'][:]) # (N, action_dim) + return { + "obs": { + "state": state_normalizer + }, + "action": action_normalizer + } + + def sample_to_data(self, sample): + state = sample['obs'].astype(np.float32) + state = self.normalizer['obs']['state'].normalize(state) + + action = sample['action'].astype(np.float32) + action = self.normalizer['action'].normalize(action) + data = { + 'obs': { + 'state': state + }, + 'action': action, + } + return data + + def __str__(self) -> str: + return f"Keys: {self.replay_buffer.keys()} Steps: {self.replay_buffer.n_steps} Episodes: {self.replay_buffer.n_episodes}" + + def __len__(self) -> int: + return len(self.sampler) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + sample = self.sampler.sample_sequence(idx) + data = self.sample_to_data(sample) + torch_data = dict_apply(data, torch.tensor) + return torch_data + + +def _data_to_obs(raw_obs, raw_actions, obs_keys, abs_action, rotation_transformer): + obs = np.concatenate([ + raw_obs[key] for key in obs_keys + ], axis=-1).astype(np.float32) + + if abs_action: + is_dual_arm = False + if raw_actions.shape[-1] == 14: + # dual arm + raw_actions = raw_actions.reshape(-1, 2, 7) + is_dual_arm = True + + pos = raw_actions[..., :3] + rot = raw_actions[..., 3:6] + gripper = raw_actions[..., 6:] + rot = rotation_transformer.forward(rot) + raw_actions = np.concatenate([ + pos, rot, gripper + ], axis=-1).astype(np.float32) + + if is_dual_arm: + raw_actions = raw_actions.reshape(-1, 20) + + data = { + 'obs': obs, + 'action': raw_actions + } + return data + + +class RobomimicImageDataset(BaseDataset): + def __init__(self, + dataset_dir, + shape_meta: dict, + n_obs_steps=None, + horizon=1, + pad_before=0, + pad_after=0, + abs_action=False, + rotation_rep='rotation_6d', + ): + super().__init__() + self.rotation_transformer = RotationTransformer( + from_rep='axis_angle', to_rep=rotation_rep) + + self.replay_buffer = _convert_robomimic_to_replay( + store=zarr.MemoryStore(), + shape_meta=shape_meta, + dataset_path=dataset_dir, + abs_action=abs_action, + rotation_transformer=self.rotation_transformer) + + rgb_keys = list() + lowdim_keys = list() + obs_shape_meta = shape_meta['obs'] + for key, attr in obs_shape_meta.items(): + type = attr.get('type', 'low_dim') + if type == 'rgb': + rgb_keys.append(key) + elif type == 'low_dim': + lowdim_keys.append(key) + + key_first_k = dict() + if n_obs_steps is not None: + # only take first k obs from images + for key in rgb_keys + lowdim_keys: + key_first_k[key] = n_obs_steps + self.sampler = SequenceSampler( + replay_buffer=self.replay_buffer, + sequence_length=horizon, + pad_before=pad_before, + pad_after=pad_after, + key_first_k=key_first_k + ) + + self.shape_meta = shape_meta + self.rgb_keys = rgb_keys + self.lowdim_keys = lowdim_keys + self.abs_action = abs_action + self.horizon = horizon + self.pad_before = pad_before + self.pad_after = pad_after + self.n_obs_steps = n_obs_steps + + self.normalizer = self.get_normalizer() + + def get_normalizer(self): + normalizer = defaultdict(dict) + for key in self.lowdim_keys: + normalizer['obs'][key] = MinMaxNormalizer(self.replay_buffer[key][:]) + for key in self.rgb_keys: + normalizer['obs'][key] = ImageNormalizer() + normalizer['action'] = MinMaxNormalizer(self.replay_buffer['action'][:]) + + return normalizer + + def __str__(self) -> str: + return f"Keys: {self.replay_buffer.keys()} Steps: {self.replay_buffer.n_steps} Episodes: {self.replay_buffer.n_episodes}" + + def __len__(self) -> int: + return len(self.sampler) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + sample = self.sampler.sample_sequence(idx) + + # obs + # to save RAM, only return first n_obs_steps of OBS + # since the rest will be discarded anyway. + # when self.n_obs_steps is None + # this slice does nothing (takes all) + T_slice = slice(self.n_obs_steps) + + obs_dict = dict() + for key in self.rgb_keys: + # move channel last to channel first + # T,H,W,C + # convert uint8 image to float32 + obs_dict[key] = np.moveaxis(sample[key][T_slice], -1, 1 + ).astype(np.float32) / 255. + # T,C,H,W + del sample[key] + obs_dict[key] = self.normalizer['obs'][key].normalize(obs_dict[key]) + + for key in self.lowdim_keys: + obs_dict[key] = sample[key][T_slice].astype(np.float32) + del sample[key] + obs_dict[key] = self.normalizer['obs'][key].normalize(obs_dict[key]) + + # action + action = sample['action'].astype(np.float32) + action = self.normalizer['action'].normalize(action) + + torch_data = { + 'obs': dict_apply(obs_dict, torch.tensor), + 'action': torch.tensor(action) + } + return torch_data + + def undo_transform_action(self, action): + raw_shape = action.shape + if raw_shape[-1] == 20: + # dual arm + action = action.reshape(-1, 2, 10) + + d_rot = action.shape[-1] - 4 + pos = action[..., :3] + rot = action[..., 3:3 + d_rot] + gripper = action[..., [-1]] + rot = self.rotation_transformer.inverse(rot) + uaction = np.concatenate([ + pos, rot, gripper + ], axis=-1) + + if raw_shape[-1] == 20: + # dual arm + uaction = uaction.reshape(*raw_shape[:-1], 14) + + return uaction + + +def _convert_actions(raw_actions, abs_action, rotation_transformer): + actions = raw_actions + if abs_action: + is_dual_arm = False + if raw_actions.shape[-1] == 14: + # dual arm + raw_actions = raw_actions.reshape(-1, 2, 7) + is_dual_arm = True + + pos = raw_actions[..., :3] + rot = raw_actions[..., 3:6] + gripper = raw_actions[..., 6:] + rot = rotation_transformer.forward(rot) + raw_actions = np.concatenate([ + pos, rot, gripper + ], axis=-1).astype(np.float32) + + if is_dual_arm: + raw_actions = raw_actions.reshape(-1, 20) + actions = raw_actions + return actions + + +def _convert_robomimic_to_replay(store, shape_meta, dataset_path, abs_action, rotation_transformer, + n_workers=None, max_inflight_tasks=None): + """ Convert Robomimic dataset to ReplayBuffer + + A ReplayBuffer is a `zarr.Group` or Dict[str, dict] that contains the following keys: + - data: zarr.Group or Dict[str, dict] + Contains the data. All data should be stored as numpy arrays with the same length. + - meta: zarr.Group or Dict[str, dict] + Contains key "episode_ends", which is a numpy array of shape (n_episodes,) that contains the + end index of each episode in the data. + + Args: + - store: zarr.Store + zarr.MemoryStore() + - shape_meta: dict + Shape metadata of the dataset. Should contain keys 'obs', 'action'. + For example: + shape_meta = { + "action": {"shape": [10, ]}, + "obs": { + "agentview_image": {"shape": [84, 84, 3], "type": "rgb"}, + "robot0_eef_pos": {"shape": [3, ], "type": "low_dim"}, + }} + - dataset_path: str + Path to the Robomimic dataset + - abs_action: bool + Whether to use position or velocity control + - rotation_transformer: RotationTransformer + Rotation transformer to convert rotation representation + """ + + """ Dataset structure of Can-PH, as an example: + - data + - demo_0 + - actions (118, 7) + - dones (118, ) + - next_obs + - agentview_image (118, 84, 84, 3) + - object (118, 14) + - robot0_eef_pos (118, 3) + - robot0_eef_quat + - robot0_eef_vel_ang + - robot0_eef_vel_lin + - robot0_eye_in_hand_image + - robot0_gripper_qpos + - robot0_gripper_qvel + - robot0_joint_pos + - robot0_joint_pos_cos + - robot0_joint_pos_sin + - robot0_joint_vel + - obs + ... + - rewards (118, ) + - states (118, 71) + - demo_1 + ...(x200 demos) + - mask + - 20_percent + - 20_percent_train + - 20_percent_valid + - 50_percent + - 50_percent_train + - 50_percent_valid + - train (180,) + - valid (20,) + + Suppose that the `shape_meta` is: + shape_meta = { + "action": {"shape": [10, ]}, + "obs": { + "agentview_image": { + "shape": [3, 84, 84], "type": "rgb", }, + "robot0_eye_in_hand_image": { + "shape": [3, 84, 84], "type": "rgb", }, + "robot0_eef_pos": { + "shape": [3, ], "type": "low_dim", }, + "robot0_eef_quat": { + "shape": [4, ], "type": "low_dim", }, + "robot0_gripper_qpos": { + "shape": [2, ], "type": "low_dim", }, }} + """ + + import multiprocessing + if n_workers is None: + n_workers = multiprocessing.cpu_count() + if max_inflight_tasks is None: + max_inflight_tasks = n_workers * 5 + + # parse shape_meta + rgb_keys = list() + lowdim_keys = list() + # construct compressors and chunks + obs_shape_meta = shape_meta['obs'] + for key, attr in obs_shape_meta.items(): + shape = attr['shape'] + type = attr.get('type', 'low_dim') + if type == 'rgb': + rgb_keys.append(key) + elif type == 'low_dim': + lowdim_keys.append(key) + # rgb_keys = ['agentview_image', 'robot0_eye_in_hand_image'] + # lowdim_keys = ['robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos'] + + # create zarr group + root = zarr.group(store) + data_group = root.require_group('data', overwrite=True) + meta_group = root.require_group('meta', overwrite=True) + + with h5py.File(dataset_path) as file: + + # count total steps + demos = file['data'] + episode_ends = list() + prev_end = 0 + for i in range(len(demos)): + demo = demos[f'demo_{i}'] + episode_length = demo['actions'].shape[0] + episode_end = prev_end + episode_length + prev_end = episode_end + episode_ends.append(episode_end) + n_steps = episode_ends[-1] + episode_starts = [0] + episode_ends[:-1] + _ = meta_group.array('episode_ends', episode_ends, + dtype=np.int64, compressor=None, overwrite=True) + + # save lowdim data + for key in tqdm(lowdim_keys + ['action'], desc="Loading lowdim data"): + data_key = 'obs/' + key + if key == 'action': + data_key = 'actions' + this_data = list() + for i in range(len(demos)): + demo = demos[f'demo_{i}'] + this_data.append(demo[data_key][:].astype(np.float32)) + this_data = np.concatenate(this_data, axis=0) + if key == 'action': + this_data = _convert_actions( + raw_actions=this_data, + abs_action=abs_action, + rotation_transformer=rotation_transformer + ) + assert this_data.shape == (n_steps,) + tuple(shape_meta['action']['shape']) + else: + assert this_data.shape == (n_steps,) + tuple(shape_meta['obs'][key]['shape']) + _ = data_group.array( + name=key, + data=this_data, + shape=this_data.shape, + chunks=this_data.shape, + compressor=None, + dtype=this_data.dtype + ) + + def img_copy(zarr_arr, zarr_idx, hdf5_arr, hdf5_idx): + try: + zarr_arr[zarr_idx] = hdf5_arr[hdf5_idx] + # make sure we can successfully decode + _ = zarr_arr[zarr_idx] + return True + except Exception as e: + return False + + with tqdm(total=n_steps * len(rgb_keys), desc="Loading image data", mininterval=1.0) as pbar: + # one chunk per thread, therefore no synchronization needed + with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as executor: + futures = set() + for key in rgb_keys: + data_key = 'obs/' + key + shape = tuple(shape_meta['obs'][key]['shape']) + c, h, w = shape + this_compressor = Jpeg2k(level=50) + img_arr = data_group.require_dataset( + name=key, + shape=(n_steps, h, w, c), + chunks=(1, h, w, c), + compressor=this_compressor, + dtype=np.uint8 + ) + for episode_idx in range(len(demos)): + demo = demos[f'demo_{episode_idx}'] + hdf5_arr = demo['obs'][key] + for hdf5_idx in range(hdf5_arr.shape[0]): + if len(futures) >= max_inflight_tasks: + # limit number of inflight tasks + completed, futures = concurrent.futures.wait(futures, + return_when=concurrent.futures.FIRST_COMPLETED) + for f in completed: + if not f.result(): + raise RuntimeError('Failed to encode image!') + pbar.update(len(completed)) + + zarr_idx = episode_starts[episode_idx] + hdf5_idx + futures.add( + executor.submit(img_copy, + img_arr, zarr_idx, hdf5_arr, hdf5_idx)) + completed, futures = concurrent.futures.wait(futures) + for f in completed: + if not f.result(): + raise RuntimeError('Failed to encode image!') + pbar.update(len(completed)) + + replay_buffer = ReplayBuffer(root) + return replay_buffer diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/robomimic_rl_dataset.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/robomimic_rl_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..970732c0015d50a3d2c91b3020d835a4426ba689 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/robomimic_rl_dataset.py @@ -0,0 +1,450 @@ +from typing import Dict +import torch +import numpy as np +import copy +import pathlib +import h5py +import zarr +from tqdm import tqdm +import concurrent.futures +from collections import defaultdict +from cleandiffuser.dataset.imagecodecs import register_codecs, Jpeg2k +from cleandiffuser.dataset.base_dataset import BaseDataset +from cleandiffuser.dataset.replay_buffer import ReplayBuffer +from cleandiffuser.dataset.dataset_utils import SequenceSampler, EmptyNormalizer, RotationTransformer, dict_apply, \ + MinMaxNormalizer, ImageNormalizer + +register_codecs() + + +class RobomimicDataset(BaseDataset): + def __init__( + self, + dataset_dir, + horizon=1, + pad_before=0, + pad_after=0, + obs_keys=['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos'], + abs_action=False, + rotation_rep='rotation_6d', + ): + super().__init__() + self.rotation_transformer = RotationTransformer( + from_rep='axis_angle', to_rep=rotation_rep) + + self.replay_buffer = ReplayBuffer.create_empty_numpy() + with h5py.File(dataset_dir) as file: + demos = file['data'] + for i in tqdm(range(len(demos)), desc="Loading hdf5 to ReplayBuffer"): + demo = demos[f'demo_{i}'] + episode = _data_to_obs( + raw_obs=demo['obs'], + raw_actions=demo['actions'][:].astype(np.float32), + rewards=demo['rewards'], + next_raw_obs=demo['next_obs'], + dones=demo['dones'], + obs_keys=obs_keys, + abs_action=abs_action, + rotation_transformer=self.rotation_transformer) + self.replay_buffer.add_episode(episode) + + self.sampler = SequenceSampler( + replay_buffer=self.replay_buffer, + sequence_length=horizon, + pad_before=pad_before, + pad_after=pad_after) + + self.horizon = horizon + self.pad_before = pad_before + self.pad_after = pad_after + self.abs_action = abs_action + self.normalizer = self.get_normalizer() + + def undo_transform_action(self, action): + raw_shape = action.shape + if raw_shape[-1] == 20: + # dual arm + action = action.reshape(-1, 2, 10) + + d_rot = action.shape[-1] - 4 + pos = action[..., :3] + rot = action[..., 3:3 + d_rot] + gripper = action[..., [-1]] + rot = self.rotation_transformer.inverse(rot) + uaction = np.concatenate([ + pos, rot, gripper + ], axis=-1) + + if raw_shape[-1] == 20: + # dual arm + uaction = uaction.reshape(*raw_shape[:-1], 14) + + return uaction + + def get_normalizer(self): + state_normalizer = MinMaxNormalizer(self.replay_buffer['obs'][:]) # (N, obs_dim) + next_state_normalizer = MinMaxNormalizer(self.replay_buffer['next_obs'][:]) # (N, obs_dim) + action_normalizer = MinMaxNormalizer(self.replay_buffer['action'][:]) # (N, action_dim) + + return { + "obs": { + "state": state_normalizer + }, + "next_obs": { + "state": next_state_normalizer + }, + "action": action_normalizer, + "reward": reward_normalizer, + } + + def sample_to_data(self, sample): + state = sample['obs'].astype(np.float32) + state = self.normalizer['obs']['state'].normalize(state) + + action = sample['action'].astype(np.float32) + action = self.normalizer['action'].normalize(action) + + next_state = sample['next_obs'].astype(np.float32) + next_state = self.normalizer['next_obs']['state'].normalize(next_state) + + reward = sample['reward'].astype(np.float32) + done = sample['done'].astype(np.float32) + + data = { + 'obs': { + 'state': state + }, + 'next_obs': { + 'state': next_state + }, + 'action': action, + 'reward': reward, + 'done': done, + } + return data + + def __str__(self) -> str: + return f"Keys: {self.replay_buffer.keys()} Steps: {self.replay_buffer.n_steps} Episodes: {self.replay_buffer.n_episodes}" + + def __len__(self) -> int: + return len(self.sampler) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + sample = self.sampler.sample_sequence(idx) + data = self.sample_to_data(sample) + torch_data = dict_apply(data, torch.from_numpy) + return torch_data + + +def _data_to_obs(raw_obs, raw_actions, rewards, next_raw_obs, dones, obs_keys, abs_action, rotation_transformer): + # for key in obs_keys: + # print(key, raw_obs[key].shape) + # exit() + obs = np.concatenate([ + raw_obs[key] for key in obs_keys + ], axis=-1).astype(np.float32) + + next_obs = np.concatenate([ + next_raw_obs[key] for key in obs_keys + ], axis=-1).astype(np.float32) + + if abs_action: + is_dual_arm = False + if raw_actions.shape[-1] == 14: + # dual arm + raw_actions = raw_actions.reshape(-1, 2, 7) + is_dual_arm = True + + pos = raw_actions[..., :3] + rot = raw_actions[..., 3:6] + gripper = raw_actions[..., 6:] + rot = rotation_transformer.forward(rot) + raw_actions = np.concatenate([ + pos, rot, gripper + ], axis=-1).astype(np.float32) + + if is_dual_arm: + raw_actions = raw_actions.reshape(-1, 20) + + data = { + 'obs': obs, + 'action': raw_actions, + 'reward': rewards, + 'next_obs': next_obs, + 'done': dones, + } + return data + + +class RobomimicImageDataset(BaseDataset): + def __init__(self, + dataset_dir, + shape_meta: dict, + n_obs_steps=None, + horizon=1, + pad_before=0, + pad_after=0, + abs_action=False, + rotation_rep='rotation_6d', + ): + super().__init__() + self.rotation_transformer = RotationTransformer( + from_rep='axis_angle', to_rep=rotation_rep) + + self.replay_buffer = _convert_robomimic_to_replay( + store=zarr.MemoryStore(), + shape_meta=shape_meta, + dataset_path=dataset_dir, + abs_action=abs_action, + rotation_transformer=self.rotation_transformer) + + rgb_keys = list() + lowdim_keys = list() + obs_shape_meta = shape_meta['obs'] + for key, attr in obs_shape_meta.items(): + type = attr.get('type', 'low_dim') + if type == 'rgb': + rgb_keys.append(key) + elif type == 'low_dim': + lowdim_keys.append(key) + + key_first_k = dict() + if n_obs_steps is not None: + # only take first k obs from images + for key in rgb_keys + lowdim_keys: + key_first_k[key] = n_obs_steps + self.sampler = SequenceSampler( + replay_buffer=self.replay_buffer, + sequence_length=horizon, + pad_before=pad_before, + pad_after=pad_after, + key_first_k=key_first_k + ) + + self.shape_meta = shape_meta + self.rgb_keys = rgb_keys + self.lowdim_keys = lowdim_keys + self.abs_action = abs_action + self.horizon = horizon + self.pad_before = pad_before + self.pad_after = pad_after + self.n_obs_steps = n_obs_steps + + self.normalizer = self.get_normalizer() + + def get_normalizer(self): + normalizer = defaultdict(dict) + for key in self.lowdim_keys: + normalizer['obs'][key] = MinMaxNormalizer(self.replay_buffer[key][:]) + for key in self.rgb_keys: + normalizer['obs'][key] = ImageNormalizer() + normalizer['action'] = MinMaxNormalizer(self.replay_buffer['action'][:]) + + return normalizer + + def __str__(self) -> str: + return f"Keys: {self.replay_buffer.keys()} Steps: {self.replay_buffer.n_steps} Episodes: {self.replay_buffer.n_episodes}" + + def __len__(self) -> int: + return len(self.sampler) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + sample = self.sampler.sample_sequence(idx) + + # obs + # to save RAM, only return first n_obs_steps of OBS + # since the rest will be discarded anyway. + # when self.n_obs_steps is None + # this slice does nothing (takes all) + T_slice = slice(self.n_obs_steps) + + obs_dict = dict() + for key in self.rgb_keys: + # move channel last to channel first + # T,H,W,C + # convert uint8 image to float32 + obs_dict[key] = np.moveaxis(sample[key][T_slice], -1, 1 + ).astype(np.float32) / 255. + # T,C,H,W + del sample[key] + obs_dict[key] = self.normalizer['obs'][key].normalize(obs_dict[key]) + + for key in self.lowdim_keys: + obs_dict[key] = sample[key][T_slice].astype(np.float32) + del sample[key] + obs_dict[key] = self.normalizer['obs'][key].normalize(obs_dict[key]) + + # action + action = sample['action'].astype(np.float32) + action = self.normalizer['action'].normalize(action) + + torch_data = { + 'obs': dict_apply(obs_dict, torch.from_numpy), + 'action': torch.from_numpy(action) + } + return torch_data + + def undo_transform_action(self, action): + raw_shape = action.shape + if raw_shape[-1] == 20: + # dual arm + action = action.reshape(-1, 2, 10) + + d_rot = action.shape[-1] - 4 + pos = action[..., :3] + rot = action[..., 3:3 + d_rot] + gripper = action[..., [-1]] + rot = self.rotation_transformer.inverse(rot) + uaction = np.concatenate([ + pos, rot, gripper + ], axis=-1) + + if raw_shape[-1] == 20: + # dual arm + uaction = uaction.reshape(*raw_shape[:-1], 14) + + return uaction + + +def _convert_actions(raw_actions, abs_action, rotation_transformer): + actions = raw_actions + if abs_action: + is_dual_arm = False + if raw_actions.shape[-1] == 14: + # dual arm + raw_actions = raw_actions.reshape(-1, 2, 7) + is_dual_arm = True + + pos = raw_actions[..., :3] + rot = raw_actions[..., 3:6] + gripper = raw_actions[..., 6:] + rot = rotation_transformer.forward(rot) + raw_actions = np.concatenate([ + pos, rot, gripper + ], axis=-1).astype(np.float32) + + if is_dual_arm: + raw_actions = raw_actions.reshape(-1, 20) + actions = raw_actions + return actions + + +def _convert_robomimic_to_replay(store, shape_meta, dataset_path, abs_action, rotation_transformer, + n_workers=None, max_inflight_tasks=None): + import multiprocessing + if n_workers is None: + n_workers = multiprocessing.cpu_count() + if max_inflight_tasks is None: + max_inflight_tasks = n_workers * 5 + + # parse shape_meta + rgb_keys = list() + lowdim_keys = list() + # construct compressors and chunks + obs_shape_meta = shape_meta['obs'] + for key, attr in obs_shape_meta.items(): + shape = attr['shape'] + type = attr.get('type', 'low_dim') + if type == 'rgb': + rgb_keys.append(key) + elif type == 'low_dim': + lowdim_keys.append(key) + + root = zarr.group(store) + data_group = root.require_group('data', overwrite=True) + meta_group = root.require_group('meta', overwrite=True) + + with h5py.File(dataset_path) as file: + # count total steps + demos = file['data'] + episode_ends = list() + prev_end = 0 + for i in range(len(demos)): + demo = demos[f'demo_{i}'] + episode_length = demo['actions'].shape[0] + episode_end = prev_end + episode_length + prev_end = episode_end + episode_ends.append(episode_end) + n_steps = episode_ends[-1] + episode_starts = [0] + episode_ends[:-1] + _ = meta_group.array('episode_ends', episode_ends, + dtype=np.int64, compressor=None, overwrite=True) + + # save lowdim data + for key in tqdm(lowdim_keys + ['action'], desc="Loading lowdim data"): + data_key = 'obs/' + key + if key == 'action': + data_key = 'actions' + this_data = list() + for i in range(len(demos)): + demo = demos[f'demo_{i}'] + this_data.append(demo[data_key][:].astype(np.float32)) + this_data = np.concatenate(this_data, axis=0) + if key == 'action': + this_data = _convert_actions( + raw_actions=this_data, + abs_action=abs_action, + rotation_transformer=rotation_transformer + ) + assert this_data.shape == (n_steps,) + tuple(shape_meta['action']['shape']) + else: + assert this_data.shape == (n_steps,) + tuple(shape_meta['obs'][key]['shape']) + _ = data_group.array( + name=key, + data=this_data, + shape=this_data.shape, + chunks=this_data.shape, + compressor=None, + dtype=this_data.dtype + ) + + def img_copy(zarr_arr, zarr_idx, hdf5_arr, hdf5_idx): + try: + zarr_arr[zarr_idx] = hdf5_arr[hdf5_idx] + # make sure we can successfully decode + _ = zarr_arr[zarr_idx] + return True + except Exception as e: + return False + + with tqdm(total=n_steps * len(rgb_keys), desc="Loading image data", mininterval=1.0) as pbar: + # one chunk per thread, therefore no synchronization needed + with concurrent.futures.ThreadPoolExecutor(max_workers=n_workers) as executor: + futures = set() + for key in rgb_keys: + data_key = 'obs/' + key + shape = tuple(shape_meta['obs'][key]['shape']) + c, h, w = shape + this_compressor = Jpeg2k(level=50) + img_arr = data_group.require_dataset( + name=key, + shape=(n_steps, h, w, c), + chunks=(1, h, w, c), + compressor=this_compressor, + dtype=np.uint8 + ) + for episode_idx in range(len(demos)): + demo = demos[f'demo_{episode_idx}'] + hdf5_arr = demo['obs'][key] + for hdf5_idx in range(hdf5_arr.shape[0]): + if len(futures) >= max_inflight_tasks: + # limit number of inflight tasks + completed, futures = concurrent.futures.wait(futures, + return_when=concurrent.futures.FIRST_COMPLETED) + for f in completed: + if not f.result(): + raise RuntimeError('Failed to encode image!') + pbar.update(len(completed)) + + zarr_idx = episode_starts[episode_idx] + hdf5_idx + futures.add( + executor.submit(img_copy, + img_arr, zarr_idx, hdf5_arr, hdf5_idx)) + completed, futures = concurrent.futures.wait(futures) + for f in completed: + if not f.result(): + raise RuntimeError('Failed to encode image!') + pbar.update(len(completed)) + + replay_buffer = ReplayBuffer(root) + return replay_buffer diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/rotation_conversions.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/rotation_conversions.py new file mode 100644 index 0000000000000000000000000000000000000000..93b10625002ef8f69936adb353f5ae54fe0efc77 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/dataset/rotation_conversions.py @@ -0,0 +1,597 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from typing import Optional, Union + +import torch +import torch.nn.functional as F + +Device = Union[str, torch.device] + +""" +The transformation matrices returned from the functions in this file assume +the points on which the transformation will be applied are column vectors. +i.e. the R matrix is structured as + + R = [ + [Rxx, Rxy, Rxz], + [Ryx, Ryy, Ryz], + [Rzx, Rzy, Rzz], + ] # (3, 3) + +This matrix can be applied to column vectors by post multiplication +by the points e.g. + + points = [[0], [1], [2]] # (3 x 1) xyz coordinates of a point + transformed_points = R * points + +To apply the same matrix to points which are row vectors, the R matrix +can be transposed and pre multiplied by the points: + +e.g. + points = [[0, 1, 2]] # (1 x 3) xyz coordinates of a point + transformed_points = points * R.transpose(1, 0) +""" + + +def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as quaternions to rotation matrices. + + Args: + quaternions: quaternions with real part first, + as tensor of shape (..., 4). + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + r, i, j, k = torch.unbind(quaternions, -1) + # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`. + two_s = 2.0 / (quaternions * quaternions).sum(-1) + + o = torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + -1, + ) + return o.reshape(quaternions.shape[:-1] + (3, 3)) + + +def _copysign(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """ + Return a tensor where each element has the absolute value taken from the, + corresponding element of a, with sign taken from the corresponding + element of b. This is like the standard copysign floating-point operation, + but is not careful about negative 0 and NaN. + + Args: + a: source tensor. + b: tensor whose signs will be used, of the same shape as a. + + Returns: + Tensor of the same shape as a with the signs of b. + """ + signs_differ = (a < 0) != (b < 0) + return torch.where(signs_differ, -a, a) + + +def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor: + """ + Returns torch.sqrt(torch.max(0, x)) + but with a zero subgradient where x is 0. + """ + ret = torch.zeros_like(x) + positive_mask = x > 0 + ret[positive_mask] = torch.sqrt(x[positive_mask]) + return ret + + +def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as rotation matrices to quaternions. + + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + + Returns: + quaternions with real part first, as tensor of shape (..., 4). + """ + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") + + batch_dim = matrix.shape[:-2] + m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( + matrix.reshape(batch_dim + (9,)), dim=-1 + ) + + q_abs = _sqrt_positive_part( + torch.stack( + [ + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ], + dim=-1, + ) + ) + + # we produce the desired quaternion multiplied by each of r, i, j, k + quat_by_rijk = torch.stack( + [ + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1), + ], + dim=-2, + ) + + # We floor here at 0.1 but the exact level is not important; if q_abs is small, + # the candidate won't be picked. + flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) + quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr)) + + # if not for numerical problems, quat_candidates[i] should be same (up to a sign), + # forall i; we pick the best-conditioned one (with the largest denominator) + out = quat_candidates[ + F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, : + ].reshape(batch_dim + (4,)) + return standardize_quaternion(out) + + +def _axis_angle_rotation(axis: str, angle: torch.Tensor) -> torch.Tensor: + """ + Return the rotation matrices for one of the rotations about an axis + of which Euler angles describe, for each value of the angle given. + + Args: + axis: Axis label "X" or "Y or "Z". + angle: any shape tensor of Euler angles in radians + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + + cos = torch.cos(angle) + sin = torch.sin(angle) + one = torch.ones_like(angle) + zero = torch.zeros_like(angle) + + if axis == "X": + R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos) + elif axis == "Y": + R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos) + elif axis == "Z": + R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one) + else: + raise ValueError("letter must be either X, Y or Z.") + + return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3)) + + +def euler_angles_to_matrix(euler_angles: torch.Tensor, convention: str) -> torch.Tensor: + """ + Convert rotations given as Euler angles in radians to rotation matrices. + + Args: + euler_angles: Euler angles in radians as tensor of shape (..., 3). + convention: Convention string of three uppercase letters from + {"X", "Y", and "Z"}. + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3: + raise ValueError("Invalid input euler angles.") + if len(convention) != 3: + raise ValueError("Convention must have 3 letters.") + if convention[1] in (convention[0], convention[2]): + raise ValueError(f"Invalid convention {convention}.") + for letter in convention: + if letter not in ("X", "Y", "Z"): + raise ValueError(f"Invalid letter {letter} in convention string.") + matrices = [ + _axis_angle_rotation(c, e) + for c, e in zip(convention, torch.unbind(euler_angles, -1)) + ] + # return functools.reduce(torch.matmul, matrices) + return torch.matmul(torch.matmul(matrices[0], matrices[1]), matrices[2]) + + +def _angle_from_tan( + axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool +) -> torch.Tensor: + """ + Extract the first or third Euler angle from the two members of + the matrix which are positive constant times its sine and cosine. + + Args: + axis: Axis label "X" or "Y or "Z" for the angle we are finding. + other_axis: Axis label "X" or "Y or "Z" for the middle axis in the + convention. + data: Rotation matrices as tensor of shape (..., 3, 3). + horizontal: Whether we are looking for the angle for the third axis, + which means the relevant entries are in the same row of the + rotation matrix. If not, they are in the same column. + tait_bryan: Whether the first and third axes in the convention differ. + + Returns: + Euler Angles in radians for each matrix in data as a tensor + of shape (...). + """ + + i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis] + if horizontal: + i2, i1 = i1, i2 + even = (axis + other_axis) in ["XY", "YZ", "ZX"] + if horizontal == even: + return torch.atan2(data[..., i1], data[..., i2]) + if tait_bryan: + return torch.atan2(-data[..., i2], data[..., i1]) + return torch.atan2(data[..., i2], -data[..., i1]) + + +def _index_from_letter(letter: str) -> int: + if letter == "X": + return 0 + if letter == "Y": + return 1 + if letter == "Z": + return 2 + raise ValueError("letter must be either X, Y or Z.") + + +def matrix_to_euler_angles(matrix: torch.Tensor, convention: str) -> torch.Tensor: + """ + Convert rotations given as rotation matrices to Euler angles in radians. + + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + convention: Convention string of three uppercase letters. + + Returns: + Euler angles in radians as tensor of shape (..., 3). + """ + if len(convention) != 3: + raise ValueError("Convention must have 3 letters.") + if convention[1] in (convention[0], convention[2]): + raise ValueError(f"Invalid convention {convention}.") + for letter in convention: + if letter not in ("X", "Y", "Z"): + raise ValueError(f"Invalid letter {letter} in convention string.") + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") + i0 = _index_from_letter(convention[0]) + i2 = _index_from_letter(convention[2]) + tait_bryan = i0 != i2 + if tait_bryan: + central_angle = torch.asin( + matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0) + ) + else: + central_angle = torch.acos(matrix[..., i0, i0]) + + o = ( + _angle_from_tan( + convention[0], convention[1], matrix[..., i2], False, tait_bryan + ), + central_angle, + _angle_from_tan( + convention[2], convention[1], matrix[..., i0, :], True, tait_bryan + ), + ) + return torch.stack(o, -1) + + +def random_quaternions( + n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None +) -> torch.Tensor: + """ + Generate random quaternions representing rotations, + i.e. versors with nonnegative real part. + + Args: + n: Number of quaternions in a batch to return. + dtype: Type to return. + device: Desired device of returned tensor. Default: + uses the current device for the default tensor type. + + Returns: + Quaternions as tensor of shape (N, 4). + """ + if isinstance(device, str): + device = torch.device(device) + o = torch.randn((n, 4), dtype=dtype, device=device) + s = (o * o).sum(1) + o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None] + return o + + +def random_rotations( + n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None +) -> torch.Tensor: + """ + Generate random rotations as 3x3 rotation matrices. + + Args: + n: Number of rotation matrices in a batch to return. + dtype: Type to return. + device: Device of returned tensor. Default: if None, + uses the current device for the default tensor type. + + Returns: + Rotation matrices as tensor of shape (n, 3, 3). + """ + quaternions = random_quaternions(n, dtype=dtype, device=device) + return quaternion_to_matrix(quaternions) + + +def random_rotation( + dtype: Optional[torch.dtype] = None, device: Optional[Device] = None +) -> torch.Tensor: + """ + Generate a single random 3x3 rotation matrix. + + Args: + dtype: Type to return + device: Device of returned tensor. Default: if None, + uses the current device for the default tensor type + + Returns: + Rotation matrix as tensor of shape (3, 3). + """ + return random_rotations(1, dtype, device)[0] + + +def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor: + """ + Convert a unit quaternion to a standard form: one in which the real + part is non negative. + + Args: + quaternions: Quaternions with real part first, + as tensor of shape (..., 4). + + Returns: + Standardized quaternions as tensor of shape (..., 4). + """ + return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions) + + +def quaternion_raw_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """ + Multiply two quaternions. + Usual torch rules for broadcasting apply. + + Args: + a: Quaternions as tensor of shape (..., 4), real part first. + b: Quaternions as tensor of shape (..., 4), real part first. + + Returns: + The product of a and b, a tensor of quaternions shape (..., 4). + """ + aw, ax, ay, az = torch.unbind(a, -1) + bw, bx, by, bz = torch.unbind(b, -1) + ow = aw * bw - ax * bx - ay * by - az * bz + ox = aw * bx + ax * bw + ay * bz - az * by + oy = aw * by - ax * bz + ay * bw + az * bx + oz = aw * bz + ax * by - ay * bx + az * bw + return torch.stack((ow, ox, oy, oz), -1) + + +def quaternion_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """ + Multiply two quaternions representing rotations, returning the quaternion + representing their composition, i.e. the versor with nonnegative real part. + Usual torch rules for broadcasting apply. + + Args: + a: Quaternions as tensor of shape (..., 4), real part first. + b: Quaternions as tensor of shape (..., 4), real part first. + + Returns: + The product of a and b, a tensor of quaternions of shape (..., 4). + """ + ab = quaternion_raw_multiply(a, b) + return standardize_quaternion(ab) + + +def quaternion_invert(quaternion: torch.Tensor) -> torch.Tensor: + """ + Given a quaternion representing rotation, get the quaternion representing + its inverse. + + Args: + quaternion: Quaternions as tensor of shape (..., 4), with real part + first, which must be versors (unit quaternions). + + Returns: + The inverse, a tensor of quaternions of shape (..., 4). + """ + + scaling = torch.tensor([1, -1, -1, -1], device=quaternion.device) + return quaternion * scaling + + +def quaternion_apply(quaternion: torch.Tensor, point: torch.Tensor) -> torch.Tensor: + """ + Apply the rotation given by a quaternion to a 3D point. + Usual torch rules for broadcasting apply. + + Args: + quaternion: Tensor of quaternions, real part first, of shape (..., 4). + point: Tensor of 3D points of shape (..., 3). + + Returns: + Tensor of rotated points of shape (..., 3). + """ + if point.size(-1) != 3: + raise ValueError(f"Points are not in 3D, {point.shape}.") + real_parts = point.new_zeros(point.shape[:-1] + (1,)) + point_as_quaternion = torch.cat((real_parts, point), -1) + out = quaternion_raw_multiply( + quaternion_raw_multiply(quaternion, point_as_quaternion), + quaternion_invert(quaternion), + ) + return out[..., 1:] + + +def axis_angle_to_matrix(axis_angle: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as axis/angle to rotation matrices. + + Args: + axis_angle: Rotations given as a vector in axis angle form, + as a tensor of shape (..., 3), where the magnitude is + the angle turned anticlockwise in radians around the + vector's direction. + + Returns: + Rotation matrices as tensor of shape (..., 3, 3). + """ + return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle)) + + +def matrix_to_axis_angle(matrix: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as rotation matrices to axis/angle. + + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + + Returns: + Rotations given as a vector in axis angle form, as a tensor + of shape (..., 3), where the magnitude is the angle + turned anticlockwise in radians around the vector's + direction. + """ + return quaternion_to_axis_angle(matrix_to_quaternion(matrix)) + + +def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as axis/angle to quaternions. + + Args: + axis_angle: Rotations given as a vector in axis angle form, + as a tensor of shape (..., 3), where the magnitude is + the angle turned anticlockwise in radians around the + vector's direction. + + Returns: + quaternions with real part first, as tensor of shape (..., 4). + """ + angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True) + half_angles = angles * 0.5 + eps = 1e-6 + small_angles = angles.abs() < eps + sin_half_angles_over_angles = torch.empty_like(angles) + sin_half_angles_over_angles[~small_angles] = ( + torch.sin(half_angles[~small_angles]) / angles[~small_angles] + ) + # for x small, sin(x/2) is about x/2 - (x/2)^3/6 + # so sin(x/2)/x is about 1/2 - (x*x)/48 + sin_half_angles_over_angles[small_angles] = ( + 0.5 - (angles[small_angles] * angles[small_angles]) / 48 + ) + quaternions = torch.cat( + [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1 + ) + return quaternions + + +def quaternion_to_axis_angle(quaternions: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as quaternions to axis/angle. + + Args: + quaternions: quaternions with real part first, + as tensor of shape (..., 4). + + Returns: + Rotations given as a vector in axis angle form, as a tensor + of shape (..., 3), where the magnitude is the angle + turned anticlockwise in radians around the vector's + direction. + """ + norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True) + half_angles = torch.atan2(norms, quaternions[..., :1]) + angles = 2 * half_angles + eps = 1e-6 + small_angles = angles.abs() < eps + sin_half_angles_over_angles = torch.empty_like(angles) + sin_half_angles_over_angles[~small_angles] = ( + torch.sin(half_angles[~small_angles]) / angles[~small_angles] + ) + # for x small, sin(x/2) is about x/2 - (x/2)^3/6 + # so sin(x/2)/x is about 1/2 - (x*x)/48 + sin_half_angles_over_angles[small_angles] = ( + 0.5 - (angles[small_angles] * angles[small_angles]) / 48 + ) + return quaternions[..., 1:] / sin_half_angles_over_angles + + +def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor: + """ + Converts 6D rotation representation by Zhou et al. [1] to rotation matrix + using Gram--Schmidt orthogonalization per Section B of [1]. + Args: + d6: 6D rotation representation, of size (*, 6) + + Returns: + batch of rotation matrices of size (*, 3, 3) + + [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. + On the Continuity of Rotation Representations in Neural Networks. + IEEE Conference on Computer Vision and Pattern Recognition, 2019. + Retrieved from http://arxiv.org/abs/1812.07035 + """ + + a1, a2 = d6[..., :3], d6[..., 3:] + b1 = F.normalize(a1, dim=-1) + b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1 + b2 = F.normalize(b2, dim=-1) + b3 = torch.cross(b1, b2, dim=-1) + return torch.stack((b1, b2, b3), dim=-2) + + +def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor: + """ + Converts rotation matrices to 6D rotation representation by Zhou et al. [1] + by dropping the last row. Note that 6D representation is not unique. + Args: + matrix: batch of rotation matrices of size (*, 3, 3) + + Returns: + 6D rotation representation, of size (*, 6) + + [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. + On the Continuity of Rotation Representations in Neural Networks. + IEEE Conference on Computer Vision and Pattern Recognition, 2019. + Retrieved from http://arxiv.org/abs/1812.07035 + """ + batch_dim = matrix.size()[:-2] + return matrix[..., :2, :].clone().reshape(batch_dim + (6,)) \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8d78325feccd6bd7fc8a7c2c8fafadd9f49e61db --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/__init__.py @@ -0,0 +1,5 @@ +from .basic import DiffusionModel +from .diffusionsde import DiscreteDiffusionSDE, ContinuousDiffusionSDE +from .newedm import ContinuousEDM +from .rectifiedflow import DiscreteRectifiedFlow, ContinuousRectifiedFlow +from .consistency_model import ContinuousConsistencyModel diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/basic.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/basic.py new file mode 100644 index 0000000000000000000000000000000000000000..b45faeff5d76c4a7451439b6c9886a871258d8f2 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/basic.py @@ -0,0 +1,106 @@ +from copy import deepcopy +from typing import Optional, Union + +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.utils import to_tensor +from cleandiffuser.classifier import BaseClassifier +from cleandiffuser.nn_condition import BaseNNCondition, IdentityCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion + + +class DiffusionModel: + + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + # NN backbone for the diffusion model + nn_diffusion: BaseNNDiffusion, + # Add a condition-process NN to enable classifier-free-guidance + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugs ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Params ---------------- # + grad_clip_norm: Optional[float] = None, + diffusion_steps: int = 1000, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + device: Union[torch.device, str] = "cpu" + ): + if optim_params is None: + optim_params = {"lr": 2e-4, "weight_decay": 1e-5} + + self.device = device + self.grad_clip_norm = grad_clip_norm + self.diffusion_steps = diffusion_steps + self.ema_rate = ema_rate + + # nn_condition is None means that the model is not conditioned on any input. + if nn_condition is None: + nn_condition = IdentityCondition() + + # In the code implementation of Diffusion models, it is common to maintain an exponential + # moving average (EMA) version of the model for inference, as it has been observed that + # this approach can result in more stable generation outcomes. + self.model = nn.ModuleDict({ + "diffusion": nn_diffusion.to(self.device), + "condition": nn_condition.to(self.device)}) + self.model_ema = deepcopy(self.model).requires_grad_(False) + + self.model.train() + self.model_ema.eval() + + self.optimizer = torch.optim.AdamW(self.model.parameters(), **optim_params) + + self.classifier = classifier + + self.fix_mask = to_tensor(fix_mask, self.device)[None, ] if fix_mask is not None else 0. + self.loss_weight = to_tensor(loss_weight, self.device)[None, ] if loss_weight is not None else 1. + + def train(self): + self.model.train() + if self.classifier is not None: + self.classifier.model.train() + + def eval(self): + self.model.eval() + if self.classifier is not None: + self.classifier.model.eval() + + def ema_update(self): + with torch.no_grad(): + for p, p_ema in zip(self.model.parameters(), self.model_ema.parameters()): + p_ema.data.mul_(self.ema_rate).add_(p.data, alpha=1. - self.ema_rate) + + def update(self, x0, condition=None, update_ema=True, **kwargs): + raise NotImplementedError + + def sample(self, *args, **kwargs): + raise NotImplementedError + + def save(self, path: str): + torch.save({ + "model": self.model.state_dict(), + "model_ema": self.model_ema.state_dict(), + }, path) + + def load(self, path: str): + checkpoint = torch.load(path, map_location=self.device) + self.model.load_state_dict(checkpoint["model"]) + self.model_ema.load_state_dict(checkpoint["model_ema"]) + + + diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/consistency_model.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/consistency_model.py new file mode 100644 index 0000000000000000000000000000000000000000..2f5e59f5b586ffd4bf9cb97fbd35fa1027bd17b2 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/consistency_model.py @@ -0,0 +1,428 @@ +from typing import Callable, List +from typing import Optional, Union + +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.diffusion import DiffusionModel +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from cleandiffuser.utils import at_least_ndim +from .newedm import ContinuousEDM + + +def erf(x): + a1 = 0.254829592 + a2 = -0.284496736 + a3 = 1.421413741 + a4 = -1.453152027 + a5 = 1.061405429 + p = 0.3275911 + + sign = np.sign(x) + x = np.abs(x) + t = 1.0 / (1.0 + p * x) + y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * np.exp(-x * x) + return sign * y + + +def compare_properties(obj1, obj2, properties: List[str]): + differences = [] + for prop in properties: + obj1_prop = getattr(obj1, prop) + obj2_prop = getattr(obj2, prop) + if isinstance(obj1_prop, torch.Tensor): + if not torch.allclose(obj1_prop, obj2_prop): + differences.append(prop) + elif isinstance(obj1_prop, np.ndarray): + if not np.allclose(obj1_prop, obj2_prop): + differences.append(prop) + else: + if obj1_prop != obj2_prop: + differences.append(prop) + return differences + + +def pseudo_huber_loss(source: torch.Tensor, target: torch.Tensor, c: float = 0.0): + return ((source - target) ** 2 + c ** 2).sqrt() - c + + +class CMCurriculumLogger: + def __init__( + self, s0: int = 10, s1: int = 1280, curriculum_cycle: int = 100_000, + sigma_min: float = 0.002, sigma_max: float = 80., rho: float = 7., + P_mean: float = -1.1, P_std: float = 2.0 + ): + self.Kprime = np.ceil(curriculum_cycle / (np.log2(np.ceil(s1 / s0)) + 1)) + self.Nk = s0 + self.s0, self.s1 = s0, s1 + self.curriculum_cycle = curriculum_cycle + self.sigma_min, self.sigma_max, self.rho = sigma_min, sigma_max, rho + self.P_mean, self.P_std = P_mean, P_std + + self.ceil_k_div_Kprime, self.k = None, None + + self.update_k(0) + + def update_k(self, k): + self.k = k + if np.ceil(k / self.Kprime) != self.ceil_k_div_Kprime: + self.ceil_k_div_Kprime = np.ceil(k / self.Kprime) + self.Nk = int(min(self.s0 * (2 ** self.ceil_k_div_Kprime), self.s1)) + + self.sigmas = ((self.sigma_min ** (1 / self.rho) + np.arange(self.Nk + 1, dtype=np.float32) + / self.Nk * (self.sigma_max ** (1 / self.rho) - self.sigma_min ** ( + 1 / self.rho))) ** self.rho) + + self.p_sigmas = (erf((np.log(self.sigmas[1:]) - self.P_mean) / (self.P_std * (2 ** 0.5))) - + erf((np.log(self.sigmas[:-1]) - self.P_mean) / (self.P_std * (2 ** 0.5)))) + self.p_sigmas = self.p_sigmas / self.p_sigmas.sum() + + def incremental_update_k(self): + self.update_k(self.k + 1) + + @property + def curriculum_process(self): + return (self.k % self.curriculum_cycle) / self.curriculum_cycle + + +class ContinuousConsistencyModel(DiffusionModel): + """**Continuous-time Consistency Model** + + The Consistency Model defines a consistency function. + A consistency function has the property of self-consistency: + its outputs are consistent for arbitrary pairs of (x_t, t) that belong to the same PF ODE trajectory. + To learn such a consistency function, the Consistency Model needs to be distilled either from a pre-trained EDM + or learned directly through consistency training loss. + This self-consistency property allows the Consistency Model in theory to achieve one-step generation. + + The current implementation of Consistency Model only supports continuous-time ODEs. + The sampling steps are required to be greater than 0. + + Args: + nn_diffusion: BaseNNDiffusion, + The neural network backbone for the Diffusion model. + nn_condition: Optional[BaseNNCondition], + The neural network backbone for the condition embedding. + + fix_mask: Union[list, np.ndarray, torch.Tensor], + Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + The mask should be in the shape of `x_shape`. + loss_weight: Union[list, np.ndarray, torch.Tensor], + Add loss weight. The weight should be in the shape of `x_shape`. + + classifier: Optional[BaseClassifier], + The Consistency Model does not support classifier guidance; please set this option to `None`. + + grad_clip_norm: Optional[float], + Gradient clipping norm. + ema_rate: float, + Exponential moving average rate. + optim_params: Optional[dict], + Optimizer parameters. + + s0: int, + The minimum number of noise levels. Default: 10. + s1: int, + The maximum number of noise levels. Default: 1280. + data_dim: int, + The dimension of the data, which affects the `pseudo_huber_constant`. + As suggested in `improved Consistency Models`, `pseudo_huber_constant` = 0.00054 * np.sqrt(data_dim). + If `data_dim` is `None`, then `pseudo_huber_constant` = 0.01 will be used. + P_mean: float, + Hyperparameter for noise sampling during training. Default: -1.1. + P_std: float, + Hyperparameter for noise sampling during training. Default: 2.0. + sigma_min: float, + The minimum standard deviation of the noise. Default: 0.002. + sigma_max: float, + The maximum standard deviation of the noise. Default: 80. + sigma_data: float, + The standard deviation of the data. Default: 0.5. + rho: float, + The power of the noise schedule. Default: 7. + curriculum_cycle: int, + The cycle of the curriculum process. + It is best to set `curriculum_cycle` to the number of model training iterations. Default: 100_000. + + x_max: Optional[torch.Tensor], + The maximum value for the input data. `None` indicates no constraint. + x_min: Optional[torch.Tensor], + The minimum value for the input data. `None` indicates no constraint. + + device: Union[torch.device, str], + The device to run the model. + """ + + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugins ---------------- # + # Do not support CG + classifier=None, + + # ------------------ Training Params ---------------- # + grad_clip_norm: Optional[float] = None, + ema_rate: float = 0.9999, + optim_params: Optional[dict] = None, + + # ------------------- Consistency Model Params ------------------- # + s0: int = 10, + s1: int = 1280, + data_dim: int = None, + P_mean: float = -1.1, + P_std: float = 2.0, + sigma_min: float = 0.002, + sigma_max: float = 80., + sigma_data: float = 0.5, + rho: float = 7.0, + curriculum_cycle: int = 100_000, + + x_max: Optional[torch.Tensor] = None, + x_min: Optional[torch.Tensor] = None, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + 0, ema_rate, optim_params, device) + + self.cur_logger = CMCurriculumLogger( + s0, s1, curriculum_cycle, sigma_min, sigma_max, rho, P_mean, P_std) + + self.pseudo_huber_constant = 0.01 if data_dim is None else 0.00054 * np.sqrt(data_dim) + + self.rho = rho + self.sigma_data, self.sigma_max, self.sigma_min = sigma_data, sigma_max, sigma_min + self.x_max = x_max.to(device) if isinstance(x_max, torch.Tensor) else x_max + self.x_min = x_min.to(device) if isinstance(x_min, torch.Tensor) else x_min + + self.edm = None + self.distillation_sigmas, self.distillation_N = None, None + + def prepare_distillation(self, edm: ContinuousEDM, distillation_N: int = 18): + checklist = [ + "sigma_data", "sigma_max", "sigma_min", "rho", "x_max", "x_min", + "fix_mask", "loss_weight", "device"] + differences = compare_properties(self, edm, checklist) + if len(differences) != 0: + raise ValueError(f"Properties {differences} are different between the EDM and the Consistency Model.") + self.edm = edm + self.model.load_state_dict(edm.model.state_dict()) + self.model_ema.load_state_dict(edm.model_ema.state_dict()) + self.distillation_N = distillation_N + self.distillation_sigmas = self.training_noise_schedule(distillation_N) + + @property + def supported_solvers(self): + return ["none"] + + @property + def clip_pred(self): + return (self.x_max is not None) or (self.x_min is not None) + + def training_noise_schedule(self, N): + sigma = ((self.sigma_min ** (1 / self.rho) + np.arange(N + 1) + / N * (self.sigma_max ** (1 / self.rho) - self.sigma_min ** (1 / self.rho))) ** self.rho) + return torch.tensor(sigma, device=self.device, dtype=torch.float32) + + # ===================== CM Pre-conditioning ======================= + def c_skip(self, sigma): + return self.sigma_data ** 2 / (self.sigma_data ** 2 + (sigma - self.sigma_min) ** 2) + + def c_out(self, sigma): + return (sigma - self.sigma_min) * self.sigma_data / (self.sigma_data ** 2 + sigma ** 2).sqrt() + + def c_in(self, sigma): + return 1 / (self.sigma_data ** 2 + sigma ** 2).sqrt() + + def c_noise(self, sigma): + return 0.25 * sigma.log() + + def f(self, x, t, condition=None, model=None): + c_skip, c_out, c_in, c_noise = self.c_skip(t), self.c_out(t), self.c_in(t), self.c_noise(t) + if model is None: + model = self.model + c_skip, c_in, c_out = ( + at_least_ndim(c_skip, x.dim()), at_least_ndim(c_in, x.dim()), at_least_ndim(c_out, x.dim())) + pred_x = c_skip * x + c_out * model["diffusion"](c_in * x, c_noise, condition) + if self.clip_pred: + pred_x = pred_x.clip(self.x_min, self.x_max) + return pred_x + + def distillation_loss(self, x0, condition=None): + + assert self.edm is not None, "Please call `prepare_distillation` before distillation." + + idx = torch.randint(self.distillation_N, (x0.shape[0],), device=self.device) + t_m, t_n = self.distillation_sigmas[idx + 1], self.distillation_sigmas[idx] + x_m, t_m, eps = self.edm.add_noise(x0, t_m, None) + + with torch.no_grad(): + condition_vec_cfg = self.edm.model_ema["condition"](condition) if condition is not None else None + pred, logp = self.edm.guided_sampling( + x_m, t_m, None, self.edm.model_ema, condition_vec_cfg, 1.0, None, 0.0, False) + dot_x = (x_m - pred) / at_least_ndim(t_m, x_m.dim()) + delta_t = t_m - t_n + x_n = x_m - dot_x * at_least_ndim(delta_t, x_m.dim()) + x_n = x_n * (1. - self.fix_mask) + x0 * self.fix_mask + + condition_vec = self.model["condition"](condition) if condition is not None else None + pred_x_m = self.f(x_m, t_m, condition_vec, self.model) + with torch.no_grad(): + condition_vec_ema = self.model_ema["condition"](condition) if condition is not None else None + pred_x_n = self.f(x_n, t_n, condition_vec_ema, self.model_ema) + + loss = (((pred_x_n - pred_x_m) ** 2) * (1 - self.fix_mask) * + self.loss_weight * at_least_ndim((1 / (t_m - t_n)), pred_x_n.dim())) + + return loss.mean(), None + + def training_loss(self, x0, condition=None): + + idx = np.random.choice(self.cur_logger.Nk, size=x0.shape[0], p=self.cur_logger.p_sigmas) + + # m = n + 1 + sigma_n = torch.tensor(self.cur_logger.sigmas[idx], device=self.device) + sigma_m = torch.tensor(self.cur_logger.sigmas[idx + 1], device=self.device) + + eps = torch.randn_like(x0) + + x_n = x0 + at_least_ndim(sigma_n, x0.dim()) * eps + x_m = x0 + at_least_ndim(sigma_m, x0.dim()) * eps + + condition = self.model["condition"](condition) if condition is not None else None + + pred_x_m = self.f(x_m, sigma_m, condition, self.model) + with torch.no_grad(): + pred_x_n = self.f(x_n, sigma_n, condition.detach(), self.model) + + loss = pseudo_huber_loss(pred_x_m, pred_x_n, self.pseudo_huber_constant) + + unweighted_loss = (loss * (1 - self.fix_mask) * self.loss_weight) + + cm_loss_weight = at_least_ndim(1 / (sigma_m - sigma_n), x0.dim()) + + return (unweighted_loss * cm_loss_weight).mean(), unweighted_loss.mean().item() + + def update(self, x0, condition=None, update_ema=True, loss_type="training", **kwargs): + """ One-step gradient update. + + Args: + x0: torch.Tensor, + Samples from the target distribution. + condition: Optional, + Condition of x0. `None` indicates no condition. + update_ema: bool, + Whether to update the exponential moving average model. + loss_type: str, + The type of loss. `training` or `distillation`. + + Returns: + log: dict, + The log dictionary. + + Examples: + >>> model = ContinuousConsistencyModel(...) + >>> x0 = torch.randn(*x_shape) + >>> condition = torch.randn(*condition_shape) + >>> log = model.update(x0, condition, loss_type="training") # training + >>> log = model.update(x0, condition, loss_type="distillation") # distillation + """ + if loss_type == "training": + loss, unweighted_loss = self.training_loss(x0, condition) + elif loss_type == "distillation": + loss, unweighted_loss = self.distillation_loss(x0, condition) + else: + raise ValueError(f"Unknown loss type: {loss_type}") + + loss.backward() + grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_norm) \ + if self.grad_clip_norm else None + self.optimizer.step() + self.optimizer.zero_grad() + + if update_ema: + self.ema_update() + + if loss_type == "training": + self.cur_logger.incremental_update_k() + + log = {"loss": loss.item(), "grad_norm": grad_norm, "unweighted_loss": unweighted_loss} + + return log + + def sample( + self, + # ---------- the known fixed portion ---------- # + prior: torch.Tensor, + # ----------------- sampling ----------------- # + solver: str = "none", + n_samples: int = 1, + sample_steps: int = 5, + sample_step_schedule: Union[str, Callable] = "uniform", + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ----------- Diffusion-X sampling ---------- + diffusion_x_sampling_steps: int = 0, + # ----------- Warm-Starting ----------- + warm_start_reference: Optional[torch.Tensor] = None, + warm_start_forward_level: float = 0.3, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + assert w_cg == 0.0 and condition_cg is None, "Consistency Distillation does not support classifier guidance." + + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + model = self.model if not use_ema else self.model_ema + + prior = prior.to(self.device) + xt = torch.randn_like(prior) * self.sigma_max * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + if preserve_history: + log["sample_history"][:, 0] = xt.cpu().numpy() + + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + + # ===================== Sampling Schedule ==================== + sample_step_schedule = ((self.sigma_min ** (1 / self.rho) + torch.arange(sample_steps + 1, device=self.device) + / sample_steps * (self.sigma_max ** (1 / self.rho) - self.sigma_min ** (1 / self.rho))) + ** self.rho) + + t = torch.full((n_samples,), sample_step_schedule[-1], dtype=torch.float32, device=self.device) + pred_x = self.f(xt, t, condition_vec_cfg, model) + pred_x = pred_x * (1. - self.fix_mask) + prior * self.fix_mask + + loop_steps = [1] * diffusion_x_sampling_steps + list(range(1, sample_steps)) + + for i in reversed(loop_steps): + t = torch.full((n_samples,), sample_step_schedule[i], dtype=torch.float32, device=self.device) + xt = pred_x + (at_least_ndim(t, xt.dim()) ** 2 - self.sigma_min ** 2).sqrt() * torch.randn_like(xt) + + pred_x = self.f(xt, t, condition_vec_cfg, model) + pred_x = pred_x * (1. - self.fix_mask) + prior * self.fix_mask + + return pred_x, log diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/ddim.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/ddim.py new file mode 100644 index 0000000000000000000000000000000000000000..d028edbeabe6fb96b7144498c8f2dbb7b55af302 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/ddim.py @@ -0,0 +1,190 @@ +from typing import Optional +from typing import Union + +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.classifier import BaseClassifier +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from cleandiffuser.utils import at_least_ndim +from .basic import DiffusionModel + + +class DDIM(DiffusionModel): + + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugs ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Params ---------------- # + grad_clip_norm: Optional[float] = None, + diffusion_steps: int = 1000, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- DPM Params ------------------- # + noise_schedule: str = "linear", # or cosine + t_eps: float = 1e-3, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + diffusion_steps, ema_rate, optim_params, device) + + self.noise_schedule = noise_schedule + self.t_eps = t_eps + + @property + def t_range(self): + if self.noise_schedule == "linear": + return self.t_eps, 1. + elif self.noise_schedule == "cosine": + return self.t_eps, 0.9946 + else: + raise ValueError(f"noise_schedule should be 'linear' or 'cosine', but got {self.noise_schedule}.") + + def alpha_schedule(self, t): + if self.noise_schedule == "linear": + beta0, beta1 = 0.1, 20 + return (-(beta1-beta0)/4*(t**2) - beta0/2*t).exp() + elif self.noise_schedule == "cosine": + s = 0.008 + return ((torch.cos(np.pi/2*(t+s)/(1+s))).log() - np.log(np.cos(np.pi/2*s/(1+s)))).exp() + else: + raise ValueError(f"noise_schedule should be 'linear' or 'cosine', but got {self.noise_schedule}.") + + # --------------------------------------------------------------------------- + # Training + + def add_noise(self, x0, t=None, eps=None): + if t is None: + t = torch.rand((x0.shape[0], ), device=self.device) + t = self.t_range[0] + t * (self.t_range[1] - self.t_range[0]) + eps = torch.randn_like(x0) if eps is None else eps + alpha = self.alpha_schedule(at_least_ndim(t, x0.dim())) + sigma = (1 - alpha ** 2).sqrt() + xt = x0 * alpha + sigma * eps + xt = xt * (1. - self.fix_mask) + x0 * self.fix_mask + return xt, t, eps + + def loss(self, x0, condition=None): + xt, t, eps = self.add_noise(x0) + condition = self.model["condition"](condition) if condition is not None else None + loss = (self.model["diffusion"](xt, t, condition) - eps) ** 2 + return (loss * self.loss_weight).mean() + + def update(self, x0, condition=None, **kwargs): + self.optimizer.zero_grad() + loss = self.loss(x0, condition) + loss.backward() + grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_norm) \ + if self.grad_clip_norm else None + self.optimizer.step() + self.ema_update() + log = {"loss": loss.item(), "grad_norm": grad_norm} + return log + + def update_classifier(self, x0, condition): + xt, t, eps = self.add_noise(x0) + log = self.classifier.update(xt, t, condition) + return log + + # --------------------------------------------------------------------------- + # Inference + + def sample( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = None, + use_ema: bool = True, + solver: str = "euler", + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + + preserve_history: bool = False, + ): + log, x_history = {}, None + model = self.model_ema if use_ema else self.model + + t = torch.linspace(self.t_range[1], self.t_range[0], sample_steps + 1, device=self.device) + alphas = self.alpha_schedule(t) + sigmas = (1 - alphas ** 2).sqrt() + logSNRs = (alphas / sigmas).log() + + if prior is None: + xt = torch.randn((n_samples, *self.default_x_shape), device=self.device) + else: + xt = torch.randn_like(prior, device=self.device) * 0.5 + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + if preserve_history: + x_history = np.empty((n_samples, sample_steps + 1, *xt.shape)) + x_history[:, 0] = xt.cpu().numpy() + + with torch.no_grad(): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + for i in range(sample_steps): + + h = logSNRs[i + 1] - logSNRs[i] + + t_batch = t[i].repeat(n_samples) + + # ----------------- CFG ----------------- # + with torch.no_grad(): + if w_cfg != 0.0 and w_cfg != 1.0: + condition_vec_cfg = torch.cat([condition_vec_cfg, torch.zeros_like(condition_vec_cfg)], 0) + eps_theta = model["diffusion"]( + torch.repeat_interleave(xt, 2, dim=0), + torch.repeat_interleave(t_batch, 2, dim=0), + condition_vec_cfg) + eps_theta = w_cfg * eps_theta[:n_samples] + (1. - w_cfg) * eps_theta[n_samples:] + else: + eps_theta = model["diffusion"](xt, t_batch, condition_vec_cfg) + # ----------------- CG ----------------- # + if self.classifier is not None and w_cg != 0.0: + log_p, grad = self.classifier.gradients(xt.clone(), t_batch, condition_vec_cg) + eps_theta = eps_theta - w_cg * sigmas[i] * grad + else: + log_p = None + + # do not change the fixed portion + xt = alphas[i + 1] / alphas[i] * xt - sigmas[i + 1] * torch.expm1(h) * eps_theta + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + if preserve_history: + x_history[:, t] = xt.cpu().numpy() + log["log_p"] = log_p + + log["sample_history"] = x_history + if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + with torch.no_grad(): + logp = self.classifier.logp(xt, t_batch, condition_vec_cg) + log["log_p"] = logp + + return xt, log + diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/ddpm.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/ddpm.py new file mode 100644 index 0000000000000000000000000000000000000000..28b10a57faec75a5778fc4c38083c18d8fb74858 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/ddpm.py @@ -0,0 +1,375 @@ +from typing import Optional, Union + +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.classifier import BaseClassifier +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from cleandiffuser.utils import ( + at_least_ndim, + cosine_beta_schedule, + linear_beta_schedule) +from .basic import DiffusionModel + + +class DDPM(DiffusionModel): + + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugs ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Params ---------------- # + grad_clip_norm: Optional[float] = None, + diffusion_steps: int = 1000, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + x_max: Optional[torch.Tensor] = None, + x_min: Optional[torch.Tensor] = None, + + # ------------------- DPM Params ------------------- # + predict_noise: bool = True, + beta_schedule: str = "cosine", # or cosine + beta_schedule_params: Optional[dict] = None, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + diffusion_steps, ema_rate, optim_params, device) + + self.predict_noise = predict_noise + + if beta_schedule_params is None: + beta_schedule_params = {} + beta_schedule_params["T"] = self.diffusion_steps + + if beta_schedule == "linear": + beta = linear_beta_schedule(**beta_schedule_params) + elif beta_schedule == "cosine": + beta = cosine_beta_schedule(**beta_schedule_params) + else: + raise ValueError(f"Unknown beta schedule: {beta_schedule}") + + self.beta = torch.tensor(beta, device=self.device, dtype=torch.float32) + self.alpha = 1 - self.beta + self.bar_alpha = torch.cumprod(self.alpha.clone(), 0) + self.x_max, self.x_min = x_max, x_min + + @property + def clip_pred(self): + return (self.x_max is not None) or (self.x_min is not None) + + # --------------------------------------------------------------------------- + # Training + + def add_noise(self, x0, t=None, eps=None): + t = torch.randint(self.diffusion_steps, (x0.shape[0],), device=self.device) if t is None else t + eps = torch.randn_like(x0) if eps is None else eps + bar_alpha = at_least_ndim(self.bar_alpha[t], x0.dim()) + xt = x0 * bar_alpha.sqrt() + eps * (1 - bar_alpha).sqrt() + xt = xt * (1. - self.fix_mask) + x0 * self.fix_mask + return xt, t, eps + + def loss(self, x0, condition=None): + xt, t, eps = self.add_noise(x0) + condition = self.model["condition"](condition) if condition is not None else None + if self.predict_noise: + loss = (self.model["diffusion"](xt, t, condition) - eps) ** 2 + else: + loss = (self.model["diffusion"](xt, t, condition) - x0) ** 2 + + return (loss * self.loss_weight * (1 - self.fix_mask)).mean() + + def update(self, x0, condition=None, update_ema=True, **kwargs): + loss = self.loss(x0, condition) + loss.backward() + grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_norm) \ + if self.grad_clip_norm else None + self.optimizer.step() + self.optimizer.zero_grad() + if update_ema: self.ema_update() + log = {"loss": loss.item(), "grad_norm": grad_norm} + return log + + def update_classifier(self, x0, condition): + xt, t, eps = self.add_noise(x0) + log = self.classifier.update(xt, t, condition) + return log + + # --------------------------------------------------------------------------- + # Inference + + def predict_function( + self, x, t, bar_alpha, + use_ema=False, requires_grad=False, + # ----------------- CFG ----------------- # + condition_vec_cfg=None, + w_cfg: float = 0.0, + # ----------------- CG ----------------- # + condition_vec_cg=None, + w_cg: float = 1.0, + ): + b = x.shape[0] + model = self.model_ema if use_ema else self.model + + # ----------------- CFG ----------------- # + with torch.set_grad_enabled(requires_grad): + if w_cfg != 0.0 and w_cfg != 1.0: + repeat_dim = [2 if i == 0 else 1 for i in range(x.dim())] + condition_vec_cfg = torch.cat([condition_vec_cfg, torch.zeros_like(condition_vec_cfg)], 0) + pred = model["diffusion"]( + x.repeat(*repeat_dim), t.repeat(2), condition_vec_cfg) + pred = w_cfg * pred[:b] + (1. - w_cfg) * pred[b:] + elif w_cfg == 0.0: + pred = model["diffusion"](x, t, None) + else: + pred = model["diffusion"](x, t, condition_vec_cfg) + + # ----------------- CG ----------------- # + if self.classifier is not None and w_cg != 0.0 and condition_vec_cg is not None: + log_p, grad = self.classifier.gradients(x.clone(), t, condition_vec_cg) + if self.predict_noise: + pred = pred - w_cg * (1 - bar_alpha).sqrt() * grad + else: + pred = pred + w_cg * (1 - bar_alpha) / bar_alpha.sqrt() * grad + else: + log_p = None + + if self.predict_noise: + if self.clip_pred: + upper_bound = (x - bar_alpha.sqrt() * self.x_min) / (1 - bar_alpha).sqrt() \ + if self.x_min is not None else None + lower_bound = (x - bar_alpha.sqrt() * self.x_max) / (1 - bar_alpha).sqrt() \ + if self.x_max is not None else None + pred = pred.clip(lower_bound, upper_bound) + pred = pred * (1 - self.fix_mask) + else: + if self.clip_pred: + pred = pred.clip(self.x_min, self.x_max) + pred = pred * (1 - self.fix_mask) + x * self.fix_mask + + return pred, {"log_p": log_p} + + def sample( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = None, + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + # initialize logger + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + # choose the model + model = self.model_ema if use_ema else self.model + + # check `sample_steps` + if sample_steps != self.diffusion_steps: + import warnings + warnings.warn(f"sample_steps != diffusion_steps, sample_steps will be set to diffusion_steps.") + sample_steps = self.diffusion_steps + + # initialize the samples + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: log["sample_history"][:, 0] = xt.cpu().numpy() + + # preprocess the conditions + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + # enter the sampling loop + for t in range(self.diffusion_steps - 1, -1, -1): + + t_batch = torch.tensor(t, device=self.device, dtype=torch.long).repeat(n_samples) + bar_alpha = self.bar_alpha[t] + bar_alpha_prev = self.bar_alpha[t - 1] if t > 0 else torch.tensor(1.0, device=self.device) + alpha = self.alpha[t] + beta = self.beta[t] + + # predict eps_theta or x_theta with CG/CFG + pred_theta, log = self.predict_function( + xt, t_batch, bar_alpha, + use_ema=use_ema, + requires_grad=requires_grad, + condition_vec_cfg=condition_vec_cfg, + condition_vec_cg=condition_vec_cg, + w_cfg=w_cfg, w_cg=w_cg) + + # one step denoise + if self.predict_noise: + + xt = 1 / alpha.sqrt() * (xt - beta / (1 - bar_alpha).sqrt() * pred_theta) + + else: + + xt = 1 / (1 - bar_alpha) * ( + alpha.sqrt() * (1 - bar_alpha_prev) * xt + + beta * bar_alpha_prev.sqrt() * pred_theta) + + if t != 0: + xt = xt + (beta * (1 - bar_alpha_prev) / (1 - bar_alpha)).sqrt() * torch.randn_like(xt) + + # Fix the known portion + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: log["sample_history"][:, 1] = xt.cpu().numpy() + + # calculate the final log_p + if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + with torch.no_grad(): + logp = self.classifier.logp(xt, t[-1].repeat(n_samples), condition_vec_cg) + log["log_p"] = logp + + return xt, log + + + def sample_x( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = None, + extra_sample_steps: int = 8, + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + # initialize logger + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + # choose the model + model = self.model_ema if use_ema else self.model + + # check `sample_steps` + if sample_steps != self.diffusion_steps: + import warnings + warnings.warn(f"sample_steps != diffusion_steps, sample_steps will be set to diffusion_steps.") + sample_steps = self.diffusion_steps + + # initialize the samples + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: log["sample_history"][:, 0] = xt.cpu().numpy() + + # preprocess the conditions + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + # enter the sampling loop + for t in range(self.diffusion_steps - 1, -1, -1): + + t_batch = torch.tensor(t, device=self.device, dtype=torch.long).repeat(n_samples) + bar_alpha = self.bar_alpha[t] + bar_alpha_prev = self.bar_alpha[t - 1] if t > 0 else torch.tensor(1.0, device=self.device) + alpha = self.alpha[t] + beta = self.beta[t] + + # predict eps_theta or x_theta with CG/CFG + pred_theta, log = self.predict_function( + xt, t_batch, bar_alpha, + use_ema=use_ema, + requires_grad=requires_grad, + condition_vec_cfg=condition_vec_cfg, + condition_vec_cg=condition_vec_cg, + w_cfg=w_cfg, w_cg=w_cg) + + # one step denoise + if self.predict_noise: + + xt = 1 / alpha.sqrt() * (xt - beta / (1 - bar_alpha).sqrt() * pred_theta) + + else: + + xt = 1 / (1 - bar_alpha) * ( + alpha.sqrt() * (1 - bar_alpha_prev) * xt + + beta * bar_alpha_prev.sqrt() * pred_theta) + + if t != 0: + xt = xt + (beta * (1 - bar_alpha_prev) / (1 - bar_alpha)).sqrt() * torch.randn_like(xt) + + # Fix the known portion + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: log["sample_history"][:, 1] = xt.cpu().numpy() + + if extra_sample_steps > 0: + + t_batch = torch.tensor(0, device=self.device, dtype=torch.long).repeat(n_samples) + bar_alpha = self.bar_alpha[0] + bar_alpha_prev = torch.tensor(1.0, device=self.device) + alpha = self.alpha[0] + beta = self.beta[0] + + for _ in range(extra_sample_steps): + + # predict eps_theta or x_theta with CG/CFG + pred_theta, log = self.predict_function( + xt, t_batch, bar_alpha, + use_ema=use_ema, + requires_grad=requires_grad, + condition_vec_cfg=condition_vec_cfg, + condition_vec_cg=condition_vec_cg, + w_cfg=w_cfg, w_cg=w_cg) + + # one step denoise + if self.predict_noise: + + xt = 1 / alpha.sqrt() * (xt - beta / (1 - bar_alpha).sqrt() * pred_theta) + + else: + + xt = 1 / (1 - bar_alpha) * ( + alpha.sqrt() * (1 - bar_alpha_prev) * xt + + beta * bar_alpha_prev.sqrt() * pred_theta) + + # Fix the known portion + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + # calculate the final log_p + if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + with torch.no_grad(): + logp = self.classifier.logp(xt, t[-1].repeat(n_samples), condition_vec_cg) + log["log_p"] = logp + + return xt, log diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/diffusionsde.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/diffusionsde.py new file mode 100644 index 0000000000000000000000000000000000000000..a5b2bf1227735fdb1bc107aa4cabe21763683cfd --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/diffusionsde.py @@ -0,0 +1,987 @@ +from typing import Optional, Union, Callable, Dict + +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.classifier import BaseClassifier +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from cleandiffuser.utils import ( + at_least_ndim, + SUPPORTED_NOISE_SCHEDULES, SUPPORTED_DISCRETIZATIONS, SUPPORTED_SAMPLING_STEP_SCHEDULE) +from .basic import DiffusionModel + +SUPPORTED_SOLVERS = [ + "ddpm", "ddim", + "ode_dpmsolver_1", "ode_dpmsolver++_1", "ode_dpmsolver++_2M", + "sde_dpmsolver_1", "sde_dpmsolver++_1", "sde_dpmsolver++_2M",] + + +def epstheta_to_xtheta(x, alpha, sigma, eps_theta): + """ + x_theta = (x - sigma * eps_theta) / alpha + """ + return (x - sigma * eps_theta) / alpha + + +def xtheta_to_epstheta(x, alpha, sigma, x_theta): + """ + eps_theta = (x - alpha * x_theta) / sigma + """ + return (x - alpha * x_theta) / sigma + + +class BaseDiffusionSDE(DiffusionModel): + + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugins ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Training Params ---------------- # + grad_clip_norm: Optional[float] = None, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- Diffusion Params ------------------- # + epsilon: float = 1e-3, + + noise_schedule: Union[str, Dict[str, Callable]] = "cosine", + noise_schedule_params: Optional[dict] = None, + + x_max: Optional[torch.Tensor] = None, + x_min: Optional[torch.Tensor] = None, + + predict_noise: bool = True, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + 0, ema_rate, optim_params, device) + + self.predict_noise = predict_noise + self.epsilon = epsilon + self.x_max = x_max.to(device) if isinstance(x_max, torch.Tensor) else x_max + self.x_min = x_min.to(device) if isinstance(x_min, torch.Tensor) else x_min + + @property + def supported_solvers(self): + return SUPPORTED_SOLVERS + + @property + def clip_pred(self): + return (self.x_max is not None) or (self.x_min is not None) + + # ==================== Training: Score Matching ====================== + + def add_noise(self, x0, t=None, eps=None): + raise NotImplementedError + + def loss(self, x0, condition=None, **kwargs): + + xt, t, eps = self.add_noise(x0) + + condition = self.model["condition"](condition) if condition is not None else None + + if self.predict_noise: + loss = (self.model["diffusion"](xt, t, condition) - eps) ** 2 + else: + loss = (self.model["diffusion"](xt, t, condition) - x0) ** 2 + + loss = loss * self.loss_weight * (1 - self.fix_mask) + + # find weighted_regression_tensor in kwargs + weighted_regression_tensor = kwargs.get("weighted_regression_tensor", None) + if weighted_regression_tensor is not None: + loss *= weighted_regression_tensor.unsqueeze(-1) + + return loss.mean() + + def update(self, x0, condition=None, update_ema=True, **kwargs): + """One-step gradient update. + Inputs: + - x0: torch.Tensor + Samples from the target distribution. + - condition: Optional + Condition of x0. `None` indicates no condition. + - update_ema: bool + Whether to update the exponential moving average model. + + Outputs: + - log: dict + The log dictionary. + """ + loss = self.loss(x0, condition, **kwargs) + + loss.backward() + grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_norm) \ + if self.grad_clip_norm else None + self.optimizer.step() + self.optimizer.zero_grad() + + if update_ema: + self.ema_update() + + log = {"loss": loss.item(), "grad_norm": grad_norm} + + return log + + def update_classifier(self, x0, condition): + + xt, t, eps = self.add_noise(x0) + + log = self.classifier.update(xt, t, condition) + + return log + + # ==================== Sampling: Solving SDE/ODE ====================== + + def classifier_guidance( + self, xt, t, alpha, sigma, + model, condition=None, w: float = 1.0, + pred=None): + """ + Guided Sampling CG: + bar_eps = eps - w * sigma * grad + bar_x0 = x0 + w * (sigma ** 2) * alpha * grad + """ + if pred is None: + pred = model["diffusion"](xt, t, None) + if self.classifier is None or w == 0.0: + return pred, None + else: + log_p, grad = self.classifier.gradients(xt.clone(), t, condition) + if self.predict_noise: + pred = pred - w * sigma * grad + else: + pred = pred + w * ((sigma ** 2) / alpha) * grad + + return pred, log_p + + def classifier_free_guidance( + self, xt, t, + model, condition=None, w: float = 1.0, + pred=None, pred_uncond=None, + requires_grad: bool = False): + """ + Guided Sampling CFG: + bar_eps = w * pred + (1 - w) * pred_uncond + bar_x0 = w * pred + (1 - w) * pred_uncond + """ + with torch.set_grad_enabled(requires_grad): + if w != 0.0 and w != 1.0: + if pred is None or pred_uncond is None: + b = xt.shape[0] + repeat_dim = [2 if i == 0 else 1 for i in range(xt.dim())] + condition = torch.cat([condition, torch.zeros_like(condition)], 0) + pred_all = model["diffusion"]( + xt.repeat(*repeat_dim), t.repeat(2), condition) + pred, pred_uncond = pred_all[:b], pred_all[b:] + elif w == 0.0: + pred = 0. + pred_uncond = model["diffusion"](xt, t, None) + else: + pred = model["diffusion"](xt, t, condition) + pred_uncond = 0. + + if self.predict_noise or not self.predict_noise: + bar_pred = w * pred + (1 - w) * pred_uncond + else: + bar_pred = pred + + return bar_pred + + def clip_prediction(self, pred, xt, alpha, sigma): + """ + Clip the prediction at each sampling step to stablize the generation. + (xt - alpha * x_max) / sigma <= eps <= (xt - alpha * x_min) / sigma + x_min <= x0 <= x_max + """ + if self.predict_noise: + if self.clip_pred: + upper_bound = (xt - alpha * self.x_min) / sigma if self.x_min is not None else None + lower_bound = (xt - alpha * self.x_max) / sigma if self.x_max is not None else None + pred = pred.clip(lower_bound, upper_bound) + else: + if self.clip_pred: + pred = pred.clip(self.x_min, self.x_max) + + return pred + + def guided_sampling( + self, xt, t, alpha, sigma, + model, + condition_cfg=None, w_cfg: float = 0.0, + condition_cg=None, w_cg: float = 0.0, + requires_grad: bool = False): + """ + One-step epsilon/x0 prediction with guidance. + """ + + pred = self.classifier_free_guidance( + xt, t, model, condition_cfg, w_cfg, None, None, requires_grad) + + pred, logp = self.classifier_guidance( + xt, t, alpha, sigma, model, condition_cg, w_cg, pred) + + return pred, logp + + def sample(self, *args, **kwargs): + raise NotImplementedError + + +class DiscreteDiffusionSDE(BaseDiffusionSDE): + """Discrete-time Diffusion SDE (VP-SDE) + + The Diffusion SDE is currently one of the most commonly used formulations of diffusion processes. + Its training process involves utilizing neural networks to estimate its scaled score function, + which is used to compute the reverse process. The Diffusion SDE has reverse processes + in both SDE and ODE forms, sharing the same marginal distribution. + The first-order discretized forms of both are equivalent to well-known models such as DDPM and DDIM. + DPM-Solvers have observed the semi-linearity of the reverse process and have computed its exact solution. + + The DiscreteDiffusionSDE is the discrete-time version of the Diffusion SDE. + It discretizes the continuous time interval into a finite number of diffusion steps + and only estimates the score function on these steps. + Therefore, in sampling, solvers can only work on these learned steps. + The sampling steps are required to be greater than 1 and less than or equal to the number of diffusion steps. + + Args: + - nn_diffusion: BaseNNDiffusion + The neural network backbone for the Diffusion model. + - nn_condition: Optional[BaseNNCondition] + The neural network backbone for the condition embedding. + + - fix_mask: Union[list, np.ndarray, torch.Tensor] + Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + The mask should be in the shape of `x_shape`. + - loss_weight: Union[list, np.ndarray, torch.Tensor] + Add loss weight. The weight should be in the shape of `x_shape`. + + - classifier: Optional[BaseClassifier] + Add a classifier to enable classifier-guidance. + + - grad_clip_norm: Optional[float] + Gradient clipping norm. + - ema_rate: float + Exponential moving average rate. + - optim_params: Optional[dict] + Optimizer parameters. + + - epsilon: float + The minimum time step for the diffusion reverse process. + In practice, using a very small value instead of `0` can avoid numerical instability. + + - diffusion_steps: int + The discretization steps for discrete-time diffusion models. + - discretization: Union[str, Callable] + The discretization method for the diffusion steps. + + - noise_schedule: Union[str, Dict[str, Callable]] + The noise schedule for the diffusion process. Can be "linear" or "cosine". + - noise_schedule_params: Optional[dict] + The parameters for the noise schedule. + + - x_max: Optional[torch.Tensor] + The maximum value for the input data. `None` indicates no constraint. + - x_min: Optional[torch.Tensor] + The minimum value for the input data. `None` indicates no constraint. + + - predict_noise: bool + Whether to predict the noise or the data. + + - device: Union[torch.device, str] + The device to run the model. + """ + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugins ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Training Params ---------------- # + grad_clip_norm: Optional[float] = None, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- Diffusion Params ------------------- # + epsilon: float = 1e-3, + + diffusion_steps: int = 1000, + discretization: Union[str, Callable] = "uniform", + + noise_schedule: Union[str, Dict[str, Callable]] = "cosine", + noise_schedule_params: Optional[dict] = None, + + x_max: Optional[torch.Tensor] = None, + x_min: Optional[torch.Tensor] = None, + + predict_noise: bool = True, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, ema_rate, optim_params, + epsilon, noise_schedule, noise_schedule_params, x_max, x_min, predict_noise, device) + + self.diffusion_steps = diffusion_steps + + if 1. / diffusion_steps < epsilon: + raise ValueError("epsilon is too large for the number of diffusion steps") + + # ================= Discretization ================= + # - Map the continuous range [epsilon, 1.] to the discrete range [0, T-1] + if isinstance(discretization, str): + if discretization in SUPPORTED_DISCRETIZATIONS.keys(): + self.t_diffusion = SUPPORTED_DISCRETIZATIONS[discretization](diffusion_steps, epsilon).to(device) + else: + Warning(f"Discretization method {discretization} is not supported. " + f"Using uniform discretization instead.") + self.t_diffusion = SUPPORTED_DISCRETIZATIONS["uniform"](diffusion_steps, epsilon).to(device) + elif callable(discretization): + self.t_diffusion = discretization(diffusion_steps, epsilon).to(device) + else: + raise ValueError("discretization must be a callable or a string") + + # ================= Noise Schedule ================= + if isinstance(noise_schedule, str): + if noise_schedule in SUPPORTED_NOISE_SCHEDULES.keys(): + self.alpha, self.sigma = SUPPORTED_NOISE_SCHEDULES[noise_schedule]["forward"]( + self.t_diffusion, **(noise_schedule_params or {})) + else: + raise ValueError(f"Noise schedule {noise_schedule} is not supported.") + elif isinstance(noise_schedule, dict): + self.alpha, self.sigma = noise_schedule["forward"](self.t_diffusion, **(noise_schedule_params or {})) + else: + raise ValueError("noise_schedule must be a callable or a string") + + self.logSNR = torch.log(self.alpha / self.sigma) + + # ==================== Training: Score Matching ====================== + + def add_noise(self, x0, t=None, eps=None): + + t = torch.randint(self.diffusion_steps, (x0.shape[0],), device=self.device) if t is None else t + eps = torch.randn_like(x0) if eps is None else eps + + alpha, sigma = at_least_ndim(self.alpha[t], x0.dim()), at_least_ndim(self.sigma[t], x0.dim()) + + xt = alpha * x0 + sigma * eps + xt = (1. - self.fix_mask) * xt + self.fix_mask * x0 + + return xt, t, eps + + # ==================== Sampling: Solving SDE/ODE ====================== + + def sample( + self, + # ---------- the known fixed portion ---------- # + prior: torch.Tensor, + # ----------------- sampling ----------------- # + solver: str = "ddpm", + n_samples: int = 1, + sample_steps: int = 5, + sample_step_schedule: Union[str, Callable] = "uniform", + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ----------- Diffusion-X sampling ---------- + diffusion_x_sampling_steps: int = 0, + # ----------- Warm-Starting ----------- + warm_start_reference: Optional[torch.Tensor] = None, + warm_start_forward_level: float = 0.3, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + # ------------------ Data collection ------------------ # + collector=None, + collect_trajectory: bool = False, + obs_dim: int = None, + initial_noise: Optional[torch.Tensor] = None, + return_initial_noise: bool = False, + **kwargs, + ): + """Sampling. + + Inputs: + - prior: torch.Tensor + The known fixed portion of the input data. Should be in the shape of generated data. + Use `torch.zeros((n_samples, *x_shape))` for non-prior sampling. + + - solver: str + The solver for the reverse process. Check `supported_solvers` property for available solvers. + - n_samples: int + The number of samples to generate. + - sample_steps: int + The number of sampling steps. Should be greater than 1 and less than or equal to the number of diffusion steps. + - sample_step_schedule: Union[str, Callable] + The schedule for the sampling steps. + - use_ema: bool + Whether to use the exponential moving average model. + - temperature: float + The temperature for sampling. + - initial_noise: Optional[torch.Tensor] + Optional standard-normal noise used to initialize the reverse process. + It is multiplied by ``temperature`` before fixed conditioning is applied. + Supplying this makes noise/trajectory distillation pairs exactly reproducible. + - return_initial_noise: bool + If true, return the exact conditioned initial state as ``log["initial_xt"]``. + + - condition_cfg: Optional + Condition for Classifier-free-guidance. + - mask_cfg: Optional + Mask for Classifier-guidance. + - w_cfg: float + Weight for Classifier-free-guidance. + - condition_cg: Optional + Condition for Classifier-guidance. + - w_cg: float + Weight for Classifier-guidance. + + - diffusion_x_sampling_steps: int + The number of diffusion steps for diffusion-x sampling. + + - warm_start_reference: Optional[torch.Tensor] + Reference data for warm-starting sampling. `None` indicates no warm-starting. + - warm_start_forward_level: float + The forward noise level to perturb the reference data. Should be in the range of `[0., 1.]`, where `1` indicates pure noise. + + - requires_grad: bool + Whether to preserve gradients. + - preserve_history: bool + Whether to preserve the sampling history. + + Outputs: + - x0: torch.Tensor + Generated samples. Be in the shape of `(n_samples, *x_shape)`. + - log: dict + The log dictionary. + """ + assert solver in SUPPORTED_SOLVERS, f"Solver {solver} is not supported." + + # ===================== Initialization ===================== + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + model = self.model if not use_ema else self.model_ema + + prior = prior.to(self.device) + if isinstance(warm_start_reference, torch.Tensor) and isinstance(initial_noise, torch.Tensor): + raise ValueError("warm_start_reference and initial_noise are mutually exclusive") + if isinstance(warm_start_reference, torch.Tensor): + diffusion_steps = int(warm_start_forward_level * self.diffusion_steps) + fwd_alpha, fwd_sigma = self.alpha[diffusion_steps], self.sigma[diffusion_steps] + xt = warm_start_reference * fwd_alpha + fwd_sigma * torch.randn_like(warm_start_reference) + else: + diffusion_steps = self.diffusion_steps + if isinstance(initial_noise, torch.Tensor): + if initial_noise.shape != prior.shape: + raise ValueError( + f"initial_noise shape {tuple(initial_noise.shape)} does not match " + f"prior shape {tuple(prior.shape)}") + xt = initial_noise.to(device=prior.device, dtype=prior.dtype) * temperature + else: + xt = torch.randn_like(prior) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if return_initial_noise: + # Clone before denoising: later solver updates must not mutate the saved pair input. + log["initial_xt"] = xt.detach().clone() + if preserve_history: + log["sample_history"][:, 0] = xt.cpu().numpy() + + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + # ===================== Sampling Schedule ==================== + if isinstance(sample_step_schedule, str): + if sample_step_schedule in SUPPORTED_SAMPLING_STEP_SCHEDULE.keys(): + sample_step_schedule = SUPPORTED_SAMPLING_STEP_SCHEDULE[sample_step_schedule]( + diffusion_steps, sample_steps) + else: + raise ValueError(f"Sampling step schedule {sample_step_schedule} is not supported.") + elif callable(sample_step_schedule): + sample_step_schedule = sample_step_schedule(diffusion_steps, sample_steps) + else: + raise ValueError("sample_step_schedule must be a callable or a string") + + alphas = self.alpha[sample_step_schedule] + sigmas = self.sigma[sample_step_schedule] + logSNRs = torch.log(alphas / sigmas) + hs = torch.zeros_like(logSNRs) + hs[1:] = logSNRs[:-1] - logSNRs[1:] # hs[0] is not correctly calculated, but it will not be used. + stds = torch.zeros((sample_steps + 1,), device=self.device) + stds[1:] = sigmas[:-1] / sigmas[1:] * (1 - (alphas[1:] / alphas[:-1]) ** 2).sqrt() + + buffer = [] + + # ===================== Denoising Loop ======================== + loop_steps = [1] * diffusion_x_sampling_steps + list(range(1, sample_steps + 1)) + for i in reversed(loop_steps): + + t = torch.full((n_samples,), sample_step_schedule[i], dtype=torch.long, device=self.device) + + xt_before = xt + + # guided sampling + pred, logp = self.guided_sampling( + xt, t, alphas[i], sigmas[i], + model, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg, requires_grad) + + # clip the prediction + pred = self.clip_prediction(pred, xt, alphas[i], sigmas[i]) + + # noise & data prediction + eps_theta = pred if self.predict_noise else xtheta_to_epstheta(xt, alphas[i], sigmas[i], pred) + x_theta = pred if not self.predict_noise else epstheta_to_xtheta(xt, alphas[i], sigmas[i], pred) + + # one-step update + if solver == "ddpm": + xt = ( + (alphas[i - 1] / alphas[i]) * (xt - sigmas[i] * eps_theta) + + (sigmas[i - 1] ** 2 - stds[i] ** 2 + 1e-8).sqrt() * eps_theta) + if i > 1: + xt += (stds[i] * torch.randn_like(xt)) + + elif solver == "ddim": + xt = (alphas[i - 1] * ((xt - sigmas[i] * eps_theta) / alphas[i]) + sigmas[i - 1] * eps_theta) + + elif solver == "ode_dpmsolver_1": + xt = (alphas[i - 1] / alphas[i]) * xt - sigmas[i - 1] * torch.expm1(hs[i]) * eps_theta + + elif solver == "ode_dpmsolver++_1": + xt = (sigmas[i - 1] / sigmas[i]) * xt - alphas[i - 1] * torch.expm1(-hs[i]) * x_theta + + elif solver == "ode_dpmsolver++_2M": + buffer.append(x_theta) + if i < sample_steps: + r = hs[i + 1] / hs[i] + D = (1 + 0.5 / r) * buffer[-1] - 0.5 / r * buffer[-2] + xt = (sigmas[i - 1] / sigmas[i]) * xt - alphas[i - 1] * torch.expm1(-hs[i]) * D + else: + xt = (sigmas[i - 1] / sigmas[i]) * xt - alphas[i - 1] * torch.expm1(-hs[i]) * x_theta + + elif solver == "sde_dpmsolver_1": + xt = ((alphas[i - 1] / alphas[i]) * xt - + 2 * sigmas[i - 1] * torch.expm1(hs[i]) * eps_theta + + sigmas[i - 1] * torch.expm1(2 * hs[i]).sqrt() * torch.randn_like(xt)) + + elif solver == "sde_dpmsolver++_1": + xt = ((sigmas[i - 1] / sigmas[i]) * (-hs[i]).exp() * xt - + alphas[i - 1] * torch.expm1(-2 * hs[i]) * x_theta + + sigmas[i - 1] * (-torch.expm1(-2 * hs[i])).sqrt() * torch.randn_like(xt)) + + elif solver == "sde_dpmsolver++_2M": + buffer.append(x_theta) + if i < sample_steps: + r = hs[i + 1] / hs[i] + D = (1 + 0.5 / r) * buffer[-1] - 0.5 / r * buffer[-2] + xt = ((sigmas[i - 1] / sigmas[i]) * (-hs[i]).exp() * xt - + alphas[i - 1] * torch.expm1(-2 * hs[i]) * D + + sigmas[i - 1] * (-torch.expm1(-2 * hs[i])).sqrt() * torch.randn_like(xt)) + else: + xt = ((sigmas[i - 1] / sigmas[i]) * (-hs[i]).exp() * xt - + alphas[i - 1] * torch.expm1(-2 * hs[i]) * x_theta + + sigmas[i - 1] * (-torch.expm1(-2 * hs[i])).sqrt() * torch.randn_like(xt)) + + # fix the known portion, and preserve the sampling history + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + # + xt_after = xt + + # === 轨迹采集（关键）=== + if collect_trajectory and collector is not None: + assert obs_dim is not None, "collect_trajectory 需要传入 obs_dim" + u = prior[:, 0, :obs_dim] # u = obs0（归一化之后） + collector.add(xt_before, xt_after, t, u, i_step=i) + + if preserve_history: + log["sample_history"][:, sample_steps - i + 1] = xt.cpu().numpy() + + # ================= Post-processing ================= + if self.classifier is not None: + with torch.no_grad(): + t = torch.zeros((n_samples,), dtype=torch.long, device=self.device) + logp = self.classifier.logp(xt, t, condition_vec_cg) + log["log_p"] = logp + + if self.clip_pred: + xt = xt.clip(self.x_min, self.x_max) + + return xt, log + + +class ContinuousDiffusionSDE(BaseDiffusionSDE): + """Continuous-time Diffusion SDE (VP-SDE) + + The Diffusion SDE is currently one of the most commonly used formulations of diffusion processes. + Its training process involves utilizing neural networks to estimate its scaled score function, + which is used to compute the reverse process. The Diffusion SDE has reverse processes + in both SDE and ODE forms, sharing the same marginal distribution. + The first-order discretized forms of both are equivalent to well-known models such as DDPM and DDIM. + DPM-Solvers have observed the semi-linearity of the reverse process and have computed its exact solution. + + The ContinuousDiffusionSDE is the continuous-time version of the Diffusion SDE. + It estimates the score function at any $t\in[0,T]$ and solves the reverse process in continuous time. + The sampling steps are required to be greater than 1. + + Args: + - nn_diffusion: BaseNNDiffusion + The neural network backbone for the Diffusion model. + - nn_condition: Optional[BaseNNCondition] + The neural network backbone for the condition embedding. + + - fix_mask: Union[list, np.ndarray, torch.Tensor] + Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + The mask should be in the shape of `x_shape`. + - loss_weight: Union[list, np.ndarray, torch.Tensor] + Add loss weight. The weight should be in the shape of `x_shape`. + + - classifier: Optional[BaseClassifier] + Add a classifier to enable classifier-guidance. + + - grad_clip_norm: Optional[float] + Gradient clipping norm. + - ema_rate: float + Exponential moving average rate. + - optim_params: Optional[dict] + Optimizer parameters. + + - epsilon: float + The minimum time step for the diffusion reverse process. + In practice, using a very small value instead of `0` can avoid numerical instability. + + - noise_schedule: Union[str, Dict[str, Callable]] + The noise schedule for the diffusion process. Can be "linear" or "cosine". + - noise_schedule_params: Optional[dict] + The parameters for the noise schedule. + + - x_max: Optional[torch.Tensor] + The maximum value for the input data. `None` indicates no constraint. + - x_min: Optional[torch.Tensor] + The minimum value for the input data. `None` indicates no constraint. + + - predict_noise: bool + Whether to predict the noise or the data. + + - device: Union[torch.device, str] + The device to run the model. + """ + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugins ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Training Params ---------------- # + grad_clip_norm: Optional[float] = None, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- Diffusion Params ------------------- # + epsilon: float = 1e-3, + + noise_schedule: Union[str, Dict[str, Callable]] = "cosine", + noise_schedule_params: Optional[dict] = None, + + x_max: Optional[torch.Tensor] = None, + x_min: Optional[torch.Tensor] = None, + + predict_noise: bool = True, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, ema_rate, optim_params, + epsilon, noise_schedule, noise_schedule_params, x_max, x_min, predict_noise, device) + + # ==================== Continuous Time-step Range ==================== + if noise_schedule == "cosine": + self.t_diffusion = [epsilon, 0.9946] + else: + self.t_diffusion = [epsilon, 1.] + + # ===================== Noise Schedule ====================== + if isinstance(noise_schedule, str): + if noise_schedule in SUPPORTED_NOISE_SCHEDULES.keys(): + self.noise_schedule_funcs = SUPPORTED_NOISE_SCHEDULES[noise_schedule] + self.noise_schedule_params = noise_schedule_params + else: + raise ValueError(f"Noise schedule {noise_schedule} is not supported.") + elif isinstance(noise_schedule, dict): + self.noise_schedule_funcs = noise_schedule + self.noise_schedule_params = noise_schedule_params + else: + raise ValueError("noise_schedule must be a callable or a string") + + # ==================== Training: Score Matching ====================== + + def add_noise(self, x0, t=None, eps=None): + + t = (torch.rand((x0.shape[0],), device=self.device) * + (self.t_diffusion[1] - self.t_diffusion[0]) + self.t_diffusion[0]) if t is None else t + + eps = torch.randn_like(x0) if eps is None else eps + + alpha, sigma = self.noise_schedule_funcs["forward"](t, **(self.noise_schedule_params or {})) + alpha = at_least_ndim(alpha, x0.dim()) + sigma = at_least_ndim(sigma, x0.dim()) + + xt = alpha * x0 + sigma * eps + xt = (1. - self.fix_mask) * xt + self.fix_mask * x0 + + return xt, t, eps + + # ==================== Sampling: Solving SDE/ODE ====================== + + def sample( + self, + # ---------- the known fixed portion ---------- # + prior: torch.Tensor, + # ----------------- sampling ----------------- # + solver: str = "ddpm", + n_samples: int = 1, + sample_steps: int = 5, + sample_step_schedule: Union[str, Callable] = "uniform_continuous", + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ----------- Diffusion-X sampling ---------- + diffusion_x_sampling_steps: int = 0, + # ----------- Warm-Starting ----------- + warm_start_reference: Optional[torch.Tensor] = None, + warm_start_forward_level: float = 0.3, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + """Sampling. + + Inputs: + - prior: torch.Tensor + The known fixed portion of the input data. Should be in the shape of generated data. + Use `torch.zeros((n_samples, *x_shape))` for non-prior sampling. + + - solver: str + The solver for the reverse process. Check `supported_solvers` property for available solvers. + - n_samples: int + The number of samples to generate. + - sample_steps: int + The number of sampling steps. Should be greater than 1. + - sample_step_schedule: Union[str, Callable] + The schedule for the sampling steps. + - use_ema: bool + Whether to use the exponential moving average model. + - temperature: float + The temperature for sampling. + + - condition_cfg: Optional + Condition for Classifier-free-guidance. + - mask_cfg: Optional + Mask for Classifier-guidance. + - w_cfg: float + Weight for Classifier-free-guidance. + - condition_cg: Optional + Condition for Classifier-guidance. + - w_cg: float + Weight for Classifier-guidance. + + - diffusion_x_sampling_steps: int + The number of diffusion steps for diffusion-x sampling. + + - warm_start_reference: Optional[torch.Tensor] + Reference data for warm-starting sampling. `None` indicates no warm-starting. + - warm_start_forward_level: float + The forward noise level to perturb the reference data. Should be in the range of `[0., 1.]`, where `1` indicates pure noise. + + - requires_grad: bool + Whether to preserve gradients. + - preserve_history: bool + Whether to preserve the sampling history. + + Outputs: + - x0: torch.Tensor + Generated samples. Be in the shape of `(n_samples, *x_shape)`. + - log: dict + The log dictionary. + """ + assert solver in SUPPORTED_SOLVERS, f"Solver {solver} is not supported." + + # ===================== Initialization ===================== + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + model = self.model if not use_ema else self.model_ema + + prior = prior.to(self.device) + if isinstance(warm_start_reference, torch.Tensor) and warm_start_forward_level > 0.: + warm_start_forward_level = self.epsilon + warm_start_forward_level * (1. - self.epsilon) + fwd_alpha, fwd_sigma = self.noise_schedule_funcs["forward"]( + torch.ones((1,), device=self.device) * warm_start_forward_level, **(self.noise_schedule_params or {})) + xt = warm_start_reference * fwd_alpha + fwd_sigma * torch.randn_like(warm_start_reference) + else: + xt = torch.randn_like(prior) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, 0] = xt.cpu().numpy() + + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + # ===================== Sampling Schedule ==================== + if isinstance(warm_start_reference, torch.Tensor) and warm_start_forward_level > 0.: + t_diffusion = [self.t_diffusion[0], warm_start_forward_level] + else: + t_diffusion = self.t_diffusion + if isinstance(sample_step_schedule, str): + if sample_step_schedule in SUPPORTED_SAMPLING_STEP_SCHEDULE.keys(): + sample_step_schedule = SUPPORTED_SAMPLING_STEP_SCHEDULE[sample_step_schedule]( + t_diffusion, sample_steps) + else: + raise ValueError(f"Sampling step schedule {sample_step_schedule} is not supported.") + elif callable(sample_step_schedule): + sample_step_schedule = sample_step_schedule(t_diffusion, sample_steps) + else: + raise ValueError("sample_step_schedule must be a callable or a string") + + alphas, sigmas = self.noise_schedule_funcs["forward"]( + sample_step_schedule, **(self.noise_schedule_params or {})) + logSNRs = torch.log(alphas / sigmas) + hs = torch.zeros_like(logSNRs) + hs[1:] = logSNRs[:-1] - logSNRs[1:] # hs[0] is not correctly calculated, but it will not be used. + stds = torch.zeros((sample_steps + 1,), device=self.device) + stds[1:] = sigmas[:-1] / sigmas[1:] * (1 - (alphas[1:] / alphas[:-1]) ** 2).sqrt() + + buffer = [] + + # ===================== Denoising Loop ======================== + loop_steps = [1] * diffusion_x_sampling_steps + list(range(1, sample_steps + 1)) + for i in reversed(loop_steps): + + t = torch.full((n_samples,), sample_step_schedule[i], dtype=torch.float32, device=self.device) + + # guided sampling + pred, logp = self.guided_sampling( + xt, t, alphas[i], sigmas[i], + model, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg, requires_grad) + + # clip the prediction + pred = self.clip_prediction(pred, xt, alphas[i], sigmas[i]) + + # transform to eps_theta + eps_theta = pred if self.predict_noise else xtheta_to_epstheta(xt, alphas[i], sigmas[i], pred) + x_theta = pred if not self.predict_noise else epstheta_to_xtheta(xt, alphas[i], sigmas[i], pred) + + # one-step update + if solver == "ddpm": + xt = ( + (alphas[i - 1] / alphas[i]) * (xt - sigmas[i] * eps_theta) + + (sigmas[i - 1] ** 2 - stds[i] ** 2 + 1e-8).sqrt() * eps_theta) + if i > 1: + xt += (stds[i] * torch.randn_like(xt)) + + elif solver == "ddim": + xt = (alphas[i - 1] * ((xt - sigmas[i] * eps_theta) / alphas[i]) + sigmas[i - 1] * eps_theta) + + elif solver == "ode_dpmsolver_1": + xt = (alphas[i - 1] / alphas[i]) * xt - sigmas[i - 1] * torch.expm1(hs[i]) * eps_theta + + elif solver == "ode_dpmsolver++_1": + xt = (sigmas[i - 1] / sigmas[i]) * xt - alphas[i - 1] * torch.expm1(-hs[i]) * x_theta + + elif solver == "ode_dpmsolver++_2M": + buffer.append(x_theta) + if i < sample_steps: + r = hs[i + 1] / hs[i] + D = (1 + 0.5 / r) * buffer[-1] - 0.5 / r * buffer[-2] + xt = (sigmas[i - 1] / sigmas[i]) * xt - alphas[i - 1] * torch.expm1(-hs[i]) * D + else: + xt = (sigmas[i - 1] / sigmas[i]) * xt - alphas[i - 1] * torch.expm1(-hs[i]) * x_theta + + elif solver == "sde_dpmsolver_1": + xt = ((alphas[i - 1] / alphas[i]) * xt - + 2 * sigmas[i - 1] * torch.expm1(hs[i]) * eps_theta + + sigmas[i - 1] * torch.expm1(2 * hs[i]).sqrt() * torch.randn_like(xt)) + + elif solver == "sde_dpmsolver++_1": + xt = ((sigmas[i - 1] / sigmas[i]) * (-hs[i]).exp() * xt - + alphas[i - 1] * torch.expm1(-2 * hs[i]) * x_theta + + sigmas[i - 1] * (-torch.expm1(-2 * hs[i])).sqrt() * torch.randn_like(xt)) + + elif solver == "sde_dpmsolver++_2M": + buffer.append(x_theta) + if i < sample_steps: + r = hs[i + 1] / hs[i] + D = (1 + 0.5 / r) * buffer[-1] - 0.5 / r * buffer[-2] + xt = ((sigmas[i - 1] / sigmas[i]) * (-hs[i]).exp() * xt - + alphas[i - 1] * torch.expm1(-2 * hs[i]) * D + + sigmas[i - 1] * (-torch.expm1(-2 * hs[i])).sqrt() * torch.randn_like(xt)) + else: + xt = ((sigmas[i - 1] / sigmas[i]) * (-hs[i]).exp() * xt - + alphas[i - 1] * torch.expm1(-2 * hs[i]) * x_theta + + sigmas[i - 1] * (-torch.expm1(-2 * hs[i])).sqrt() * torch.randn_like(xt)) + + # fix the known portion, and preserve the sampling history + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, sample_steps - i + 1] = xt.cpu().numpy() + + # ================= Post-processing ================= + if self.classifier is not None and w_cg != 0.: + with torch.no_grad(): + t = torch.zeros((n_samples,), dtype=torch.long, device=self.device) + logp = self.classifier.logp(xt, t, condition_vec_cg) + log["log_p"] = logp + + if self.clip_pred: + xt = xt.clip(self.x_min, self.x_max) + + return xt, log diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/dpmsolver.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/dpmsolver.py new file mode 100644 index 0000000000000000000000000000000000000000..9c302d5a99568db962e86336a3dbf12285a87f84 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/dpmsolver.py @@ -0,0 +1,889 @@ +from typing import Optional, Union + +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.classifier import BaseClassifier +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from cleandiffuser.utils import at_least_ndim +from .basic import DiffusionModel + +SAMPLER_CONFIG = { + + "ode_dpm_1": { + "predict_noise": True, + "order": 1, }, + + "ddim": { + "predict_noise": True, + "order": 1, }, + + "sde_dpm_1": { + "predict_noise": True, + "order": 1, }, + + "ode_dpmpp_1": { + "predict_noise": False, + "order": 1, }, + + "sde_dpmpp_1": { + "predict_noise": False, + "order": 1, }, + + "ode_dpm_2": { + "predict_noise": True, + "order": 2, }, + + "sde_dpm_2": { + "predict_noise": True, + "order": 2, }, + + "ode_dpmpp_2": { + "predict_noise": False, + "order": 2, }, + + "sde_dpmpp_2": { + "predict_noise": False, + "order": 2, }, } + + +def epstheta_to_xtheta(x, alpha, sigma, eps_theta): + """ + x_theta = (x - sigma * eps_theta) / alpha + """ + return (x - sigma * eps_theta) / alpha + + +def xtheta_to_epstheta(x, alpha, sigma, x_theta): + """ + eps_theta = (x - alpha * x_theta) / sigma + """ + return (x - alpha * x_theta) / sigma + + +def onestep_estimate(xt, pred_theta, i, alphas, sigmas, h, sampler): + + if sampler[:-2] == "ode_dpm" or sampler == "ddim": + xt_next = ( + alphas[i] / alphas[i - 1] * xt - + sigmas[i] * torch.expm1(h[i]) * pred_theta) + elif sampler[:-2] == "sde_dpm": + xt_next = ( + alphas[i] / alphas[i - 1] * xt - + 2. * sigmas[i] * torch.expm1(h[i]) * pred_theta + + sigmas[i] * torch.expm1(2. * h[i]).sqrt() * torch.randn_like(xt)) + elif sampler[:-2] == "ode_dpmpp": + xt_next = ( + sigmas[i] / sigmas[i - 1] * xt - + alphas[i] * torch.expm1(-h[i]) * pred_theta) + elif sampler[:-2] == "sde_dpmpp": + xt_next = ( + sigmas[i] / sigmas[i - 1] * (-h[i]).exp() * xt - + alphas[i] * torch.expm1(-2. * h[i]) * pred_theta + + sigmas[i] * (-1. * torch.expm1(-2. * h[i])).sqrt() * torch.randn_like(xt)) + else: + raise ValueError(f'Unknown sampler: {sampler}.') + + return xt_next + + +class DPMSolver(DiffusionModel): + + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugs ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Params ---------------- # + grad_clip_norm: Optional[float] = None, + diffusion_steps: int = 1000, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + x_max: Optional[torch.Tensor] = None, + x_min: Optional[torch.Tensor] = None, + + # ------------------- DPM Params ------------------- # + predict_noise: bool = False, + noise_schedule: str = "linear", # or cosine + t_eps: float = 1e-3, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + diffusion_steps, ema_rate, optim_params, device) + + self.predict_noise = predict_noise + self.noise_schedule = noise_schedule + self.t_eps = t_eps + self.x_max, self.x_min = x_max, x_min + + @property + def clip_pred(self): + return (self.x_max is not None) or (self.x_min is not None) + + @property + def supported_samplers(self): + return list(SAMPLER_CONFIG.keys()) + + @property + def t_range(self): + if self.noise_schedule == "linear": + return self.t_eps, 1. + elif self.noise_schedule == "cosine": + return self.t_eps, 0.9946 + else: + raise ValueError(f"noise_schedule should be 'linear' or 'cosine', but got {self.noise_schedule}.") + + def alpha_schedule(self, t): + if self.noise_schedule == "linear": + beta0, beta1 = 0.1, 20 + return (-(beta1 - beta0) / 4 * (t ** 2) - beta0 / 2 * t).exp() + elif self.noise_schedule == "cosine": + s = 0.008 + return ((torch.cos(np.pi / 2 * (t + s) / (1 + s))).log() - np.log(np.cos(np.pi / 2 * s / (1 + s)))).exp() + else: + raise ValueError(f"noise_schedule should be 'linear' or 'cosine', but got {self.noise_schedule}.") + + # --------------------------------------------------------------------------- + # Training + + def add_noise(self, x0, t=None, eps=None): + if t is None: + t = torch.rand((x0.shape[0],), device=self.device) + t = self.t_range[0] + t * (self.t_range[1] - self.t_range[0]) + eps = torch.randn_like(x0) if eps is None else eps + alpha = self.alpha_schedule(at_least_ndim(t, x0.dim())) + sigma = (1 - alpha ** 2).sqrt() + xt = x0 * alpha + sigma * eps + xt = xt * (1. - self.fix_mask) + x0 * self.fix_mask + return xt, t, eps + + def loss(self, x0, condition=None): + xt, t, eps = self.add_noise(x0) + condition = self.model["condition"](condition) if condition is not None else None + + if self.predict_noise: + loss = (self.model["diffusion"](xt, t, condition) - eps) ** 2 + else: + loss = (self.model["diffusion"](xt, t, condition) - x0) ** 2 + + return (loss * self.loss_weight * (1. - self.fix_mask)).mean() + + def update(self, x0, condition=None, update_ema=True, **kwargs): + self.optimizer.zero_grad() + loss = self.loss(x0, condition) + loss.backward() + grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_norm) \ + if self.grad_clip_norm else None + self.optimizer.step() + if update_ema: self.ema_update() + log = {"loss": loss.item(), "grad_norm": grad_norm} + return log + + def update_classifier(self, x0, condition): + xt, t, eps = self.add_noise(x0) + log = self.classifier.update(xt, t, condition) + return log + + # --------------------------------------------------------------------------- + # Inference + + def predict_function( + self, x, t, alpha, sigma, + use_ema=False, requires_grad=False, predict_noise=False, + # ----------------- CFG ----------------- # + condition_vec_cfg=None, + w_cfg: float = 0.0, + # ----------------- CG ----------------- # + condition_vec_cg=None, + w_cg: float = 1.0, + ): + b = x.shape[0] + model = self.model_ema if use_ema else self.model + + # ----------------- CFG ----------------- # + with torch.set_grad_enabled(requires_grad): + if w_cfg != 0.0 and w_cfg != 1.0: + repeat_dim = [2 if i == 0 else 1 for i in range(x.dim())] + condition_vec_cfg = torch.cat([condition_vec_cfg, torch.zeros_like(condition_vec_cfg)], 0) + pred = model["diffusion"]( + x.repeat(*repeat_dim), t.repeat(2), condition_vec_cfg) + pred = w_cfg * pred[:b] + (1. - w_cfg) * pred[b:] + elif w_cfg == 0.0: + pred = model["diffusion"](x, t, None) + else: + pred = model["diffusion"](x, t, condition_vec_cfg) + + if self.predict_noise and predict_noise: + # The network predicts noise, and we need to predict noise here + pass + elif self.predict_noise and not predict_noise: + # The network predicts noise, and we need to predict x0 here + pred = epstheta_to_xtheta(x, alpha, sigma, pred) + elif not self.predict_noise and predict_noise: + # The network predicts x0, and we need to predict noise here + pred = xtheta_to_epstheta(x, alpha, sigma, pred) + else: + # The network predicts x0, and we need to predict x0 here + pass + + # ----------------- CG ----------------- # + if self.classifier is not None and w_cg != 0.0 and condition_vec_cg is not None: + log_p, grad = self.classifier.gradients(x.clone(), t, condition_vec_cg) + if predict_noise: + pred = pred - w_cg * sigma * grad + else: + pred = pred + w_cg * ((sigma ** 2) / alpha) * grad + else: + log_p = None + + # If the bounds are not None, clip `pred` + if predict_noise: + if self.clip_pred: + upper_bound = (x - alpha * self.x_min) / sigma if self.x_min is not None else None + lower_bound = (x - alpha * self.x_max) / sigma if self.x_max is not None else None + pred = pred.clip(lower_bound, upper_bound) + pred = pred * (1 - self.fix_mask) + else: + if self.clip_pred: + pred = pred.clip(self.x_min, self.x_max) + pred = pred * (1 - self.fix_mask) + x * self.fix_mask + + return pred, {"log_p": log_p} + + # def first_order_sample( + # self, + # sampler: str = "ode_dpm_1", + # # ---------- the known fixed portion ---------- # + # prior: Optional[torch.Tensor] = None, + # # ----------------- sampling ----------------- # + # n_samples: int = 1, + # sample_steps: int = None, + # use_ema: bool = True, + # kappa: float = 1.0, + # temperature: float = 1.0, + # # ------------------ guidance ------------------ # + # condition_cfg=None, + # mask_cfg=None, + # w_cfg: float = 0.0, + # condition_cg=None, + # w_cg: float = 0.0, + # + # requires_grad: bool = False, + # preserve_history: bool = False, + # ): + # log = { + # "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + # + # model = self.model_ema if use_ema else self.model + # + # idx = torch.arange(sample_steps + 1) + # t = (((sample_steps - idx) / sample_steps * self.t_range[1] ** (1 / kappa) + + # idx / sample_steps * self.t_range[0] ** (1 / kappa)) ** kappa).to(self.device) + # + # alphas = self.alpha_schedule(t) + # sigmas = (1 - alphas ** 2).sqrt() + # logSNRs = (alphas / sigmas).log() + # h = torch.zeros_like(logSNRs) + # h[1:] = logSNRs[1:] - logSNRs[:-1] + # + # xt = torch.randn_like(prior, device=self.device) * temperature + # xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + # if preserve_history: log["sample_history"][:, 0] = xt.cpu().numpy() + # + # with torch.set_grad_enabled(requires_grad): + # condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + # condition_vec_cg = condition_cg + # + # for i in range(1, sample_steps + 1): + # + # pred_theta, log = self.predict_function( + # xt, t[i - 1].repeat(n_samples), alphas[i - 1], sigmas[i - 1], + # use_ema=use_ema, + # requires_grad=requires_grad, + # predict_noise=SAMPLER_CONFIG[sampler]["predict_noise"], + # condition_vec_cfg=condition_vec_cfg, + # condition_vec_cg=condition_vec_cg, + # w_cfg=w_cfg, w_cg=w_cg) + # + # xt = onestep_estimate(xt, pred_theta, i, alphas, sigmas, h, sampler) + # xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + # if preserve_history: log["sample_history"][:, 0] = xt.cpu().numpy() + # + # if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + # with torch.no_grad(): + # logp = self.classifier.logp(xt, t[-1].repeat(n_samples), condition_vec_cg) + # log["log_p"] = logp + # + # return xt, log + # + # def second_order_sample_multistep( + # self, + # sampler: str = "ode_dpmpp_2", + # # ---------- the known fixed portion ---------- # + # prior: Optional[torch.Tensor] = None, + # # ----------------- sampling ----------------- # + # n_samples: int = 1, + # sample_steps: int = None, + # use_ema: bool = True, + # kappa: float = 1.0, + # temperature: float = 1.0, + # # ------------------ guidance ------------------ # + # condition_cfg=None, + # mask_cfg=None, + # w_cfg: float = 0.0, + # condition_cg=None, + # w_cg: float = 0.0, + # + # requires_grad: bool = False, + # preserve_history: bool = False, + # ): + # log = { + # "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + # + # model = self.model_ema if use_ema else self.model + # + # idx = torch.arange(sample_steps + 1) + # t = (((sample_steps - idx) / sample_steps * self.t_range[1] ** (1 / kappa) + + # idx / sample_steps * self.t_range[0] ** (1 / kappa)) ** kappa).to(self.device) + # + # alphas = self.alpha_schedule(t) + # sigmas = (1 - alphas ** 2).sqrt() + # logSNRs = (alphas / sigmas).log() + # h = torch.zeros_like(logSNRs) + # h[1:] = logSNRs[1:] - logSNRs[:-1] + # + # buffer = [] + # + # xt = torch.randn_like(prior, device=self.device) * temperature + # xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + # if preserve_history: log["sample_history"][:, 0] = xt.cpu().numpy() + # + # with torch.set_grad_enabled(requires_grad): + # condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + # condition_vec_cg = condition_cg + # + # # pre-update + # + # pred_theta, log = self.predict_function( + # xt, t[0].repeat(n_samples), alphas[0], sigmas[0], + # use_ema=use_ema, + # requires_grad=requires_grad, + # predict_noise=SAMPLER_CONFIG[sampler]["predict_noise"], + # condition_vec_cfg=condition_vec_cfg, + # condition_vec_cg=condition_vec_cg, + # w_cfg=w_cfg, w_cg=w_cg) + # buffer.append(pred_theta.clone()) + # + # xt = onestep_estimate(xt, pred_theta, 1, alphas, sigmas, h, sampler) + # xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + # if preserve_history: log["sample_history"][:, 0] = xt.cpu().numpy() + # + # pred_theta, log = self.predict_function( + # xt, t[1].repeat(n_samples), alphas[1], sigmas[1], + # use_ema=use_ema, + # requires_grad=requires_grad, + # predict_noise=SAMPLER_CONFIG[sampler]["predict_noise"], + # condition_vec_cfg=condition_vec_cfg, + # condition_vec_cg=condition_vec_cg, + # w_cfg=w_cfg, w_cg=w_cg) + # buffer.append(pred_theta.clone()) + # + # for i in range(2, sample_steps + 1): + # + # r = h[i - 1] / h[i] + # + # D = (1 + 0.5 / r) * buffer[-1] - 0.5 / r * buffer[-2] + # + # xt = onestep_estimate(xt, D, i, alphas, sigmas, h, sampler) + # xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + # if preserve_history: log["sample_history"][:, 1] = xt.cpu().numpy() + # + # if i < sample_steps + 1: + # pred_theta, log = self.predict_function( + # xt, t[i].repeat(n_samples), alphas[i], sigmas[i], + # use_ema=use_ema, + # requires_grad=requires_grad, + # predict_noise=SAMPLER_CONFIG[sampler]["predict_noise"], + # condition_vec_cfg=condition_vec_cfg, + # condition_vec_cg=condition_vec_cg, + # w_cfg=w_cfg, w_cg=w_cg) + # buffer.append(pred_theta.clone()) + # + # if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + # with torch.no_grad(): + # logp = self.classifier.logp(xt, t[-1].repeat(n_samples), condition_vec_cg) + # log["log_p"] = logp + # + # return xt, log + + def sample( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + use_ema: bool = True, + temperature: float = 1.0, + kappa: float = 1.0, + sampler: str = "ddim", + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + assert sampler in self.supported_samplers, f"Sampler '{sampler}' is not supported." + + # initialize logger + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + # choose the model + model = self.model_ema if use_ema else self.model + + # schedule the sampling steps + idx = torch.arange(sample_steps + 1) + t = (((sample_steps - idx) / sample_steps * self.t_range[1] ** (1 / kappa) + + idx / sample_steps * self.t_range[0] ** (1 / kappa)) ** kappa).to(self.device) + + alphas = self.alpha_schedule(t) + sigmas = (1 - alphas ** 2).sqrt() + logSNRs = (alphas / sigmas).log() + h = torch.zeros_like(logSNRs) + h[1:] = logSNRs[1:] - logSNRs[:-1] + + buffer = [] # only for 2nd order samplers + + # initialize the samples + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: log["sample_history"][:, 0] = xt.cpu().numpy() + + # preprocess the conditions + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + # enter the sampling loop + for i in range(1, sample_steps + 1): + + # predict eps_theta or x_theta with CG/CFG + pred_theta, log = self.predict_function( + xt, t[i - 1].repeat(n_samples), alphas[i - 1], sigmas[i - 1], + use_ema=use_ema, + requires_grad=requires_grad, + predict_noise=SAMPLER_CONFIG[sampler]["predict_noise"], + condition_vec_cfg=condition_vec_cfg, + condition_vec_cg=condition_vec_cg, + w_cfg=w_cfg, w_cg=w_cg) + if SAMPLER_CONFIG[sampler]["order"] == 2: buffer.append(pred_theta.clone()) + + # one step denoise + if SAMPLER_CONFIG[sampler]["order"] == 2 and i > 1: + + r = h[i - 1] / h[i] + D = (1 + 0.5 / r) * buffer[-1] - 0.5 / r * buffer[-2] + + xt = onestep_estimate(xt, D, i, alphas, sigmas, h, sampler) + + else: + + xt = onestep_estimate(xt, pred_theta, i, alphas, sigmas, h, sampler) + + # fix the known portion, and preserve the sampling history + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: log["sample_history"][:, 1] = xt.cpu().numpy() + + # calculate the final log_p + if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + with torch.no_grad(): + logp = self.classifier.logp(xt, t[-1].repeat(n_samples), condition_vec_cg) + log["log_p"] = logp + + # clip xt if bound is not None + if self.clip_pred: + xt = xt.clip(self.x_min, self.x_max) + + return xt, log + + def sample_x( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + extra_sample_steps: int = 8, + use_ema: bool = True, + temperature: float = 1.0, + kappa: float = 1.0, + sampler: str = "ddim", + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + assert sampler in self.supported_samplers, f"Sampler '{sampler}' is not supported." + + # initialize logger + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + # choose the model + model = self.model_ema if use_ema else self.model + + # schedule the sampling steps + idx = torch.arange(sample_steps + 1) + t = (((sample_steps - idx) / sample_steps * self.t_range[1] ** (1 / kappa) + + idx / sample_steps * self.t_range[0] ** (1 / kappa)) ** kappa).to(self.device) + + alphas = self.alpha_schedule(t) + sigmas = (1 - alphas ** 2).sqrt() + logSNRs = (alphas / sigmas).log() + h = torch.zeros_like(logSNRs) + h[1:] = logSNRs[1:] - logSNRs[:-1] + + buffer = [] # only for 2nd order samplers + + # initialize the samples + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: log["sample_history"][:, 0] = xt.cpu().numpy() + + # preprocess the conditions + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + # enter the sampling loop + for i in range(1, sample_steps + 1): + + # predict eps_theta or x_theta with CG/CFG + pred_theta, log = self.predict_function( + xt, t[i - 1].repeat(n_samples), alphas[i - 1], sigmas[i - 1], + use_ema=use_ema, + requires_grad=requires_grad, + predict_noise=SAMPLER_CONFIG[sampler]["predict_noise"], + condition_vec_cfg=condition_vec_cfg, + condition_vec_cg=condition_vec_cg, + w_cfg=w_cfg, w_cg=w_cg) + if SAMPLER_CONFIG[sampler]["order"] == 2: buffer.append(pred_theta.clone()) + + # one step denoise + if SAMPLER_CONFIG[sampler]["order"] == 2 and i > 1: + + r = h[i - 1] / h[i] + D = (1 + 0.5 / r) * buffer[-1] - 0.5 / r * buffer[-2] + + xt = onestep_estimate(xt, D, i, alphas, sigmas, h, sampler) + + else: + + xt = onestep_estimate(xt, pred_theta, i, alphas, sigmas, h, sampler) + + # fix the known portion, and preserve the sampling history + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: log["sample_history"][:, 1] = xt.cpu().numpy() # TODO: BUG logging + + # diffusion-X extra sampling steps, only for 1st order solvers + if SAMPLER_CONFIG[sampler]["order"] == 1 and extra_sample_steps > 0: + for i in range(extra_sample_steps): + + # predict eps_theta or x_theta with CG/CFG + pred_theta, log = self.predict_function( + xt, t[sample_steps - 1].repeat(n_samples), alphas[sample_steps - 1], sigmas[sample_steps - 1], + use_ema=use_ema, + requires_grad=requires_grad, + predict_noise=SAMPLER_CONFIG[sampler]["predict_noise"], + condition_vec_cfg=condition_vec_cfg, + condition_vec_cg=condition_vec_cg, + w_cfg=w_cfg, w_cg=w_cg) + + xt = onestep_estimate(xt, pred_theta, sample_steps, alphas, sigmas, h, sampler) + + # fix the known portion, and preserve the sampling history + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + # calculate the final log_p + if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + with torch.no_grad(): + logp = self.classifier.logp(xt, t[-1].repeat(n_samples), condition_vec_cg) + log["log_p"] = logp + + # clip xt if bound is not None + if self.clip_pred: + xt = xt.clip(self.x_min, self.x_max) + + return xt, log + + + + # def sample_with_ode_dpm_1( + # self, + # # ---------- the known fixed portion ---------- # + # prior: Optional[torch.Tensor] = None, + # # ----------------- sampling ----------------- # + # n_samples: int = 1, + # sample_steps: int = None, + # use_ema: bool = True, + # kappa: float = 1.0, + # temperature: float = 1.0, + # # ------------------ guidance ------------------ # + # condition_cfg=None, + # mask_cfg=None, + # w_cfg: float = 0.0, + # condition_cg=None, + # w_cg: float = 0.0, + # + # preserve_history: bool = False, + # ): + # log = { + # "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + # + # model = self.model_ema if use_ema else self.model + # + # idx = torch.arange(sample_steps + 1) + # t = (((sample_steps - idx) / sample_steps * self.t_range[1] ** (1 / kappa) + + # idx / sample_steps * self.t_range[0] ** (1 / kappa)) ** kappa).to(self.device) + # + # alphas = self.alpha_schedule(t) + # sigmas = (1 - alphas ** 2).sqrt() + # logSNRs = (alphas / sigmas).log() + # h = torch.zeros_like(logSNRs) + # h[1:] = logSNRs[1:] - logSNRs[:-1] + # + # xt = torch.randn_like(prior, device=self.device) * temperature + # xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + # if preserve_history: log["sample_history"][:, 0] = xt.cpu().numpy() + # + # with torch.no_grad(): + # condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + # condition_vec_cg = condition_cg + # + # for i in range(1, sample_steps + 1): + # + # eps_theta, log = self.predict_eps_theta( + # xt, t[i - 1].repeat(n_samples), alphas[i - 1], sigmas[i - 1], + # use_ema, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg) + # + # xt = alphas[i] / alphas[i - 1] * xt - sigmas[i] * torch.expm1(h[i]) * eps_theta + # xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + # + # if preserve_history: + # x_history[:, i] = xt.cpu().numpy() + # + # log["sample_history"] = x_history + # if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + # with torch.no_grad(): + # logp = self.classifier.logp(xt, t[-1].repeat(n_samples), condition_vec_cg) + # log["log_p"] = logp + # + # return xt, log + # + # def sample_with_sde_dpm_1( + # self, + # # ---------- the known fixed portion ---------- # + # prior: Optional[torch.Tensor] = None, + # # ----------------- sampling ----------------- # + # n_samples: int = 1, + # sample_steps: int = None, + # use_ema: bool = True, + # kappa: float = 1.0, + # temperature: float = 1.0, + # # ------------------ guidance ------------------ # + # condition_cfg=None, + # mask_cfg=None, + # w_cfg: float = 0.0, + # condition_cg=None, + # w_cg: float = 0.0, + # + # preserve_history: bool = False, + # ): + # log, x_history = {}, None + # model = self.model_ema if use_ema else self.model + # + # idx = torch.arange(sample_steps + 1) + # t = (((sample_steps - idx) / sample_steps * self.t_range[1] ** (1 / kappa) + + # idx / sample_steps * self.t_range[0] ** (1 / kappa)) ** kappa).to(self.device) + # + # alphas = self.alpha_schedule(t) + # sigmas = (1 - alphas ** 2).sqrt() + # logSNRs = (alphas / sigmas).log() + # h = torch.zeros_like(logSNRs) + # h[1:] = logSNRs[1:] - logSNRs[:-1] + # + # if prior is None: + # xt = torch.randn((n_samples, *self.default_x_shape), device=self.device) * temperature + # else: + # xt = torch.randn_like(prior, device=self.device) * temperature + # xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + # + # if preserve_history: + # x_history = np.empty((n_samples, sample_steps + 1, *xt.shape)) + # x_history[:, 0] = xt.cpu().numpy() + # + # with torch.no_grad(): + # condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + # condition_vec_cg = condition_cg + # + # for i in range(1, sample_steps + 1): + # + # eps_theta, log = self.predict_eps_theta( + # xt, t[i - 1].repeat(n_samples), alphas[i - 1], sigmas[i - 1], + # use_ema, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg) + # + # xt = ( + # alphas[i] / alphas[i - 1] * xt - + # 2. * sigmas[i] * torch.expm1(h[i]) * eps_theta + + # sigmas[i] * torch.expm1(2. * h[i]).sqrt() * torch.randn_like(xt)) + # xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + # + # if preserve_history: + # x_history[:, i] = xt.cpu().numpy() + # + # log["sample_history"] = x_history + # if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + # with torch.no_grad(): + # logp = self.classifier.logp(xt, t[-1].repeat(n_samples), condition_vec_cg) + # log["log_p"] = logp + # + # return xt, log + # + # def sample_with_sde_dpmpp_1( + # self, + # # ---------- the known fixed portion ---------- # + # prior: Optional[torch.Tensor] = None, + # # ----------------- sampling ----------------- # + # n_samples: int = 1, + # sample_steps: int = None, + # use_ema: bool = True, + # kappa: float = 1.0, + # temperature: float = 1.0, + # # ------------------ guidance ------------------ # + # condition_cfg=None, + # mask_cfg=None, + # w_cfg: float = 0.0, + # condition_cg=None, + # w_cg: float = 0.0, + # + # preserve_history: bool = False, + # ): + # log, x_history = {}, None + # model = self.model_ema if use_ema else self.model + # + # idx = torch.arange(sample_steps + 1) + # t = (((sample_steps - idx) / sample_steps * self.t_range[1] ** (1 / kappa) + + # idx / sample_steps * self.t_range[0] ** (1 / kappa)) ** kappa).to(self.device) + # + # alphas = self.alpha_schedule(t) + # sigmas = (1 - alphas ** 2).sqrt() + # logSNRs = (alphas / sigmas).log() + # h = torch.zeros_like(logSNRs) + # h[1:] = logSNRs[1:] - logSNRs[:-1] + # + # if prior is None: + # xt = torch.randn((n_samples, *self.default_x_shape), device=self.device) * temperature + # else: + # xt = torch.randn_like(prior, device=self.device) * temperature + # xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + # + # if preserve_history: + # x_history = np.empty((n_samples, sample_steps + 1, *xt.shape)) + # x_history[:, 0] = xt.cpu().numpy() + # + # with torch.no_grad(): + # condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + # condition_vec_cg = condition_cg + # + # for i in range(1, sample_steps + 1): + # + # x_theta, log = self.predict_x_theta( + # xt, t[i - 1].repeat(n_samples), alphas[i - 1], sigmas[i - 1], + # use_ema, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg) + # + # xt = ( + # sigmas[i] / sigmas[i - 1] * (-h[i]).exp() * xt + + # alphas[i] * (1. - (-2. * h[i]).exp()) * x_theta + + # sigmas[i] * (1. - (-2. * h[i]).exp()).sqrt() * torch.randn_like(xt)) + # xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + # + # if preserve_history: + # x_history[:, i] = xt.cpu().numpy() + # + # log["sample_history"] = x_history + # if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + # with torch.no_grad(): + # logp = self.classifier.logp(xt, t[-1].repeat(n_samples), condition_vec_cg) + # log["log_p"] = logp + # + # return xt, log + + # def sample_with_ode_dpm_2( + # self, + # # ---------- the known fixed portion ---------- # + # prior: Optional[torch.Tensor] = None, + # # ----------------- sampling ----------------- # + # n_samples: int = 1, + # sample_steps: int = None, + # use_ema: bool = True, + # kappa: float = 1.0, + # temperature: float = 1.0, + # # ------------------ guidance ------------------ # + # condition_cfg=None, + # mask_cfg=None, + # w_cfg: float = 0.0, + # condition_cg=None, + # w_cg: float = 0.0, + # + # preserve_history: bool = False, + # ): + # raise NotImplementedError + + # def sample_with_ode_dpmpp_1( + # self, + # # ---------- the known fixed portion ---------- # + # prior: Optional[torch.Tensor] = None, + # # ----------------- sampling ----------------- # + # n_samples: int = 1, + # sample_steps: int = 5, + # use_ema: bool = True, + # kappa: float = 1.0, + # temperature: float = 1.0, + # # ------------------ guidance ------------------ # + # condition_cfg=None, + # mask_cfg=None, + # w_cfg: float = 0.0, + # condition_cg=None, + # w_cg: float = 0.0, + # + # preserve_history: bool = False, + # ): + # raise NotImplementedError diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/dpmsolver_discrete.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/dpmsolver_discrete.py new file mode 100644 index 0000000000000000000000000000000000000000..45baecaddae8a23540f6a0cf50c56d1b54f5368d --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/dpmsolver_discrete.py @@ -0,0 +1,598 @@ +from typing import Optional +from typing import Union + +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.classifier import BaseClassifier +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from cleandiffuser.utils import at_least_ndim +from .basic import DiffusionModel + + +def epstheta_to_xtheta(x, alpha, sigma, eps_theta): + """ + x_theta = (x - sigma * eps_theta) / alpha + """ + return (x - sigma * eps_theta) / alpha + + +def xtheta_to_epstheta(x, alpha, sigma, x_theta): + """ + eps_theta = (x - alpha * x_theta) / sigma + """ + return (x - alpha * x_theta) / sigma + + +class DPMSolverDiscrete(DiffusionModel): + + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugs ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Params ---------------- # + grad_clip_norm: Optional[float] = None, + diffusion_steps: int = 1000, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- DPM Params ------------------- # + predict_noise: bool = False, + noise_schedule: str = "linear", # or cosine + t_eps: float = 1e-3, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + diffusion_steps, ema_rate, optim_params, device) + + self.predict_noise = predict_noise + self.noise_schedule = noise_schedule + self.t_eps = t_eps + + if self.noise_schedule == "linear": + beta0, beta1 = 0.1, 20 + self.t_cont = torch.linspace(t_eps, 1., diffusion_steps, device=device) + self.alphas = (-(beta1 - beta0) / 4 * (self.t_cont ** 2) - beta0 / 2 * self.t_cont).exp() + elif self.noise_schedule == "cosine": + s = 0.008 + self.t_cont = torch.linspace(t_eps, 0.9946, diffusion_steps, device=device) + self.alphas = ((torch.cos(np.pi / 2 * (self.t_cont + s) / (1 + s))).log() + - np.log(np.cos(np.pi / 2 * s / (1 + s)))).exp() + else: + raise ValueError(f"noise_schedule should be 'linear' or 'cosine', but got {self.noise_schedule}.") + + self.sigmas = (1 - self.alphas ** 2).sqrt() + self.logSNRs = (self.alphas / self.sigmas).log() + + # --------------------------------------------------------------------------- + # Training + + def add_noise(self, x0, idx=None, eps=None): + if idx is None: + idx = torch.randint(self.diffusion_steps, (x0.shape[0],), device=self.device) + eps = torch.randn_like(x0) if eps is None else eps + alpha = at_least_ndim(self.alphas[idx], x0.dim()) + sigma = at_least_ndim(self.sigmas[idx], x0.dim()) + xt = x0 * alpha + sigma * eps + xt = xt * (1. - self.fix_mask) + x0 * self.fix_mask + return xt, idx, eps + + def loss(self, x0, condition=None): + xt, idx, eps = self.add_noise(x0) + condition = self.model["condition"](condition) if condition is not None else None + + if self.predict_noise: + loss = (self.model["diffusion"](xt, idx, condition) - eps) ** 2 + else: + loss = (self.model["diffusion"](xt, idx, condition) - x0) ** 2 + + return (loss * self.loss_weight * (1. - self.fix_mask)).mean() + + def update(self, x0, condition=None, **kwargs): + self.optimizer.zero_grad() + loss = self.loss(x0, condition) + loss.backward() + grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_norm) \ + if self.grad_clip_norm else None + self.optimizer.step() + self.ema_update() + log = {"loss": loss.item(), "grad_norm": grad_norm} + return log + + def update_classifier(self, x0, condition): + xt, idx, eps = self.add_noise(x0) + log = self.classifier.update(xt, idx, condition) + return log + + # --------------------------------------------------------------------------- + # Inference + + def predict_eps_theta( + self, x, idx, alpha, sigma, use_ema=False, + # ----------------- CFG ----------------- # + condition_vec_cfg=None, + w_cfg: float = 0.0, + # ----------------- CG ----------------- # + condition_vec_cg=None, + w_cg: float = 1.0, + ): + b = x.shape[0] + model = self.model_ema if use_ema else self.model + + # ----------------- CFG ----------------- # + with torch.no_grad(): + if w_cfg != 0.0 and w_cfg != 1.0: + repeat_dim = [2 if i == 0 else 1 for i in range(x.dim())] + condition_vec_cfg = torch.cat([condition_vec_cfg, torch.zeros_like(condition_vec_cfg)], 0) + pred = model["diffusion"]( + x.repeat(*repeat_dim), idx.repeat(2), condition_vec_cfg) + pred = w_cfg * pred[:b] + (1. - w_cfg) * pred[b:] + else: + pred = model["diffusion"](x, idx, condition_vec_cfg) + + if self.predict_noise: + eps_theta = pred + else: + eps_theta = xtheta_to_epstheta(x, alpha, sigma, pred) + + # ----------------- CG ----------------- # + if self.classifier is not None and w_cg != 0.0 and condition_vec_cg is not None: + log_p, grad = self.classifier.gradients(x.clone(), idx, condition_vec_cg) + eps_theta = eps_theta - w_cg * sigma * grad + else: + log_p = None + + eps_theta = eps_theta * (1 - self.fix_mask) + + return eps_theta, {"log_p": log_p} + + def predict_x_theta( + self, x, idx, alpha, sigma, use_ema=False, + # ----------------- CFG ----------------- # + condition_vec_cfg=None, + w_cfg: float = 0.0, + # ----------------- CG ----------------- # + condition_vec_cg=None, + w_cg: float = 1.0, + ): + b = x.shape[0] + model = self.model_ema if use_ema else self.model + + # ----------------- CFG ----------------- # + with torch.no_grad(): + if w_cfg != 0.0 and w_cfg != 1.0: + repeat_dim = [2 if i == 0 else 1 for i in range(x.dim())] + condition_vec_cfg = torch.cat([condition_vec_cfg, torch.zeros_like(condition_vec_cfg)], 0) + pred = model["diffusion"]( + x.repeat(*repeat_dim), idx.repeat(2), condition_vec_cfg) + pred = w_cfg * pred[:b] + (1. - w_cfg) * pred[b:] + else: + pred = model["diffusion"](x, idx, condition_vec_cfg) + + if self.predict_noise: + x_theta = epstheta_to_xtheta(x, alpha, sigma, pred) + else: + x_theta = pred + + # ----------------- CG ----------------- # + if self.classifier is not None and w_cg != 0.0 and condition_vec_cg is not None: + log_p, grad = self.classifier.gradients(x.clone(), idx, condition_vec_cg) + x_theta = x_theta + w_cg * ((sigma ** 2) / alpha) * grad + else: + log_p = None + + x_theta = x_theta * (1 - self.fix_mask) + x * self.fix_mask + + return x_theta, {"log_p": log_p} + + def sample_with_ode_dpm_1( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = None, + use_ema: bool = True, + kappa: float = 1.0, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + + preserve_history: bool = False, + ): + log, x_history = {}, None + model = self.model_ema if use_ema else self.model + + idx = torch.arange(sample_steps + 1) + t_cond = (((sample_steps - idx) / sample_steps * 1. ** (1 / kappa) + + idx / sample_steps * self.t_eps ** (1 / kappa)) ** kappa).to(self.device) + t_sample = (torch.max(t_cond-1/self.diffusion_steps, torch.zeros_like(t_cond))*self.diffusion_steps).long() + + alphas = self.alphas[t_sample] + sigmas = self.sigmas[t_sample] + logSNRs = self.logSNRs[t_sample] + h = torch.zeros_like(logSNRs) + h[1:] = logSNRs[1:] - logSNRs[:-1] + + if prior is None: + xt = torch.randn((n_samples, *self.default_x_shape), device=self.device) * temperature + else: + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + if preserve_history: + x_history = np.empty((n_samples, sample_steps + 1, *xt.shape)) + x_history[:, 0] = xt.cpu().numpy() + + with torch.no_grad(): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + for i in range(1, sample_steps + 1): + + eps_theta, log = self.predict_eps_theta( + xt, t_sample[i - 1].repeat(n_samples), alphas[i - 1], sigmas[i - 1], + use_ema, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg) + + xt = alphas[i] / alphas[i - 1] * xt - sigmas[i] * torch.expm1(h[i]) * eps_theta + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + if preserve_history: + x_history[:, i] = xt.cpu().numpy() + + log["sample_history"] = x_history + if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + with torch.no_grad(): + logp = self.classifier.logp(xt, t_sample[-1].repeat(n_samples), condition_vec_cg) + log["log_p"] = logp + + return xt, log + + def sample_with_sde_dpm_1( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = None, + use_ema: bool = True, + kappa: float = 1.0, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + + preserve_history: bool = False, + ): + log, x_history = {}, None + model = self.model_ema if use_ema else self.model + + idx = torch.arange(sample_steps + 1) + t_cond = (((sample_steps - idx) / sample_steps * 1. ** (1 / kappa) + + idx / sample_steps * self.t_eps ** (1 / kappa)) ** kappa).to(self.device) + t_sample = (torch.max(t_cond - 1 / self.diffusion_steps, + torch.zeros_like(t_cond)) * self.diffusion_steps).long() + + alphas = self.alphas[t_sample] + sigmas = self.sigmas[t_sample] + logSNRs = self.logSNRs[t_sample] + h = torch.zeros_like(logSNRs) + h[1:] = logSNRs[1:] - logSNRs[:-1] + + if prior is None: + xt = torch.randn((n_samples, *self.default_x_shape), device=self.device) * temperature + else: + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + if preserve_history: + x_history = np.empty((n_samples, sample_steps + 1, *xt.shape)) + x_history[:, 0] = xt.cpu().numpy() + + with torch.no_grad(): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + for i in range(1, sample_steps + 1): + + eps_theta, log = self.predict_eps_theta( + xt, t_sample[i - 1].repeat(n_samples), alphas[i - 1], sigmas[i - 1], + use_ema, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg) + + xt = ( + alphas[i] / alphas[i - 1] * xt - + 2. * sigmas[i] * torch.expm1(h[i]) * eps_theta + + sigmas[i] * torch.expm1(2. * h[i]).sqrt() * torch.randn_like(xt)) + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + if preserve_history: + x_history[:, i] = xt.cpu().numpy() + + log["sample_history"] = x_history + if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + with torch.no_grad(): + logp = self.classifier.logp(xt, t_sample[-1].repeat(n_samples), condition_vec_cg) + log["log_p"] = logp + + return xt, log + + def sample_with_sde_dpmpp_1( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = None, + use_ema: bool = True, + kappa: float = 1.0, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + + preserve_history: bool = False, + ): + log, x_history = {}, None + model = self.model_ema if use_ema else self.model + + idx = torch.arange(sample_steps + 1) + t_cond = (((sample_steps - idx) / sample_steps * 1. ** (1 / kappa) + + idx / sample_steps * self.t_eps ** (1 / kappa)) ** kappa).to(self.device) + t_sample = (torch.max(t_cond - 1 / self.diffusion_steps, + torch.zeros_like(t_cond)) * self.diffusion_steps).long() + + alphas = self.alphas[t_sample] + sigmas = self.sigmas[t_sample] + logSNRs = self.logSNRs[t_sample] + h = torch.zeros_like(logSNRs) + h[1:] = logSNRs[1:] - logSNRs[:-1] + + if prior is None: + xt = torch.randn((n_samples, *self.default_x_shape), device=self.device) * temperature + else: + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + if preserve_history: + x_history = np.empty((n_samples, sample_steps + 1, *xt.shape)) + x_history[:, 0] = xt.cpu().numpy() + + with torch.no_grad(): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + for i in range(1, sample_steps + 1): + + x_theta, log = self.predict_x_theta( + xt, t_sample[i - 1].repeat(n_samples), alphas[i - 1], sigmas[i - 1], + use_ema, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg) + + xt = ( + sigmas[i] / sigmas[i - 1] * (-h[i]).exp() * xt + + alphas[i] * (1. - (-2. * h[i]).exp()) * x_theta + + sigmas[i] * (1. - (-2. * h[i]).exp()).sqrt() * torch.randn_like(xt)) + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + if preserve_history: + x_history[:, i] = xt.cpu().numpy() + + log["sample_history"] = x_history + if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + with torch.no_grad(): + logp = self.classifier.logp(xt, t_sample[-1].repeat(n_samples), condition_vec_cg) + log["log_p"] = logp + + return xt, log + + def sample_with_ode_dpm_2( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = None, + use_ema: bool = True, + kappa: float = 1.0, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + + preserve_history: bool = False, + ): + raise NotImplementedError + + def sample_with_ode_dpmpp_1( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + use_ema: bool = True, + kappa: float = 1.0, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + + preserve_history: bool = False, + ): + raise NotImplementedError + + def sample_with_ode_dpmpp_2( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + use_ema: bool = True, + kappa: float = 1.0, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + + preserve_history: bool = False, + ): + log, x_history = {}, None + model = self.model_ema if use_ema else self.model + + idx = torch.arange(sample_steps + 1) + t_cond = (((sample_steps - idx) / sample_steps * 1. ** (1 / kappa) + + idx / sample_steps * self.t_eps ** (1 / kappa)) ** kappa).to(self.device) + t_sample = (torch.max(t_cond - 1 / self.diffusion_steps, + torch.zeros_like(t_cond)) * self.diffusion_steps).long() + + alphas = self.alphas[t_sample] + sigmas = self.sigmas[t_sample] + logSNRs = self.logSNRs[t_sample] + h = torch.zeros_like(logSNRs) + h[1:] = logSNRs[1:] - logSNRs[:-1] + buffer = [] + + if prior is None: + xt = torch.randn((n_samples, *self.default_x_shape), device=self.device) * temperature + else: + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + if preserve_history: + x_history = np.empty((n_samples, sample_steps + 1, *xt.shape)) + x_history[:, 0] = xt.cpu().numpy() + + with torch.no_grad(): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + x_theta_t0, log = self.predict_x_theta( + xt, t_sample[0].repeat(n_samples), alphas[0], sigmas[0], use_ema, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg) + x_theta_t0 = x_theta_t0 * (1. - self.fix_mask) + prior * self.fix_mask + buffer.append(x_theta_t0) + + xt = sigmas[1] / sigmas[0] * xt - alphas[1] * torch.expm1(-h[1]) * buffer[-1] + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + x_history[:, 1] = xt.cpu().numpy() + + x_theta_t1, log = self.predict_x_theta( + xt, t_sample[1].repeat(n_samples), alphas[1], sigmas[1], use_ema, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg) + x_theta_t1 = x_theta_t1 * (1. - self.fix_mask) + prior * self.fix_mask + buffer.append(x_theta_t1) + + for i in range(2, sample_steps + 1): + + r = h[i - 1] / h[i] + + D = (1 + 0.5 / r) * buffer[-1] - 0.5 / r * buffer[-2] + + xt = sigmas[i] / sigmas[i - 1] * xt - alphas[i] * torch.expm1(-h[i]) * D + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + if i < sample_steps + 1: + x_theta_ti, log = self.predict_x_theta( + xt, t_sample[i].repeat(n_samples), alphas[i], sigmas[i], + use_ema, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg) + x_theta_ti = x_theta_ti * (1. - self.fix_mask) + prior * self.fix_mask + buffer.append(x_theta_ti.clone()) + + if preserve_history: + x_history[:, i] = xt.cpu().numpy() + + log["sample_history"] = x_history + if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + with torch.no_grad(): + logp = self.classifier.logp(xt, t_sample[-1].repeat(n_samples), condition_vec_cg) + log["log_p"] = logp + + return xt, log + + def sample( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + use_ema: bool = True, + kappa: float = 1.0, + temperature: float = 1.0, + solver: str = "dpmsolver++", # or ddim + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + + preserve_history: bool = False, + ): + if solver == "ode_dpm_1" or solver == "ddim": + return self.sample_with_ode_dpm_1( + prior, n_samples, sample_steps, use_ema, kappa, temperature, condition_cfg, mask_cfg, w_cfg, + condition_cg, w_cg, + preserve_history) + elif solver == "sde_dpm_1": + return self.sample_with_sde_dpm_1( + prior, n_samples, sample_steps, use_ema, kappa, temperature, condition_cfg, mask_cfg, w_cfg, + condition_cg, w_cg, + preserve_history) + elif solver == "sde_dpmpp_1": + return self.sample_with_sde_dpmpp_1( + prior, n_samples, sample_steps, use_ema, kappa, temperature, condition_cfg, mask_cfg, w_cfg, + condition_cg, w_cg, + preserve_history) + elif solver == "ode_dpm_2": + return self.sample_with_ode_dpm_2( + prior, n_samples, sample_steps, use_ema, kappa, temperature, condition_cfg, mask_cfg, w_cfg, + condition_cg, w_cg, + preserve_history) + elif solver == "ode_dpmpp_1": + return self.sample_with_ode_dpmpp_1( + prior, n_samples, sample_steps, use_ema, kappa, temperature, condition_cfg, mask_cfg, w_cfg, + condition_cg, w_cg, + preserve_history) + elif solver == "ode_dpmpp_2": + return self.sample_with_ode_dpmpp_2( + prior, n_samples, sample_steps, use_ema, kappa, temperature, condition_cfg, mask_cfg, w_cfg, + condition_cg, w_cg, + preserve_history) + else: + raise ValueError(f"Solver: {solver} has not been implemented yet.") diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/edm.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/edm.py new file mode 100644 index 0000000000000000000000000000000000000000..51bb63b41160bbe48ab3f486ebeb839488808383 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/edm.py @@ -0,0 +1,426 @@ +from typing import Optional +from typing import Union + +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.utils import at_least_ndim +from cleandiffuser.classifier import BaseClassifier +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from .basic import DiffusionModel + + +class EDMArchetecture(DiffusionModel): + + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugs ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Params ---------------- # + grad_clip_norm: Optional[float] = None, + diffusion_steps: int = 1000, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + diffusion_steps, ema_rate, optim_params, device) + + self.dot_scale_s = None + self.dot_sigma_s = None + self.scale_s = None + self.t_s = None + self.sigma_s = None + self.x_weight_s, self.D_weight_s = None, None + self.sample_steps = None + + def set_sample_steps(self, N: int): + raise NotImplementedError + + def c_skip(self, sigma): + raise NotImplementedError + + def c_out(self, sigma): + raise NotImplementedError + + def c_in(self, sigma): + raise NotImplementedError + + def c_noise(self, sigma): + raise NotImplementedError + + def loss_weighting(self, sigma): + raise NotImplementedError + + def sample_noise_distribution(self, N): + raise NotImplementedError + + def sample_scale_distribution(self, N): + raise NotImplementedError + + def D(self, x, sigma, condition=None, use_ema=False): + """ Prepositioning in EDM """ + c_skip, c_out, c_in, c_noise = self.c_skip(sigma), self.c_out(sigma), self.c_in(sigma), self.c_noise(sigma) + F = self.model_ema["diffusion"] if use_ema else self.model["diffusion"] + c_noise = at_least_ndim(c_noise.squeeze(), 1) + return c_skip * x + c_out * F(c_in * x, c_noise, condition) + + # --------------------------------------------------------------------------- + # Training + + def loss(self, x0, condition=None): + sigma = self.sample_noise_distribution(x0.shape[0]) + sigma = at_least_ndim(sigma, x0.dim()) + eps = torch.randn_like(x0) * sigma * (1. - self.fix_mask) + condition = self.model["condition"](condition) if condition is not None else None + loss = (self.loss_weighting(sigma) * (self.D(x0 + eps, sigma, condition) - x0) ** 2) + return (loss * self.loss_weight).mean() + + def update(self, x0, condition=None, **kwargs): + self.optimizer.zero_grad() + loss = self.loss(x0, condition) + loss.backward() + grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_norm) \ + if self.grad_clip_norm else None + self.optimizer.step() + self.ema_update() + log = {"loss": loss.item(), "grad_norm": grad_norm} + return log + + def update_classifier(self, x0, condition): + sigma = self.sample_noise_distribution(x0.shape[0]) + sigma = at_least_ndim(sigma, x0.dim()) + noise = self.c_noise(sigma) + eps = torch.randn_like(x0) * sigma * (1. - self.fix_mask) + log = self.classifier.update(x0 + eps, at_least_ndim(noise.squeeze(), 1), condition) + return log + + # --------------------------------------------------------------------------- + # Inference + + def dot_x( + self, x, i, use_ema=False, + # ----------------- CFG ----------------- # + condition_vec_cfg=None, + w_cfg: float = 0.0, + # ----------------- CG ----------------- # + condition_vec_cg=None, + w_cg: float = 1.0, + ): + b = x.shape[0] + sigma = at_least_ndim(self.sigma_s[i].repeat(b), x.dim()) + noise = self.c_noise(sigma) + unscale = 1. / self.scale_s[i] * (1. - self.fix_mask) + self.fix_mask + # ----------------- CFG ----------------- # + with torch.no_grad(): + if w_cfg != 0.0 and w_cfg != 1.0: + repeat_dim = [2 if i == 0 else 1 for i in range(x.dim())] + condition_vec_cfg = torch.cat([condition_vec_cfg, torch.zeros_like(condition_vec_cfg)], 0) + D = self.D( + (x * unscale).repeat(*repeat_dim), + sigma.repeat(*repeat_dim), + condition_vec_cfg, use_ema) + D = w_cfg * D[:b] + (1. - w_cfg) * D[b:] + elif w_cfg == 0.0: + D = self.D(x * unscale, sigma, None, use_ema) + else: + D = self.D(x * unscale, sigma, condition_vec_cfg, use_ema) + # ----------------- CG ----------------- # + if self.classifier is not None and w_cg != 0.0 and condition_vec_cg is not None: + log_p, grad = self.classifier.gradients(x * unscale, + at_least_ndim(noise.squeeze(), 1), condition_vec_cg) + D = D + w_cg * self.scale_s[i] * (sigma ** 2) * grad + else: + log_p = None + + # do not change the fixed portion + dot_x = self.x_weight_s[i] * x - self.D_weight_s[i] * D + dot_x = dot_x * (1. - self.fix_mask) + + return dot_x, {"log_p": log_p} + + def sample( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + use_ema: bool = True, + solver: str = "euler", + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + + preserve_history: bool = False, + **kwargs + ): + """ + Sample from the diffusion model. + --- + Input: + - prior: Optional[torch.Tensor] = None + The known fixed portion of the input data. Should be in the shape of `(n_samples, *x_shape)`. + Leave the unknown part as `0`. If `None`, which means `fix_mask` is `None`, the model takes no prior. + + - sample_steps: int = 5 + Number of sampling steps. + + - use_ema: bool = True + Whether to use the EMA model. If `False`, you should `eval` the model. + + - solver: str = "euler" + The solver to use. Can be either "euler" or "heun". + + - condition_cfg: Optional[torch.Tensor] = None + The condition for the CFG. Should be in the shape of `(n_samples, *shape_of_nn_condition_input)`. + If `None`, the model takes no condition. + + - mask_cfg: Optional[torch.Tensor] = None + The mask for the CFG. Should be in the shape of `(n_samples, *shape_of_nn_condition_mask)`. + Model will ignore the `mask_cfg==0` parts in `condition_cfg`. If `None`, the model takes no mask. + + - w_cfg: float = 0.0 + The weight for the CFG. If `0.0`, the model takes no CFG. + + - condition_cg: Optional[torch.Tensor] = None + The condition for the CG. Should be in the shape of `(n_samples, 1)`. + If `None`, the model takes no condition. + + - w_cg: float = 0.0 + The weight for the CG. If `0.0`, the model takes no CG. + + - preserve_history: bool = False + Whether to preserve the history of the sampling process. If `True`, the model will return the history. + + Output: + - xt: torch.Tensor + The sampled data. Should be in the shape of `(n_samples, *x_shape)`. + + - log: dict + The log of the sampling process. Contains the following keys: + - "sample_history": np.ndarray + The history of the sampling process. Should be in the shape of `(n_samples, N + 1, *x_shape)`. + If `preserve_history` is `False`, this key will not exist. + - "log_p": torch.Tensor + The log probability of the sampled data estimated by CG. + Should be in the shape of `(n_samples,)`. + """ + if sample_steps != self.sample_steps: + self.set_sample_steps(sample_steps) + + N = self.sample_steps + model = self.model_ema if use_ema else self.model + log, x_history = {}, None + + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + + if prior is None: + xt = torch.randn((n_samples, *self.default_x_shape), + device=self.device) * self.sigma_s[0] * self.scale_s[0] + else: + xt = torch.randn_like(prior, device=self.device) * self.sigma_s[0] * self.scale_s[0] + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + if preserve_history: + x_history = np.empty((n_samples, N + 1, *xt.shape)) + x_history[:, 0] = xt.cpu().numpy() + + for i in range(N): + # tmp_w_cg = w_cg if i < 1 else 0.0 + delta_x, log = self.dot_x(xt, i, use_ema, condition_vec_cfg, w_cfg, condition_cg, w_cg) + delta_t = self.t_s[i] - self.t_s[i + 1] + x_tp1 = xt - delta_x * delta_t + if prior is not None: + x_tp1 = x_tp1 * (1. - self.fix_mask) + prior * self.fix_mask + if solver == "heun": + if i != N - 1 and self.sigma_s[i + 1] > 0.005: + delta_x_2, log = self.dot_x(x_tp1, i + 1, use_ema, condition_vec_cfg, w_cfg, condition_cg, w_cg) + x_tp1 = xt - (delta_x + delta_x_2) / 2. * delta_t + if prior is not None: + x_tp1 = x_tp1 * (1. - self.fix_mask) + prior * self.fix_mask + + xt = x_tp1 + + if preserve_history: + x_history[:, i + 1] = xt.cpu().numpy() + + log["sample_history"] = x_history + if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + with torch.no_grad(): + logp = self.classifier.logp( + xt, at_least_ndim(self.c_noise(self.sigma_s[-1]).squeeze(), 1), condition_cg) + log["log_p"] = logp + + return xt, log + + def sample_x( + self, + # ---------- the known fixed portion ---------- # + prior: Optional[torch.Tensor] = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + extra_sample_steps: int = 8, + use_ema: bool = True, + solver: str = "euler", + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + + preserve_history: bool = False, + ): + if sample_steps != self.sample_steps: + self.set_sample_steps(sample_steps) + + N = self.sample_steps + model = self.model_ema if use_ema else self.model + log, x_history = {}, None + + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + + if prior is None: + xt = torch.randn((n_samples, *self.default_x_shape), + device=self.device) * self.sigma_s[0] * self.scale_s[0] + else: + xt = torch.randn_like(prior, device=self.device) * self.sigma_s[0] * self.scale_s[0] + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + + if preserve_history: + x_history = np.empty((n_samples, N + 1, *xt.shape)) + x_history[:, 0] = xt.cpu().numpy() + + for i in range(N): + # tmp_w_cg = w_cg if i < 1 else 0.0 + delta_x, log = self.dot_x(xt, i, use_ema, condition_vec_cfg, w_cfg, condition_cg, w_cg) + delta_t = self.t_s[i] - self.t_s[i + 1] + x_tp1 = xt - delta_x * delta_t + if prior is not None: + x_tp1 = x_tp1 * (1. - self.fix_mask) + prior * self.fix_mask + if solver == "heun": + if i != N - 1 and self.sigma_s[i + 1] > 0.005: + delta_x_2, log = self.dot_x(x_tp1, i + 1, use_ema, condition_vec_cfg, w_cfg, condition_cg, w_cg) + x_tp1 = xt - (delta_x + delta_x_2) / 2. * delta_t + if prior is not None: + x_tp1 = x_tp1 * (1. - self.fix_mask) + prior * self.fix_mask + + xt = x_tp1 + + if preserve_history: + x_history[:, i + 1] = xt.cpu().numpy() + + if extra_sample_steps > 0: + + delta_t = self.t_s[N - 1] - self.t_s[N] + + for _ in range(extra_sample_steps): + + delta_x, log = self.dot_x(xt, N - 1, use_ema, condition_vec_cfg, w_cfg, condition_cg, w_cg) + x_tp1 = xt - delta_x * delta_t + if prior is not None: + x_tp1 = x_tp1 * (1. - self.fix_mask) + prior * self.fix_mask + + xt = x_tp1 + + log["sample_history"] = x_history + if log["log_p"] is None and self.classifier is not None and condition_cg is not None: + with torch.no_grad(): + logp = self.classifier.logp( + xt, at_least_ndim(self.c_noise(self.sigma_s[-1]).squeeze(), 1), condition_cg) + log["log_p"] = logp + + return xt, log + + +class EDM(EDMArchetecture): + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugs ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Params ---------------- # + grad_clip_norm: Optional[float] = None, + diffusion_steps: int = 1000, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- EDM Params ------------------- # + sigma_data: float = 0.5, + sigma_min: float = 0.002, + sigma_max: float = 80., + rho: float = 7., + P_mean: float = -1.2, + P_std: float = 1.2, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + diffusion_steps, ema_rate, optim_params, device) + + self.sigma_data = sigma_data + self.sigma_min, self.sigma_max, self.rho = sigma_min, sigma_max, rho + self.P_mean, self.P_std = P_mean, P_std + + def set_sample_steps(self, N: int): + self.sample_steps = N + self.sigma_s = (self.sigma_max ** (1 / self.rho) + torch.arange(N + 1, device=self.device) / N * + (self.sigma_min ** (1 / self.rho) - self.sigma_max ** (1 / self.rho))) ** self.rho + self.t_s = self.sigma_s + self.scale_s = torch.ones_like(self.sigma_s) + self.dot_sigma_s = torch.ones_like(self.sigma_s) + self.dot_scale_s = torch.zeros_like(self.sigma_s) + self.x_weight_s = (self.dot_sigma_s / self.sigma_s + self.dot_scale_s / self.scale_s) + self.D_weight_s = self.dot_sigma_s / self.sigma_s * self.scale_s + + def c_skip(self, sigma): return self.sigma_data ** 2 / (self.sigma_data ** 2 + sigma ** 2) + + def c_out(self, sigma): return sigma * self.sigma_data / (self.sigma_data ** 2 + sigma ** 2).sqrt() + + def c_in(self, sigma): return 1 / (self.sigma_data ** 2 + sigma ** 2).sqrt() + + def c_noise(self, sigma): return 0.25 * sigma.log() + + def loss_weighting(self, sigma): return (self.sigma_data ** 2 + sigma ** 2) / ((sigma * self.sigma_data) ** 2) + + def sample_noise_distribution(self, N): + log_sigma = torch.randn(N, device=self.device) * self.P_std + self.P_mean + return log_sigma.exp() + + def sample_scale_distribution(self, N): + return torch.ones(N, device=self.device) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/edmddim.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/edmddim.py new file mode 100644 index 0000000000000000000000000000000000000000..d0facfbcf8ccb9996f1507d113e4d9608fa2d706 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/edmddim.py @@ -0,0 +1,84 @@ +from typing import Union, Optional + +import numpy as np +import torch + +from cleandiffuser.classifier import BaseClassifier +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from .edm import EDMArchetecture + + +class EDMDDIM(EDMArchetecture): + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugs ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Params ---------------- # + grad_clip_norm: Optional[float] = None, + diffusion_steps: int = 1000, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- EDM Params ------------------- # + C1: float = 0.001, + C2: float = 0.008, + j0: float = 8, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + diffusion_steps, ema_rate, optim_params, device) + + self.C1, self.C2, self.j0 = C1, C2, j0 + self.u = None + + def set_sample_steps(self, N: int): + self.sample_steps = N + bar_alpha = torch.sin( + torch.arange(self.diffusion_steps + 1, device=self.device) / + (self.diffusion_steps * (self.C2 + 1)) * np.pi / 2) ** 2 + tmp = torch.max(bar_alpha[:-1] / bar_alpha[1:], torch.tensor(self.C1, device=self.device)) + self.u = torch.empty_like(bar_alpha[:-1]) + for i in range(self.diffusion_steps): + if i == 0: + self.u[-1 - i] = (1 / tmp[-1 - i] - 1).sqrt() + else: + self.u[-1 - i] = ((self.u[-i] ** 2 + 1) / tmp[-1 - i] - 1).sqrt() + idx = torch.arange(N, device=self.device) + self.t_s = self.u[torch.floor(self.j0 + (self.diffusion_steps - 1 - self.j0) / + (N - 1) * idx + 0.5).long()] + self.sigma_s = self.t_s + self.scale_s = torch.ones_like(self.sigma_s) * 1.0 + self.dot_sigma_s = torch.ones_like(self.sigma_s) * 1.0 + self.dot_scale_s = torch.zeros_like(self.sigma_s) + self.x_weight_s = (self.dot_sigma_s / self.sigma_s + self.dot_scale_s / self.scale_s) + self.D_weight_s = self.dot_sigma_s / self.sigma_s * self.scale_s + + def c_skip(self, sigma): return torch.ones_like(sigma) + + def c_out(self, sigma): return -sigma + + def c_in(self, sigma): return 1 / (1 + sigma ** 2).sqrt() + + def c_noise(self, sigma): return sigma + + def loss_weighting(self, sigma): return 1 / (sigma ** 2) + + def sample_noise_distribution(self, N): + j = torch.randint(0, self.diffusion_steps, (N, 1), device=self.device) + return self.u[j] diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/koopman_net.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/koopman_net.py new file mode 100644 index 0000000000000000000000000000000000000000..1650983b2cc30fb1f513564e9e3dfca86e38c0b0 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/koopman_net.py @@ -0,0 +1,427 @@ +import math +from typing import Any, Dict, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# ========================= +# Utilities +# ========================= + + +def load_koopman_checkpoint(path: str, map_location: Any = "cpu") -> Tuple[Dict[str, Any], Dict[str, Any]]: + """ + Loads a Koopman student checkpoint. Supports legacy flat state_dict or + wrapped dicts: {"state_dict": ..., "gamma_mode": "z"|"z0"|"obs0", ...}. + """ + obj = torch.load(path, map_location=map_location) + if isinstance(obj, dict) and "state_dict" in obj: + meta = {k: v for k, v in obj.items() if k != "state_dict"} + return obj["state_dict"], meta + return obj, {} + + +def orthogonal_init_(tensor, scale=1.0): + with torch.no_grad(): + nn.init.orthogonal_(tensor) + tensor.mul_(scale) + +def build_encoder(input_dim: int, latent_dim: int) -> nn.Module: + return nn.Sequential( + nn.Conv1d(input_dim, 128, kernel_size=5, padding=2), + nn.ReLU(), + nn.Conv1d(128, 256, kernel_size=5, padding=2), + nn.ReLU(), + nn.AdaptiveAvgPool1d(1), + nn.Flatten(), + nn.Linear(256, latent_dim), + ) + +def build_decoder(latent_dim: int, horizon: int, input_dim: int) -> nn.Module: + return nn.Sequential( + nn.Linear(latent_dim, 1024), + nn.ReLU(), + nn.Linear(1024, horizon * input_dim), + ) + +class BaseDistiller(nn.Module): + def get_spectral_penalty(self) -> torch.Tensor: + p = next(self.parameters()) + return torch.tensor(0.0, device=p.device) + + def get_pinv_loss(self) -> torch.Tensor: + p = next(self.parameters()) + return torch.tensor(0.0, device=p.device) + +# ========================= +# KDM: static linear +# ========================= +class KDMKoopmanDistiller(BaseDistiller): + def __init__(self, input_dim, horizon, latent_dim=512, num_iters=1, + detach_clean=False, shared_encoder=False, noisy_latent=0.2): + super().__init__() + self.horizon = horizon + self.input_dim = input_dim + self.latent_dim = latent_dim + self.num_iters = num_iters + self.detach_clean = detach_clean + self.shared_encoder = shared_encoder + self.noisy_latent = noisy_latent + + if shared_encoder: + self.encoder = build_encoder(input_dim, latent_dim) + else: + self.encoder_clean = build_encoder(input_dim, latent_dim) + self.encoder_noise = build_encoder(input_dim, latent_dim) + + self.koopman = nn.Linear(latent_dim, latent_dim, bias=False) + orthogonal_init_(self.koopman.weight, scale=0.9) + + self.decoder = build_decoder(latent_dim, horizon, input_dim) + + def encode_clean(self, x): + return self.encoder(x.permute(0, 2, 1)) if self.shared_encoder else self.encoder_clean(x.permute(0, 2, 1)) + + def encode_noise(self, x): + return self.encoder(x.permute(0, 2, 1)) if self.shared_encoder else self.encoder_noise(x.permute(0, 2, 1)) + + def decode(self, z): + return self.decoder(z).view(-1, self.horizon, self.input_dim) + + def forward(self, x_noise, x_clean=None): + z_noise = self.encode_noise(x_noise) + if self.training and self.noisy_latent > 0: + z_noise = z_noise + torch.randn_like(z_noise) * self.noisy_latent + + z_pred = z_noise + for _ in range(self.num_iters): + z_pred = self.koopman(z_pred) + + if self.training: + assert x_clean is not None + z_clean = self.encode_clean(x_clean) + z_clean_noisy = z_clean + torch.randn_like(z_clean) * self.noisy_latent if self.noisy_latent > 0 else z_clean + z_clean_target = z_clean.detach() if self.detach_clean else z_clean + + x_recon = self.decode(z_clean_noisy) + x_pred = self.decode(z_pred) + return z_clean_target, z_pred, x_recon, x_pred + else: + return self.decode(z_pred) + + +# ========================= +# KDM-F: static factorized (Aligned with NeurIPS 2025 FastKoopmanMatrix) +# ========================= +class StaticFactorizedKoopman(nn.Module): + """ + KDM-F 风格：严格对齐 NeurIPS 2025 的 FastKoopmanMatrix + 1. 采用 QR 分解进行正交初始化 + 2. 采用双重指数 exp(-exp(nu)) 严格约束特征值模长在 (0, 1) 内 + """ + def __init__(self, latent_dim: int): + super().__init__() + self.d = latent_dim + scale = 1.0 / math.sqrt(2 * self.d) + + # 1. QR 分解正交初始化 P_re 和 P_im + P_re_init = self.random_orthonormal_matrix(self.d) * scale + P_im_init = self.random_orthonormal_matrix(self.d) * scale + self.P_re = nn.Parameter(P_re_init) + self.P_im = nn.Parameter(P_im_init) + + # P_inv 原版使用标准高斯初始化，依赖 P_inv Loss 在训练中对齐 + self.P_inv_re = nn.Parameter(torch.randn(self.d, self.d) * scale) + self.P_inv_im = nn.Parameter(torch.randn(self.d, self.d) * scale) + + # 2. 严格收缩的特征值参数化 (防止动力学发散) + rmin, rmax = 0.0, 1.0 + max_phase = 6.283 # 2 * pi + + u1 = torch.rand(self.d) + u2 = torch.rand(self.d) + + # nu_log 映射，使得 lambda_mod = exp(-exp(nu_log)) + init_nu = torch.log(-0.5 * torch.log(u1 * (rmax + rmin) * (rmax - rmin) + rmin**2)) + self.nu_log = nn.Parameter(init_nu) + + # theta_log 映射，使得 phase = exp(theta_log) + self.theta_log = nn.Parameter(torch.log(max_phase * u2)) + + def random_orthonormal_matrix(self, size): + """生成随机正交矩阵用于初始化""" + q, _ = torch.linalg.qr(torch.randn(size, size)) + return q + + def normalize_rows(self, x, eps=1e-8): + """行归一化，原版前向传播时使用""" + return x / (x.norm(dim=-1, keepdim=True) + eps) + + def to_real_block_matrix(self, re, im): + """将复数矩阵重组为实数块对角矩阵""" + top = torch.cat([re, -im], dim=1) + bottom = torch.cat([im, re], dim=1) + return torch.cat([top, bottom], dim=0) + + def forward(self, z): + # 归一化特征向量 + P_re_norm = self.normalize_rows(self.P_re) + P_im_norm = self.normalize_rows(self.P_im) + + P = self.to_real_block_matrix(P_re_norm, P_im_norm) + P_inv = self.to_real_block_matrix(self.P_inv_re, self.P_inv_im) + + # 还原严格有界的特征值 + lambda_mod = torch.exp(-torch.exp(self.nu_log)) + phase = torch.exp(self.theta_log) + lambda_re = lambda_mod * torch.cos(phase) + lambda_im = lambda_mod * torch.sin(phase) + + Lambda = self.to_real_block_matrix(torch.diag(lambda_re), torch.diag(lambda_im)) + + # 状态增广并演化 + z_aug = torch.cat([z, torch.zeros_like(z)], dim=-1) + z_pushed = z_aug @ P_inv.T @ Lambda.T @ P.T + + # 截断回原维度 + return z_pushed[..., :self.d] + + def get_spectral_penalty(self): + # 对应原版 EigenSpecKoopmanLossTypes.OnTheCircle + # 因为 lambda_mod 已经被双重指数限制在 (0, 1) 之间 + # 谱惩罚的作用是将它们推向单位圆边缘，防止过度收缩导致特征消失 + lambda_mod = torch.exp(-torch.exp(self.nu_log)) + return torch.abs(lambda_mod - 1.0).mean() + + def get_pinv_loss(self): + # 逆矩阵对齐损失 + P_re_norm = self.normalize_rows(self.P_re) + P_im_norm = self.normalize_rows(self.P_im) + P = self.to_real_block_matrix(P_re_norm, P_im_norm) + P_inv = self.to_real_block_matrix(self.P_inv_re, self.P_inv_im) + + I = torch.eye(2 * self.d, device=P.device) + return F.mse_loss(P_inv @ P, I) + F.mse_loss(P @ P_inv, I) + +class KDMFactorizedDistiller(BaseDistiller): + def __init__(self, input_dim, horizon, latent_dim=512, num_iters=1, + spectral_penalty_coef=0.01, detach_clean=False, + shared_encoder=False, noisy_latent=0.2): + super().__init__() + self.horizon = horizon + self.input_dim = input_dim + self.latent_dim = latent_dim + self.num_iters = num_iters + self.detach_clean = detach_clean + self.spectral_penalty_coef = spectral_penalty_coef + self.shared_encoder = shared_encoder + self.noisy_latent = noisy_latent + + if shared_encoder: + self.encoder = build_encoder(input_dim, latent_dim) + else: + self.encoder_clean = build_encoder(input_dim, latent_dim) + self.encoder_noise = build_encoder(input_dim, latent_dim) + + self.koopman = StaticFactorizedKoopman(latent_dim) + self.decoder = build_decoder(latent_dim, horizon, input_dim) + + def encode_clean(self, x): + return self.encoder(x.permute(0, 2, 1)) if self.shared_encoder else self.encoder_clean(x.permute(0, 2, 1)) + + def encode_noise(self, x): + return self.encoder(x.permute(0, 2, 1)) if self.shared_encoder else self.encoder_noise(x.permute(0, 2, 1)) + + def decode(self, z): + return self.decoder(z).view(-1, self.horizon, self.input_dim) + + def forward(self, x_noise, x_clean=None): + z_noise = self.encode_noise(x_noise) + if self.training and self.noisy_latent > 0: + z_noise = z_noise + torch.randn_like(z_noise) * self.noisy_latent + + z_pred = z_noise + for _ in range(self.num_iters): + z_pred = self.koopman(z_pred) + + if self.training: + assert x_clean is not None + z_clean = self.encode_clean(x_clean) + z_clean_noisy = z_clean + torch.randn_like(z_clean) * self.noisy_latent if self.noisy_latent > 0 else z_clean + z_clean_target = z_clean.detach() if self.detach_clean else z_clean + + x_recon = self.decode(z_clean_noisy) + x_pred = self.decode(z_pred) + return z_clean_target, z_pred, x_recon, x_pred + else: + return self.decode(z_pred) + + def get_spectral_penalty(self): + return self.koopman.get_spectral_penalty() * self.spectral_penalty_coef + + def get_pinv_loss(self): + return self.koopman.get_pinv_loss() + +# ========================= +# Regression baseline +# ========================= +class ResidualMLPEvolver(nn.Module): + def __init__(self, latent_dim: int, hidden_dim: int = 1024, depth: int = 2): + super().__init__() + layers = [] + d_in = latent_dim + for i in range(depth): + d_out = hidden_dim if i < depth - 1 else latent_dim + layers.append(nn.Linear(d_in, d_out)) + if i < depth - 1: + layers.append(nn.ReLU()) + d_in = d_out + self.net = nn.Sequential(*layers) + + def forward(self, z): + return z + self.net(z) + +class RegressionDistiller(BaseDistiller): + def __init__(self, input_dim, horizon, latent_dim=512, num_iters=1, detach_clean=False): + super().__init__() + self.horizon = horizon + self.input_dim = input_dim + self.latent_dim = latent_dim + self.num_iters = num_iters + self.detach_clean = detach_clean + + self.encoder = build_encoder(input_dim, latent_dim) + self.decoder = build_decoder(latent_dim, horizon, input_dim) + self.evolver = ResidualMLPEvolver(latent_dim, hidden_dim=1024, depth=2) + + def forward(self, x_noise, x_clean=None): + z_noise = self.encoder(x_noise.permute(0, 2, 1)) + + z_pred = z_noise + for _ in range(self.num_iters): + z_pred = self.evolver(z_pred) + + if self.training: + assert x_clean is not None + z_clean = self.encoder(x_clean.permute(0, 2, 1)) + z_clean_target = z_clean.detach() if self.detach_clean else z_clean + + x_recon = self.decoder(z_clean).view(-1, self.horizon, self.input_dim) + x_pred = self.decoder(z_pred).view(-1, self.horizon, self.input_dim) + return z_clean_target, z_pred, x_recon, x_pred + else: + return self.decoder(z_pred).view(-1, self.horizon, self.input_dim) + +# ========================= +# Ours: dynamic factorized Koopman +# ========================= +class FactorizedDynamicKoopman(nn.Module): + """ + Modal gains gamma: eigen_net input is either current latent z (default), + fixed initial latent z0 (gamma_mode='z0'), or first-timestep obs slice (gamma_mode='obs0'). + """ + + def __init__(self, latent_dim: int, gamma_mode: str = "z", obs_dim: Optional[int] = None): + super().__init__() + self.latent_dim = latent_dim + self.gamma_mode = gamma_mode + if gamma_mode not in ("z", "z0", "obs0"): + raise ValueError(f"gamma_mode must be 'z', 'z0', or 'obs0', got {gamma_mode!r}") + if gamma_mode == "obs0" and obs_dim is None: + raise ValueError("obs0 gamma_mode requires obs_dim") + + self.obs_dim = obs_dim + cond_dim = obs_dim if gamma_mode == "obs0" else latent_dim + + self.P = nn.Linear(latent_dim, latent_dim, bias=False) + self.P_inv = nn.Linear(latent_dim, latent_dim, bias=False) + + nn.init.orthogonal_(self.P.weight) + self.P_inv.weight.data = self.P.weight.data.T.clone() + + self.eigen_net = nn.Sequential( + nn.Linear(cond_dim, 128), + nn.ReLU(), + nn.Linear(128, latent_dim) + ) + + def forward(self, z: torch.Tensor, c_context: Optional[torch.Tensor] = None): + if self.gamma_mode == "z": + gamma_in = z + else: + if c_context is None: + raise ValueError(f"gamma_mode={self.gamma_mode!r} requires c_context (z0 or obs0).") + gamma_in = c_context + eigen_values = self.eigen_net(gamma_in) + z_projected = self.P_inv(z) + z_evolved = z_projected * eigen_values + z_next = self.P(z_evolved) + return z_next, eigen_values + + def get_pinv_loss(self): + I = torch.eye(self.latent_dim, device=self.P.weight.device) + return F.mse_loss(self.P.weight @ self.P_inv.weight, I) + +class KoopmanDistiller(BaseDistiller): + def __init__(self, input_dim, horizon, latent_dim=512, num_iters=1, + spectral_penalty_coef=0.01, detach_clean=True, + gamma_mode: str = "z", obs_dim: Optional[int] = None): + super().__init__() + self.horizon = horizon + self.input_dim = input_dim + self.latent_dim = latent_dim + self.num_iters = num_iters + self.spectral_penalty_coef = spectral_penalty_coef + self.detach_clean = detach_clean + self.gamma_mode = gamma_mode + if gamma_mode == "obs0" and obs_dim is None: + raise ValueError("KoopmanDistiller: obs0 gamma_mode requires obs_dim") + self.obs_dim = obs_dim + + self.encoder = build_encoder(input_dim, latent_dim) + self.decoder = build_decoder(latent_dim, horizon, input_dim) + self.koopman = FactorizedDynamicKoopman(latent_dim, gamma_mode=gamma_mode, obs_dim=obs_dim) + + def gamma_context(self, x_noise: torch.Tensor) -> Optional[torch.Tensor]: + """Conditioning vector for eigen_net when gamma_mode is z0 or obs0; else None.""" + if self.gamma_mode == "z": + return None + if self.gamma_mode == "z0": + return self.encoder(x_noise.permute(0, 2, 1)) + return x_noise[:, 0, : self.obs_dim] + + def forward(self, x_noise, x_clean=None): + z_noise = self.encoder(x_noise.permute(0, 2, 1)) + c_ctx = self.gamma_context(x_noise) + + all_eigen_values = [] + z_pred = z_noise + for _ in range(self.num_iters): + z_pred, eigen_vals = self.koopman(z_pred, c_ctx) + all_eigen_values.append(eigen_vals) + + self.current_eigen_values = torch.cat(all_eigen_values, dim=0) + + if self.training: + assert x_clean is not None + z_clean = self.encoder(x_clean.permute(0, 2, 1)) + z_clean_target = z_clean.detach() if self.detach_clean else z_clean + + x_recon = self.decoder(z_clean).view(-1, self.horizon, self.input_dim) + x_pred = self.decoder(z_pred).view(-1, self.horizon, self.input_dim) + return z_clean_target, z_pred, x_recon, x_pred + else: + return self.decoder(z_pred).view(-1, self.horizon, self.input_dim) + + def get_spectral_penalty(self): + if not hasattr(self, "current_eigen_values"): + p = next(self.parameters()) + return torch.tensor(0.0, device=p.device) + + radii = torch.abs(self.current_eigen_values) + penalty = F.relu(radii - 1.0).mean() + return penalty * self.spectral_penalty_coef + + def get_pinv_loss(self): + return self.koopman.get_pinv_loss() \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/newedm.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/newedm.py new file mode 100644 index 0000000000000000000000000000000000000000..e43b4f1ef9d3b23b39371dfbd470d9b17daa9ea5 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/newedm.py @@ -0,0 +1,438 @@ +from typing import Optional, Union + +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.classifier import BaseClassifier +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from cleandiffuser.utils import ( + at_least_ndim) +from .basic import DiffusionModel + + +class ContinuousEDM(DiffusionModel): + """Continuous-time EDM + + EDM posits that the concepts of `t` in the diffusion process and the noise schedule are equivalent. + Previous noise schedules can be interpreted as perturbing the data with Gaussian noise + followed by scaling. EDM sets the standard deviation of noise as `t`, the scale as 1, + and devises a series of preconditioning steps to aid in model learning. + + The current implementation of EDM only supports continuous-time ODEs. + The sampling steps are required to be greater than 1. + + Args: + - nn_diffusion: BaseNNDiffusion + The neural network backbone for the Diffusion model. + - nn_condition: Optional[BaseNNCondition] + The neural network backbone for the condition embedding. + + - fix_mask: Union[list, np.ndarray, torch.Tensor] + Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + The mask should be in the shape of `x_shape`. + - loss_weight: Union[list, np.ndarray, torch.Tensor] + Add loss weight. The weight should be in the shape of `x_shape`. + + - classifier: Optional[BaseClassifier] + Add a classifier to enable classifier-guidance. + + - grad_clip_norm: Optional[float] + Gradient clipping norm. + - ema_rate: float + Exponential moving average rate. + - optim_params: Optional[dict] + Optimizer parameters. + + - sigma_data: float + The standard deviation of the data. Default: 0.5. + - sigma_min: float + The minimum standard deviation of the noise. Default: 0.002. + - sigma_max: float + The maximum standard deviation of the noise. Default: 80. + - rho: float + The power of the noise schedule. Default: 7. + - P_mean: float + Hyperparameter for noise sampling during training. Default: -1.2. + - P_std: float + Hyperparameter for noise sampling during training. Default: 1.2. + + - x_max: Optional[torch.Tensor] + The maximum value for the input data. `None` indicates no constraint. + - x_min: Optional[torch.Tensor] + The minimum value for the input data. `None` indicates no constraint. + + - device: Union[torch.device, str] + The device to run the model. + """ + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugins ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Training Params ---------------- # + grad_clip_norm: Optional[float] = None, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- Diffusion Params ------------------- # + sigma_data: float = 0.5, + sigma_min: float = 0.002, + sigma_max: float = 80., + rho: float = 7., + P_mean: float = -1.2, + P_std: float = 1.2, + + x_max: Optional[torch.Tensor] = None, + x_min: Optional[torch.Tensor] = None, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + 0, ema_rate, optim_params, device) + + self.sigma_data, self.sigma_min, self.sigma_max = sigma_data, sigma_min, sigma_max + self.rho, self.P_mean, self.P_std = rho, P_mean, P_std + + self.x_max = x_max.to(device) if isinstance(x_max, torch.Tensor) else x_max + self.x_min = x_min.to(device) if isinstance(x_min, torch.Tensor) else x_min + + # ==================== Continuous Time-step Range ==================== + self.t_diffusion = [sigma_min, sigma_max] + + # ======================= Noise Schedule ========================= + # scale(t) = 1., dot_scale(t) = 0. + # sigma(t) = t , dot_sigma(t) = 1. + + @property + def supported_solvers(self): + return ["euler", "heun"] + + @property + def clip_pred(self): + return (self.x_max is not None) or (self.x_min is not None) + + # ===================== EDM Pre-conditioning ========================= + def c_skip(self, sigma): + return self.sigma_data ** 2 / (self.sigma_data ** 2 + sigma ** 2) + + def c_out(self, sigma): + return sigma * self.sigma_data / (self.sigma_data ** 2 + sigma ** 2).sqrt() + + def c_in(self, sigma): + return 1 / (self.sigma_data ** 2 + sigma ** 2).sqrt() + + def c_noise(self, sigma): + return 0.25 * sigma.log() + + def D(self, x, sigma, condition=None, model=None): + """ Prepositioning in EDM """ + c_skip, c_out, c_in, c_noise = self.c_skip(sigma), self.c_out(sigma), self.c_in(sigma), self.c_noise(sigma) + if model is None: + model = self.model + c_skip, c_in, c_out = at_least_ndim(c_skip, x.dim()), at_least_ndim(c_in, x.dim()), at_least_ndim(c_out, x.dim()) + return c_skip * x + c_out * model["diffusion"](c_in * x, c_noise, condition) + + # ==================== Training: Score Matching ====================== + + def add_noise(self, x0, t=None, eps=None): + + t = (torch.randn((x0.shape[0], ), device=self.device) * self.P_std + self.P_mean).exp() if t is None else t + + eps = torch.randn_like(x0) if eps is None else eps + + scale = 1. + sigma = at_least_ndim(t, x0.dim()) + + xt = scale * x0 + sigma * eps + xt = (1. - self.fix_mask) * xt + self.fix_mask * x0 + + return xt, t, eps + + def loss(self, x0, condition=None): + + xt, t, eps = self.add_noise(x0) + + condition = self.model["condition"](condition) if condition is not None else None + + loss = (self.D(xt, t, condition) - x0) ** 2 + + edm_loss_weight = at_least_ndim((t ** 2 + self.sigma_data ** 2) / ((t * self.sigma_data) ** 2), x0.dim()) + + return (loss * self.loss_weight * (1 - self.fix_mask) * edm_loss_weight).mean() + + def update(self, x0, condition=None, update_ema=True, **kwargs): + """One-step gradient update. + Inputs: + - x0: torch.Tensor + Samples from the target distribution. + - condition: Optional + Condition of x0. `None` indicates no condition. + - update_ema: bool + Whether to update the exponential moving average model. + + Outputs: + - log: dict + The log dictionary. + """ + loss = self.loss(x0, condition) + + loss.backward() + grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_norm) \ + if self.grad_clip_norm else None + self.optimizer.step() + self.optimizer.zero_grad() + + if update_ema: + self.ema_update() + + log = {"loss": loss.item(), "grad_norm": grad_norm} + + return log + + def update_classifier(self, x0, condition): + + xt, t, eps = self.add_noise(x0) + + log = self.classifier.update(xt, t.log() / 4., condition) + + return log + + # ==================== Sampling: Solving SDE/ODE ====================== + + def classifier_guidance( + self, xt, t, sigma, + model, condition=None, w: float = 1.0, + pred=None): + """ + Guided Sampling CG: + bar_eps = eps - w * sigma * grad + bar_x0 = x0 + w * (sigma ** 2) * alpha * grad + """ + if pred is None: + pred = self.D(xt, t, None, model) + if self.classifier is None or w == 0.0 or condition is None: + return pred, None + else: + log_p, grad = self.classifier.gradients( + xt.clone(), t.log() / 4., condition) + pred = pred + w * (at_least_ndim(sigma, pred.dim()) ** 2) * grad + + return pred, log_p + + def classifier_free_guidance( + self, xt, t, + model, condition=None, w: float = 1.0, + pred=None, pred_uncond=None, + requires_grad: bool = False): + """ + Guided Sampling CFG: + bar_eps = w * pred + (1 - w) * pred_uncond + bar_x0 = w * pred + (1 - w) * pred_uncond + """ + with torch.set_grad_enabled(requires_grad): + if w != 0.0 and w != 1.0: + if pred is None or pred_uncond is None: + b = xt.shape[0] + repeat_dim = [2 if i == 0 else 1 for i in range(xt.dim())] + condition = torch.cat([condition, torch.zeros_like(condition)], 0) + pred_all = self.D( + xt.repeat(*repeat_dim), + t.repeat(2), condition, model) + pred, pred_uncond = pred_all[:b], pred_all[b:] + elif w == 0.0: + pred = 0. + pred_uncond = self.D(xt, t, None, model) + else: + pred = self.D(xt, t, condition, model) + pred_uncond = 0. + + bar_pred = w * pred + (1 - w) * pred_uncond + + return bar_pred + + def guided_sampling( + self, xt, t, sigma, + model, + condition_cfg=None, w_cfg: float = 0.0, + condition_cg=None, w_cg: float = 0.0, + requires_grad: bool = False): + """ + One-step epsilon/x0 prediction with guidance. + """ + + pred = self.classifier_free_guidance( + xt, t, model, condition_cfg, w_cfg, None, None, requires_grad) + + pred, logp = self.classifier_guidance( + xt, t, sigma, model, condition_cg, w_cg, pred) + + return pred, logp + + def sample( + self, + # ---------- the known fixed portion ---------- # + prior: torch.Tensor, + # ----------------- sampling ----------------- # + solver: str = "euler", # euler or heun + n_samples: int = 1, + sample_steps: int = 5, + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ----------- Diffusion-X sampling ---------- + diffusion_x_sampling_steps: int = 0, + # ----------- Warm-Starting ----------- + warm_start_reference: Optional[torch.Tensor] = None, + warm_start_forward_level: float = 0.3, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + """Sampling. + + Inputs: + - prior: torch.Tensor + The known fixed portion of the input data. Should be in the shape of generated data. + Use `torch.zeros((n_samples, *x_shape))` for non-prior sampling. + + - solver: str + The solver for the reverse process. Check `supported_solvers` property for available solvers. + - n_samples: int + The number of samples to generate. + - sample_steps: int + The number of sampling steps. Should be greater than 1 and less than or equal to the number of diffusion steps. + - sample_step_schedule: Union[str, Callable] + The schedule for the sampling steps. + - use_ema: bool + Whether to use the exponential moving average model. + - temperature: float + The temperature for sampling. + + - condition_cfg: Optional + Condition for Classifier-free-guidance. + - mask_cfg: Optional + Mask for Classifier-guidance. + - w_cfg: float + Weight for Classifier-free-guidance. + - condition_cg: Optional + Condition for Classifier-guidance. + - w_cg: float + Weight for Classifier-guidance. + + - diffusion_x_sampling_steps: int + The number of diffusion steps for diffusion-x sampling. + + - warm_start_reference: Optional[torch.Tensor] + Reference data for warm-starting sampling. `None` indicates no warm-starting. + - warm_start_forward_level: float + The forward noise level to perturb the reference data. Should be in the range of `[0., 1.]`, where `1` indicates pure noise. + + - requires_grad: bool + Whether to preserve gradients. + - preserve_history: bool + Whether to preserve the sampling history. + + Outputs: + - x0: torch.Tensor + Generated samples. Be in the shape of `(n_samples, *x_shape)`. + - log: dict + The log dictionary. + """ + assert solver in ["euler", "heun"], f"Solver {solver} is not supported. Use 'euler' or 'heun' instead." + + # ===================== Initialization ===================== + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + model = self.model if not use_ema else self.model_ema + + prior = prior.to(self.device) + if isinstance(warm_start_reference, torch.Tensor) and warm_start_forward_level > 0.: + fwd_sigma = self.sigma_min + (self.sigma_max - self.sigma_min) * warm_start_forward_level + xt = warm_start_reference + fwd_sigma * torch.randn_like(warm_start_reference) + else: + fwd_sigma = self.sigma_max + xt = torch.randn_like(prior) * self.sigma_max * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, 0] = xt.cpu().numpy() + + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + # ===================== Sampling Schedule ==================== + sample_step_schedule = ((self.sigma_min ** (1 / self.rho) + torch.arange(sample_steps + 1, device=self.device) + / sample_steps * (fwd_sigma ** (1 / self.rho) - self.sigma_min ** (1 / self.rho))) + ** self.rho) + + sigmas = sample_step_schedule + + # ===================== Denoising Loop ======================== + loop_steps = [1] * diffusion_x_sampling_steps + list(range(1, sample_steps + 1)) + for i in reversed(loop_steps): + + t = torch.full((n_samples,), sample_step_schedule[i], dtype=torch.float32, device=self.device) + + # guided sampling + pred, logp = self.guided_sampling( + xt, t, sigmas[i], + model, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg, requires_grad) + + # clip the prediction + if self.clip_pred: + pred = pred.clip(self.x_min, self.x_max) + + # one-step update + dot_x = (xt - pred) / at_least_ndim(sigmas[i], xt.dim()) + delta_t = sample_step_schedule[i] - sample_step_schedule[i - 1] + x_next = xt - dot_x * delta_t + x_next = x_next * (1. - self.fix_mask) + prior * self.fix_mask + + if solver == "heun" and i > 1: + pred, logp = self.guided_sampling( + x_next, t / sample_step_schedule[i] * sample_step_schedule[i - 1], sigmas[i - 1], + model, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg, requires_grad) + if self.clip_pred: + pred = pred.clip(self.x_min, self.x_max) + dot_x_prime = (x_next - pred) / at_least_ndim(sigmas[i - 1], xt.dim()) + x_next = xt - (dot_x + dot_x_prime) / 2. * delta_t + x_next = x_next * (1. - self.fix_mask) + prior * self.fix_mask + + xt = x_next + + if preserve_history: + log["sample_history"][:, sample_steps - i + 1] = xt.cpu().numpy() + + # ================= Post-processing ================= + if self.classifier is not None: + with torch.no_grad(): + t = torch.ones((n_samples,), dtype=torch.long, device=self.device) * self.sigma_min + logp = self.classifier.logp(xt, t.log() / 4., condition_vec_cg) + log["log_p"] = logp + + if self.clip_pred: + xt = xt.clip(self.x_min, self.x_max) + + return xt, log diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/rectifiedflow.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/rectifiedflow.py new file mode 100644 index 0000000000000000000000000000000000000000..a48db1dc1ba862281461e818a8fb842b0667e1e5 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/rectifiedflow.py @@ -0,0 +1,632 @@ +from typing import Optional, Union, Callable + +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.classifier import BaseClassifier +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from cleandiffuser.utils import ( + at_least_ndim, + SUPPORTED_DISCRETIZATIONS, SUPPORTED_SAMPLING_STEP_SCHEDULE) +from .basic import DiffusionModel + + +class DiscreteRectifiedFlow(DiffusionModel): + """Discrete-time Rectified Flow + + Rectified flow learns a straight flow for transporting between two probability distributions. + The straight characteristic of this flow enables ODEs to achieve high-quality data generation + with very few sampling steps (achieving one-step sampling when the ODE is completely straight). + Rectified flow can continuously optimize model performance by refining the learned ODE through multiple reflow procedures. + Since it does not require Gaussian to be the source distribution, our implementation supports transformations + between two arbitrary distributions. However, if not specifically specified during training, + the model will default to learning the transport from a Gaussian distribution to the target distribution. + + The DiscreteRectifiedFlow is the discrete-time version of Rectified flow. + It discretizes the continuous time interval into a finite number of diffusion steps + and only estimates the score function on these steps. + Therefore, in sampling, solvers can only work on these learned steps. + The sampling steps are required to be greater than 1 and less than or equal to the number of diffusion steps. + + Args: + - nn_diffusion: BaseNNDiffusion + The neural network backbone for the Diffusion model. + - nn_condition: Optional[BaseNNCondition] + The neural network backbone for the condition embedding. + + - fix_mask: Union[list, np.ndarray, torch.Tensor] + Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + The mask should be in the shape of `x_shape`. + - loss_weight: Union[list, np.ndarray, torch.Tensor] + Add loss weight. The weight should be in the shape of `x_shape`. + + - classifier: Optional[BaseClassifier] + Add a classifier to enable classifier-guidance. + + - grad_clip_norm: Optional[float] + Gradient clipping norm. + - ema_rate: float + Exponential moving average rate. + - optim_params: Optional[dict] + Optimizer parameters. + + - diffusion_steps: int + The discretization steps for discrete-time diffusion models. + - discretization: Union[str, Callable] + The discretization method for the diffusion steps. + + - x_max: Optional[torch.Tensor] + The maximum value for the input data. `None` indicates no constraint. + - x_min: Optional[torch.Tensor] + The minimum value for the input data. `None` indicates no constraint. + + - device: Union[torch.device, str] + The device to run the model. + """ + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugins ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Training Params ---------------- # + grad_clip_norm: Optional[float] = None, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- Diffusion Params ------------------- # + diffusion_steps: int = 1000, + + discretization: Union[str, Callable] = "uniform", + + x_max: Optional[torch.Tensor] = None, + x_min: Optional[torch.Tensor] = None, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + diffusion_steps, ema_rate, optim_params, device) + + assert classifier is None, "Rectified Flow does not support classifier-guidance." + + self.x_max, self.x_min = x_max, x_min + + # ================= Discretization ================= + # - Map the continuous range [0., 1.] to the discrete range [0, T-1] + if isinstance(discretization, str): + if discretization in SUPPORTED_DISCRETIZATIONS.keys(): + self.t_diffusion = SUPPORTED_DISCRETIZATIONS[discretization](diffusion_steps, 0.).to(device) + else: + Warning(f"Discretization method {discretization} is not supported. " + f"Using uniform discretization instead.") + self.t_diffusion = SUPPORTED_DISCRETIZATIONS["uniform"](diffusion_steps, 0.).to(device) + elif callable(discretization): + self.t_diffusion = discretization(diffusion_steps, 0.).to(device) + else: + raise ValueError("discretization must be a callable or a string") + + @property + def supported_solvers(self): + return ["euler"] + + @property + def clip_pred(self): + return (self.x_max is not None) or (self.x_min is not None) + + # ==================== Training: Straighten Flow ====================== + + def loss(self, x0, x1=None, condition=None): + + # x1 is the samples of source distribution. + # If x1 is None, then we assume x1 is from a standard Gaussian distribution. + if x1 is None: + x1 = torch.randn_like(x0) + else: + assert x0.shape == x1.shape, "x0 and x1 must have the same shape" + + t = torch.randint(self.diffusion_steps, (x0.shape[0],), device=self.device) + t_c = at_least_ndim(self.t_diffusion[t], x0.dim()) + + xt = t_c * x1 + (1 - t_c) * x0 + xt = xt * (1. - self.fix_mask) + x0 * self.fix_mask + + condition = self.model["condition"](condition) if condition is not None else None + + loss = (self.model["diffusion"](xt, t, condition) - (x0 - x1)) ** 2 + + return (loss * self.loss_weight * (1 - self.fix_mask)).mean() + + def update(self, x0, condition=None, update_ema=True, x1=None, **kwargs): + """One-step gradient update. + Inputs: + - x0: torch.Tensor + Samples from the target distribution. + - condition: Optional + Condition of x0. `None` indicates no condition. + - update_ema: bool + Whether to update the exponential moving average model. + - x1: torch.Tensor + Samples from the source distribution. `None` indicates standard Gaussian samples. + + Outputs: + - log: dict + The log dictionary. + """ + loss = self.loss(x0, x1, condition) + + loss.backward() + grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_norm) \ + if self.grad_clip_norm else None + self.optimizer.step() + self.optimizer.zero_grad() + + if update_ema: + self.ema_update() + + log = {"loss": loss.item(), "grad_norm": grad_norm} + + return log + + # ==================== Sampling: Solving a straight ODE flow ====================== + + def sample( + self, + # ---------- the known fixed portion ---------- # + prior: torch.Tensor, + x1: torch.Tensor = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + sample_step_schedule: Union[str, Callable] = "uniform", + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ----------- Diffusion-X sampling ---------- + diffusion_x_sampling_steps: int = 0, + # ----------- Warm-Starting ----------- + warm_start_reference: Optional[torch.Tensor] = None, + warm_start_forward_level: float = 0.3, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + """Sampling. + + Inputs: + - prior: torch.Tensor + The known fixed portion of the input data. Should be in the shape of generated data. + Use `torch.zeros((n_samples, *x_shape))` for non-prior sampling. + - x1: torch.Tensor + The samples from the source distribution. `None` indicates standard Gaussian samples. + + - n_samples: int + The number of samples to generate. + - sample_steps: int + The number of sampling steps. Should be greater than 1 and less than or equal to the number of diffusion steps. + - sample_step_schedule: Union[str, Callable] + The schedule for the sampling steps. + - use_ema: bool + Whether to use the exponential moving average model. + - temperature: float + The temperature for sampling. + + - condition_cfg: Optional + Condition for Classifier-free-guidance. + - mask_cfg: Optional + Mask for Classifier-guidance. + - w_cfg: float + Weight for Classifier-free-guidance. + - condition_cg: Optional + Condition for Classifier-guidance. + - w_cg: float + Weight for Classifier-guidance. + + - diffusion_x_sampling_steps: int + The number of diffusion steps for diffusion-x sampling. + + - warm_start_reference: Optional[torch.Tensor] + Reference data for warm-starting sampling. `None` indicates no warm-starting. + - warm_start_forward_level: float + The forward noise level to perturb the reference data. Should be in the range of `[0., 1.]`, where `1` indicates pure noise. + + - requires_grad: bool + Whether to preserve gradients. + - preserve_history: bool + Whether to preserve the sampling history. + + Outputs: + - x0: torch.Tensor + Generated samples. Be in the shape of `(n_samples, *x_shape)`. + - log: dict + The log dictionary. + """ + assert w_cg == 0.0 and condition_cg is None, "Rectified Flow does not support classifier-guidance." + + prior = prior.to(self.device) + if isinstance(warm_start_reference, torch.Tensor): + diffusion_steps = int(warm_start_forward_level * self.diffusion_steps) + t_c = at_least_ndim(self.t_diffusion[diffusion_steps], prior.dim()) + x1 = torch.randn_like(prior) * t_c + warm_start_reference * (1 - t_c) + else: + diffusion_steps = self.diffusion_steps + if x1 is None: + x1 = torch.randn_like(prior) * temperature + else: + assert prior.shape == x1.shape, "prior and x1 must have the same shape" + + # ===================== Initialization ===================== + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + model = self.model if not use_ema else self.model_ema + + xt = x1.clone() + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, 0] = xt.cpu().numpy() + + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + + # ===================== Sampling Schedule ==================== + if isinstance(sample_step_schedule, str): + if sample_step_schedule in SUPPORTED_SAMPLING_STEP_SCHEDULE.keys(): + sample_step_schedule = SUPPORTED_SAMPLING_STEP_SCHEDULE[sample_step_schedule]( + diffusion_steps, sample_steps) + else: + raise ValueError(f"Sampling step schedule {sample_step_schedule} is not supported.") + elif callable(sample_step_schedule): + sample_step_schedule = sample_step_schedule(diffusion_steps, sample_steps) + else: + raise ValueError("sample_step_schedule must be a callable or a string") + + # ===================== Denoising Loop ======================== + loop_steps = [1] * diffusion_x_sampling_steps + list(range(1, sample_steps + 1)) + for i in reversed(loop_steps): + + t = torch.full((n_samples,), sample_step_schedule[i], dtype=torch.long, device=self.device) + + start_t, end_t = self.t_diffusion[sample_step_schedule[i]], self.t_diffusion[sample_step_schedule[i - 1]] + delta_t = start_t - end_t + + # velocity + if w_cfg != 0.0 and w_cfg != 1.0 and condition_vec_cfg is not None: + repeat_dim = [2 if i == 0 else 1 for i in range(xt.dim())] + vel_all = model["diffusion"]( + xt.repeat(*repeat_dim), t.repeat(2), + torch.cat([condition_vec_cfg, torch.zeros_like(condition_vec_cfg)], 0)) + vel_cond, vel_uncond = vel_all.chunk(2, dim=0) + vel = w_cfg * vel_cond + (1 - w_cfg) * vel_uncond + elif w_cfg == 0.0 or condition_vec_cfg is None: + vel = model["diffusion"](xt, t, None) + else: + vel = model["diffusion"](xt, t, condition_vec_cfg) + + # one-step update + xt = xt + delta_t * vel + + # fix the known portion, and preserve the sampling history + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, sample_steps - i + 1] = xt.cpu().numpy() + + # ================= Post-processing ================= + if self.clip_pred: + xt = xt.clip(self.x_min, self.x_max) + + return xt, log + + +class ContinuousRectifiedFlow(DiffusionModel): + """Continuous-time Rectified Flow + + Rectified flow learns a straight flow for transporting between two probability distributions. + The straight characteristic of this flow enables ODEs to achieve high-quality data generation + with very few sampling steps (achieving one-step sampling when the ODE is completely straight). + Rectified flow can continuously optimize model performance by refining the learned ODE through multiple reflow procedures. + Since it does not require Gaussian to be the source distribution, our implementation supports transformations + between two arbitrary distributions. However, if not specifically specified during training, + the model will default to learning the transport from a Gaussian distribution to the target distribution. + + The ContinuousRectifiedFlow is the continuous-time version of Rectified flow. + It estimates the score function at any $t\in[0,T]$ and solves the reverse process in continuous time. + The sampling steps are required to be greater than 1. + + Args: + - nn_diffusion: BaseNNDiffusion + The neural network backbone for the Diffusion model. + - nn_condition: Optional[BaseNNCondition] + The neural network backbone for the condition embedding. + + - fix_mask: Union[list, np.ndarray, torch.Tensor] + Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + The mask should be in the shape of `x_shape`. + - loss_weight: Union[list, np.ndarray, torch.Tensor] + Add loss weight. The weight should be in the shape of `x_shape`. + + - classifier: Optional[BaseClassifier] + Add a classifier to enable classifier-guidance. + + - grad_clip_norm: Optional[float] + Gradient clipping norm. + - ema_rate: float + Exponential moving average rate. + - optim_params: Optional[dict] + Optimizer parameters. + + - x_max: Optional[torch.Tensor] + The maximum value for the input data. `None` indicates no constraint. + - x_min: Optional[torch.Tensor] + The minimum value for the input data. `None` indicates no constraint. + + - device: Union[torch.device, str] + The device to run the model. + """ + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugins ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Training Params ---------------- # + grad_clip_norm: Optional[float] = None, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- Diffusion Params ------------------- # + x_max: Optional[torch.Tensor] = None, + x_min: Optional[torch.Tensor] = None, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + 0, ema_rate, optim_params, device) + + assert classifier is None, "Rectified Flow does not support classifier-guidance." + + self.x_max, self.x_min = x_max, x_min + + @property + def supported_solvers(self): + return ["euler"] + + @property + def clip_pred(self): + return (self.x_max is not None) or (self.x_min is not None) + + # ==================== Training: Straighten Flow ====================== + + def loss(self, x0, x1=None, condition=None): + + # x1 is the samples of source distribution. + # If x1 is None, then we assume x1 is from a standard Gaussian distribution. + if x1 is None: + x1 = torch.randn_like(x0) + else: + assert x0.shape == x1.shape, "x0 and x1 must have the same shape" + + t = torch.rand((x0.shape[0],), device=self.device) + + xt = x0 + at_least_ndim(t, x0.dim()) * (x1 - x0) + + xt = xt * (1. - self.fix_mask) + x0 * self.fix_mask + + condition = self.model["condition"](condition) if condition is not None else None + + loss = (self.model["diffusion"](xt, t, condition) - (x0 - x1)) ** 2 + + return (loss * self.loss_weight * (1 - self.fix_mask)).mean() + + def update(self, x0, condition=None, update_ema=True, x1=None, **kwargs): + """One-step gradient update. + Inputs: + - x0: torch.Tensor + Samples from the target distribution. + - condition: Optional + Condition of x0. `None` indicates no condition. + - update_ema: bool + Whether to update the exponential moving average model. + - x1: torch.Tensor + Samples from the source distribution. `None` indicates standard Gaussian samples. + + Outputs: + - log: dict + The log dictionary. + """ + loss = self.loss(x0, x1, condition) + + loss.backward() + grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_norm) \ + if self.grad_clip_norm else None + self.optimizer.step() + self.optimizer.zero_grad() + + if update_ema: + self.ema_update() + + log = {"loss": loss.item(), "grad_norm": grad_norm} + + return log + + # ==================== Sampling: Solving a straight ODE flow ====================== + + def sample( + self, + # ---------- the known fixed portion ---------- # + prior: torch.Tensor, + x1: torch.Tensor = None, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + sample_step_schedule: Union[str, Callable] = "uniform_continuous", + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ----------- Diffusion-X sampling ---------- + diffusion_x_sampling_steps: int = 0, + # ----------- Warm-Starting ----------- + warm_start_reference: Optional[torch.Tensor] = None, + warm_start_forward_level: float = 0.3, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + """Sampling. + + Inputs: + - prior: torch.Tensor + The known fixed portion of the input data. Should be in the shape of generated data. + Use `torch.zeros((n_samples, *x_shape))` for non-prior sampling. + - x1: torch.Tensor + The samples from the source distribution. `None` indicates standard Gaussian samples. + + - n_samples: int + The number of samples to generate. + - sample_steps: int + The number of sampling steps. Should be greater than 1 and less than or equal to the number of diffusion steps. + - sample_step_schedule: Union[str, Callable] + The schedule for the sampling steps. + - use_ema: bool + Whether to use the exponential moving average model. + - temperature: float + The temperature for sampling. + + - condition_cfg: Optional + Condition for Classifier-free-guidance. + - mask_cfg: Optional + Mask for Classifier-guidance. + - w_cfg: float + Weight for Classifier-free-guidance. + - condition_cg: Optional + Condition for Classifier-guidance. + - w_cg: float + Weight for Classifier-guidance. + + - diffusion_x_sampling_steps: int + The number of diffusion steps for diffusion-x sampling. + + - requires_grad: bool + Whether to preserve gradients. + - preserve_history: bool + Whether to preserve the sampling history. + + Outputs: + - x0: torch.Tensor + Generated samples. Be in the shape of `(n_samples, *x_shape)`. + - log: dict + The log dictionary. + """ + assert w_cg == 0.0 and condition_cg is None, "Rectified Flow does not support classifier-guidance." + + prior = prior.to(self.device) + if isinstance(warm_start_reference, torch.Tensor): + t_c = torch.ones_like(prior) * warm_start_forward_level + x1 = torch.randn_like(prior) * t_c + warm_start_reference * (1 - t_c) + else: + if x1 is None: + x1 = torch.randn_like(prior) * temperature + else: + assert prior.shape == x1.shape, "prior and x1 must have the same shape" + + # ===================== Initialization ===================== + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + model = self.model if not use_ema else self.model_ema + + xt = x1.clone() + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, 0] = xt.cpu().numpy() + + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + + # ===================== Sampling Schedule ==================== + if isinstance(warm_start_reference, torch.Tensor) and warm_start_forward_level > 0.: + final_t = warm_start_forward_level + else: + final_t = 1. + if isinstance(sample_step_schedule, str): + if sample_step_schedule in SUPPORTED_SAMPLING_STEP_SCHEDULE.keys(): + sample_step_schedule = SUPPORTED_SAMPLING_STEP_SCHEDULE[sample_step_schedule]( + [0., final_t], sample_steps) + else: + raise ValueError(f"Sampling step schedule {sample_step_schedule} is not supported.") + elif callable(sample_step_schedule): + sample_step_schedule = sample_step_schedule([0., final_t], sample_steps) + else: + raise ValueError("sample_step_schedule must be a callable or a string") + + # ===================== Denoising Loop ======================== + loop_steps = [1] * diffusion_x_sampling_steps + list(range(1, sample_steps + 1)) + for i in reversed(loop_steps): + + t = torch.full((n_samples,), sample_step_schedule[i], dtype=torch.float32, device=self.device) + + delta_t = sample_step_schedule[i] - sample_step_schedule[i - 1] + + # velocity + if w_cfg != 0.0 and w_cfg != 1.0 and condition_vec_cfg is not None: + repeat_dim = [2 if i == 0 else 1 for i in range(xt.dim())] + vel_all = model["diffusion"]( + xt.repeat(*repeat_dim), t.repeat(2), + torch.cat([condition_vec_cfg, torch.zeros_like(condition_vec_cfg)], 0)) + vel_cond, vel_uncond = vel_all.chunk(2, dim=0) + vel = w_cfg * vel_cond + (1 - w_cfg) * vel_uncond + elif w_cfg == 0.0 or condition_vec_cfg is None: + vel = model["diffusion"](xt, t, None) + else: + vel = model["diffusion"](xt, t, condition_vec_cfg) + + # one-step update + xt = xt + delta_t * vel + + # fix the known portion, and preserve the sampling history + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, sample_steps - i + 1] = xt.cpu().numpy() + + # ================= Post-processing ================= + if self.clip_pred: + xt = xt.clip(self.x_min, self.x_max) + + return xt, log diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/veode.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/veode.py new file mode 100644 index 0000000000000000000000000000000000000000..9450e90b79498c8c7c3cb951f017594b332a9f76 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/veode.py @@ -0,0 +1,73 @@ +from typing import Optional +from typing import Union + +import numpy as np +import torch + +from cleandiffuser.classifier import BaseClassifier +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from .edm import EDMArchetecture + + +class VEODE(EDMArchetecture): + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugs ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Params ---------------- # + grad_clip_norm: Optional[float] = None, + diffusion_steps: int = 1000, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- EDM Params ------------------- # + sigma_min: float = 0.02, + sigma_max: float = 100., + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + diffusion_steps, ema_rate, optim_params, device) + + self.sigma_min, self.sigma_max = sigma_min, sigma_max + + def set_sample_steps(self, N: int): + self.sample_steps = N + self.sigma_s = self.sigma_max * (self.sigma_min / self.sigma_max) ** ( + torch.arange(N, device=self.device) / (N - 1)) + self.t_s = self.sigma_s ** 2 + self.scale_s = torch.ones_like(self.sigma_s) * 1.0 + self.dot_sigma_s = 1 / (2 * self.sigma_s) + self.dot_scale_s = torch.zeros_like(self.sigma_s) + self.x_weight_s = (self.dot_sigma_s / self.sigma_s + self.dot_scale_s / self.scale_s) + self.D_weight_s = self.dot_sigma_s / self.sigma_s * self.scale_s + + def c_skip(self, sigma): return torch.ones_like(sigma) + + def c_out(self, sigma): return sigma + + def c_in(self, sigma): return torch.ones_like(sigma) + + def c_noise(self, sigma): return (0.5 * sigma).log() + + def loss_weighting(self, sigma): return 1 / (sigma ** 2) + + def sample_noise_distribution(self, N): + log_sigma = torch.rand((N, 1), device=self.device) * np.log(self.sigma_max / self.sigma_min) + np.log( + self.sigma_min) + return log_sigma.exp() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/vpode.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/vpode.py new file mode 100644 index 0000000000000000000000000000000000000000..fba6eeb294562a075bbcd99933aca14faf286a82 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/vpode.py @@ -0,0 +1,78 @@ +from typing import Optional +from typing import Union + +import numpy as np +import torch + +from cleandiffuser.classifier import BaseClassifier +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from .edm import EDMArchetecture + + +class VPODE(EDMArchetecture): + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugs ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Params ---------------- # + grad_clip_norm: Optional[float] = None, + diffusion_steps: int = 1000, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- EDM Params ------------------- # + beta_min: float = 0.1, + beta_max: float = 20., + eps_s: float = 1e-3, + eps_t: float = 1e-5, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + diffusion_steps, ema_rate, optim_params, device) + + self.beta_min, self.beta_max = beta_min, beta_max + self.eps_s, self.eps_t = eps_s, eps_t + self.beta_d = beta_max - beta_min + + def set_sample_steps(self, N: int): + self.sample_steps = N + self.t_s = torch.arange(N, device=self.device) / (N - 1) * (1e-3 - 1) + 1 + self.sigma_s = ((0.5 * self.beta_d * (self.t_s ** 2) + self.beta_min * self.t_s).exp() - 1.).sqrt() + self.scale_s = 1 / (1 + self.sigma_s ** 2).sqrt() + self.dot_sigma_s = (0.5 * (self.sigma_s ** 2 + 1) * (self.beta_d * self.t_s + self.beta_min) / self.sigma_s) + self.dot_scale_s = -self.sigma_s / (1 + self.sigma_s ** 2) ** 1.5 * self.dot_sigma_s + self.x_weight_s = (self.dot_sigma_s / self.sigma_s + self.dot_scale_s / self.scale_s) + self.D_weight_s = self.dot_sigma_s / self.sigma_s * self.scale_s + + def c_skip(self, sigma): return torch.ones_like(sigma) + + def c_out(self, sigma): return -sigma + + def c_in(self, sigma): return 1 / (1 + sigma ** 2).sqrt() + + def c_noise(self, sigma): + scale = 1 / (1 + sigma ** 2).sqrt() + t = ((self.beta_min ** 2 - 4 * self.beta_d * scale.log()).sqrt() - self.beta_min) / self.beta_d + return ((self.diffusion_steps - 1) * t).long() + + def loss_weighting(self, sigma): return 1 / (sigma ** 2) + + def sample_noise_distribution(self, N): + t = torch.rand((N, 1), device=self.device) * (1 - self.eps_t) + self.eps_t + return ((0.5 * self.beta_d * t ** 2 + self.beta_min * t).exp() - 1).sqrt() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/vpsde.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/vpsde.py new file mode 100644 index 0000000000000000000000000000000000000000..d36be9bc98e2dd99de4bae1070697e74bf9cb4af --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/diffusion/vpsde.py @@ -0,0 +1,1099 @@ +from typing import Optional, Union, Callable, Dict + +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.classifier import BaseClassifier +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from cleandiffuser.utils import ( + at_least_ndim, + SUPPORTED_NOISE_SCHEDULES, SUPPORTED_DISCRETIZATIONS, SUPPORTED_SAMPLING_STEP_SCHEDULE) +from .basic import DiffusionModel + +SUPPORTED_SOLVERS = [ + "ddpm", "ddim", "ode_dpmsolver++_1", "sde_dpmsolver++_1", "ode_dpmsolver++_2M" +] + + +def epstheta_to_xtheta(x, alpha, sigma, eps_theta): + """ + x_theta = (x - sigma * eps_theta) / alpha + """ + return (x - sigma * eps_theta) / alpha + + +def xtheta_to_epstheta(x, alpha, sigma, x_theta): + """ + eps_theta = (x - alpha * x_theta) / sigma + """ + return (x - alpha * x_theta) / sigma + + +class DiscreteVPSDE(DiffusionModel): + + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugins ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Training Params ---------------- # + grad_clip_norm: Optional[float] = None, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- Diffusion Params ------------------- # + epsilon: float = 1e-3, + diffusion_steps: int = 1000, + + discretization: Union[str, Callable] = "uniform", + noise_schedule: Union[str, Dict[str, Callable]] = "cosine", + noise_schedule_params: Optional[dict] = None, + + x_max: Optional[torch.Tensor] = None, + x_min: Optional[torch.Tensor] = None, + + predict_noise: bool = True, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + diffusion_steps, ema_rate, optim_params, device) + + self.predict_noise = predict_noise + self.epsilon = epsilon + self.x_max, self.x_min = x_max, x_min + + if 1. / diffusion_steps < epsilon: + raise ValueError("epsilon is too large for the number of diffusion steps") + + # ================= Discretization ================= + # - Map the continuous range [epsilon, 1.] to the discrete range [0, T-1] + if isinstance(discretization, str): + if discretization in SUPPORTED_DISCRETIZATIONS.keys(): + self.t_diffusion = SUPPORTED_DISCRETIZATIONS[discretization](diffusion_steps, epsilon).to(device) + else: + Warning(f"Discretization method {discretization} is not supported. " + f"Using uniform discretization instead.") + self.t_diffusion = SUPPORTED_DISCRETIZATIONS["uniform"](diffusion_steps, epsilon).to(device) + elif callable(discretization): + self.t_diffusion = discretization(diffusion_steps, epsilon).to(device) + else: + raise ValueError("discretization must be a callable or a string") + + # ================= Noise Schedule ================= + if isinstance(noise_schedule, str): + if noise_schedule in SUPPORTED_NOISE_SCHEDULES.keys(): + self.alpha, self.sigma = SUPPORTED_NOISE_SCHEDULES[noise_schedule]["forward"]( + self.t_diffusion, **(noise_schedule_params or {})) + else: + raise ValueError(f"Noise schedule {noise_schedule} is not supported.") + elif isinstance(noise_schedule, dict): + self.alpha, self.sigma = noise_schedule["forward"](self.t_diffusion, **(noise_schedule_params or {})) + else: + raise ValueError("noise_schedule must be a callable or a string") + + self.logSNR = torch.log(self.alpha / self.sigma) + + # # ================= Extra Steps for DPM-Solver 2 ================== + # self.inter_logSNR = torch.zeros_like(self.logSNR) + # self.inter_logSNR[1:] = (self.logSNR[1:] + self.logSNR[:-1]) / 2. + # # inter_logSNR[0] is not correctly calculated, but it will not be used. + # if isinstance(noise_schedule, str): + # if noise_schedule in SUPPORTED_NOISE_SCHEDULES.keys(): + # self.inter_t_diffusion = SUPPORTED_NOISE_SCHEDULES[noise_schedule]["reverse"]( + # logSNR=self.inter_logSNR, **(noise_schedule_params or {})) + # else: + # raise ValueError(f"Noise schedule {noise_schedule} is not supported.") + # elif isinstance(noise_schedule, dict): + # self.inter_t_diffusion = noise_schedule["reverse"]( + # logSNR=self.inter_logSNR, **(noise_schedule_params or {})) + # else: + # raise ValueError("noise_schedule must be a callable or a string") + + @property + def clip_pred(self): + return (self.x_max is not None) or (self.x_min is not None) + + # ==================== Training: Score Matching ====================== + + def add_noise(self, x0, t=None, eps=None): + + t = torch.randint(self.diffusion_steps, (x0.shape[0],), device=self.device) if t is None else t + eps = torch.randn_like(x0) if eps is None else eps + + alpha, sigma = at_least_ndim(self.alpha[t], x0.dim()), at_least_ndim(self.sigma[t], x0.dim()) + + xt = alpha * x0 + sigma * eps + xt = (1. - self.fix_mask) * xt + self.fix_mask * x0 + + return xt, t, eps + + def loss(self, x0, condition=None): + + xt, t, eps = self.add_noise(x0) + + condition = self.model["condition"](condition) if condition is not None else None + + if self.predict_noise: + loss = (self.model["diffusion"](xt, t, condition) - eps) ** 2 + else: + loss = (self.model["diffusion"](xt, t, condition) - x0) ** 2 + + return (loss * self.loss_weight * (1 - self.fix_mask)).mean() + + def update(self, x0, condition=None, update_ema=True, **kwargs): + + loss = self.loss(x0, condition) + + loss.backward() + grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_norm) \ + if self.grad_clip_norm else None + self.optimizer.step() + self.optimizer.zero_grad() + + if update_ema: + self.ema_update() + + log = {"loss": loss.item(), "grad_norm": grad_norm} + + return log + + def update_classifier(self, x0, condition): + + xt, t, eps = self.add_noise(x0) + + log = self.classifier.update(xt, t, condition) + + return log + + # ==================== Sampling: Solving SDE/ODE ====================== + + def classifier_guidance( + self, xt, t, alpha, sigma, + model, condition=None, w: float = 1.0, + pred=None): + """ + Guided Sampling CG: + bar_eps = eps - w * sigma * grad + bar_x0 = x0 + w * (sigma ** 2) * alpha * grad + """ + if pred is None: + pred = model["diffusion"](xt, t, None) + if self.classifier is None or w == 0.0 or condition is None: + return pred, None + else: + log_p, grad = self.classifier.gradients(xt.clone(), t, condition) + if self.predict_noise: + pred = pred - w * sigma * grad + else: + pred = pred + w * ((sigma ** 2) / alpha) * grad + + return pred, log_p + + def classifier_free_guidance( + self, xt, t, + model, condition=None, w: float = 1.0, + pred=None, pred_uncond=None, + requires_grad: bool = False): + """ + Guided Sampling CFG: + bar_eps = w * pred + (1 - w) * pred_uncond + bar_x0 = w * pred + (1 - w) * pred_uncond + """ + with torch.set_grad_enabled(requires_grad): + if w != 0.0 and w != 1.0: + if pred is None or pred_uncond is None: + b = xt.shape[0] + repeat_dim = [2 if i == 0 else 1 for i in range(xt.dim())] + condition = torch.cat([condition, torch.zeros_like(condition)], 0) + pred_all = model["diffusion"]( + xt.repeat(*repeat_dim), t.repeat(2), condition) + pred, pred_uncond = pred_all[:b], pred_all[b:] + elif w == 0.0: + pred = 0. + pred_uncond = model["diffusion"](xt, t, None) + else: + pred = model["diffusion"](xt, t, condition) + pred_uncond = 0. + + if self.predict_noise or not self.predict_noise: + bar_pred = w * pred + (1 - w) * pred_uncond + else: + bar_pred = pred + + return bar_pred + + def clip_prediction(self, pred, xt, alpha, sigma): + """ + Clip the prediction at each sampling step to stablize the generation. + (xt - alpha * x_max) / sigma <= eps <= (xt - alpha * x_min) / sigma + x_min <= x0 <= x_max + """ + if self.predict_noise: + if self.clip_pred: + upper_bound = (xt - alpha * self.x_min) / sigma if self.x_min is not None else None + lower_bound = (xt - alpha * self.x_max) / sigma if self.x_max is not None else None + pred = pred.clip(lower_bound, upper_bound) + else: + if self.clip_pred: + pred = pred.clip(self.x_min, self.x_max) + + return pred + + def guided_sampling( + self, xt, t, alpha, sigma, + model, + condition_cfg=None, w_cfg: float = 0.0, + condition_cg=None, w_cg: float = 0.0, + requires_grad: bool = False): + """ + One-step epsilon/x0 prediction with guidance. + """ + + pred = self.classifier_free_guidance( + xt, t, model, condition_cfg, w_cfg, None, None, requires_grad) + + pred, logp = self.classifier_guidance( + xt, t, alpha, sigma, model, condition_cg, w_cg, pred) + + return pred, logp + + def ddpm_solver( + self, + # ---------- the known fixed portion ---------- # + prior: torch.Tensor, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + sample_step_schedule: Union[str, Callable] = "uniform", + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + + # ===================== Initialization ===================== + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + model = self.model if not use_ema else self.model_ema + + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, 0] = xt.cpu().numpy() + + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + # ===================== Sampling Schedule ==================== + if isinstance(sample_step_schedule, str): + if sample_step_schedule in SUPPORTED_SAMPLING_STEP_SCHEDULE.keys(): + sample_step_schedule = SUPPORTED_SAMPLING_STEP_SCHEDULE[sample_step_schedule]( + self.diffusion_steps, sample_steps) + else: + raise ValueError(f"Sampling step schedule {sample_step_schedule} is not supported.") + elif callable(sample_step_schedule): + sample_step_schedule = sample_step_schedule(self.diffusion_steps, sample_steps) + else: + raise ValueError("sample_step_schedule must be a callable or a string") + + alphas = self.alpha[sample_step_schedule] + sigmas = self.sigma[sample_step_schedule] + logSNRs = torch.log(alphas / sigmas) + hs = torch.zeros_like(logSNRs) + hs[1:] = logSNRs[:-1] - logSNRs[1:] # hs[0] is not correctly calculated, but it will not be used. + stds = torch.zeros((sample_steps + 1,), device=self.device) + stds[1:] = sigmas[:-1] / sigmas[1:] * (1 - (alphas[1:] / alphas[:-1]) ** 2).sqrt() + + # ===================== Denoising Loop ======================== + + for i in reversed(range(1, sample_steps + 1)): + + t = torch.full((n_samples,), sample_step_schedule[i], dtype=torch.long, device=self.device) + + # guided sampling + pred, logp = self.guided_sampling( + xt, t, alphas[i], sigmas[i], + model, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg, requires_grad) + + # clip the prediction + pred = self.clip_prediction(pred, xt, alphas[i], sigmas[i]) + + # transform to eps_theta + eps_theta = pred if self.predict_noise else xtheta_to_epstheta(xt, alphas[i], sigmas[i], pred) + + # one-step update + xt = ( + (alphas[i - 1] / alphas[i]) * (xt - sigmas[i] * eps_theta) + + (sigmas[i - 1] ** 2 - stds[i] ** 2 + 1e-8).sqrt() * eps_theta) + if i > 1: + xt += (stds[i] * torch.randn_like(xt)) + + # fix the known portion, and preserve the sampling history + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, sample_steps - i + 1] = xt.cpu().numpy() + + # ================= Post-processing ================= + if self.classifier is not None: + with torch.no_grad(): + t = torch.zeros((n_samples,), dtype=torch.long, device=self.device) + logp = self.classifier.logp(xt, t, condition_vec_cg) + log["log_p"] = logp + + if self.clip_pred: + xt = xt.clip(self.x_min, self.x_max) + + return xt, log + + def ddim_solver( + self, + # ---------- the known fixed portion ---------- # + prior: torch.Tensor, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + sample_step_schedule: Union[str, Callable] = "uniform", + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + + # ===================== Initialization ===================== + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + model = self.model if not use_ema else self.model_ema + + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, 0] = xt.cpu().numpy() + + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + # ===================== Sampling Schedule ==================== + if isinstance(sample_step_schedule, str): + if sample_step_schedule in SUPPORTED_SAMPLING_STEP_SCHEDULE.keys(): + sample_step_schedule = SUPPORTED_SAMPLING_STEP_SCHEDULE[sample_step_schedule]( + self.diffusion_steps, sample_steps) + else: + raise ValueError(f"Sampling step schedule {sample_step_schedule} is not supported.") + elif callable(sample_step_schedule): + sample_step_schedule = sample_step_schedule(self.diffusion_steps, sample_steps) + else: + raise ValueError("sample_step_schedule must be a callable or a string") + + alphas = self.alpha[sample_step_schedule] + sigmas = self.sigma[sample_step_schedule] + logSNRs = torch.log(alphas / sigmas) + hs = torch.zeros_like(logSNRs) + hs[1:] = logSNRs[:-1] - logSNRs[1:] # hs[0] is not correctly calculated, but it will not be used. + stds = torch.zeros((sample_steps + 1,), device=self.device) + stds[1:] = sigmas[:-1] / sigmas[1:] * (1 - (alphas[1:] / alphas[:-1]) ** 2).sqrt() + + # ===================== Denoising Loop ======================== + + for i in reversed(range(1, sample_steps + 1)): + + t = torch.full((n_samples,), sample_step_schedule[i], dtype=torch.long, device=self.device) + + # guided sampling + pred, logp = self.guided_sampling( + xt, t, alphas[i], sigmas[i], + model, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg, requires_grad) + + # clip the prediction + pred = self.clip_prediction(pred, xt, alphas[i], sigmas[i]) + + # transform to eps_theta + eps_theta = pred if self.predict_noise else xtheta_to_epstheta(xt, alphas[i], sigmas[i], pred) + + # one-step update + xt = (alphas[i - 1] * ((xt - sigmas[i] * eps_theta) / alphas[i]) + sigmas[i - 1] * eps_theta) + + # fix the known portion, and preserve the sampling history + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, sample_steps - i + 1] = xt.cpu().numpy() + + # ================= Post-processing ================= + if self.classifier is not None: + with torch.no_grad(): + t = torch.zeros((n_samples,), dtype=torch.long, device=self.device) + logp = self.classifier.logp(xt, t, condition_vec_cg) + log["log_p"] = logp + + if self.clip_pred: + xt = xt.clip(self.x_min, self.x_max) + + return xt, log + + def ode_dpm_solver_pp_1( + self, + # ---------- the known fixed portion ---------- # + prior: torch.Tensor, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + sample_step_schedule: Union[str, Callable] = "uniform", + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + + # ===================== Initialization ===================== + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + model = self.model if not use_ema else self.model_ema + + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, 0] = xt.cpu().numpy() + + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + # ===================== Sampling Schedule ==================== + if isinstance(sample_step_schedule, str): + if sample_step_schedule in SUPPORTED_SAMPLING_STEP_SCHEDULE.keys(): + sample_step_schedule = SUPPORTED_SAMPLING_STEP_SCHEDULE[sample_step_schedule]( + self.diffusion_steps, sample_steps) + else: + raise ValueError(f"Sampling step schedule {sample_step_schedule} is not supported.") + elif callable(sample_step_schedule): + sample_step_schedule = sample_step_schedule(self.diffusion_steps, sample_steps) + else: + raise ValueError("sample_step_schedule must be a callable or a string") + + alphas = self.alpha[sample_step_schedule] + sigmas = self.sigma[sample_step_schedule] + logSNRs = torch.log(alphas / sigmas) + hs = torch.zeros_like(logSNRs) + hs[1:] = logSNRs[:-1] - logSNRs[1:] # hs[0] is not correctly calculated, but it will not be used. + stds = torch.zeros((sample_steps + 1,), device=self.device) + stds[1:] = sigmas[:-1] / sigmas[1:] * (1 - (alphas[1:] / alphas[:-1]) ** 2).sqrt() + + # ===================== Denoising Loop ======================== + + for i in reversed(range(1, sample_steps + 1)): + + t = torch.full((n_samples,), sample_step_schedule[i], dtype=torch.long, device=self.device) + + # guided sampling + pred, logp = self.guided_sampling( + xt, t, alphas[i], sigmas[i], + model, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg, requires_grad) + + # clip the prediction + pred = self.clip_prediction(pred, xt, alphas[i], sigmas[i]) + + # transform to eps_theta + x_theta = pred if not self.predict_noise else epstheta_to_xtheta(xt, alphas[i], sigmas[i], pred) + + # one-step update + xt = (sigmas[i - 1] / sigmas[i]) * xt - alphas[i - 1] * torch.expm1(-hs[i]) * x_theta + + # fix the known portion, and preserve the sampling history + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, sample_steps - i + 1] = xt.cpu().numpy() + + # ================= Post-processing ================= + if self.classifier is not None: + with torch.no_grad(): + t = torch.zeros((n_samples,), dtype=torch.long, device=self.device) + logp = self.classifier.logp(xt, t, condition_vec_cg) + log["log_p"] = logp + + if self.clip_pred: + xt = xt.clip(self.x_min, self.x_max) + + return xt, log + + def sde_dpm_solver_pp_1( + self, + # ---------- the known fixed portion ---------- # + prior: torch.Tensor, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + sample_step_schedule: Union[str, Callable] = "uniform", + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + + # ===================== Initialization ===================== + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + model = self.model if not use_ema else self.model_ema + + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, 0] = xt.cpu().numpy() + + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + # ===================== Sampling Schedule ==================== + if isinstance(sample_step_schedule, str): + if sample_step_schedule in SUPPORTED_SAMPLING_STEP_SCHEDULE.keys(): + sample_step_schedule = SUPPORTED_SAMPLING_STEP_SCHEDULE[sample_step_schedule]( + self.diffusion_steps, sample_steps) + else: + raise ValueError(f"Sampling step schedule {sample_step_schedule} is not supported.") + elif callable(sample_step_schedule): + sample_step_schedule = sample_step_schedule(self.diffusion_steps, sample_steps) + else: + raise ValueError("sample_step_schedule must be a callable or a string") + + alphas = self.alpha[sample_step_schedule] + sigmas = self.sigma[sample_step_schedule] + logSNRs = torch.log(alphas / sigmas) + hs = torch.zeros_like(logSNRs) + hs[1:] = logSNRs[:-1] - logSNRs[1:] # hs[0] is not correctly calculated, but it will not be used. + stds = torch.zeros((sample_steps + 1,), device=self.device) + stds[1:] = sigmas[:-1] / sigmas[1:] * (1 - (alphas[1:] / alphas[:-1]) ** 2).sqrt() + + # ===================== Denoising Loop ======================== + + for i in reversed(range(1, sample_steps + 1)): + + t = torch.full((n_samples,), sample_step_schedule[i], dtype=torch.long, device=self.device) + + # guided sampling + pred, logp = self.guided_sampling( + xt, t, alphas[i], sigmas[i], + model, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg, requires_grad) + + # clip the prediction + pred = self.clip_prediction(pred, xt, alphas[i], sigmas[i]) + + # transform to eps_theta + x_theta = pred if not self.predict_noise else epstheta_to_xtheta(xt, alphas[i], sigmas[i], pred) + + # one-step update + xt = ((sigmas[i - 1] / sigmas[i]) * (-hs[i]).exp() * xt - + alphas[i - 1] * torch.expm1(-2 * hs[i]) * x_theta + + sigmas[i - 1] * (-torch.expm1(-2 * hs[i])).sqrt() * torch.randn_like(xt)) + + # fix the known portion, and preserve the sampling history + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, sample_steps - i + 1] = xt.cpu().numpy() + + # ================= Post-processing ================= + if self.classifier is not None: + with torch.no_grad(): + t = torch.zeros((n_samples,), dtype=torch.long, device=self.device) + logp = self.classifier.logp(xt, t, condition_vec_cg) + log["log_p"] = logp + + if self.clip_pred: + xt = xt.clip(self.x_min, self.x_max) + + return xt, log + + def ode_dpm_solver_pp_2M( + self, + # ---------- the known fixed portion ---------- # + prior: torch.Tensor, + # ----------------- sampling ----------------- # + n_samples: int = 1, + sample_steps: int = 5, + sample_step_schedule: Union[str, Callable] = "uniform", + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + + # ===================== Initialization ===================== + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + model = self.model if not use_ema else self.model_ema + + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, 0] = xt.cpu().numpy() + + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + # ===================== Sampling Schedule ==================== + if isinstance(sample_step_schedule, str): + if sample_step_schedule in SUPPORTED_SAMPLING_STEP_SCHEDULE.keys(): + sample_step_schedule = SUPPORTED_SAMPLING_STEP_SCHEDULE[sample_step_schedule]( + self.diffusion_steps, sample_steps) + else: + raise ValueError(f"Sampling step schedule {sample_step_schedule} is not supported.") + elif callable(sample_step_schedule): + sample_step_schedule = sample_step_schedule(self.diffusion_steps, sample_steps) + else: + raise ValueError("sample_step_schedule must be a callable or a string") + + alphas = self.alpha[sample_step_schedule] + sigmas = self.sigma[sample_step_schedule] + logSNRs = torch.log(alphas / sigmas) + hs = torch.zeros_like(logSNRs) + hs[1:] = logSNRs[:-1] - logSNRs[1:] # hs[0] is not correctly calculated, but it will not be used. + stds = torch.zeros((sample_steps + 1,), device=self.device) + stds[1:] = sigmas[:-1] / sigmas[1:] * (1 - (alphas[1:] / alphas[:-1]) ** 2).sqrt() + + buffer = [] + + # ===================== Denoising Loop ======================== + + for i in reversed(range(1, sample_steps + 1)): + + t = torch.full((n_samples,), sample_step_schedule[i], dtype=torch.long, device=self.device) + + # guided sampling + pred, logp = self.guided_sampling( + xt, t, alphas[i], sigmas[i], + model, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg, requires_grad) + + # clip the prediction + pred = self.clip_prediction(pred, xt, alphas[i], sigmas[i]) + + # transform to x_theta + x_theta = pred if not self.predict_noise else epstheta_to_xtheta(xt, alphas[i], sigmas[i], pred) + + buffer.append(x_theta) + + # one-step update + if i < sample_steps: + r = hs[i + 1] / hs[i] + D = (1 + 0.5 / r) * buffer[-1] - 0.5 / r * buffer[-2] + xt = (sigmas[i - 1] / sigmas[i]) * xt - alphas[i - 1] * torch.expm1(-hs[i]) * D + else: + xt = (sigmas[i - 1] / sigmas[i]) * xt - alphas[i - 1] * torch.expm1(-hs[i]) * x_theta + + # fix the known portion, and preserve the sampling history + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, sample_steps - i + 1] = xt.cpu().numpy() + + # ================= Post-processing ================= + if self.classifier is not None: + with torch.no_grad(): + t = torch.zeros((n_samples,), dtype=torch.long, device=self.device) + logp = self.classifier.logp(xt, t, condition_vec_cg) + log["log_p"] = logp + + if self.clip_pred: + xt = xt.clip(self.x_min, self.x_max) + + return xt, log + + +class ContinuousVPSDE(DiffusionModel): + + def __init__( + self, + + # ----------------- Neural Networks ----------------- # + nn_diffusion: BaseNNDiffusion, + nn_condition: Optional[BaseNNCondition] = None, + + # ----------------- Masks ----------------- # + # Fix some portion of the input data, and only allow the diffusion model to complete the rest part. + fix_mask: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + # Add loss weight + loss_weight: Union[list, np.ndarray, torch.Tensor] = None, # be in the shape of `x_shape` + + # ------------------ Plugins ---------------- # + # Add a classifier to enable classifier-guidance + classifier: Optional[BaseClassifier] = None, + + # ------------------ Training Params ---------------- # + grad_clip_norm: Optional[float] = None, + ema_rate: float = 0.995, + optim_params: Optional[dict] = None, + + # ------------------- Diffusion Params ------------------- # + epsilon: float = 1e-3, + + discretization: Union[str, Callable] = "uniform", + noise_schedule: Union[str, Dict[str, Callable]] = "cosine", + noise_schedule_params: Optional[dict] = None, + + x_max: Optional[torch.Tensor] = None, + x_min: Optional[torch.Tensor] = None, + + predict_noise: bool = True, + + device: Union[torch.device, str] = "cpu" + ): + super().__init__( + nn_diffusion, nn_condition, fix_mask, loss_weight, classifier, grad_clip_norm, + 0, ema_rate, optim_params, device) + + self.predict_noise = predict_noise + self.epsilon = epsilon + self.x_max, self.x_min = x_max, x_min + + # ==================== Continuous Time-step Range ==================== + if noise_schedule == "cosine": + self.t_diffusion = [epsilon, 0.9946] + else: + self.t_diffusion = [epsilon, 1.] + + # ===================== Noise Schedule ====================== + if isinstance(noise_schedule, str): + if noise_schedule in SUPPORTED_NOISE_SCHEDULES.keys(): + self.noise_schedule_funcs = SUPPORTED_NOISE_SCHEDULES[noise_schedule] + self.noise_schedule_params = noise_schedule_params + else: + raise ValueError(f"Noise schedule {noise_schedule} is not supported.") + elif isinstance(noise_schedule, dict): + self.noise_schedule_funcs = noise_schedule + self.noise_schedule_params = noise_schedule_params + else: + raise ValueError("noise_schedule must be a callable or a string") + + @property + def clip_pred(self): + return (self.x_max is not None) or (self.x_min is not None) + + # ==================== Training: Score Matching ====================== + + def add_noise(self, x0, t=None, eps=None): + + t = (torch.rand((x0.shape[0],), device=self.device) * + (self.t_diffusion[1] - self.t_diffusion[0]) + self.t_diffusion[0]) if t is None else t + + eps = torch.randn_like(x0) if eps is None else eps + + alpha, sigma = self.noise_schedule_funcs["forward"](t, **(self.noise_schedule_params or {})) + alpha = at_least_ndim(alpha, x0.dim()) + sigma = at_least_ndim(sigma, x0.dim()) + + xt = alpha * x0 + sigma * eps + xt = (1. - self.fix_mask) * xt + self.fix_mask * x0 + + return xt, t, eps + + def loss(self, x0, condition=None): + + xt, t, eps = self.add_noise(x0) + + condition = self.model["condition"](condition) if condition is not None else None + + if self.predict_noise: + loss = (self.model["diffusion"](xt, t, condition) - eps) ** 2 + else: + loss = (self.model["diffusion"](xt, t, condition) - x0) ** 2 + + return (loss * self.loss_weight * (1 - self.fix_mask)).mean() + + def update(self, x0, condition=None, update_ema=True, **kwargs): + + loss = self.loss(x0, condition) + + loss.backward() + grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_norm) \ + if self.grad_clip_norm else None + self.optimizer.step() + self.optimizer.zero_grad() + + if update_ema: + self.ema_update() + + log = {"loss": loss.item(), "grad_norm": grad_norm} + + return log + + def update_classifier(self, x0, condition): + + xt, t, eps = self.add_noise(x0) + + log = self.classifier.update(xt, t, condition) + + return log + + # ==================== Sampling: Solving SDE/ODE ====================== + + def classifier_guidance( + self, xt, t, alpha, sigma, + model, condition=None, w: float = 1.0, + pred=None): + """ + Guided Sampling CG: + bar_eps = eps - w * sigma * grad + bar_x0 = x0 + w * (sigma ** 2) * alpha * grad + """ + if pred is None: + pred = model["diffusion"](xt, t, None) + if self.classifier is None or w == 0.0 or condition is None: + return pred, None + else: + log_p, grad = self.classifier.gradients(xt.clone(), t, condition) + if self.predict_noise: + pred = pred - w * sigma * grad + else: + pred = pred + w * ((sigma ** 2) / alpha) * grad + + return pred, log_p + + def classifier_free_guidance( + self, xt, t, + model, condition=None, w: float = 1.0, + pred=None, pred_uncond=None, + requires_grad: bool = False): + """ + Guided Sampling CFG: + bar_eps = w * pred + (1 - w) * pred_uncond + bar_x0 = w * pred + (1 - w) * pred_uncond + """ + with torch.set_grad_enabled(requires_grad): + if w != 0.0 and w != 1.0: + if pred is None or pred_uncond is None: + b = xt.shape[0] + repeat_dim = [2 if i == 0 else 1 for i in range(xt.dim())] + condition = torch.cat([condition, torch.zeros_like(condition)], 0) + pred_all = model["diffusion"]( + xt.repeat(*repeat_dim), t.repeat(2), condition) + pred, pred_uncond = pred_all[:b], pred_all[b:] + elif w == 0.0: + pred = 0. + pred_uncond = model["diffusion"](xt, t, None) + else: + pred = model["diffusion"](xt, t, condition) + pred_uncond = 0. + + if self.predict_noise or not self.predict_noise: + bar_pred = w * pred + (1 - w) * pred_uncond + else: + bar_pred = pred + + return bar_pred + + def clip_prediction(self, pred, xt, alpha, sigma): + """ + Clip the prediction at each sampling step to stablize the generation. + (xt - alpha * x_max) / sigma <= eps <= (xt - alpha * x_min) / sigma + x_min <= x0 <= x_max + """ + if self.clip_pred: + if self.predict_noise: + upper_bound = (xt - alpha * self.x_min) / sigma if self.x_min is not None else None + lower_bound = (xt - alpha * self.x_max) / sigma if self.x_max is not None else None + pred = pred.clip(lower_bound, upper_bound) + else: + pred = pred.clip(self.x_min, self.x_max) + return pred + + def guided_sampling( + self, xt, t, alpha, sigma, + model, + condition_cfg=None, w_cfg: float = 0.0, + condition_cg=None, w_cg: float = 0.0, + requires_grad: bool = False): + """ + One-step epsilon/x0 prediction with guidance. + """ + + pred = self.classifier_free_guidance( + xt, t, model, condition_cfg, w_cfg, None, None, requires_grad) + + pred, logp = self.classifier_guidance( + xt, t, alpha, sigma, model, condition_cg, w_cg, pred) + + return pred, logp + + def sample( + self, + # ---------- the known fixed portion ---------- # + prior: torch.Tensor, + # ----------------- sampling ----------------- # + solver: str = "ddpm", + n_samples: int = 1, + sample_steps: int = 5, + sample_step_schedule: Union[str, Callable] = "uniform_continuous", + use_ema: bool = True, + temperature: float = 1.0, + # ------------------ guidance ------------------ # + condition_cfg=None, + mask_cfg=None, + w_cfg: float = 0.0, + condition_cg=None, + w_cg: float = 0.0, + # ------------------ others ------------------ # + requires_grad: bool = False, + preserve_history: bool = False, + **kwargs, + ): + assert solver in SUPPORTED_SOLVERS, f"Solver {solver} is not supported." + + # ===================== Initialization ===================== + log = { + "sample_history": np.empty((n_samples, sample_steps + 1, *prior.shape)) if preserve_history else None, } + + model = self.model if not use_ema else self.model_ema9999 + + xt = torch.randn_like(prior, device=self.device) * temperature + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, 0] = xt.cpu().numpy() + + with torch.set_grad_enabled(requires_grad): + condition_vec_cfg = model["condition"](condition_cfg, mask_cfg) if condition_cfg is not None else None + condition_vec_cg = condition_cg + + # ===================== Sampling Schedule ==================== + if isinstance(sample_step_schedule, str): + if sample_step_schedule in SUPPORTED_SAMPLING_STEP_SCHEDULE.keys(): + sample_step_schedule = SUPPORTED_SAMPLING_STEP_SCHEDULE[sample_step_schedule]( + self.t_diffusion, sample_steps) + else: + raise ValueError(f"Sampling step schedule {sample_step_schedule} is not supported.") + elif callable(sample_step_schedule): + sample_step_schedule = sample_step_schedule(self.t_diffusion, sample_steps) + else: + raise ValueError("sample_step_schedule must be a callable or a string") + + alphas, sigmas = self.noise_schedule_funcs["forward"]( + sample_step_schedule, **(self.noise_schedule_params or {})) + logSNRs = torch.log(alphas / sigmas) + hs = torch.zeros_like(logSNRs) + hs[1:] = logSNRs[:-1] - logSNRs[1:] # hs[0] is not correctly calculated, but it will not be used. + stds = torch.zeros((sample_steps + 1,), device=self.device) + stds[1:] = sigmas[:-1] / sigmas[1:] * (1 - (alphas[1:] / alphas[:-1]) ** 2).sqrt() + + buffer = [] + + # ===================== Denoising Loop ======================== + + for i in reversed(range(1, sample_steps + 1)): + + t = torch.full((n_samples,), sample_step_schedule[i], dtype=torch.float32, device=self.device) + + # guided sampling + pred, logp = self.guided_sampling( + xt, t, alphas[i], sigmas[i], + model, condition_vec_cfg, w_cfg, condition_vec_cg, w_cg, requires_grad) + + # clip the prediction + pred = self.clip_prediction(pred, xt, alphas[i], sigmas[i]) + + # transform to eps_theta + eps_theta = pred if self.predict_noise else xtheta_to_epstheta(xt, alphas[i], sigmas[i], pred) + x_theta = pred if not self.predict_noise else epstheta_to_xtheta(xt, alphas[i], sigmas[i], pred) + + # one-step update + if solver == "ddpm": + xt = ( + (alphas[i - 1] / alphas[i]) * (xt - sigmas[i] * eps_theta) + + (sigmas[i - 1] ** 2 - stds[i] ** 2 + 1e-8).sqrt() * eps_theta) + if i > 1: + xt += (stds[i] * torch.randn_like(xt)) + + elif solver == "ddim": + xt = (alphas[i - 1] * ((xt - sigmas[i] * eps_theta) / alphas[i]) + sigmas[i - 1] * eps_theta) + + elif solver == "ode_dpmsolver++_1": + xt = (sigmas[i - 1] / sigmas[i]) * xt - alphas[i - 1] * torch.expm1(-hs[i]) * x_theta + + elif solver == "sde_dpmsolver++_1": + xt = ((sigmas[i - 1] / sigmas[i]) * (-hs[i]).exp() * xt - + alphas[i - 1] * torch.expm1(-2 * hs[i]) * x_theta + + sigmas[i - 1] * (-torch.expm1(-2 * hs[i])).sqrt() * torch.randn_like(xt)) + + elif solver == "ode_dpmsolver++_2M": + buffer.append(x_theta) + if i < sample_steps: + r = hs[i + 1] / hs[i] + D = (1 + 0.5 / r) * buffer[-1] - 0.5 / r * buffer[-2] + xt = (sigmas[i - 1] / sigmas[i]) * xt - alphas[i - 1] * torch.expm1(-hs[i]) * D + else: + xt = (sigmas[i - 1] / sigmas[i]) * xt - alphas[i - 1] * torch.expm1(-hs[i]) * x_theta + + # fix the known portion, and preserve the sampling history + xt = xt * (1. - self.fix_mask) + prior * self.fix_mask + if preserve_history: + log["sample_history"][:, sample_steps - i + 1] = xt.cpu().numpy() + + # ================= Post-processing ================= + if self.classifier is not None: + with torch.no_grad(): + t = torch.zeros((n_samples,), dtype=torch.long, device=self.device) + logp = self.classifier.logp(xt, t, condition_vec_cg) + log["log_p"] = logp + + if self.clip_pred: + xt = xt.clip(self.x_min, self.x_max) + + return xt, log diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/async_vector_env.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/async_vector_env.py new file mode 100644 index 0000000000000000000000000000000000000000..dfb0f620ac3d39861e0dfc1a3a58fb22d443f244 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/async_vector_env.py @@ -0,0 +1,671 @@ +""" +Back ported methods: call, set_attr from v0.26 +Disabled auto-reset after done +Added render method. +""" + + +import numpy as np +import multiprocessing as mp +import time +import sys +from enum import Enum +from copy import deepcopy + +from gym import logger +from gym.vector.vector_env import VectorEnv +from gym.error import ( + AlreadyPendingCallError, + NoAsyncCallError, + ClosedEnvironmentError, + CustomSpaceError, +) +from gym.vector.utils import ( + create_shared_memory, + create_empty_array, + write_to_shared_memory, + read_from_shared_memory, + concatenate, + CloudpickleWrapper, + clear_mpi_env_vars, +) + +__all__ = ["AsyncVectorEnv"] + + +class AsyncState(Enum): + DEFAULT = "default" + WAITING_RESET = "reset" + WAITING_STEP = "step" + WAITING_CALL = "call" + + +class AsyncVectorEnv(VectorEnv): + """Vectorized environment that runs multiple environments in parallel. It + uses `multiprocessing` processes, and pipes for communication. + Parameters + ---------- + env_fns : iterable of callable + Functions that create the environments. + observation_space : `gym.spaces.Space` instance, optional + Observation space of a single environment. If `None`, then the + observation space of the first environment is taken. + action_space : `gym.spaces.Space` instance, optional + Action space of a single environment. If `None`, then the action space + of the first environment is taken. + shared_memory : bool (default: `True`) + If `True`, then the observations from the worker processes are + communicated back through shared variables. This can improve the + efficiency if the observations are large (e.g. images). + copy : bool (default: `True`) + If `True`, then the `reset` and `step` methods return a copy of the + observations. + context : str, optional + Context for multiprocessing. If `None`, then the default context is used. + Only available in Python 3. + daemon : bool (default: `True`) + If `True`, then subprocesses have `daemon` flag turned on; that is, they + will quit if the head process quits. However, `daemon=True` prevents + subprocesses to spawn children, so for some environments you may want + to have it set to `False` + worker : function, optional + WARNING - advanced mode option! If set, then use that worker in a subprocess + instead of a default one. Can be useful to override some inner vector env + logic, for instance, how resets on done are handled. Provides high + degree of flexibility and a high chance to shoot yourself in the foot; thus, + if you are writing your own worker, it is recommended to start from the code + for `_worker` (or `_worker_shared_memory`) method below, and add changes + """ + + def __init__( + self, + env_fns, + dummy_env_fn=None, + observation_space=None, + action_space=None, + shared_memory=True, + copy=True, + context=None, + daemon=True, + worker=None, + ): + ctx = mp.get_context(context) + self.env_fns = env_fns + self.shared_memory = shared_memory + self.copy = copy + + # Added dummy_env_fn to fix OpenGL error in Mujoco + # disable any OpenGL rendering in dummy_env_fn, since it + # will conflict with OpenGL context in the forked child process + if dummy_env_fn is None: + dummy_env_fn = env_fns[0] + dummy_env = dummy_env_fn() + self.metadata = dummy_env.metadata + + if (observation_space is None) or (action_space is None): + observation_space = observation_space or dummy_env.observation_space + action_space = action_space or dummy_env.action_space + dummy_env.close() + del dummy_env + super(AsyncVectorEnv, self).__init__( + num_envs=len(env_fns), + observation_space=observation_space, + action_space=action_space, + ) + + if self.shared_memory: + try: + _obs_buffer = create_shared_memory( + self.single_observation_space, n=self.num_envs, ctx=ctx + ) + self.observations = read_from_shared_memory( + _obs_buffer, self.single_observation_space, n=self.num_envs + ) + except CustomSpaceError: + raise ValueError( + "Using `shared_memory=True` in `AsyncVectorEnv` " + "is incompatible with non-standard Gym observation spaces " + "(i.e. custom spaces inheriting from `gym.Space`), and is " + "only compatible with default Gym spaces (e.g. `Box`, " + "`Tuple`, `Dict`) for batching. Set `shared_memory=False` " + "if you use custom observation spaces." + ) + else: + _obs_buffer = None + self.observations = create_empty_array( + self.single_observation_space, n=self.num_envs, fn=np.zeros + ) + + self.parent_pipes, self.processes = [], [] + self.error_queue = ctx.Queue() + target = _worker_shared_memory if self.shared_memory else _worker + target = worker or target + with clear_mpi_env_vars(): + for idx, env_fn in enumerate(self.env_fns): + parent_pipe, child_pipe = ctx.Pipe() + process = ctx.Process( + target=target, + name="Worker<{0}>-{1}".format(type(self).__name__, idx), + args=( + idx, + CloudpickleWrapper(env_fn), + child_pipe, + parent_pipe, + _obs_buffer, + self.error_queue, + ), + ) + + self.parent_pipes.append(parent_pipe) + self.processes.append(process) + + process.daemon = daemon + process.start() + child_pipe.close() + + self._state = AsyncState.DEFAULT + self._check_observation_spaces() + + def seed(self, seeds=None): + self._assert_is_running() + if seeds is None: + seeds = [None for _ in range(self.num_envs)] + if isinstance(seeds, int): + seeds = [seeds + i for i in range(self.num_envs)] + assert len(seeds) == self.num_envs + + if self._state != AsyncState.DEFAULT: + raise AlreadyPendingCallError( + "Calling `seed` while waiting " + "for a pending call to `{0}` to complete.".format(self._state.value), + self._state.value, + ) + + for pipe, seed in zip(self.parent_pipes, seeds): + pipe.send(("seed", seed)) + _, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) + self._raise_if_errors(successes) + + def reset_async(self): + self._assert_is_running() + if self._state != AsyncState.DEFAULT: + raise AlreadyPendingCallError( + "Calling `reset_async` while waiting " + "for a pending call to `{0}` to complete".format(self._state.value), + self._state.value, + ) + + for pipe in self.parent_pipes: + pipe.send(("reset", None)) + self._state = AsyncState.WAITING_RESET + + def reset_wait(self, timeout=None): + """ + Parameters + ---------- + timeout : int or float, optional + Number of seconds before the call to `reset_wait` times out. If + `None`, the call to `reset_wait` never times out. + Returns + ------- + observations : sample from `observation_space` + A batch of observations from the vectorized environment. + """ + self._assert_is_running() + if self._state != AsyncState.WAITING_RESET: + raise NoAsyncCallError( + "Calling `reset_wait` without any prior " "call to `reset_async`.", + AsyncState.WAITING_RESET.value, + ) + + if not self._poll(timeout): + self._state = AsyncState.DEFAULT + raise mp.TimeoutError( + "The call to `reset_wait` has timed out after " + "{0} second{1}.".format(timeout, "s" if timeout > 1 else "") + ) + + results, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) + self._raise_if_errors(successes) + self._state = AsyncState.DEFAULT + + if not self.shared_memory: + self.observations = concatenate( + results, self.observations, self.single_observation_space + ) + + return deepcopy(self.observations) if self.copy else self.observations + + def step_async(self, actions): + """ + Parameters + ---------- + actions : iterable of samples from `action_space` + List of actions. + """ + self._assert_is_running() + if self._state != AsyncState.DEFAULT: + raise AlreadyPendingCallError( + "Calling `step_async` while waiting " + "for a pending call to `{0}` to complete.".format(self._state.value), + self._state.value, + ) + + for pipe, action in zip(self.parent_pipes, actions): + pipe.send(("step", action)) + self._state = AsyncState.WAITING_STEP + + def step_wait(self, timeout=None): + """ + Parameters + ---------- + timeout : int or float, optional + Number of seconds before the call to `step_wait` times out. If + `None`, the call to `step_wait` never times out. + Returns + ------- + observations : sample from `observation_space` + A batch of observations from the vectorized environment. + rewards : `np.ndarray` instance (dtype `np.float_`) + A vector of rewards from the vectorized environment. + dones : `np.ndarray` instance (dtype `np.bool_`) + A vector whose entries indicate whether the episode has ended. + infos : list of dict + A list of auxiliary diagnostic information. + """ + self._assert_is_running() + if self._state != AsyncState.WAITING_STEP: + raise NoAsyncCallError( + "Calling `step_wait` without any prior call " "to `step_async`.", + AsyncState.WAITING_STEP.value, + ) + + if not self._poll(timeout): + self._state = AsyncState.DEFAULT + raise mp.TimeoutError( + "The call to `step_wait` has timed out after " + "{0} second{1}.".format(timeout, "s" if timeout > 1 else "") + ) + + results, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) + self._raise_if_errors(successes) + self._state = AsyncState.DEFAULT + observations_list, rewards, dones, infos = zip(*results) + + if not self.shared_memory: + self.observations = concatenate( + observations_list, self.observations, self.single_observation_space + ) + + return ( + deepcopy(self.observations) if self.copy else self.observations, + np.array(rewards), + np.array(dones, dtype=np.bool_), + infos, + ) + + def close_extras(self, timeout=None, terminate=False): + """ + Parameters + ---------- + timeout : int or float, optional + Number of seconds before the call to `close` times out. If `None`, + the call to `close` never times out. If the call to `close` times + out, then all processes are terminated. + terminate : bool (default: `False`) + If `True`, then the `close` operation is forced and all processes + are terminated. + """ + timeout = 0 if terminate else timeout + try: + if self._state != AsyncState.DEFAULT: + logger.warn( + "Calling `close` while waiting for a pending " + "call to `{0}` to complete.".format(self._state.value) + ) + function = getattr(self, "{0}_wait".format(self._state.value)) + function(timeout) + except mp.TimeoutError: + terminate = True + + if terminate: + for process in self.processes: + if process.is_alive(): + process.terminate() + else: + for pipe in self.parent_pipes: + if (pipe is not None) and (not pipe.closed): + pipe.send(("close", None)) + for pipe in self.parent_pipes: + if (pipe is not None) and (not pipe.closed): + pipe.recv() + + for pipe in self.parent_pipes: + if pipe is not None: + pipe.close() + for process in self.processes: + process.join() + + def _poll(self, timeout=None): + self._assert_is_running() + if timeout is None: + return True + end_time = time.perf_counter() + timeout + delta = None + for pipe in self.parent_pipes: + delta = max(end_time - time.perf_counter(), 0) + if pipe is None: + return False + if pipe.closed or (not pipe.poll(delta)): + return False + return True + + def _check_observation_spaces(self): + self._assert_is_running() + for pipe in self.parent_pipes: + pipe.send(("_check_observation_space", self.single_observation_space)) + same_spaces, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) + self._raise_if_errors(successes) + if not all(same_spaces): + raise RuntimeError( + "Some environments have an observation space " + "different from `{0}`. In order to batch observations, the " + "observation spaces from all environments must be " + "equal.".format(self.single_observation_space) + ) + + def _assert_is_running(self): + if self.closed: + raise ClosedEnvironmentError( + "Trying to operate on `{0}`, after a " + "call to `close()`.".format(type(self).__name__) + ) + + def _raise_if_errors(self, successes): + if all(successes): + return + + num_errors = self.num_envs - sum(successes) + assert num_errors > 0 + for _ in range(num_errors): + index, exctype, value = self.error_queue.get() + logger.error( + "Received the following error from Worker-{0}: " + "{1}: {2}".format(index, exctype.__name__, value) + ) + logger.error("Shutting down Worker-{0}.".format(index)) + self.parent_pipes[index].close() + self.parent_pipes[index] = None + + logger.error("Raising the last exception back to the main process.") + raise exctype(value) + + def call_async(self, name: str, *args, **kwargs): + """Calls the method with name asynchronously and apply args and kwargs to the method. + + Args: + name: Name of the method or property to call. + *args: Arguments to apply to the method call. + **kwargs: Keyword arguments to apply to the method call. + + Raises: + ClosedEnvironmentError: If the environment was closed (if :meth:`close` was previously called). + AlreadyPendingCallError: Calling `call_async` while waiting for a pending call to complete + """ + self._assert_is_running() + if self._state != AsyncState.DEFAULT: + raise AlreadyPendingCallError( + "Calling `call_async` while waiting " + f"for a pending call to `{self._state.value}` to complete.", + self._state.value, + ) + + for pipe in self.parent_pipes: + pipe.send(("_call", (name, args, kwargs))) + self._state = AsyncState.WAITING_CALL + + def call_wait(self, timeout = None) -> list: + """Calls all parent pipes and waits for the results. + + Args: + timeout: Number of seconds before the call to `step_wait` times out. + If `None` (default), the call to `step_wait` never times out. + + Returns: + List of the results of the individual calls to the method or property for each environment. + + Raises: + NoAsyncCallError: Calling `call_wait` without any prior call to `call_async`. + TimeoutError: The call to `call_wait` has timed out after timeout second(s). + """ + self._assert_is_running() + if self._state != AsyncState.WAITING_CALL: + raise NoAsyncCallError( + "Calling `call_wait` without any prior call to `call_async`.", + AsyncState.WAITING_CALL.value, + ) + + if not self._poll(timeout): + self._state = AsyncState.DEFAULT + raise mp.TimeoutError( + f"The call to `call_wait` has timed out after {timeout} second(s)." + ) + + results, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) + self._raise_if_errors(successes) + self._state = AsyncState.DEFAULT + + return results + + def call(self, name: str, *args, **kwargs): + """Call a method, or get a property, from each parallel environment. + + Args: + name (str): Name of the method or property to call. + *args: Arguments to apply to the method call. + **kwargs: Keyword arguments to apply to the method call. + + Returns: + List of the results of the individual calls to the method or property for each environment. + """ + self.call_async(name, *args, **kwargs) + return self.call_wait() + + + def call_each(self, name: str, + args_list: list=None, + kwargs_list: list=None, + timeout = None): + n_envs = len(self.parent_pipes) + if args_list is None: + args_list = [[]] * n_envs + assert len(args_list) == n_envs + + if kwargs_list is None: + kwargs_list = [dict()] * n_envs + assert len(kwargs_list) == n_envs + + # send + self._assert_is_running() + if self._state != AsyncState.DEFAULT: + raise AlreadyPendingCallError( + "Calling `call_async` while waiting " + f"for a pending call to `{self._state.value}` to complete.", + self._state.value, + ) + + for i, pipe in enumerate(self.parent_pipes): + pipe.send(("_call", (name, args_list[i], kwargs_list[i]))) + self._state = AsyncState.WAITING_CALL + + # receive + self._assert_is_running() + if self._state != AsyncState.WAITING_CALL: + raise NoAsyncCallError( + "Calling `call_wait` without any prior call to `call_async`.", + AsyncState.WAITING_CALL.value, + ) + + if not self._poll(timeout): + self._state = AsyncState.DEFAULT + raise mp.TimeoutError( + f"The call to `call_wait` has timed out after {timeout} second(s)." + ) + + results, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) + self._raise_if_errors(successes) + self._state = AsyncState.DEFAULT + + return results + + + def set_attr(self, name: str, values): + """Sets an attribute of the sub-environments. + + Args: + name: Name of the property to be set in each individual environment. + values: Values of the property to be set to. If ``values`` is a list or + tuple, then it corresponds to the values for each individual + environment, otherwise a single value is set for all environments. + + Raises: + ValueError: Values must be a list or tuple with length equal to the number of environments. + AlreadyPendingCallError: Calling `set_attr` while waiting for a pending call to complete. + """ + self._assert_is_running() + if not isinstance(values, (list, tuple)): + values = [values for _ in range(self.num_envs)] + if len(values) != self.num_envs: + raise ValueError( + "Values must be a list or tuple with length equal to the " + f"number of environments. Got `{len(values)}` values for " + f"{self.num_envs} environments." + ) + + if self._state != AsyncState.DEFAULT: + raise AlreadyPendingCallError( + "Calling `set_attr` while waiting " + f"for a pending call to `{self._state.value}` to complete.", + self._state.value, + ) + + for pipe, value in zip(self.parent_pipes, values): + pipe.send(("_setattr", (name, value))) + _, successes = zip(*[pipe.recv() for pipe in self.parent_pipes]) + self._raise_if_errors(successes) + + def render(self, *args, **kwargs): + return self.call('render', *args, **kwargs) + + + +def _worker(index, env_fn, pipe, parent_pipe, shared_memory, error_queue): + assert shared_memory is None + env = env_fn() + parent_pipe.close() + try: + while True: + command, data = pipe.recv() + if command == "reset": + observation = env.reset() + pipe.send((observation, True)) + elif command == "step": + observation, reward, done, info = env.step(data) + # if done: + # observation = env.reset() + pipe.send(((observation, reward, done, info), True)) + elif command == "seed": + env.seed(data) + pipe.send((None, True)) + elif command == "close": + pipe.send((None, True)) + break + elif command == "_call": + name, args, kwargs = data + if name in ["reset", "step", "seed", "close"]: + raise ValueError( + f"Trying to call function `{name}` with " + f"`_call`. Use `{name}` directly instead." + ) + function = getattr(env, name) + if callable(function): + pipe.send((function(*args, **kwargs), True)) + else: + pipe.send((function, True)) + elif command == "_setattr": + name, value = data + setattr(env, name, value) + pipe.send((None, True)) + + elif command == "_check_observation_space": + pipe.send((data == env.observation_space, True)) + else: + raise RuntimeError( + "Received unknown command `{0}`. Must " + "be one of {`reset`, `step`, `seed`, `close`, " + "`_check_observation_space`}.".format(command) + ) + except (KeyboardInterrupt, Exception): + error_queue.put((index,) + sys.exc_info()[:2]) + pipe.send((None, False)) + finally: + env.close() + + +def _worker_shared_memory(index, env_fn, pipe, parent_pipe, shared_memory, error_queue): + assert shared_memory is not None + env = env_fn() + observation_space = env.observation_space + parent_pipe.close() + try: + while True: + command, data = pipe.recv() + if command == "reset": + observation = env.reset() + write_to_shared_memory( + index, observation, shared_memory, observation_space + ) + pipe.send((None, True)) + elif command == "step": + observation, reward, done, info = env.step(data) + # if done: + # observation = env.reset() + write_to_shared_memory( + index, observation, shared_memory, observation_space + ) + pipe.send(((None, reward, done, info), True)) + elif command == "seed": + env.seed(data) + pipe.send((None, True)) + elif command == "close": + pipe.send((None, True)) + break + elif command == "_call": + name, args, kwargs = data + if name in ["reset", "step", "seed", "close"]: + raise ValueError( + f"Trying to call function `{name}` with " + f"`_call`. Use `{name}` directly instead." + ) + function = getattr(env, name) + if callable(function): + pipe.send((function(*args, **kwargs), True)) + else: + pipe.send((function, True)) + elif command == "_setattr": + name, value = data + setattr(env, name, value) + pipe.send((None, True)) + elif command == "_check_observation_space": + pipe.send((data == observation_space, True)) + else: + raise RuntimeError( + "Received unknown command `{0}`. Must " + "be one of {`reset`, `step`, `seed`, `close`, " + "`_check_observation_space`}.".format(command) + ) + except (KeyboardInterrupt, Exception): + error_queue.put((index,) + sys.exc_info()[:2]) + pipe.send((None, False)) + finally: + env.close() \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..307ec5372581a0ca14dd9d1db063e4997c01c229 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/__init__.py @@ -0,0 +1,7 @@ +from gym.envs.registration import register + +register( + id='block-push-multi-modal-v0', + entry_point='cleandiffuser.env.block_pushing.block_pushing_multimodal:BlockPushMultimodal', + max_episode_steps=350, +) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/block_pushing.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/block_pushing.py new file mode 100644 index 0000000000000000000000000000000000000000..6f78f2f5595abc44cd436cb4b9beb78be2233300 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/block_pushing.py @@ -0,0 +1,1092 @@ +# coding=utf-8 +# Copyright 2022 The Reach ML Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Simple block environments for the XArm.""" + +import collections +import enum +import math +import time +from typing import Dict, List, Optional, Tuple, Union + +import gym +from gym import spaces +from gym.envs import registration +from cleandiffuser.env.block_pushing.utils import utils_pybullet +from cleandiffuser.env.block_pushing.utils import xarm_sim_robot +from cleandiffuser.env.block_pushing.utils.pose3d import Pose3d +from cleandiffuser.env.block_pushing.utils.utils_pybullet import ObjState +from cleandiffuser.env.block_pushing.utils.utils_pybullet import XarmState +import numpy as np +from scipy.spatial import transform +import pybullet +import pybullet_utils.bullet_client as bullet_client + +import matplotlib.pyplot as plt + +BLOCK_URDF_PATH = "third_party/py/envs/assets/block.urdf" +PLANE_URDF_PATH = "third_party/bullet/examples/pybullet/gym/pybullet_data/" "plane.urdf" +WORKSPACE_URDF_PATH = "third_party/py/envs/assets/workspace.urdf" +ZONE_URDF_PATH = "third_party/py/envs/assets/zone.urdf" +INSERT_URDF_PATH = "third_party/py/envs/assets/insert.urdf" + +EFFECTOR_HEIGHT = 0.06 +EFFECTOR_DOWN_ROTATION = transform.Rotation.from_rotvec([0, math.pi, 0]) + +WORKSPACE_BOUNDS = np.array(((0.15, -0.5), (0.7, 0.5))) + +# Min/max bounds calculated from oracle data using: +# ibc/environments/board2d_dataset_statistics.ipynb +# to calculate [mean - 3 * std, mean + 3 * std] using the oracle data. +# pylint: disable=line-too-long +ACTION_MIN = np.array([-0.02547718, -0.02090043], np.float32) +ACTION_MAX = np.array([0.02869084, 0.04272365], np.float32) +EFFECTOR_TARGET_TRANSLATION_MIN = np.array( + [0.1774151772260666, -0.6287994794547558], np.float32 +) +EFFECTOR_TARGET_TRANSLATION_MAX = np.array( + [0.5654461532831192, 0.5441607423126698], np.float32 +) +EFFECTOR_TARGET_TO_BLOCK_TRANSLATION_MIN = np.array( + [-0.07369826920330524, -0.11395704373717308], np.float32 +) +EFFECTOR_TARGET_TO_BLOCK_TRANSLATION_MAX = np.array( + [0.10131562314927578, 0.19391131028532982], np.float32 +) +EFFECTOR_TARGET_TO_TARGET_TRANSLATION_MIN = np.array( + [-0.17813862301409245, -0.3309651017189026], np.float32 +) +EFFECTOR_TARGET_TO_TARGET_TRANSLATION_MAX = np.array( + [0.23726161383092403, 0.8404090404510498], np.float32 +) +BLOCK_ORIENTATION_COS_SIN_MIN = np.array( + [-2.0649861991405487, -0.6154364347457886], np.float32 +) +BLOCK_ORIENTATION_COS_SIN_MAX = np.array( + [1.6590178310871124, 1.8811014890670776], np.float32 +) +TARGET_ORIENTATION_COS_SIN_MIN = np.array( + [-1.0761439241468906, -0.8846937336493284], np.float32 +) +TARGET_ORIENTATION_COS_SIN_MAX = np.array( + [-0.8344330154359341, 0.8786859593819827], np.float32 +) + +# Hardcoded Pose joints to make sure we don't have surprises from using the +# IK solver on reset. The joint poses correspond to the Pose with: +# rotation = rotation3.Rotation3.from_axis_angle([0, 1, 0], math.pi) +# translation = np.array([0.3, -0.4, 0.07]) +INITIAL_JOINT_POSITIONS = np.array( + [ + -0.9254632489674508, + 0.6990770671568564, + -1.106629064060494, + 0.0006653351931553931, + 0.3987969742311386, + -4.063402065624296, + ] +) + +DEFAULT_CAMERA_POSE = (1.0, 0, 0.75) +DEFAULT_CAMERA_ORIENTATION = (np.pi / 4, np.pi, -np.pi / 2) +IMAGE_WIDTH = 320 +IMAGE_HEIGHT = 240 +CAMERA_INTRINSICS = ( + 0.803 * IMAGE_WIDTH, # fx + 0, + IMAGE_WIDTH / 2.0, # cx + 0, + 0.803 * IMAGE_WIDTH, # fy + IMAGE_HEIGHT / 2.0, # cy + 0, + 0, + 1, +) + +# "Realistic" visuals. +X_MIN_REAL = 0.15 +X_MAX_REAL = 0.6 +Y_MIN_REAL = -0.3048 +Y_MAX_REAL = 0.3048 +WORKSPACE_BOUNDS_REAL = np.array(((X_MIN_REAL, Y_MIN_REAL), (X_MAX_REAL, Y_MAX_REAL))) +WORKSPACE_URDF_PATH_REAL = "third_party/py/ibc/environments/assets/workspace_real.urdf" +CAMERA_POSE_REAL = (0.75, 0, 0.5) +CAMERA_ORIENTATION_REAL = (np.pi / 5, np.pi, -np.pi / 2) + +IMAGE_WIDTH_REAL = 320 +IMAGE_HEIGHT_REAL = 180 +CAMERA_INTRINSICS_REAL = ( + 0.803 * IMAGE_WIDTH_REAL, # fx + 0, + IMAGE_WIDTH_REAL / 2.0, # cx + 0, + 0.803 * IMAGE_WIDTH_REAL, # fy + IMAGE_HEIGHT_REAL / 2.0, # cy + 0, + 0, + 1, +) +# pylint: enable=line-too-long + + +def build_env_name(task, shared_memory, use_image_obs, use_normalized_env=False): + """Construct the env name from parameters.""" + if isinstance(task, str): + task = BlockTaskVariant[task] + env_name = "Block" + task.value + + if use_image_obs: + env_name = env_name + "Rgb" + if use_normalized_env: + env_name = env_name + "Normalized" + if shared_memory: + env_name = "Shared" + env_name + + env_name = env_name + "-v0" + + return env_name + + +class BlockTaskVariant(enum.Enum): + REACH = "Reach" + REACH_NORMALIZED = "ReachNormalized" + PUSH = "Push" + PUSH_NORMALIZED = "PushNormalized" + INSERT = "Insert" + + +def sleep_spin(sleep_time_sec): + """Spin wait sleep. Avoids time.sleep accuracy issues on Windows.""" + if sleep_time_sec <= 0: + return + t0 = time.perf_counter() + while time.perf_counter() - t0 < sleep_time_sec: + pass + + +class BlockPush(gym.Env): + """Simple XArm environment for block pushing.""" + + def __init__( + self, + control_frequency=10.0, + task=BlockTaskVariant.PUSH, + image_size=None, + shared_memory=False, + seed=None, + goal_dist_tolerance=0.01, + effector_height=None, + visuals_mode="default", + abs_action=False + ): + """Creates an env instance. + + Args: + control_frequency: Control frequency for the arm. Each env step will + advance the simulation by 1/control_frequency seconds. + task: enum for which task, see BlockTaskVariant enum. + image_size: Optional image size (height, width). If None, no image + observations will be used. + shared_memory: If True `pybullet.SHARED_MEMORY` is used to connect to + pybullet. Useful to debug. + seed: Optional seed for the environment. + goal_dist_tolerance: float, how far away from the goal to terminate. + effector_height: float, custom height for end effector. + visuals_mode: 'default' or 'real'. + """ + # pybullet.connect(pybullet.GUI) + # pybullet.resetDebugVisualizerCamera( + # cameraDistance=1.5, + # cameraYaw=0, + # cameraPitch=-40, + # cameraTargetPosition=[0.55, -0.35, 0.2], + # ) + if visuals_mode != "default" and visuals_mode != "real": + raise ValueError("visuals_mode must be `real` or `default`.") + self._task = task + self._connection_mode = pybullet.DIRECT + if shared_memory: + self._connection_mode = pybullet.SHARED_MEMORY + + self.goal_dist_tolerance = goal_dist_tolerance + + self.effector_height = effector_height or EFFECTOR_HEIGHT + + self._visuals_mode = visuals_mode + if visuals_mode == "default": + self._camera_pose = DEFAULT_CAMERA_POSE + self._camera_orientation = DEFAULT_CAMERA_ORIENTATION + self.workspace_bounds = WORKSPACE_BOUNDS + self._image_size = image_size + self._camera_instrinsics = CAMERA_INTRINSICS + self._workspace_urdf_path = WORKSPACE_URDF_PATH + else: + self._camera_pose = CAMERA_POSE_REAL + self._camera_orientation = CAMERA_ORIENTATION_REAL + self.workspace_bounds = WORKSPACE_BOUNDS_REAL + self._image_size = image_size + self._camera_instrinsics = CAMERA_INTRINSICS_REAL + self._workspace_urdf_path = WORKSPACE_URDF_PATH_REAL + + self.action_space = spaces.Box(low=-0.1, high=0.1, shape=(2,)) # x, y + self.observation_space = self._create_observation_space(image_size) + + self._rng = np.random.RandomState(seed=seed) + self._block_ids = None + self._previous_state = None + self._robot = None + self._workspace_uid = None + self._target_id = None + self._target_pose = None + self._target_effector_pose = None + self._pybullet_client = None + self.reach_target_translation = None + self._setup_pybullet_scene() + self._saved_state = None + + assert isinstance(self._pybullet_client, bullet_client.BulletClient) + self._control_frequency = control_frequency + self._step_frequency = ( + 1 / self._pybullet_client.getPhysicsEngineParameters()["fixedTimeStep"] + ) + + self._last_loop_time = None + self._last_loop_frame_sleep_time = None + if self._step_frequency % self._control_frequency != 0: + raise ValueError( + "Control frequency should be a multiple of the " + "configured Bullet TimeStep." + ) + self._sim_steps_per_step = int(self._step_frequency / self._control_frequency) + + self.rendered_img = None + self._abs_action = abs_action + + # Use saved_state and restore to make reset safe as no simulation state has + # been updated at this state, but the assets are now loaded. + self.save_state() + self.reset() + + @property + def pybullet_client(self): + return self._pybullet_client + + @property + def robot(self): + return self._robot + + @property + def workspace_uid(self): + return self._workspace_uid + + @property + def target_effector_pose(self): + return self._target_effector_pose + + @property + def target_pose(self): + return self._target_pose + + @property + def control_frequency(self): + return self._control_frequency + + @property + def connection_mode(self): + return self._connection_mode + + def save_state(self): + self._saved_state = self._pybullet_client.saveState() + + def set_goal_dist_tolerance(self, val): + self.goal_dist_tolerance = val + + def get_control_frequency(self): + return self._control_frequency + + def compute_state(self): + return self._compute_state() + + def get_goal_translation(self): + """Return the translation component of the goal (2D).""" + if self._task == BlockTaskVariant.REACH: + return np.concatenate([self.reach_target_translation, [0]]) + else: + return self._target_pose.translation if self._target_pose else None + + def get_obj_ids(self): + return self._block_ids + + def _setup_workspace_and_robot(self, end_effector="suction"): + self._pybullet_client.resetSimulation() + self._pybullet_client.configureDebugVisualizer(pybullet.COV_ENABLE_GUI, 0) + self._pybullet_client.setPhysicsEngineParameter(enableFileCaching=0) + self._pybullet_client.setGravity(0, 0, -9.8) + + utils_pybullet.load_urdf( + self._pybullet_client, PLANE_URDF_PATH, basePosition=[0, 0, -0.001] + ) + self._workspace_uid = utils_pybullet.load_urdf( + self._pybullet_client, + self._workspace_urdf_path, + basePosition=[0.35, 0, 0.0], + ) + + self._robot = xarm_sim_robot.XArmSimRobot( + self._pybullet_client, + initial_joint_positions=INITIAL_JOINT_POSITIONS, + end_effector=end_effector, + color="white" if self._visuals_mode == "real" else "default", + ) + + def _setup_pybullet_scene(self): + self._pybullet_client = bullet_client.BulletClient(self._connection_mode) + + # Temporarily disable rendering to speed up loading URDFs. + pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 0) + + self._setup_workspace_and_robot() + + if self._task == BlockTaskVariant.INSERT: + target_urdf_path = INSERT_URDF_PATH + else: + target_urdf_path = ZONE_URDF_PATH + + self._target_id = utils_pybullet.load_urdf( + self._pybullet_client, target_urdf_path, useFixedBase=True + ) + self._block_ids = [ + utils_pybullet.load_urdf( + self._pybullet_client, BLOCK_URDF_PATH, useFixedBase=False + ) + ] + + # Re-enable rendering. + pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 1) + + self.step_simulation_to_stabilize() + + def step_simulation_to_stabilize(self, nsteps=100): + for _ in range(nsteps): + self._pybullet_client.stepSimulation() + + def seed(self, seed=None): + self._rng = np.random.RandomState(seed=seed) + + def _set_robot_target_effector_pose(self, pose): + self._target_effector_pose = pose + self._robot.set_target_effector_pose(pose) + + def reset(self, reset_poses=True): + workspace_center_x = 0.4 + + if reset_poses: + self._pybullet_client.restoreState(self._saved_state) + + rotation = transform.Rotation.from_rotvec([0, math.pi, 0]) + translation = np.array([0.3, -0.4, self.effector_height]) + starting_pose = Pose3d(rotation=rotation, translation=translation) + self._set_robot_target_effector_pose(starting_pose) + + # Reset block pose. + block_x = workspace_center_x + self._rng.uniform(low=-0.1, high=0.1) + block_y = -0.2 + self._rng.uniform(low=-0.15, high=0.15) + block_translation = np.array([block_x, block_y, 0]) + block_sampled_angle = self._rng.uniform(math.pi) + block_rotation = transform.Rotation.from_rotvec([0, 0, block_sampled_angle]) + + self._pybullet_client.resetBasePositionAndOrientation( + self._block_ids[0], + block_translation.tolist(), + block_rotation.as_quat().tolist(), + ) + + # Reset target pose. + target_x = workspace_center_x + self._rng.uniform(low=-0.10, high=0.10) + target_y = 0.2 + self._rng.uniform(low=-0.15, high=0.15) + target_translation = np.array([target_x, target_y, 0.020]) + + target_sampled_angle = math.pi + self._rng.uniform( + low=-math.pi / 6, high=math.pi / 6 + ) + target_rotation = transform.Rotation.from_rotvec( + [0, 0, target_sampled_angle] + ) + + self._pybullet_client.resetBasePositionAndOrientation( + self._target_id, + target_translation.tolist(), + target_rotation.as_quat().tolist(), + ) + else: + ( + target_translation, + target_orientation_quat, + ) = self._pybullet_client.getBasePositionAndOrientation(self._target_id) + target_rotation = transform.Rotation.from_quat(target_orientation_quat) + target_translation = np.array(target_translation) + + self._target_pose = Pose3d( + rotation=target_rotation, translation=target_translation + ) + + if reset_poses: + self.step_simulation_to_stabilize() + + state = self._compute_state() + self._previous_state = state + + if self._task == BlockTaskVariant.REACH: + self._compute_reach_target(state) + + self._init_goal_distance = self._compute_goal_distance(state) + init_goal_eps = 1e-7 + assert self._init_goal_distance > init_goal_eps + self.best_fraction_reduced_goal_dist = 0.0 + + return state + + def _compute_goal_distance(self, state): + goal_translation = self.get_goal_translation() + if self._task != BlockTaskVariant.REACH: + goal_distance = np.linalg.norm( + state["block_translation"] - goal_translation[0:2] + ) + else: + goal_distance = np.linalg.norm( + state["effector_translation"] - goal_translation[0:2] + ) + return goal_distance + + def _compute_reach_target(self, state): + xy_block = state["block_translation"] + xy_target = state["target_translation"] + + xy_block_to_target = xy_target - xy_block + xy_dir_block_to_target = (xy_block_to_target) / np.linalg.norm( + xy_block_to_target + ) + self.reach_target_translation = xy_block + -1 * xy_dir_block_to_target * 0.05 + + def _compute_state(self): + effector_pose = self._robot.forward_kinematics() + block_position_and_orientation = ( + self._pybullet_client.getBasePositionAndOrientation(self._block_ids[0]) + ) + block_pose = Pose3d( + rotation=transform.Rotation.from_quat(block_position_and_orientation[1]), + translation=block_position_and_orientation[0], + ) + + def _yaw_from_pose(pose): + return np.array([pose.rotation.as_euler("xyz", degrees=False)[-1]]) + + obs = collections.OrderedDict( + block_translation=block_pose.translation[0:2], + block_orientation=_yaw_from_pose(block_pose), + effector_translation=effector_pose.translation[0:2], + effector_target_translation=self._target_effector_pose.translation[0:2], + target_translation=self._target_pose.translation[0:2], + target_orientation=_yaw_from_pose(self._target_pose), + ) + if self._image_size is not None: + obs["rgb"] = self._render_camera(self._image_size) + return obs + + def _step_robot_and_sim(self, action): + """Steps the robot and pybullet sim.""" + # Compute target_effector_pose by shifting the effector's pose by the + # action. + if self._abs_action: + target_effector_translation = np.array([action[0], action[1], 0]) + else: + target_effector_translation = np.array( + self._target_effector_pose.translation + ) + np.array([action[0], action[1], 0]) + + target_effector_translation[0:2] = np.clip( + target_effector_translation[0:2], + self.workspace_bounds[0], + self.workspace_bounds[1], + ) + target_effector_translation[-1] = self.effector_height + target_effector_pose = Pose3d( + rotation=EFFECTOR_DOWN_ROTATION, translation=target_effector_translation + ) + + self._set_robot_target_effector_pose(target_effector_pose) + + # Update sleep time dynamically to stay near real-time. + frame_sleep_time = 0 + if self._connection_mode == pybullet.SHARED_MEMORY: + cur_time = time.time() + if self._last_loop_time is not None: + # Calculate the total, non-sleeping time from the previous frame, this + # includes the actual step as well as any compute that happens in the + # caller thread (model inference, etc). + compute_time = ( + cur_time + - self._last_loop_time + - self._last_loop_frame_sleep_time * self._sim_steps_per_step + ) + # Use this to calculate the current frame's total sleep time to ensure + # that env.step runs at policy rate. This is an estimate since the + # previous frame's compute time may not match the current frame. + total_sleep_time = max((1 / self._control_frequency) - compute_time, 0) + # Now spread this out over the inner sim steps. This doesn't change + # control in any way, but makes the animation appear smooth. + frame_sleep_time = total_sleep_time / self._sim_steps_per_step + else: + # No estimate of the previous frame's compute, assume it is zero. + frame_sleep_time = 1 / self._step_frequency + + # Cache end of this loop time, to compute sleep time on next iteration. + self._last_loop_time = cur_time + self._last_loop_frame_sleep_time = frame_sleep_time + + for _ in range(self._sim_steps_per_step): + if self._connection_mode == pybullet.SHARED_MEMORY: + sleep_spin(frame_sleep_time) + self._pybullet_client.stepSimulation() + + def step(self, action): + self._step_robot_and_sim(action) + + state = self._compute_state() + + goal_distance = self._compute_goal_distance(state) + fraction_reduced_goal_distance = 1.0 - ( + goal_distance / self._init_goal_distance + ) + if fraction_reduced_goal_distance > self.best_fraction_reduced_goal_dist: + self.best_fraction_reduced_goal_dist = fraction_reduced_goal_distance + + done = False + reward = self.best_fraction_reduced_goal_dist + + # Terminate the episode if the block is close enough to the target. + if goal_distance < self.goal_dist_tolerance: + reward = 1.0 + done = True + + return state, reward, done, {} + + @property + def succeeded(self): + state = self._compute_state() + goal_distance = self._compute_goal_distance(state) + if goal_distance < self.goal_dist_tolerance: + return True + return False + + @property + def goal_distance(self): + state = self._compute_state() + return self._compute_goal_distance(state) + + def render(self, mode="rgb_array"): + if self._image_size is not None: + image_size = self._image_size + else: + # This allows rendering even for state-only obs, + # for visualization. + image_size = (IMAGE_HEIGHT, IMAGE_WIDTH) + + data = self._render_camera(image_size=(image_size[0], image_size[1])) + if mode == "human": + if self.rendered_img is None: + self.rendered_img = plt.imshow( + np.zeros((image_size[0], image_size[1], 4)) + ) + else: + self.rendered_img.set_data(data) + plt.draw() + plt.pause(0.00001) + return data + + def close(self): + self._pybullet_client.disconnect() + + def calc_camera_params(self, image_size): + # Mimic RealSense D415 camera parameters. + intrinsics = self._camera_instrinsics + + # Set default camera poses. + front_position = self._camera_pose + front_rotation = self._camera_orientation + front_rotation = self._pybullet_client.getQuaternionFromEuler(front_rotation) + # Default camera configs. + zrange = (0.01, 10.0) + + # OpenGL camera settings. + lookdir = np.float32([0, 0, 1]).reshape(3, 1) + updir = np.float32([0, -1, 0]).reshape(3, 1) + rotation = self._pybullet_client.getMatrixFromQuaternion(front_rotation) + rotm = np.float32(rotation).reshape(3, 3) + lookdir = (rotm @ lookdir).reshape(-1) + updir = (rotm @ updir).reshape(-1) + lookat = front_position + lookdir + focal_len = intrinsics[0] + znear, zfar = zrange + viewm = self._pybullet_client.computeViewMatrix(front_position, lookat, updir) + fovh = (image_size[0] / 2) / focal_len + fovh = 180 * np.arctan(fovh) * 2 / np.pi + + # Notes: 1) FOV is vertical FOV 2) aspect must be float + aspect_ratio = image_size[1] / image_size[0] + projm = self._pybullet_client.computeProjectionMatrixFOV( + fovh, aspect_ratio, znear, zfar + ) + + return viewm, projm, front_position, lookat, updir + + def _render_camera(self, image_size): + """Render RGB image with RealSense configuration.""" + viewm, projm, _, _, _ = self.calc_camera_params(image_size) + + # Render with OpenGL camera settings. + _, _, color, _, _ = self._pybullet_client.getCameraImage( + width=image_size[1], + height=image_size[0], + viewMatrix=viewm, + projectionMatrix=projm, + flags=pybullet.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX, + renderer=pybullet.ER_BULLET_HARDWARE_OPENGL, + ) + + # Get color image. + color_image_size = (image_size[0], image_size[1], 4) + color = np.array(color, dtype=np.uint8).reshape(color_image_size) + color = color[:, :, :3] # remove alpha channel + + return color.astype(np.uint8) + + def _create_observation_space(self, image_size): + pi2 = math.pi * 2 + + obs_dict = collections.OrderedDict( + block_translation=spaces.Box(low=-5, high=5, shape=(2,)), # x,y + block_orientation=spaces.Box(low=-pi2, high=pi2, shape=(1,)), # phi + effector_translation=spaces.Box( + low=self.workspace_bounds[0] - 0.1, # Small buffer for to IK noise. + high=self.workspace_bounds[1] + 0.1, + ), # x,y + effector_target_translation=spaces.Box( + low=self.workspace_bounds[0] - 0.1, # Small buffer for to IK noise. + high=self.workspace_bounds[1] + 0.1, + ), # x,y + target_translation=spaces.Box(low=-5, high=5, shape=(2,)), # x,y + target_orientation=spaces.Box( + low=-pi2, + high=pi2, + shape=(1,), + ), # theta + ) + if image_size is not None: + obs_dict["rgb"] = spaces.Box( + low=0, high=255, shape=(image_size[0], image_size[1], 3), dtype=np.uint8 + ) + return spaces.Dict(obs_dict) + + def get_pybullet_state(self): + """Save pybullet state of the scene. + + Returns: + dict containing 'robots', 'robot_end_effectors', 'targets', 'objects', + each containing a list of ObjState. + """ + state: Dict[str, List[ObjState]] = {} + + state["robots"] = [ + XarmState.get_bullet_state( + self._pybullet_client, + self.robot.xarm, + target_effector_pose=self._target_effector_pose, + goal_translation=self.get_goal_translation(), + ) + ] + + state["robot_end_effectors"] = [] + if self.robot.end_effector: + state["robot_end_effectors"].append( + ObjState.get_bullet_state( + self._pybullet_client, self.robot.end_effector + ) + ) + + state["targets"] = [] + if self._target_id: + state["targets"].append( + ObjState.get_bullet_state(self._pybullet_client, self._target_id) + ) + + state["objects"] = [] + for obj_id in self.get_obj_ids(): + state["objects"].append( + ObjState.get_bullet_state(self._pybullet_client, obj_id) + ) + + return state + + def set_pybullet_state(self, state): + """Restore pyullet state. + + WARNING: py_environment wrapper assumes environments aren't reset in their + constructor and will often reset the environment unintentionally. It is + always recommended that you call env.reset on the tfagents wrapper before + playback (replaying pybullet_state). + + Args: + state: dict containing 'robots', 'robot_end_effectors', 'targets', + 'objects', each containing a list of ObjState. + """ + + assert isinstance(state["robots"][0], XarmState) + xarm_state: XarmState = state["robots"][0] + xarm_state.set_bullet_state(self._pybullet_client, self.robot.xarm) + self._set_robot_target_effector_pose(xarm_state.target_effector_pose) + + def _set_state_safe(obj_state, obj_id): + if obj_state is not None: + assert obj_id is not None, "Cannot set state for missing object." + obj_state.set_bullet_state(self._pybullet_client, obj_id) + else: + assert obj_id is None, f"No state found for obj_id {obj_id}" + + robot_end_effectors = state["robot_end_effectors"] + _set_state_safe( + None if not robot_end_effectors else robot_end_effectors[0], + self.robot.end_effector, + ) + + targets = state["targets"] + _set_state_safe(None if not targets else targets[0], self._target_id) + + obj_ids = self.get_obj_ids() + assert len(state["objects"]) == len(obj_ids), "State length mismatch" + for obj_state, obj_id in zip(state["objects"], obj_ids): + _set_state_safe(obj_state, obj_id) + + self.reset(reset_poses=False) + + +class BlockPushNormalized(gym.Env): + """Simple XArm environment for block pushing, normalized state and actions.""" + + def __init__( + self, + control_frequency=10.0, + task=BlockTaskVariant.PUSH_NORMALIZED, + image_size=None, + shared_memory=False, + seed=None, + ): + """Creates an env instance. + + Args: + control_frequency: Control frequency for the arm. Each env step will + advance the simulation by 1/control_frequency seconds. + task: enum for which task, see BlockTaskVariant enum. + image_size: Optional image size (height, width). If None, no image + observations will be used. + shared_memory: If True `pybullet.SHARED_MEMORY` is used to connect to + pybullet. Useful to debug. + seed: Optional seed for the environment. + """ + # Map normalized task to unnormalized task. + if task == BlockTaskVariant.PUSH_NORMALIZED: + env_task = BlockTaskVariant.PUSH + elif task == BlockTaskVariant.REACH_NORMALIZED: + env_task = BlockTaskVariant.REACH + else: + raise ValueError("Unsupported task %s" % str(task)) + self._env = BlockPush( + control_frequency, env_task, image_size, shared_memory, seed + ) + self.action_space = spaces.Box(low=-1, high=1, shape=(2,)) + self.observation_space = spaces.Dict( + collections.OrderedDict( + effector_target_translation=spaces.Box(low=-1, high=1, shape=(2,)), + effector_target_to_block_translation=spaces.Box( + low=-1, high=1, shape=(2,) + ), + block_orientation_cos_sin=spaces.Box(low=-1, high=1, shape=(2,)), + effector_target_to_target_translation=spaces.Box( + low=-1, high=1, shape=(2,) + ), + target_orientation_cos_sin=spaces.Box(low=-1, high=1, shape=(2,)), + ) + ) + self.reset() + + def get_control_frequency(self): + return self._env.get_control_frequency() + + @property + def reach_target_translation(self): + return self._env.reach_target_translation + + def seed(self, seed=None): + self._env.seed(seed) + + def reset(self): + state = self._env.reset() + return self.calc_normalized_state(state) + + def step(self, action): + # The environment is normalized [mean-3*std, mean+3*std] -> [-1, 1]. + action = np.clip(action, a_min=-1.0, a_max=1.0) + state, reward, done, info = self._env.step( + self.calc_unnormalized_action(action) + ) + state = self.calc_normalized_state(state) + reward = reward * 100 # Keep returns in [0, 100] + return state, reward, done, info + + def render(self, mode="rgb_array"): + return self._env.render(mode) + + def close(self): + self._env.close() + + @staticmethod + def _normalize(values, values_min, values_max): + offset = (values_max + values_min) * 0.5 + scale = (values_max - values_min) * 0.5 + return (values - offset) / scale # [min, max] -> [-1, 1] + + @staticmethod + def _unnormalize(values, values_min, values_max): + offset = (values_max + values_min) * 0.5 + scale = (values_max - values_min) * 0.5 + return values * scale + offset # [-1, 1] -> [min, max] + + @classmethod + def calc_normalized_action(cls, action): + return cls._normalize(action, ACTION_MIN, ACTION_MAX) + + @classmethod + def calc_unnormalized_action(cls, norm_action): + return cls._unnormalize(norm_action, ACTION_MIN, ACTION_MAX) + + @classmethod + def calc_normalized_state(cls, state): + + effector_target_translation = cls._normalize( + state["effector_target_translation"], + EFFECTOR_TARGET_TRANSLATION_MIN, + EFFECTOR_TARGET_TRANSLATION_MAX, + ) + + effector_target_to_block_translation = cls._normalize( + state["block_translation"] - state["effector_target_translation"], + EFFECTOR_TARGET_TO_BLOCK_TRANSLATION_MIN, + EFFECTOR_TARGET_TO_BLOCK_TRANSLATION_MAX, + ) + ori = state["block_orientation"][0] + block_orientation_cos_sin = cls._normalize( + np.array([math.cos(ori), math.sin(ori)], np.float32), + BLOCK_ORIENTATION_COS_SIN_MIN, + BLOCK_ORIENTATION_COS_SIN_MAX, + ) + + effector_target_to_target_translation = cls._normalize( + state["target_translation"] - state["effector_target_translation"], + EFFECTOR_TARGET_TO_TARGET_TRANSLATION_MIN, + EFFECTOR_TARGET_TO_TARGET_TRANSLATION_MAX, + ) + ori = state["target_orientation"][0] + target_orientation_cos_sin = cls._normalize( + np.array([math.cos(ori), math.sin(ori)], np.float32), + TARGET_ORIENTATION_COS_SIN_MIN, + TARGET_ORIENTATION_COS_SIN_MAX, + ) + + # Note: We do not include effector_translation in the normalized state. + # This means the unnormalized -> normalized mapping is not invertable. + return collections.OrderedDict( + effector_target_translation=effector_target_translation, + effector_target_to_block_translation=effector_target_to_block_translation, + block_orientation_cos_sin=block_orientation_cos_sin, + effector_target_to_target_translation=effector_target_to_target_translation, + target_orientation_cos_sin=target_orientation_cos_sin, + ) + + @classmethod + def calc_unnormalized_state(cls, norm_state): + + effector_target_translation = cls._unnormalize( + norm_state["effector_target_translation"], + EFFECTOR_TARGET_TRANSLATION_MIN, + EFFECTOR_TARGET_TRANSLATION_MAX, + ) + # Note: normalized state does not include effector_translation state, this + # means this component will be missing (and is marked nan). + effector_translation = np.array([np.nan, np.nan], np.float32) + + effector_target_to_block_translation = cls._unnormalize( + norm_state["effector_target_to_block_translation"], + EFFECTOR_TARGET_TO_BLOCK_TRANSLATION_MIN, + EFFECTOR_TARGET_TO_BLOCK_TRANSLATION_MAX, + ) + block_translation = ( + effector_target_to_block_translation + effector_target_translation + ) + ori_cos_sin = cls._unnormalize( + norm_state["block_orientation_cos_sin"], + BLOCK_ORIENTATION_COS_SIN_MIN, + BLOCK_ORIENTATION_COS_SIN_MAX, + ) + block_orientation = np.array( + [math.atan2(ori_cos_sin[1], ori_cos_sin[0])], np.float32 + ) + + effector_target_to_target_translation = cls._unnormalize( + norm_state["effector_target_to_target_translation"], + EFFECTOR_TARGET_TO_TARGET_TRANSLATION_MIN, + EFFECTOR_TARGET_TO_TARGET_TRANSLATION_MAX, + ) + target_translation = ( + effector_target_to_target_translation + effector_target_translation + ) + ori_cos_sin = cls._unnormalize( + norm_state["target_orientation_cos_sin"], + TARGET_ORIENTATION_COS_SIN_MIN, + TARGET_ORIENTATION_COS_SIN_MAX, + ) + target_orientation = np.array( + [math.atan2(ori_cos_sin[1], ori_cos_sin[0])], np.float32 + ) + + return collections.OrderedDict( + block_translation=block_translation, + block_orientation=block_orientation, + effector_translation=effector_translation, + effector_target_translation=effector_target_translation, + target_translation=target_translation, + target_orientation=target_orientation, + ) + + def get_pybullet_state(self): + return self._env.get_pybullet_state() + + def set_pybullet_state(self, state): + return self._env.set_pybullet_state(state) + + @property + def pybullet_client(self): + return self._env.pybullet_client + + def calc_camera_params(self, image_size): + return self._env.calc_camera_params(image_size) + + def _compute_state(self): + return self.calc_normalized_state( + self._env._compute_state() + ) # pylint: disable=protected-access + + +# Make sure we only register once to allow us to reload the module in colab for +# debugging. +if "BlockPush-v0" in registration.registry.env_specs: + del registration.registry.env_specs["BlockInsert-v0"] + del registration.registry.env_specs["BlockPush-v0"] + del registration.registry.env_specs["BlockPushNormalized-v0"] + del registration.registry.env_specs["BlockPushRgbNormalized-v0"] + del registration.registry.env_specs["BlockReach-v0"] + del registration.registry.env_specs["BlockReachNormalized-v0"] + del registration.registry.env_specs["BlockReachRgbNormalized-v0"] + del registration.registry.env_specs["SharedBlockInsert-v0"] + del registration.registry.env_specs["SharedBlockPush-v0"] + del registration.registry.env_specs["SharedBlockReach-v0"] + +registration.register( + id="BlockInsert-v0", + entry_point=BlockPush, + kwargs=dict(task=BlockTaskVariant.INSERT), + max_episode_steps=50, +) +registration.register(id="BlockPush-v0", entry_point=BlockPush, max_episode_steps=100) +registration.register( + id="BlockPushNormalized-v0", + entry_point=BlockPushNormalized, + kwargs=dict(task=BlockTaskVariant.PUSH_NORMALIZED), + max_episode_steps=100, +) +registration.register( + id="BlockPushRgb-v0", + entry_point=BlockPush, + max_episode_steps=100, + kwargs=dict(image_size=(IMAGE_HEIGHT, IMAGE_WIDTH)), +) +registration.register( + id="BlockPushRgbNormalized-v0", + entry_point=BlockPushNormalized, + kwargs=dict( + task=BlockTaskVariant.PUSH_NORMALIZED, image_size=(IMAGE_HEIGHT, IMAGE_WIDTH) + ), + max_episode_steps=100, +) +registration.register( + id="BlockReach-v0", + entry_point=BlockPush, + kwargs=dict(task=BlockTaskVariant.REACH), + max_episode_steps=50, +) +registration.register( + id="BlockReachRgb-v0", + entry_point=BlockPush, + max_episode_steps=100, + kwargs=dict(task=BlockTaskVariant.REACH, image_size=(IMAGE_HEIGHT, IMAGE_WIDTH)), +) +registration.register( + id="BlockReachNormalized-v0", + entry_point=BlockPushNormalized, + kwargs=dict(task=BlockTaskVariant.REACH_NORMALIZED), + max_episode_steps=50, +) +registration.register( + id="BlockReachRgbNormalized-v0", + entry_point=BlockPushNormalized, + kwargs=dict( + task=BlockTaskVariant.REACH_NORMALIZED, image_size=(IMAGE_HEIGHT, IMAGE_WIDTH) + ), + max_episode_steps=50, +) + +registration.register( + id="SharedBlockInsert-v0", + entry_point=BlockPush, + kwargs=dict(task=BlockTaskVariant.INSERT, shared_memory=True), + max_episode_steps=50, +) +registration.register( + id="SharedBlockPush-v0", + entry_point=BlockPush, + kwargs=dict(shared_memory=True), + max_episode_steps=100, +) +registration.register( + id="SharedBlockPushNormalized-v0", + entry_point=BlockPushNormalized, + kwargs=dict(task=BlockTaskVariant.PUSH_NORMALIZED, shared_memory=True), + max_episode_steps=100, +) +registration.register( + id="SharedBlockReach-v0", + entry_point=BlockPush, + kwargs=dict(task=BlockTaskVariant.REACH, shared_memory=True), + max_episode_steps=50, +) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/block_pushing_discontinuous.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/block_pushing_discontinuous.py new file mode 100644 index 0000000000000000000000000000000000000000..38aaef381868a7c313c54a56a497105c047b4a03 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/block_pushing_discontinuous.py @@ -0,0 +1,338 @@ +# coding=utf-8 +# Copyright 2022 The Reach ML Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Discontinuous block pushing.""" +import collections +import enum +import math +from typing import List, Optional + +from gym import spaces +from gym.envs import registration +from cleandiffuser.env.block_pushing import block_pushing +from cleandiffuser.env.block_pushing.utils import utils_pybullet +from cleandiffuser.env.block_pushing.utils.pose3d import Pose3d +import numpy as np +from scipy.spatial import transform +import pybullet +import pybullet_utils.bullet_client as bullet_client + +ZONE2_URDF_PATH = "third_party/py/envs/assets/zone2.urdf" + +MIN_TARGET_DIST = 0.15 +NUM_RESET_ATTEMPTS = 1000 + + +def build_env_name(task, shared_memory, use_image_obs): + """Construct the env name from parameters.""" + del task + env_name = "BlockPushDiscontinuous" + + if use_image_obs: + env_name = env_name + "Rgb" + + if shared_memory: + env_name = "Shared" + env_name + + env_name = env_name + "-v0" + + return env_name + + +class BlockTaskVariant(enum.Enum): + REACH = "Reach" + REACH_NORMALIZED = "ReachNormalized" + PUSH = "Push" + PUSH_NORMALIZED = "PushNormalized" + INSERT = "Insert" + + +# pytype: skip-file +class BlockPushDiscontinuous(block_pushing.BlockPush): + """Discontinuous block pushing.""" + + def __init__( + self, + control_frequency=10.0, + task=BlockTaskVariant.PUSH, + image_size=None, + shared_memory=False, + seed=None, + goal_dist_tolerance=0.04, + ): + super(BlockPushDiscontinuous, self).__init__( + control_frequency=control_frequency, + task=task, + image_size=image_size, + shared_memory=shared_memory, + seed=seed, + goal_dist_tolerance=goal_dist_tolerance, + ) + + @property + def target_poses(self): + return self._target_poses + + def get_goal_translation(self): + """Return the translation component of the goal (2D).""" + if self._target_poses: + return [i.translation for i in self._target_poses] + else: + return None + + def _setup_pybullet_scene(self): + self._pybullet_client = bullet_client.BulletClient(self._connection_mode) + + # Temporarily disable rendering to speed up loading URDFs. + pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 0) + + self._setup_workspace_and_robot() + target_urdf_path = block_pushing.ZONE_URDF_PATH + + self._target_ids = [] + for _ in [block_pushing.ZONE_URDF_PATH, ZONE2_URDF_PATH]: + self._target_ids.append( + utils_pybullet.load_urdf( + self._pybullet_client, target_urdf_path, useFixedBase=True + ) + ) + self._block_ids = [ + utils_pybullet.load_urdf( + self._pybullet_client, block_pushing.BLOCK_URDF_PATH, useFixedBase=False + ) + ] + + # Re-enable rendering. + pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 1) + + self.step_simulation_to_stabilize() + + def _reset_target_poses(self, workspace_center_x): + """Resets target poses.""" + self._target_poses = [None for _ in range(len(self._target_ids))] + + def _reset_target_pose(idx, avoid=None): + def _get_random_translation(): + # Choose x,y randomly. + target_x = workspace_center_x + self._rng.uniform(low=-0.10, high=0.10) + # Fix ys for this environment. + if idx == 0: + target_y = 0 + else: + target_y = 0.4 + target_translation = np.array([target_x, target_y, 0.020]) + return target_translation + + if avoid is None: + target_translation = _get_random_translation() + else: + # Reject targets too close to `avoid`. + for _ in range(NUM_RESET_ATTEMPTS): + target_translation = _get_random_translation() + dist = np.linalg.norm(target_translation[0] - avoid[0]) + if dist > MIN_TARGET_DIST: + break + target_sampled_angle = math.pi + self._rng.uniform( + low=-math.pi / 6, high=math.pi / 6 + ) + target_rotation = transform.Rotation.from_rotvec( + [0, 0, target_sampled_angle] + ) + self._pybullet_client.resetBasePositionAndOrientation( + self._target_ids[idx], + target_translation.tolist(), + target_rotation.as_quat().tolist(), + ) + self._target_poses[idx] = Pose3d( + rotation=target_rotation, translation=target_translation + ) + + try_idx = 0 + while True: + # Choose the first target. + _reset_target_pose(0) + # Choose the second target, avoiding the first. + _reset_target_pose(1, avoid=self._target_poses[0].translation) + dist = np.linalg.norm( + self._target_poses[0].translation[0] + - self._target_poses[1].translation[0] + ) + if dist > MIN_TARGET_DIST: + break + try_idx += 1 + if try_idx >= NUM_RESET_ATTEMPTS: + raise ValueError("could not find matching target") + assert dist > MIN_TARGET_DIST + + def reset(self): + self._pybullet_client.restoreState(self._saved_state) + + rotation = transform.Rotation.from_rotvec([0, math.pi, 0]) + translation = np.array([0.3, -0.4, block_pushing.EFFECTOR_HEIGHT]) + starting_pose = Pose3d(rotation=rotation, translation=translation) + self._set_robot_target_effector_pose(starting_pose) + + workspace_center_x = 0.4 + + # Reset block pose. + block_x = workspace_center_x + self._rng.uniform(low=-0.1, high=0.1) + block_y = -0.2 + self._rng.uniform(low=-0.15, high=0.15) + block_translation = np.array([block_x, block_y, 0]) + block_sampled_angle = self._rng.uniform(math.pi) + block_rotation = transform.Rotation.from_rotvec([0, 0, block_sampled_angle]) + + self._pybullet_client.resetBasePositionAndOrientation( + self._block_ids[0], + block_translation.tolist(), + block_rotation.as_quat().tolist(), + ) + + # Reset target pose. + self._reset_target_poses(workspace_center_x) + + self.step_simulation_to_stabilize() + state = self._compute_state() + self._previous_state = state + self.min_dist_to_first_goal = np.inf + self.min_dist_to_second_goal = np.inf + self.steps = 0 + return state + + def _compute_goal_distance(self, state): + # Reward is 1. blocks is inside any target. + return np.mean([self.min_dist_to_first_goal, self.min_dist_to_second_goal]) + + def _compute_state(self): + effector_pose = self._robot.forward_kinematics() + block_position_and_orientation = ( + self._pybullet_client.getBasePositionAndOrientation(self._block_ids[0]) + ) + block_pose = Pose3d( + rotation=transform.Rotation.from_quat(block_position_and_orientation[1]), + translation=block_position_and_orientation[0], + ) + + def _yaw_from_pose(pose): + return np.array([pose.rotation.as_euler("xyz", degrees=False)[-1]]) + + obs = collections.OrderedDict( + block_translation=block_pose.translation[0:2], + block_orientation=_yaw_from_pose(block_pose), + effector_translation=effector_pose.translation[0:2], + effector_target_translation=self._target_effector_pose.translation[0:2], + target_translation=self._target_poses[0].translation[0:2], + target_orientation=_yaw_from_pose(self._target_poses[0]), + target2_translation=self._target_poses[1].translation[0:2], + target2_orientation=_yaw_from_pose(self._target_poses[1]), + ) + if self._image_size is not None: + obs["rgb"] = self._render_camera(self._image_size) + return obs + + def step(self, action): + self._step_robot_and_sim(action) + state = self._compute_state() + reward = self._get_reward(state) + done = False + if reward > 0.0: + done = True + # Cache so we can compute success. + self.state = state + return state, reward, done, {} + + def dist(self, state, target): + # Reward is 1. blocks is inside any target. + return np.linalg.norm( + state["block_translation"] - state["%s_translation" % target] + ) + + def _get_reward(self, state): + """Reward is 1.0 if agent hits both goals and stays at second.""" + # This also statefully updates these values. + self.min_dist_to_first_goal = min( + self.dist(state, "target"), self.min_dist_to_first_goal + ) + self.min_dist_to_second_goal = min( + self.dist(state, "target2"), self.min_dist_to_second_goal + ) + + def _reward(thresh): + reward_first = True if self.min_dist_to_first_goal < thresh else False + reward_second = True if self.min_dist_to_second_goal < thresh else False + return 1.0 if (reward_first and reward_second) else 0.0 + + reward = _reward(self.goal_dist_tolerance) + return reward + + @property + def succeeded(self): + thresh = self.goal_dist_tolerance + hit_first = True if self.min_dist_to_first_goal < thresh else False + hit_second = True if self.min_dist_to_first_goal < thresh else False + current_distance_to_second = self.dist(self.state, "target2") + still_at_second = True if current_distance_to_second < thresh else False + return hit_first and hit_second and still_at_second + + def _create_observation_space(self, image_size): + pi2 = math.pi * 2 + + obs_dict = collections.OrderedDict( + block_translation=spaces.Box(low=-5, high=5, shape=(2,)), # x,y + block_orientation=spaces.Box(low=-pi2, high=pi2, shape=(1,)), # phi + effector_translation=spaces.Box( + # Small buffer for to IK noise. + low=block_pushing.WORKSPACE_BOUNDS[0] - 0.1, + high=block_pushing.WORKSPACE_BOUNDS[1] + 0.1, + ), # x,y + effector_target_translation=spaces.Box( + # Small buffer for to IK noise. + low=block_pushing.WORKSPACE_BOUNDS[0] - 0.1, + high=block_pushing.WORKSPACE_BOUNDS[1] + 0.1, + ), # x,y + target_translation=spaces.Box(low=-5, high=5, shape=(2,)), # x,y + target_orientation=spaces.Box( + low=-pi2, + high=pi2, + shape=(1,), + ), # theta + target2_translation=spaces.Box(low=-5, high=5, shape=(2,)), # x,y + target2_orientation=spaces.Box( + low=-pi2, + high=pi2, + shape=(1,), + ), # theta + ) + if image_size is not None: + obs_dict["rgb"] = spaces.Box( + low=0, high=255, shape=(image_size[0], image_size[1], 3), dtype=np.uint8 + ) + return spaces.Dict(obs_dict) + + +if "BlockPushDiscontinuous-v0" in registration.registry.env_specs: + del registration.registry.env_specs["BlockPushDiscontinuous-v0"] + +registration.register( + id="BlockPushDiscontinuous-v0", + entry_point=BlockPushDiscontinuous, + max_episode_steps=200, +) + +registration.register( + id="BlockPushDiscontinuousRgb-v0", + entry_point=BlockPushDiscontinuous, + max_episode_steps=200, + kwargs=dict(image_size=(block_pushing.IMAGE_HEIGHT, block_pushing.IMAGE_WIDTH)), +) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/block_pushing_multimodal.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/block_pushing_multimodal.py new file mode 100644 index 0000000000000000000000000000000000000000..f4dbba54ecdeb9840f7e6af261908f7ffcca3348 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/block_pushing_multimodal.py @@ -0,0 +1,802 @@ +# coding=utf-8 +# Copyright 2022 The Reach ML Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Multimodal block environments for the XArm.""" + +import collections +import logging +import math +from typing import Dict, List, Optional, Union +import copy +import time + +from gym import spaces +from gym.envs import registration +from cleandiffuser.env.block_pushing import block_pushing +from cleandiffuser.env.block_pushing.utils import utils_pybullet +from cleandiffuser.env.block_pushing.utils.pose3d import Pose3d +from cleandiffuser.env.block_pushing.utils.utils_pybullet import ObjState +from cleandiffuser.env.block_pushing.utils.utils_pybullet import XarmState +import numpy as np +from scipy.spatial import transform +import pybullet +import pybullet_utils.bullet_client as bullet_client + +# pytype: skip-file +BLOCK2_URDF_PATH = "third_party/py/envs/assets/block2.urdf" +ZONE2_URDF_PATH = "third_party/py/envs/assets/zone2.urdf" + +# When resetting multiple targets, they should all be this far apart. +MIN_BLOCK_DIST = 0.1 +MIN_TARGET_DIST = 0.12 +# pylint: enable=line-too-long +NUM_RESET_ATTEMPTS = 1000 + +# Random movement of blocks +RANDOM_X_SHIFT = 0.1 +RANDOM_Y_SHIFT = 0.15 + +logging.basicConfig( + level="INFO", + format="%(asctime)s [%(levelname)s] %(message)s", + filemode="w", +) +logger = logging.getLogger() + + +def build_env_name(task, shared_memory, use_image_obs): + """Construct the env name from parameters.""" + del task + env_name = "BlockPushMultimodal" + + if use_image_obs: + env_name = env_name + "Rgb" + + if shared_memory: + env_name = "Shared" + env_name + + env_name = env_name + "-v0" + + return env_name + + +class BlockPushEventManager: + def __init__(self): + self.event_steps = { + 'REACH_0': -1, + 'REACH_1': -1, + 'TARGET_0_0': -1, + 'TARGET_0_1': -1, + 'TARGET_1_0': -1, + 'TARGET_1_1': -1 + } + + def reach(self, step, block_id): + key = f'REACH_{block_id}' + if self.event_steps[key] < 0: + self.event_steps[key] = step + + def target(self, step, block_id, target_id): + key = f'TARGET_{block_id}_{target_id}' + if self.event_steps[key] < 0: + self.event_steps[key] = step + + def reset(self): + for key in list(self.event_steps): + self.event_steps[key] = -1 + + def get_info(self): + return copy.deepcopy(self.event_steps) + +class BlockPushMultimodal(block_pushing.BlockPush): + """2 blocks, 2 targets.""" + + def __init__( + self, + control_frequency=10.0, + task=block_pushing.BlockTaskVariant.PUSH, + image_size=None, + shared_memory=False, + seed=None, + goal_dist_tolerance=0.05, + abs_action=False + ): + """Creates an env instance. + + Args: + control_frequency: Control frequency for the arm. Each env step will + advance the simulation by 1/control_frequency seconds. + task: enum for which task, see BlockTaskVariant enum. + image_size: Optional image size (height, width). If None, no image + observations will be used. + shared_memory: If True `pybullet.SHARED_MEMORY` is used to connect to + pybullet. Useful to debug. + seed: Optional seed for the environment. + goal_dist_tolerance: float, how far away from the goal to terminate. + """ + self._target_ids = None + self._target_poses = None + self._event_manager = BlockPushEventManager() + super(BlockPushMultimodal, self).__init__( + control_frequency=control_frequency, + task=task, + image_size=image_size, + shared_memory=shared_memory, + seed=seed, + goal_dist_tolerance=goal_dist_tolerance, + ) + self._init_distance = [-1.0, -1.0] + self._in_target = [[-1.0, -1.0], [-1.0, -1.0]] + self._first_move = [-1, -1] + self._step_num = 0 + self._abs_action = abs_action + + @property + def target_poses(self): + return self._target_poses + + def get_goal_translation(self): + """Return the translation component of the goal (2D).""" + if self._target_poses: + return [i.translation for i in self._target_poses] + else: + return None + + def _setup_pybullet_scene(self): + self._pybullet_client = bullet_client.BulletClient(self._connection_mode) + + # Temporarily disable rendering to speed up loading URDFs. + pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 0) + + self._setup_workspace_and_robot() + + self._target_ids = [ + utils_pybullet.load_urdf(self._pybullet_client, i, useFixedBase=True) + for i in [block_pushing.ZONE_URDF_PATH, ZONE2_URDF_PATH] + ] + self._block_ids = [] + for i in [block_pushing.BLOCK_URDF_PATH, BLOCK2_URDF_PATH]: + self._block_ids.append( + utils_pybullet.load_urdf(self._pybullet_client, i, useFixedBase=False) + ) + + # Re-enable rendering. + pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 1) + + self.step_simulation_to_stabilize() + + def _reset_block_poses(self, workspace_center_x): + """Resets block poses.""" + + # Helper for choosing random block position. + def _reset_block_pose(idx, add=0.0, avoid=None): + def _get_random_translation(): + block_x = ( + workspace_center_x + + add + + self._rng.uniform(low=-RANDOM_X_SHIFT, high=RANDOM_X_SHIFT) + ) + block_y = -0.2 + self._rng.uniform( + low=-RANDOM_Y_SHIFT, high=RANDOM_Y_SHIFT + ) + block_translation = np.array([block_x, block_y, 0]) + return block_translation + + if avoid is None: + block_translation = _get_random_translation() + else: + # Reject targets too close to `avoid`. + for _ in range(NUM_RESET_ATTEMPTS): + block_translation = _get_random_translation() + dist = np.linalg.norm(block_translation[0] - avoid[0]) + # print('block inner try_idx %d, dist %.3f' % (try_idx, dist)) + if dist > MIN_BLOCK_DIST: + break + block_sampled_angle = self._rng.uniform(math.pi) + block_rotation = transform.Rotation.from_rotvec([0, 0, block_sampled_angle]) + self._pybullet_client.resetBasePositionAndOrientation( + self._block_ids[idx], + block_translation.tolist(), + block_rotation.as_quat().tolist(), + ) + return block_translation + + # Reject targets too close to `avoid`. + for _ in range(NUM_RESET_ATTEMPTS): + # Reset first block. + b0_translation = _reset_block_pose(0) + # Reset second block away from first block. + b1_translation = _reset_block_pose(1, avoid=b0_translation) + dist = np.linalg.norm(b0_translation[0] - b1_translation[0]) + if dist > MIN_BLOCK_DIST: + break + else: + raise ValueError("could not find matching block") + assert dist > MIN_BLOCK_DIST + + def _reset_target_poses(self, workspace_center_x): + """Resets target poses.""" + + def _reset_target_pose(idx, add=0.0, avoid=None): + def _get_random_translation(): + # Choose x,y randomly. + target_x = ( + workspace_center_x + + add + + self._rng.uniform( + low=-0.05 * RANDOM_X_SHIFT, high=0.05 * RANDOM_X_SHIFT + ) + ) + target_y = 0.2 + self._rng.uniform( + low=-0.05 * RANDOM_Y_SHIFT, high=0.05 * RANDOM_Y_SHIFT + ) + target_translation = np.array([target_x, target_y, 0.020]) + return target_translation + + if avoid is None: + target_translation = _get_random_translation() + else: + # Reject targets too close to `avoid`. + for _ in range(NUM_RESET_ATTEMPTS): + target_translation = _get_random_translation() + dist = np.linalg.norm(target_translation[0] - avoid[0]) + # print('target inner try_idx %d, dist %.3f' % (try_idx, dist)) + if dist > MIN_TARGET_DIST: + break + target_sampled_angle = math.pi + self._rng.uniform( + low=-math.pi / 30, high=math.pi / 30 + ) + target_rotation = transform.Rotation.from_rotvec( + [0, 0, target_sampled_angle] + ) + self._pybullet_client.resetBasePositionAndOrientation( + self._target_ids[idx], + target_translation.tolist(), + target_rotation.as_quat().tolist(), + ) + self._target_poses[idx] = Pose3d( + rotation=target_rotation, translation=target_translation + ) + + if self._target_poses is None: + self._target_poses = [None for _ in range(len(self._target_ids))] + + for _ in range(NUM_RESET_ATTEMPTS): + # Choose the first target. + add = 0.12 * self._rng.choice([-1, 1]) + # Randomly flip the location of the targets. + _reset_target_pose(0, add=add) + _reset_target_pose(1, add=-add, avoid=self._target_poses[0].translation) + dist = np.linalg.norm( + self._target_poses[0].translation[0] + - self._target_poses[1].translation[0] + ) + if dist > MIN_TARGET_DIST: + break + else: + raise ValueError("could not find matching target") + assert dist > MIN_TARGET_DIST + + def _reset_object_poses(self, workspace_center_x, workspace_center_y): + # Reset block poses. + self._reset_block_poses(workspace_center_x) + + # Reset target poses. + self._reset_target_poses(workspace_center_x) + + self._init_distance = [-1.0, -1.0] + self._in_target = [[-1.0, -1.0], [-1.0, -1.0]] + self._step_num = 0 + + def reset(self, reset_poses=True): + workspace_center_x = 0.4 + workspace_center_y = 0.0 + + if reset_poses: + self._pybullet_client.restoreState(self._saved_state) + + rotation = transform.Rotation.from_rotvec([0, math.pi, 0]) + translation = np.array([0.3, -0.4, block_pushing.EFFECTOR_HEIGHT]) + starting_pose = Pose3d(rotation=rotation, translation=translation) + self._set_robot_target_effector_pose(starting_pose) + self._reset_object_poses(workspace_center_x, workspace_center_y) + + # else: + self._target_poses = [ + self._get_target_pose(idx) for idx in self._target_ids + ] + + if reset_poses: + self.step_simulation_to_stabilize() + + state = self._compute_state() + self._previous_state = state + self._event_manager.reset() + return state + + def _get_target_pose(self, idx): + ( + target_translation, + target_orientation_quat, + ) = self._pybullet_client.getBasePositionAndOrientation(idx) + target_rotation = transform.Rotation.from_quat(target_orientation_quat) + target_translation = np.array(target_translation) + return Pose3d(rotation=target_rotation, translation=target_translation) + + def _compute_reach_target(self, state): + xy_block = state["block_translation"] + xy_target = state["target_translation"] + + xy_block_to_target = xy_target - xy_block + xy_dir_block_to_target = (xy_block_to_target) / np.linalg.norm( + xy_block_to_target + ) + self.reach_target_translation = xy_block + -1 * xy_dir_block_to_target * 0.05 + + def _compute_state(self): + effector_pose = self._robot.forward_kinematics() + + def _get_block_pose(idx): + block_position_and_orientation = ( + self._pybullet_client.getBasePositionAndOrientation( + self._block_ids[idx] + ) + ) + block_pose = Pose3d( + rotation=transform.Rotation.from_quat( + block_position_and_orientation[1] + ), + translation=block_position_and_orientation[0], + ) + return block_pose + + block_poses = [_get_block_pose(i) for i in range(len(self._block_ids))] + + def _yaw_from_pose(pose): + return np.array([pose.rotation.as_euler("xyz", degrees=False)[-1] % np.pi]) + + obs = collections.OrderedDict( + block_translation=block_poses[0].translation[0:2], + block_orientation=_yaw_from_pose(block_poses[0]), + block2_translation=block_poses[1].translation[0:2], + block2_orientation=_yaw_from_pose(block_poses[1]), + effector_translation=effector_pose.translation[0:2], + effector_target_translation=self._target_effector_pose.translation[0:2], + target_translation=self._target_poses[0].translation[0:2], + target_orientation=_yaw_from_pose(self._target_poses[0]), + target2_translation=self._target_poses[1].translation[0:2], + target2_orientation=_yaw_from_pose(self._target_poses[1]), + ) + + for i in range(2): + new_distance = np.linalg.norm( + block_poses[i].translation[0:2] + ) # + np.linalg.norm(_yaw_from_pose(block_poses[i])) + if self._init_distance[i] == -1: + self._init_distance[i] = new_distance + else: + if self._init_distance[i] != 100: + if np.abs(new_distance - self._init_distance[i]) > 1e-3: + logger.info(f"Block {i} moved on step {self._step_num}") + self._event_manager.reach(step=self._step_num, block_id=i) + self._init_distance[i] = 100 + + self._step_num += 1 + if self._image_size is not None: + obs["rgb"] = self._render_camera(self._image_size) + return obs + + def step(self, action): + self._step_robot_and_sim(action) + + state = self._compute_state() + done = False + reward = self._get_reward(state) + if reward >= 0.5: + # Terminate the episode if both blocks are close enough to the targets. + done = True + + info = self._event_manager.get_info() + return state, reward, done, info + + def _step_robot_and_sim(self, action): + """Steps the robot and pybullet sim.""" + # Compute target_effector_pose by shifting the effector's pose by the + # action. + if self._abs_action: + target_effector_translation = np.array([action[0], action[1], 0]) + else: + target_effector_translation = np.array( + self._target_effector_pose.translation + ) + np.array([action[0], action[1], 0]) + + target_effector_translation[0:2] = np.clip( + target_effector_translation[0:2], + self.workspace_bounds[0], + self.workspace_bounds[1], + ) + target_effector_translation[-1] = self.effector_height + target_effector_pose = Pose3d( + rotation=block_pushing.EFFECTOR_DOWN_ROTATION, translation=target_effector_translation + ) + + self._set_robot_target_effector_pose(target_effector_pose) + + # Update sleep time dynamically to stay near real-time. + frame_sleep_time = 0 + if self._connection_mode == pybullet.SHARED_MEMORY: + cur_time = time.time() + if self._last_loop_time is not None: + # Calculate the total, non-sleeping time from the previous frame, this + # includes the actual step as well as any compute that happens in the + # caller thread (model inference, etc). + compute_time = ( + cur_time + - self._last_loop_time + - self._last_loop_frame_sleep_time * self._sim_steps_per_step + ) + # Use this to calculate the current frame's total sleep time to ensure + # that env.step runs at policy rate. This is an estimate since the + # previous frame's compute time may not match the current frame. + total_sleep_time = max((1 / self._control_frequency) - compute_time, 0) + # Now spread this out over the inner sim steps. This doesn't change + # control in any way, but makes the animation appear smooth. + frame_sleep_time = total_sleep_time / self._sim_steps_per_step + else: + # No estimate of the previous frame's compute, assume it is zero. + frame_sleep_time = 1 / self._step_frequency + + # Cache end of this loop time, to compute sleep time on next iteration. + self._last_loop_time = cur_time + self._last_loop_frame_sleep_time = frame_sleep_time + + for _ in range(self._sim_steps_per_step): + if self._connection_mode == pybullet.SHARED_MEMORY: + block_pushing.sleep_spin(frame_sleep_time) + self._pybullet_client.stepSimulation() + + def _get_reward(self, state): + # Reward is 1. if both blocks are inside targets, but not the same target. + targets = ["target", "target2"] + + def _block_target_dist(block, target): + return np.linalg.norm( + state["%s_translation" % block] - state["%s_translation" % target] + ) + + def _closest_target(block): + # Distances to all targets. + dists = [_block_target_dist(block, t) for t in targets] + # Which is closest. + closest_target = targets[np.argmin(dists)] + closest_dist = np.min(dists) + # Is it in the closest target? + in_target = closest_dist < self.goal_dist_tolerance + return closest_target, in_target + + blocks = ["block", "block2"] + + reward = 0.0 + + for t_i, t in enumerate(targets): + for b_i, b in enumerate(blocks): + if self._in_target[t_i][b_i] == -1: + dist = _block_target_dist(b, t) + if dist < self.goal_dist_tolerance: + self._in_target[t_i][b_i] = 0 + logger.info( + f"Block {b_i} entered target {t_i} on step {self._step_num}" + ) + self._event_manager.target(step=self._step_num, block_id=b_i, target_id=t_i) + reward += 0.49 + + b0_closest_target, b0_in_target = _closest_target("block") + b1_closest_target, b1_in_target = _closest_target("block2") + # reward = 0.0 + if b0_in_target and b1_in_target and (b0_closest_target != b1_closest_target): + reward = 0.51 + return reward + + def _compute_goal_distance(self, state): + blocks = ["block", "block2"] + + def _target_block_dist(target, block): + return np.linalg.norm( + state["%s_translation" % block] - state["%s_translation" % target] + ) + + def _closest_block_dist(target): + dists = [_target_block_dist(target, b) for b in blocks] + closest_dist = np.min(dists) + return closest_dist + + t0_closest_dist = _closest_block_dist("target") + t1_closest_dist = _closest_block_dist("target2") + return np.mean([t0_closest_dist, t1_closest_dist]) + + @property + def succeeded(self): + state = self._compute_state() + reward = self._get_reward(state) + if reward >= 0.5: + return True + return False + + def _create_observation_space(self, image_size): + pi2 = math.pi * 2 + + obs_dict = collections.OrderedDict( + block_translation=spaces.Box(low=-5, high=5, shape=(2,)), # x,y + block_orientation=spaces.Box(low=-pi2, high=pi2, shape=(1,)), # phi + block2_translation=spaces.Box(low=-5, high=5, shape=(2,)), # x,y + block2_orientation=spaces.Box(low=-pi2, high=pi2, shape=(1,)), # phi + effector_translation=spaces.Box( + low=block_pushing.WORKSPACE_BOUNDS[0] - 0.1, + high=block_pushing.WORKSPACE_BOUNDS[1] + 0.1, + ), # x,y + effector_target_translation=spaces.Box( + low=block_pushing.WORKSPACE_BOUNDS[0] - 0.1, + high=block_pushing.WORKSPACE_BOUNDS[1] + 0.1, + ), # x,y + target_translation=spaces.Box(low=-5, high=5, shape=(2,)), # x,y + target_orientation=spaces.Box( + low=-pi2, + high=pi2, + shape=(1,), + ), # theta + target2_translation=spaces.Box(low=-5, high=5, shape=(2,)), # x,y + target2_orientation=spaces.Box( + low=-pi2, + high=pi2, + shape=(1,), + ), # theta + ) + if image_size is not None: + obs_dict["rgb"] = spaces.Box( + low=0, high=255, shape=(image_size[0], image_size[1], 3), dtype=np.uint8 + ) + return spaces.Dict(obs_dict) + + def get_pybullet_state(self): + """Save pybullet state of the scene. + + Returns: + dict containing 'robots', 'robot_end_effectors', 'targets', 'objects', + each containing a list of ObjState. + """ + state: Dict[str, List[ObjState]] = {} + + state["robots"] = [ + XarmState.get_bullet_state( + self._pybullet_client, + self.robot.xarm, + target_effector_pose=self._target_effector_pose, + goal_translation=None, + ) + ] + + state["robot_end_effectors"] = [] + if self.robot.end_effector: + state["robot_end_effectors"].append( + ObjState.get_bullet_state( + self._pybullet_client, self.robot.end_effector + ) + ) + + state["targets"] = [] + if self._target_ids: + for target_id in self._target_ids: + state["targets"].append( + ObjState.get_bullet_state(self._pybullet_client, target_id) + ) + + state["objects"] = [] + for obj_id in self.get_obj_ids(): + state["objects"].append( + ObjState.get_bullet_state(self._pybullet_client, obj_id) + ) + + return state + + def set_pybullet_state(self, state): + """Restore pyullet state. + + WARNING: py_environment wrapper assumes environments aren't reset in their + constructor and will often reset the environment unintentionally. It is + always recommended that you call env.reset on the tfagents wrapper before + playback (replaying pybullet_state). + + Args: + state: dict containing 'robots', 'robot_end_effectors', 'targets', + 'objects', each containing a list of ObjState. + """ + + assert isinstance(state["robots"][0], XarmState) + xarm_state: XarmState = state["robots"][0] + xarm_state.set_bullet_state(self._pybullet_client, self.robot.xarm) + self._set_robot_target_effector_pose(xarm_state.target_effector_pose) + + def _set_state_safe(obj_state, obj_id): + if obj_state is not None: + assert obj_id is not None, "Cannot set state for missing object." + obj_state.set_bullet_state(self._pybullet_client, obj_id) + else: + assert obj_id is None, f"No state found for obj_id {obj_id}" + + robot_end_effectors = state["robot_end_effectors"] + _set_state_safe( + None if not robot_end_effectors else robot_end_effectors[0], + self.robot.end_effector, + ) + + for target_state, target_id in zip(state["targets"], self._target_ids): + _set_state_safe(target_state, target_id) + + obj_ids = self.get_obj_ids() + assert len(state["objects"]) == len(obj_ids), "State length mismatch" + for obj_state, obj_id in zip(state["objects"], obj_ids): + _set_state_safe(obj_state, obj_id) + + self.reset(reset_poses=False) + + +class BlockPushHorizontalMultimodal(BlockPushMultimodal): + def _reset_object_poses(self, workspace_center_x, workspace_center_y): + # Reset block poses. + self._reset_block_poses(workspace_center_y) + + # Reset target poses. + self._reset_target_poses(workspace_center_y) + + def _reset_block_poses(self, workspace_center_y): + """Resets block poses.""" + + # Helper for choosing random block position. + def _reset_block_pose(idx, add=0.0, avoid=None): + def _get_random_translation(): + block_x = 0.35 + 0.5 * self._rng.uniform( + low=-RANDOM_X_SHIFT, high=RANDOM_X_SHIFT + ) + block_y = ( + workspace_center_y + + add + + 0.5 * self._rng.uniform(low=-RANDOM_Y_SHIFT, high=RANDOM_Y_SHIFT) + ) + block_translation = np.array([block_x, block_y, 0]) + return block_translation + + if avoid is None: + block_translation = _get_random_translation() + else: + # Reject targets too close to `avoid`. + for _ in range(NUM_RESET_ATTEMPTS): + block_translation = _get_random_translation() + dist = np.linalg.norm(block_translation[0] - avoid[0]) + # print('block inner try_idx %d, dist %.3f' % (try_idx, dist)) + if dist > MIN_BLOCK_DIST: + break + block_sampled_angle = self._rng.uniform(math.pi) + block_rotation = transform.Rotation.from_rotvec([0, 0, block_sampled_angle]) + self._pybullet_client.resetBasePositionAndOrientation( + self._block_ids[idx], + block_translation.tolist(), + block_rotation.as_quat().tolist(), + ) + return block_translation + + # Reject targets too close to `avoid`. + for _ in range(NUM_RESET_ATTEMPTS): + # Reset first block. + add = 0.2 * self._rng.choice([-1, 1]) + b0_translation = _reset_block_pose(0, add=add) + # Reset second block away from first block. + b1_translation = _reset_block_pose(1, add=-add, avoid=b0_translation) + dist = np.linalg.norm(b0_translation[0] - b1_translation[0]) + if dist > MIN_BLOCK_DIST: + break + else: + raise ValueError("could not find matching block") + assert dist > MIN_BLOCK_DIST + + def _reset_target_poses(self, workspace_center_y): + """Resets target poses.""" + + def _reset_target_pose(idx, add=0.0, avoid=None): + def _get_random_translation(): + # Choose x,y randomly. + target_x = 0.5 + self._rng.uniform( + low=-0.05 * RANDOM_X_SHIFT, high=0.05 * RANDOM_X_SHIFT + ) + target_y = ( + workspace_center_y + + add + + self._rng.uniform( + low=-0.05 * RANDOM_Y_SHIFT, high=0.05 * RANDOM_Y_SHIFT + ) + ) + target_translation = np.array([target_x, target_y, 0.020]) + return target_translation + + if avoid is None: + target_translation = _get_random_translation() + else: + # Reject targets too close to `avoid`. + for _ in range(NUM_RESET_ATTEMPTS): + target_translation = _get_random_translation() + dist = np.linalg.norm(target_translation[0] - avoid[0]) + # print('target inner try_idx %d, dist %.3f' % (try_idx, dist)) + if dist > MIN_TARGET_DIST: + break + target_sampled_angle = math.pi + self._rng.uniform( + low=-math.pi / 30, high=math.pi / 30 + ) + target_rotation = transform.Rotation.from_rotvec( + [0, 0, target_sampled_angle] + ) + self._pybullet_client.resetBasePositionAndOrientation( + self._target_ids[idx], + target_translation.tolist(), + target_rotation.as_quat().tolist(), + ) + self._target_poses[idx] = Pose3d( + rotation=target_rotation, translation=target_translation + ) + + if self._target_poses is None: + self._target_poses = [None for _ in range(len(self._target_ids))] + + for _ in range(NUM_RESET_ATTEMPTS): + # Choose the first target. + add = 0.2 * self._rng.choice([-1, 1]) + # Randomly flip the location of the targets. + _reset_target_pose(0, add=add) + _reset_target_pose(1, add=-add, avoid=self._target_poses[0].translation) + dist = np.linalg.norm( + self._target_poses[0].translation[0] + - self._target_poses[1].translation[0] + ) + break + # if dist > MIN_TARGET_DIST: + # break + else: + raise ValueError("could not find matching target") + # assert dist > MIN_TARGET_DIST + + +if "BlockPushMultimodal-v0" in registration.registry.env_specs: + del registration.registry.env_specs["BlockPushMultimodal-v0"] + +registration.register( + id="BlockPushMultimodal-v0", entry_point=BlockPushMultimodal, max_episode_steps=350 +) + +registration.register( + id="BlockPushMultimodalFlipped-v0", + entry_point=BlockPushHorizontalMultimodal, + max_episode_steps=25, +) + +registration.register( + id="SharedBlockPushMultimodal-v0", + entry_point=BlockPushMultimodal, + kwargs=dict(shared_memory=True), + max_episode_steps=350, +) +registration.register( + id="BlockPushMultimodalRgb-v0", + entry_point=BlockPushMultimodal, + max_episode_steps=350, + kwargs=dict(image_size=(block_pushing.IMAGE_HEIGHT, block_pushing.IMAGE_WIDTH)), +) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/discontinuous_push_oracle.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/discontinuous_push_oracle.py new file mode 100644 index 0000000000000000000000000000000000000000..4d775e1b40f189232e8e5991b2486eddeff19631 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/discontinuous_push_oracle.py @@ -0,0 +1,70 @@ +# coding=utf-8 +# Copyright 2022 The Reach ML Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pushes to first target, waits, then pushes to second target.""" + +import cleandiffuser.env.block_pushing.oracles.oriented_push_oracle as oriented_push_oracle_module +import numpy as np +from tf_agents.trajectories import policy_step +from tf_agents.trajectories import time_step as ts +from tf_agents.typing import types + +# Only used for debug visualization. +import pybullet # pylint: disable=unused-import + + +class DiscontinuousOrientedPushOracle(oriented_push_oracle_module.OrientedPushOracle): + """Pushes to first target, waits, then pushes to second target.""" + + def __init__(self, env, goal_tolerance=0.04, wait=0): + super(DiscontinuousOrientedPushOracle, self).__init__(env) + self._countdown = 0 + self._wait = wait + self._goal_dist_tolerance = goal_tolerance + + def reset(self): + self.phase = "move_to_pre_block" + self._countdown = 0 + + def _action(self, time_step, policy_state): + if time_step.is_first(): + self.reset() + # Move to first target first. + self._current_target = "target" + self._has_switched = False + + def _block_target_dist(block, target): + dist = np.linalg.norm( + time_step.observation["%s_translation" % block] + - time_step.observation["%s_translation" % target] + ) + return dist + + d1 = _block_target_dist("block", "target") + if d1 < self._goal_dist_tolerance and not self._has_switched: + self._countdown = self._wait + # If first block has been pushed to first target, switch to second block. + self._has_switched = True + self._current_target = "target2" + + xy_delta = self._get_action_for_block_target( + time_step, block="block", target=self._current_target + ) + + if self._countdown > 0: + xy_delta = np.zeros_like(xy_delta) + self._countdown -= 1 + + return policy_step.PolicyStep(action=np.asarray(xy_delta, dtype=np.float32)) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/multimodal_push_oracle.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/multimodal_push_oracle.py new file mode 100644 index 0000000000000000000000000000000000000000..360e4c11afd5e730746c57f6627cb202d994d7aa --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/multimodal_push_oracle.py @@ -0,0 +1,187 @@ +# coding=utf-8 +# Copyright 2022 The Reach ML Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Oracle for multimodal pushing task.""" +import cleandiffuser.env.block_pushing.oracles.oriented_push_oracle as oriented_push_oracle_module +import numpy as np +from tf_agents.trajectories import policy_step +from tf_agents.trajectories import time_step as ts +from tf_agents.typing import types + +# Only used for debug visualization. +import pybullet # pylint: disable=unused-import + + +class MultimodalOrientedPushOracle(oriented_push_oracle_module.OrientedPushOracle): + """Oracle for multimodal pushing task.""" + + def __init__(self, env, goal_dist_tolerance=0.04, action_noise_std=0.0): + super(MultimodalOrientedPushOracle, self).__init__(env) + self._goal_dist_tolerance = goal_dist_tolerance + self._action_noise_std = action_noise_std + + def reset(self): + self.origin = None + self.first_preblock = None + self.phase = "move_to_pre_block" + + def _get_move_to_preblock(self, xy_pre_block, xy_ee): + max_step_velocity = 0.3 + # Go 5 cm away from the block, on the line between the block and target. + xy_delta_to_preblock = xy_pre_block - xy_ee + diff = np.linalg.norm(xy_delta_to_preblock) + if diff < 0.001: + self.phase = "move_to_block" + if self.first_preblock is None: + self.first_preblock = np.copy(xy_pre_block) + xy_delta = xy_delta_to_preblock + return xy_delta, max_step_velocity + + def _get_action_for_block_target(self, time_step, block="block", target="target"): + # Specifying this as velocity makes it independent of control frequency. + max_step_velocity = 0.35 + + info = self._get_action_info(time_step, block, target) + + if self.origin is None: + self.origin = np.copy(info.xy_ee) + + if self.phase == "move_to_pre_block": + xy_delta, max_step_velocity = self._get_move_to_preblock( + info.xy_pre_block, info.xy_ee + ) + + if self.phase == "return_to_first_preblock": + max_step_velocity = 0.3 + if self.first_preblock is None: + self.first_preblock = self.origin + # Return to the first preblock. + xy_delta_to_origin = self.first_preblock - info.xy_ee + diff = np.linalg.norm(xy_delta_to_origin) + if diff < 0.001: + self.phase = "return_to_origin" + xy_delta = xy_delta_to_origin + + if self.phase == "return_to_origin": + max_step_velocity = 0.3 + # Go 5 cm away from the block, on the line between the block and target. + xy_delta_to_origin = self.origin - info.xy_ee + diff = np.linalg.norm(xy_delta_to_origin) + if diff < 0.001: + self.phase = "move_to_pre_block" + xy_delta = xy_delta_to_origin + + if self.phase == "move_to_block": + xy_delta = self._get_move_to_block( + info.xy_delta_to_nexttoblock, + info.theta_threshold_to_orient, + info.theta_error, + ) + + if self.phase == "push_block": + xy_delta = self._get_push_block( + info.theta_error, + info.theta_threshold_to_orient, + info.xy_delta_to_touchingblock, + ) + + orient_circle_diameter = 0.025 + + if self.phase == "orient_block_left" or self.phase == "orient_block_right": + max_step_velocity = 0.15 + + if self.phase == "orient_block_left": + xy_delta = self._get_orient_block_left( + info.xy_dir_block_to_ee, + orient_circle_diameter, + info.xy_block, + info.xy_ee, + info.theta_error, + info.theta_threshold_flat_enough, + ) + + if self.phase == "orient_block_right": + xy_delta = self._get_orient_block_right( + info.xy_dir_block_to_ee, + orient_circle_diameter, + info.xy_block, + info.xy_ee, + info.theta_error, + info.theta_threshold_flat_enough, + ) + + if self._action_noise_std != 0.0: + xy_delta += self._np_random_state.randn(2) * self._action_noise_std + + max_step_distance = max_step_velocity * (1 / self._env.get_control_frequency()) + length = np.linalg.norm(xy_delta) + if length > max_step_distance: + xy_direction = xy_delta / length + xy_delta = xy_direction * max_step_distance + return xy_delta + + def _choose_goal_order(self): + """Chooses block->target order for multimodal pushing.""" + # Define all possible ((first_block, first_target), + # (second_block, second_target)). + possible_orders = [ + (("block", "target"), ("block2", "target2")), + (("block", "target2"), ("block2", "target")), + (("block2", "target"), ("block", "target2")), + (("block2", "target2"), ("block", "target")), + ] + # import pdb; pdb.set_trace() + # result = random.choice(possible_orders) + result = possible_orders[self._env._rng.choice(len(possible_orders))] + return result + + def _action(self, time_step, policy_state): + if time_step.is_first(): + self.reset() + ( + (self._first_block, self._first_target), + (self._second_block, self._second_target), + ) = self._choose_goal_order() + self._current_block, self._current_target = ( + self._first_block, + self._first_target, + ) + self._has_switched = False + + def _block_target_dist(block, target): + dist = np.linalg.norm( + time_step.observation["%s_translation" % block] + - time_step.observation["%s_translation" % target] + ) + return dist + + if ( + _block_target_dist(self._first_block, self._first_target) + < self._goal_dist_tolerance + and not self._has_switched + ): + # If first block has been pushed to first target, switch to second block. + self._current_block, self._current_target = ( + self._second_block, + self._second_target, + ) + self._has_switched = True + self.phase = "return_to_first_preblock" + + xy_delta = self._get_action_for_block_target( + time_step, block=self._current_block, target=self._current_target + ) + + return policy_step.PolicyStep(action=np.asarray(xy_delta, dtype=np.float32)) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/oriented_push_oracle.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/oriented_push_oracle.py new file mode 100644 index 0000000000000000000000000000000000000000..abc4621092fe5bcb71cd7889506e870d9228606c --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/oriented_push_oracle.py @@ -0,0 +1,258 @@ +# coding=utf-8 +# Copyright 2022 The Reach ML Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Oracle for pushing task which orients the block then pushes it.""" + +import cleandiffuser.env.block_pushing.oracles.pushing_info as pushing_info_module +import numpy as np +from tf_agents.policies import py_policy +from tf_agents.trajectories import policy_step +from tf_agents.trajectories import time_step as ts +from tf_agents.typing import types + +# Only used for debug visualization. +import pybullet # pylint: disable=unused-import + + +class OrientedPushOracle(py_policy.PyPolicy): + """Oracle for pushing task which orients the block then pushes it.""" + + def __init__(self, env, action_noise_std=0.0): + super(OrientedPushOracle, self).__init__( + env.time_step_spec(), env.action_spec() + ) + self._env = env + self._np_random_state = np.random.RandomState(0) + self.phase = "move_to_pre_block" + self._action_noise_std = action_noise_std + + def reset(self): + self.phase = "move_to_pre_block" + + def get_theta_from_vector(self, vector): + return np.arctan2(vector[1], vector[0]) + + def theta_to_rotation2d(self, theta): + r = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]) + return r + + def rotate(self, theta, xy_dir_block_to_ee): + rot_2d = self.theta_to_rotation2d(theta) + return rot_2d @ xy_dir_block_to_ee + + def _get_action_info(self, time_step, block, target): + xy_block = time_step.observation["%s_translation" % block][:2] + theta_block = time_step.observation["%s_orientation" % block] + xy_target = time_step.observation["%s_translation" % target][:2] + xy_ee = time_step.observation["effector_target_translation"][:2] + + xy_block_to_target = xy_target - xy_block + xy_dir_block_to_target = (xy_block_to_target) / np.linalg.norm( + xy_block_to_target + ) + theta_to_target = self.get_theta_from_vector(xy_dir_block_to_target) + + theta_error = theta_to_target - theta_block + # Block has 4-way symmetry. + while theta_error > np.pi / 4: + theta_error -= np.pi / 2.0 + while theta_error < -np.pi / 4: + theta_error += np.pi / 2.0 + + xy_pre_block = xy_block + -xy_dir_block_to_target * 0.05 + xy_nexttoblock = xy_block + -xy_dir_block_to_target * 0.03 + xy_touchingblock = xy_block + -xy_dir_block_to_target * 0.01 + xy_delta_to_nexttoblock = xy_nexttoblock - xy_ee + xy_delta_to_touchingblock = xy_touchingblock - xy_ee + + xy_block_to_ee = xy_ee - xy_block + xy_dir_block_to_ee = xy_block_to_ee / np.linalg.norm(xy_block_to_ee) + + theta_threshold_to_orient = 0.2 + theta_threshold_flat_enough = 0.03 + return pushing_info_module.PushingInfo( + xy_block=xy_block, + xy_ee=xy_ee, + xy_pre_block=xy_pre_block, + xy_delta_to_nexttoblock=xy_delta_to_nexttoblock, + xy_delta_to_touchingblock=xy_delta_to_touchingblock, + xy_dir_block_to_ee=xy_dir_block_to_ee, + theta_threshold_to_orient=theta_threshold_to_orient, + theta_threshold_flat_enough=theta_threshold_flat_enough, + theta_error=theta_error, + ) + + def _get_move_to_preblock(self, xy_pre_block, xy_ee): + max_step_velocity = 0.3 + # Go 5 cm away from the block, on the line between the block and target. + xy_delta_to_preblock = xy_pre_block - xy_ee + diff = np.linalg.norm(xy_delta_to_preblock) + if diff < 0.001: + self.phase = "move_to_block" + xy_delta = xy_delta_to_preblock + return xy_delta, max_step_velocity + + def _get_move_to_block( + self, xy_delta_to_nexttoblock, theta_threshold_to_orient, theta_error + ): + diff = np.linalg.norm(xy_delta_to_nexttoblock) + if diff < 0.001: + self.phase = "push_block" + # If need to re-oorient, then re-orient. + if theta_error > theta_threshold_to_orient: + self.phase = "orient_block_left" + if theta_error < -theta_threshold_to_orient: + self.phase = "orient_block_right" + # Otherwise, push into the block. + xy_delta = xy_delta_to_nexttoblock + return xy_delta + + def _get_push_block( + self, theta_error, theta_threshold_to_orient, xy_delta_to_touchingblock + ): + # If need to reorient, go back to move_to_pre_block, move_to_block first. + if theta_error > theta_threshold_to_orient: + self.phase = "move_to_pre_block" + if theta_error < -theta_threshold_to_orient: + self.phase = "move_to_pre_block" + xy_delta = xy_delta_to_touchingblock + return xy_delta + + def _get_orient_block_left( + self, + xy_dir_block_to_ee, + orient_circle_diameter, + xy_block, + xy_ee, + theta_error, + theta_threshold_flat_enough, + ): + xy_dir_block_to_ee = self.rotate(0.2, xy_dir_block_to_ee) + xy_block_to_ee = xy_dir_block_to_ee * orient_circle_diameter + xy_push_left_spot = xy_block + xy_block_to_ee + xy_delta = xy_push_left_spot - xy_ee + if theta_error < theta_threshold_flat_enough: + self.phase = "move_to_pre_block" + return xy_delta + + def _get_orient_block_right( + self, + xy_dir_block_to_ee, + orient_circle_diameter, + xy_block, + xy_ee, + theta_error, + theta_threshold_flat_enough, + ): + xy_dir_block_to_ee = self.rotate(-0.2, xy_dir_block_to_ee) + xy_block_to_ee = xy_dir_block_to_ee * orient_circle_diameter + xy_push_left_spot = xy_block + xy_block_to_ee + xy_delta = xy_push_left_spot - xy_ee + if theta_error > -theta_threshold_flat_enough: + self.phase = "move_to_pre_block" + return xy_delta + + def _get_action_for_block_target(self, time_step, block="block", target="target"): + # Specifying this as velocity makes it independent of control frequency. + max_step_velocity = 0.35 + info = self._get_action_info(time_step, block, target) + + if self.phase == "move_to_pre_block": + xy_delta, max_step_velocity = self._get_move_to_preblock( + info.xy_pre_block, info.xy_ee + ) + + if self.phase == "move_to_block": + xy_delta = self._get_move_to_block( + info.xy_delta_to_nexttoblock, + info.theta_threshold_to_orient, + info.theta_error, + ) + + if self.phase == "push_block": + xy_delta = self._get_push_block( + info.theta_error, + info.theta_threshold_to_orient, + info.xy_delta_to_touchingblock, + ) + + orient_circle_diameter = 0.025 + + if self.phase == "orient_block_left" or self.phase == "orient_block_right": + max_step_velocity = 0.15 + + if self.phase == "orient_block_left": + xy_delta = self._get_orient_block_left( + info.xy_dir_block_to_ee, + orient_circle_diameter, + info.xy_block, + info.xy_ee, + info.theta_error, + info.theta_threshold_flat_enough, + ) + + if self.phase == "orient_block_right": + xy_delta = self._get_orient_block_right( + info.xy_dir_block_to_ee, + orient_circle_diameter, + info.xy_block, + info.xy_ee, + info.theta_error, + info.theta_threshold_flat_enough, + ) + + if self._action_noise_std != 0.0: + xy_delta += self._np_random_state.randn(2) * self._action_noise_std + + max_step_distance = max_step_velocity * (1 / self._env.get_control_frequency()) + length = np.linalg.norm(xy_delta) + if length > max_step_distance: + xy_direction = xy_delta / length + xy_delta = xy_direction * max_step_distance + return xy_delta + + def _action(self, time_step, policy_state): + if time_step.is_first(): + self.reset() + xy_delta = self._get_action_for_block_target( + time_step, block="block", target="target" + ) + return policy_step.PolicyStep(action=np.asarray(xy_delta, dtype=np.float32)) + + +class OrientedPushNormalizedOracle(py_policy.PyPolicy): + """Oracle for pushing task which orients the block then pushes it.""" + + def __init__(self, env): + super(OrientedPushNormalizedOracle, self).__init__( + env.time_step_spec(), env.action_spec() + ) + self._oracle = OrientedPushOracle(env) + self._env = env + + def reset(self): + self._oracle.reset() + + def _action(self, time_step, policy_state): + time_step = time_step._asdict() + time_step["observation"] = self._env.calc_unnormalized_state( + time_step["observation"] + ) + step = self._oracle._action( + ts.TimeStep(**time_step), policy_state + ) # pylint: disable=protected-access + return policy_step.PolicyStep( + action=self._env.calc_normalized_action(step.action) + ) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/pushing_info.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/pushing_info.py new file mode 100644 index 0000000000000000000000000000000000000000..0c01cbb0fb6e345d7bc47675d06cfe00a3362f51 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/pushing_info.py @@ -0,0 +1,35 @@ +# coding=utf-8 +# Copyright 2022 The Reach ML Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Dataclass holding info needed for pushing oracles.""" +import dataclasses +from typing import Any + + +@dataclasses.dataclass +class PushingInfo: + """Holds onto info necessary for pushing state machine.""" + + xy_block: Any = None + xy_ee: Any = None + xy_pre_block: Any = None + xy_delta_to_nexttoblock: Any = None + xy_delta_to_touchingblock: Any = None + xy_dir_block_to_ee: Any = None + theta_threshold_to_orient: Any = None + theta_threshold_flat_enough: Any = None + theta_error: Any = None + obstacle_poses: Any = None + distance_to_target: Any = None diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/reach_oracle.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/reach_oracle.py new file mode 100644 index 0000000000000000000000000000000000000000..8b7746cfddc4eec5694286da4593c64b8ac331f0 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/oracles/reach_oracle.py @@ -0,0 +1,61 @@ +# coding=utf-8 +# Copyright 2022 The Reach ML Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reach oracle.""" +import numpy as np +from tf_agents.policies import py_policy +from tf_agents.trajectories import policy_step +from tf_agents.trajectories import time_step as ts +from tf_agents.typing import types + +# Only used for debug visualization. +import pybullet # pylint: disable=unused-import + + +class ReachOracle(py_policy.PyPolicy): + """Oracle for moving to a specific spot relative to the block and target.""" + + def __init__(self, env, block_pushing_oracles_action_std=0.0): + super(ReachOracle, self).__init__(env.time_step_spec(), env.action_spec()) + self._env = env + self._np_random_state = np.random.RandomState(0) + self._block_pushing_oracles_action_std = block_pushing_oracles_action_std + + def _action(self, time_step, policy_state): + + # Specifying this as velocity makes it independent of control frequency. + max_step_velocity = 0.2 + + xy_ee = time_step.observation["effector_target_translation"] + + # This should be observable from block and target translation, + # but re-using the computation from the env so that it's only done once, and + # used for reward / completion computation. + xy_pre_block = self._env.reach_target_translation + + xy_delta = xy_pre_block - xy_ee + + if self._block_pushing_oracles_action_std != 0.0: + xy_delta += ( + self._np_random_state.randn(2) * self._block_pushing_oracles_action_std + ) + + max_step_distance = max_step_velocity * (1 / self._env.get_control_frequency()) + length = np.linalg.norm(xy_delta) + if length > max_step_distance: + xy_direction = xy_delta / length + xy_delta = xy_direction * max_step_distance + + return policy_step.PolicyStep(action=np.asarray(xy_delta, dtype=np.float32)) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/utils/pose3d.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/utils/pose3d.py new file mode 100644 index 0000000000000000000000000000000000000000..a25beeb18b90d0204b36861bce10df44d632c025 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/utils/pose3d.py @@ -0,0 +1,70 @@ +# coding=utf-8 +# Copyright 2022 The Reach ML Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A simple 6DOF pose container. +""" + +import dataclasses +import numpy as np +from scipy.spatial import transform + + +class NoCopyAsDict(object): + """Base class for dataclasses. Avoids a copy in the asdict() call.""" + + def asdict(self): + """Replacement for dataclasses.asdict. + + TF Dataset does not handle dataclasses.asdict, which uses copy.deepcopy when + setting values in the output dict. This causes issues with tf.Dataset. + Instead, shallow copy contents. + + Returns: + dict containing contents of dataclass. + """ + return {k.name: getattr(self, k.name) for k in dataclasses.fields(self)} + + +@dataclasses.dataclass +class Pose3d(NoCopyAsDict): + """Simple container for translation and rotation.""" + + rotation: transform.Rotation + translation: np.ndarray + + @property + def vec7(self): + return np.concatenate([self.translation, self.rotation.as_quat()]) + + def serialize(self): + return { + "rotation": self.rotation.as_quat().tolist(), + "translation": self.translation.tolist(), + } + + @staticmethod + def deserialize(data): + return Pose3d( + rotation=transform.Rotation.from_quat(data["rotation"]), + translation=np.array(data["translation"]), + ) + + def __eq__(self, other): + return np.array_equal( + self.rotation.as_quat(), other.rotation.as_quat() + ) and np.array_equal(self.translation, other.translation) + + def __ne__(self, other): + return not self.__eq__(other) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/utils/utils_pybullet.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/utils/utils_pybullet.py new file mode 100644 index 0000000000000000000000000000000000000000..0a1894b9ea2e57ee4defa6fe2ef766bfb948013d --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/utils/utils_pybullet.py @@ -0,0 +1,451 @@ +# coding=utf-8 +# Copyright 2022 The Reach ML Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Assortment of utilities to interact with bullet within g3.""" +import dataclasses +import datetime +import getpass +import gzip +import json +import os +import time +from typing import Any, Dict, List, Optional, Tuple + +from absl import logging +from cleandiffuser.env.block_pushing.utils.pose3d import Pose3d +import numpy as np +from scipy.spatial import transform +import six + + +import pybullet +import pybullet_data +import pybullet_utils.bullet_client as bullet_client + +Vec3 = Tuple[float, float, float] +Vec4 = Tuple[float, float, float, float] +PYBULLET_STATE_VERSION = 2 # Basic versioning of serialized pybullet state. + + +# Note about rotation_to_matrix and matrix_to_rotation below: +# The abstractions below allow us to use older versions of scipy. +def rotation_to_matrix(rotation): + if hasattr(rotation, "as_dcm"): + return rotation.as_dcm() + else: + assert hasattr(rotation, "as_matrix") + return rotation.as_matrix() + + +def matrix_to_rotation(matrix): + if hasattr(transform.Rotation, "from_dcm"): + return transform.Rotation.from_dcm(matrix) + else: + assert hasattr(transform.Rotation, "from_matrix") + return transform.Rotation.from_matrix(matrix) + + +def load_urdf(pybullet_client, file_path, *args, **kwargs): + """Loads the given URDF filepath.""" + + # Handles most general file open case. + try: + if os.path.exists(file_path): + return pybullet_client.loadURDF(file_path, *args, **kwargs) + except pybullet_client.error: + pass + + try: + import pathlib + asset_path = str(pathlib.Path(__file__).parent.parent.joinpath('assets')) + if file_path.startswith("third_party/py/envs/assets/"): + pybullet_client.setAdditionalSearchPath(asset_path) + file_path = file_path[len("third_party/py/envs/assets/") :] + if file_path.startswith( + "third_party/bullet/examples/pybullet/gym/pybullet_data/" + ): + pybullet_client.setAdditionalSearchPath(pybullet_data.getDataPath()) + file_path = file_path[55:] + # logging.info("Loading URDF %s", file_path) + return pybullet_client.loadURDF(file_path, *args, **kwargs) + except pybullet.error: + raise FileNotFoundError("Cannot load the URDF file {}".format(file_path)) + + +def add_visual_sphere(client, center=(0, 0, 0), radius=0.1, rgba=(0.5, 0.5, 0.5, 0.5)): + """Add a sphere to bullet scene (visual only, no physics). + + Args: + client: pybullet client (or pybullet library handle). + center: Center of sphere. + radius: Sphere radius. + rgba: rgba color of sphere. + + Returns: + Unique integer bullet id of constructed object. + """ + vis_obj_id = client.createVisualShape( + client.GEOM_SPHERE, radius=radius, rgbaColor=rgba + ) + obj_id = client.createMultiBody( + baseCollisionShapeIndex=-1, baseVisualShapeIndex=vis_obj_id, basePosition=center + ) + return obj_id + + +def pybullet_mat_to_numpy_4x4(pybullet_matrix): + assert len(pybullet_matrix) == 16, "pybullet matrix should be len 16" + return np.transpose(np.reshape(np.array(pybullet_matrix, dtype=np.float64), (4, 4))) + + +def decompose_view_matrix(pybullet_view_matrix): + """Decompose view matrix into pos + quat format (assumes mat is rigid!).""" + # It would be MUCH better to use something from bullet, however pybullet does + # not expose all of the linear algebra library. + mat = pybullet_mat_to_numpy_4x4(pybullet_view_matrix) + + # View matrix is now: + # | R_11 R_12 R_13 t_1 | + # | R_21 R_22 R_23 t_2 | + # | R_31 R_32 R_33 t_3 | + # | 0 0 0 1 | + + # R is the inverse eye to target at orientation, and t is R * eye. + mat_view_to_world = np.linalg.inv(mat) + + # mat_view_to_world is the view to world transform, therefore the translation + # component of this matrix is simply the world space position (since mat * + # (0, 0, 0, 1)) is just copying the right column. + world_xyz_view = np.copy(mat_view_to_world[0:3, 3]) + + mat_view_to_world[0:3, 3] = 0 # Zero out the position change. + world_quat_view = matrix_to_rotation(mat_view_to_world).as_quat() + + return world_xyz_view, world_quat_view + + +def world_obj_to_view(world_xyz_obj, world_quat_obj, camera_view, client): + """Transform object into view space.""" + world_xyz_view, world_quat_view = decompose_view_matrix(camera_view) + view_xyz_world, view_quat_world = client.invertTransform( + world_xyz_view, world_quat_view + ) + view_xyz_obj, view_quat_obj = client.multiplyTransforms( + view_xyz_world, view_quat_world, world_xyz_obj, world_quat_obj + ) + + return view_xyz_obj, view_quat_obj + + +def image_xy_to_view_ray(xy, cam_width, cam_height, proj_mat_inv): + """Calculate view-space ray from pixel location.""" + # Recall (from http://www.songho.ca/opengl/gl_projectionmatrix.html): + # xyzw_clip = M_proj * xyzw_eye, and + # xyz_ndc = xyzw_clip[0:3] / xwzw_clip[3]. + xyz_ndc = np.array( + [2.0 * xy[0] / cam_width - 1.0, -(2.0 * xy[1] / cam_height - 1.0), 0] + ) # in [-1, 1] + xyzw_clip = np.concatenate([xyz_ndc, [1]]) + xyzw_eye = proj_mat_inv @ xyzw_clip + origin = np.zeros(3) + vec = xyzw_eye[:3] / max(np.linalg.norm(xyzw_eye[:3]), 1e-6) + return origin, vec + + +def view_ray_to_world_ray(origin, vec, view_mat_inv): + """Transform view-space ray into world space.""" + origin = view_mat_inv @ np.concatenate([origin, [1]]) + vec = view_mat_inv @ np.concatenate([vec, [0]]) + + return origin[:3], vec[:3] + + +def ray_to_plane_test(ray_origin, ray_vec, plane_origin, plane_normal): + """Perform a ray-plane intersection test.""" + ln = np.dot(plane_normal, ray_vec) + if abs(ln) < np.finfo(np.float32).eps: + return None + + # Solve for the intersection fraction t. + t = np.dot(plane_normal, plane_origin - ray_origin) / ln + if t >= 0: + return ray_origin + ray_vec * t + else: + return None + + +def get_workspace(env): + ( + workspace_origin, + workspace_quat, + ) = env.pybullet_client.getBasePositionAndOrientation(env.workspace_uid) + workspace_normal = rotation_to_matrix(transform.Rotation.from_quat(workspace_quat))[ + 2, 0:3 + ] + + return workspace_origin, workspace_normal + + +def reset_camera_pose(env, view_type): + """Reset camera pose to canonical frame.""" + p = env.pybullet_client + + if view_type == "POLICY": + camera_info = p.getDebugVisualizerCamera() + image_size = (camera_info[0], camera_info[1]) + + viewm, _, front_position, lookat, _ = env.calc_camera_params(image_size) + + euler = matrix_to_rotation(pybullet_mat_to_numpy_4x4(viewm)[0:3, 0:3]).as_euler( + "xyz", degrees=False + ) + pitch = euler[1] + yaw = -euler[2] + # The distance is a bit far away (the GL view has higher FOV). + distance = np.linalg.norm(front_position - lookat) * 0.6 + elif view_type == "TOP_DOWN": + workspace_origin, _ = get_workspace(env) + distance = 0.5 + lookat = workspace_origin + yaw = np.pi / 2 + # Note: pi/2 pitch results in gimble lock and pybullet doesn't support it. + pitch = -(np.pi / 2 - 1e-5) + else: + raise ValueError("unsupported view_type %s" % view_type) + p.resetDebugVisualizerCamera( + cameraDistance=distance, + cameraYaw=360 * yaw / (2.0 * np.pi), + cameraPitch=360 * pitch / (2.0 * np.pi), + cameraTargetPosition=lookat, + ) + + +def _lists_to_tuple(obj): + if isinstance(obj, list): + return tuple([_lists_to_tuple(v) for v in obj]) + else: + return obj + + +@dataclasses.dataclass +class ObjState: + """A container for storing pybullet object state.""" + + obj_id: int + + # base_pose: (xyz, quat). + base_pose: Tuple[Vec3, Vec4] + # base_vel: (vel, ang_vel). + base_vel: Tuple[Vec3, Vec3] + joint_info: Any + joint_state: Any + + @staticmethod + def get_bullet_state(client, obj_id): + """Read Pybullet internal state.""" + base_pose = client.getBasePositionAndOrientation(obj_id) + base_vel = client.getBaseVelocity(obj_id) + + joint_info = [] + joint_state = [] + for i in range(client.getNumJoints(obj_id)): + joint_state.append(client.getJointState(obj_id, i)) + joint_info.append(ObjState._get_joint_info(client, obj_id, i)) + + return ObjState( + obj_id=obj_id, + base_pose=base_pose, + base_vel=base_vel, + joint_info=tuple(joint_info), + joint_state=tuple(joint_state), + ) + + @staticmethod + def _get_joint_info(client, obj_id, joint_index): + ji = client.getJointInfo(obj_id, joint_index) + return tuple([v if not isinstance(v, bytes) else v.decode("utf-8") for v in ji]) + + def set_bullet_state(self, client, obj_id): + """Hard set the current bullet state.""" + xyz, quat = self.base_pose + client.resetBasePositionAndOrientation(obj_id, xyz, quat) + vel, ang_vel = self.base_vel + client.resetBaseVelocity(obj_id, vel, ang_vel) + + njoints = client.getNumJoints(obj_id) + if njoints != len(self.joint_info) or njoints != len(self.joint_state): + raise ValueError("Incorrect number of joint info state pairs.") + + for i, (joint_info, joint_state) in enumerate( + zip(self.joint_info, self.joint_state) + ): + joint_index = joint_info[0] + if joint_index != i: + raise ValueError("Joint index mismatch.") + + # Check that the current joint we're trying to restore state for has the + # same info as the state joint. + cur_joint_info = ObjState._get_joint_info(client, obj_id, joint_index) + if cur_joint_info != joint_info: + raise ValueError( + "joint_info mismatch %s vs %s (expected)" + % (str(cur_joint_info), str(joint_info)) + ) + joint_position = joint_state[0] + joint_velocity = joint_state[1] + client.resetJointState( + obj_id, i, targetValue=joint_position, targetVelocity=joint_velocity + ) + + def serialize(self): + return { + "obj_id": self.obj_id, + "base_pose": self.base_pose, + "base_vel": self.base_vel, + "joint_info": self.joint_info, + "joint_state": self.joint_state, + } + + @staticmethod + def deserialize(data): + return ObjState( + obj_id=_lists_to_tuple(data["obj_id"]), + base_pose=_lists_to_tuple(data["base_pose"]), + base_vel=_lists_to_tuple(data["base_vel"]), + joint_info=_lists_to_tuple(data["joint_info"]), + joint_state=_lists_to_tuple(data["joint_state"]), + ) + + +@dataclasses.dataclass +class XarmState(ObjState): + """A container for storing pybullet robot state.""" + + # The set point of the robot's controller. + target_effector_pose: Pose3d + goal_translation: Optional[Vec3] + + @staticmethod + def get_bullet_state(client, obj_id, target_effector_pose, goal_translation): + if goal_translation is not None: + goal_translation = tuple(goal_translation.tolist()) + return XarmState( + **dataclasses.asdict(ObjState.get_bullet_state(client, obj_id)), + target_effector_pose=target_effector_pose, + goal_translation=goal_translation + ) + + def serialize(self): + data = ObjState.serialize(self) + data["target_effector_pose"] = self.target_effector_pose.serialize() + if self.goal_translation is not None: + data["goal_translation"] = self.goal_translation + else: + data["goal_translation"] = [] + return data + + @staticmethod + def deserialize(data): + goal_translation = ( + None + if not data["goal_translation"] + else _lists_to_tuple(data["goal_translation"]) + ) + return XarmState( + obj_id=data["obj_id"], + base_pose=_lists_to_tuple(data["base_pose"]), + base_vel=_lists_to_tuple(data["base_vel"]), + joint_info=_lists_to_tuple(data["joint_info"]), + joint_state=_lists_to_tuple(data["joint_state"]), + goal_translation=goal_translation, + target_effector_pose=Pose3d.deserialize(data["target_effector_pose"]), + ) + + +def _serialize_pybullet_state(pybullet_state): + """Convert data to POD types.""" + if isinstance(pybullet_state, list): + return [_serialize_pybullet_state(entry) for entry in pybullet_state] + elif isinstance(pybullet_state, dict): + assert "_serialized_obj_name" not in pybullet_state + return { + key: _serialize_pybullet_state(value) + for key, value in pybullet_state.items() + } + elif isinstance(pybullet_state, (XarmState, ObjState)): + return { + "_serialized_obj_name": type(pybullet_state).__name__, + "_serialized_data": pybullet_state.serialize(), + } + elif isinstance(pybullet_state, int): + return pybullet_state + else: + raise ValueError( + "Unhandled type for object %s, type %s" + % (str(pybullet_state), type(pybullet_state)) + ) + + +def _deserialize_pybullet_state(state): + """Parse data from POD types.""" + if isinstance(state, list): + return [_deserialize_pybullet_state(item) for item in state] + elif isinstance(state, dict): + if "_serialized_obj_name" in state: + if state["_serialized_obj_name"] == XarmState.__name__: + return XarmState.deserialize(state["_serialized_data"]) + elif state["_serialized_obj_name"] == ObjState.__name__: + return ObjState.deserialize(state["_serialized_data"]) + else: + raise ValueError("Unsupported: %s" % state["_serialized_obj_name"]) + else: + return { + key: _deserialize_pybullet_state(value) for key, value in state.items() + } + elif isinstance(state, int): + return state + else: + raise ValueError("Unhandled type for object %s" % str(state)) + + +def write_pybullet_state(filename, pybullet_state, task, actions=None): + """Serialize pybullet state to json file.""" + import torch + data = { + "pybullet_state": _serialize_pybullet_state(pybullet_state), + "state_version": PYBULLET_STATE_VERSION, + "ts_ms": int(time.mktime(datetime.datetime.now().timetuple())) * 1000, + "user": getpass.getuser(), + "task": task, + "actions": actions if actions is not None else [], + } + torch.save(data, filename) + + +def read_pybullet_state(filename): + """Deserialize pybullet state from json file.""" + import torch + data = torch.load(filename) + + assert isinstance(data, dict) + + if data["state_version"] != PYBULLET_STATE_VERSION: + raise ValueError( + "incompatible state data (version %d, expected %d)" + % (data["state_version"], PYBULLET_STATE_VERSION) + ) + + data["pybullet_state"] = _deserialize_pybullet_state(data["pybullet_state"]) + return data diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/utils/xarm_sim_robot.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/utils/xarm_sim_robot.py new file mode 100644 index 0000000000000000000000000000000000000000..eb15a9f02dad6f3441a072c2ca2e59adb802a441 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/block_pushing/utils/xarm_sim_robot.py @@ -0,0 +1,230 @@ +# coding=utf-8 +# Copyright 2022 The Reach ML Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""XArm Robot Kinematics.""" +from cleandiffuser.env.block_pushing.utils import utils_pybullet +from cleandiffuser.env.block_pushing.utils.pose3d import Pose3d +import numpy as np +from scipy.spatial import transform +import pybullet + +XARM_URDF_PATH = ( + "third_party/bullet/examples/pybullet/gym/pybullet_data/" "xarm/xarm6_robot.urdf" +) +SUCTION_URDF_PATH = "third_party/py/envs/assets/suction/" "suction-head-long.urdf" +CYLINDER_URDF_PATH = "third_party/py/envs/assets/suction/" "cylinder.urdf" +CYLINDER_REAL_URDF_PATH = "third_party/py/envs/assets/suction/" "cylinder_real.urdf" +HOME_JOINT_POSITIONS = np.deg2rad([0, -20, -80, 0, 100, -30]) + + +class XArmSimRobot: + """A simulated PyBullet XArm robot, mostly for forward/inverse kinematics.""" + + def __init__( + self, + pybullet_client, + initial_joint_positions=HOME_JOINT_POSITIONS, + end_effector="none", + color="default", + ): + self._pybullet_client = pybullet_client + self.initial_joint_positions = initial_joint_positions + + if color == "default": + self.xarm = utils_pybullet.load_urdf( + pybullet_client, XARM_URDF_PATH, [0, 0, 0] + ) + else: + raise ValueError("Unrecognized xarm color %s" % color) + + # Get revolute joints of robot (skip fixed joints). + joints = [] + joint_indices = [] + for i in range(self._pybullet_client.getNumJoints(self.xarm)): + joint_info = self._pybullet_client.getJointInfo(self.xarm, i) + if joint_info[2] == pybullet.JOINT_REVOLUTE: + joints.append(joint_info[0]) + joint_indices.append(i) + # Note examples in pybullet do this, but it is not clear what the + # benefits are. + self._pybullet_client.changeDynamics( + self.xarm, i, linearDamping=0, angularDamping=0 + ) + + self._n_joints = len(joints) + self._joints = tuple(joints) + self._joint_indices = tuple(joint_indices) + + # Move robot to home joint configuration + self.reset_joints(self.initial_joint_positions) + self.effector_link = 6 + + if ( + end_effector == "suction" + or end_effector == "cylinder" + or end_effector == "cylinder_real" + ): + self.end_effector = self._setup_end_effector(end_effector) + else: + if end_effector != "none": + raise ValueError('end_effector "%s" is not supported.' % end_effector) + self.end_effector = None + + def _setup_end_effector(self, end_effector): + """Adds a suction or cylinder end effector.""" + pose = self.forward_kinematics() + if end_effector == "suction": + body = utils_pybullet.load_urdf( + self._pybullet_client, + SUCTION_URDF_PATH, + pose.translation, + pose.rotation.as_quat(), + ) + elif end_effector == "cylinder": + body = utils_pybullet.load_urdf( + self._pybullet_client, + CYLINDER_URDF_PATH, + pose.translation, + pose.rotation.as_quat(), + ) + elif end_effector == "cylinder_real": + body = utils_pybullet.load_urdf( + self._pybullet_client, + CYLINDER_REAL_URDF_PATH, + pose.translation, + pose.rotation.as_quat(), + ) + else: + raise ValueError('end_effector "%s" is not supported.' % end_effector) + + constraint_id = self._pybullet_client.createConstraint( + parentBodyUniqueId=self.xarm, + parentLinkIndex=6, + childBodyUniqueId=body, + childLinkIndex=-1, + jointType=pybullet.JOINT_FIXED, + jointAxis=(0, 0, 0), + parentFramePosition=(0, 0, 0), + childFramePosition=(0, 0, 0), + ) + self._pybullet_client.changeConstraint(constraint_id, maxForce=50) + + return body + + def reset_joints(self, joint_values): + """Sets the position of the Robot's joints. + + *Note*: This should only be used at the start while not running the + simulation resetJointState overrides all physics simulation. + + Args: + joint_values: Iterable with desired joint positions. + """ + for i in range(self._n_joints): + self._pybullet_client.resetJointState( + self.xarm, self._joints[i], joint_values[i] + ) + + def get_joints_measured(self): + joint_states = self._pybullet_client.getJointStates( + self.xarm, self._joint_indices + ) + joint_positions = np.array([state[0] for state in joint_states]) + joint_velocities = np.array([state[1] for state in joint_states]) + joint_torques = np.array([state[3] for state in joint_states]) + return joint_positions, joint_velocities, joint_torques + + def get_joint_positions(self): + joint_states = self._pybullet_client.getJointStates( + self.xarm, self._joint_indices + ) + joint_positions = np.array([state[0] for state in joint_states]) + return joint_positions + + def forward_kinematics(self): + """Forward kinematics.""" + effector_state = self._pybullet_client.getLinkState( + self.xarm, self.effector_link + ) + return Pose3d( + translation=np.array(effector_state[0]), + rotation=transform.Rotation.from_quat(effector_state[1]), + ) + + def inverse_kinematics( + self, world_effector_pose, max_iterations=100, residual_threshold=1e-10 + ): + """Inverse kinematics. + + Args: + world_effector_pose: Target Pose3d for the robot's end effector. + max_iterations: Refine the IK solution until the distance between target + and actual end effector position is below this threshold, or the + maxNumIterations is reached. Default is 20 iterations. + residual_threshold: Refine the IK solution until the distance between + target and actual end effector position is below this threshold, or the + maxNumIterations is reached. + + Returns: + Numpy array with required joint angles to reach the requested pose. + """ + return np.array( + self._pybullet_client.calculateInverseKinematics( + self.xarm, + self.effector_link, + world_effector_pose.translation, + world_effector_pose.rotation.as_quat(), # as_quat returns xyzw. + lowerLimits=[-17] * 6, + upperLimits=[17] * 6, + jointRanges=[17] * 6, + restPoses=[0, 0] + self.get_joint_positions()[2:].tolist(), + maxNumIterations=max_iterations, + residualThreshold=residual_threshold, + ) + ) + + def set_target_effector_pose(self, world_effector_pose): + target_joint_positions = self.inverse_kinematics(world_effector_pose) + self.set_target_joint_positions(target_joint_positions) + + def set_target_joint_velocities(self, target_joint_velocities): + self._pybullet_client.setJointMotorControlArray( + self.xarm, + self._joint_indices, + pybullet.VELOCITY_CONTROL, + targetVelocities=target_joint_velocities, + forces=[5 * 240.0] * 6, + ) + + def set_target_joint_positions(self, target_joint_positions): + self._pybullet_client.setJointMotorControlArray( + self.xarm, + self._joint_indices, + pybullet.POSITION_CONTROL, + targetPositions=target_joint_positions, + forces=[5 * 240.0] * 6, + ) + + def set_alpha_transparency(self, alpha): + visual_shape_data = self._pybullet_client.getVisualShapeData(self.xarm) + + for i in range(self._pybullet_client.getNumJoints(self.xarm)): + object_id, link_index, _, _, _, _, _, rgba_color = visual_shape_data[i] + assert object_id == self.xarm, "xarm id mismatch." + assert link_index == i, "Link visual data was returned out of order." + rgba_color = list(rgba_color[0:3]) + [alpha] + self._pybullet_client.changeVisualShape( + self.xarm, linkIndex=i, rgbaColor=rgba_color + ) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..be6fd3d494398a465d8b3c903985f439a8d8d739 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/__init__.py @@ -0,0 +1,30 @@ +"""Environments using kitchen and Franka robot.""" +from gym.envs.registration import register + +register( + id="kitchen-microwave-kettle-light-slider-v0", + entry_point="cleandiffuser.env.kitchen.v0:KitchenMicrowaveKettleLightSliderV0", + max_episode_steps=280, + reward_threshold=1.0, +) + +register( + id="kitchen-microwave-kettle-burner-light-v0", + entry_point="cleandiffuser.env.kitchen.v0:KitchenMicrowaveKettleBottomBurnerLightV0", + max_episode_steps=280, + reward_threshold=1.0, +) + +register( + id="kitchen-kettle-microwave-light-slider-v0", + entry_point="cleandiffuser.env.kitchen.v0:KitchenKettleMicrowaveLightSliderV0", + max_episode_steps=280, + reward_threshold=1.0, +) + +register( + id="kitchen-all-v0", + entry_point="cleandiffuser.env.kitchen.v0:KitchenAllV0", + max_episode_steps=280, + reward_threshold=1.0, +) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/base.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/base.py new file mode 100644 index 0000000000000000000000000000000000000000..a0c5e1aaec6bbd19326a6cb1481fb6730c8a2510 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/base.py @@ -0,0 +1,156 @@ +import sys +import os +# hack to import adept envs +ADEPT_DIR = os.path.join(os.path.dirname(__file__), 'relay_policy_learning', 'adept_envs') +sys.path.append(ADEPT_DIR) + +import logging +import numpy as np +import adept_envs +from adept_envs.franka.kitchen_multitask_v0 import KitchenTaskRelaxV1 + +OBS_ELEMENT_INDICES = { + "bottom burner": np.array([11, 12]), + "top burner": np.array([15, 16]), + "light switch": np.array([17, 18]), + "slide cabinet": np.array([19]), + "hinge cabinet": np.array([20, 21]), + "microwave": np.array([22]), + "kettle": np.array([23, 24, 25, 26, 27, 28, 29]), +} +OBS_ELEMENT_GOALS = { + "bottom burner": np.array([-0.88, -0.01]), + "top burner": np.array([-0.92, -0.01]), + "light switch": np.array([-0.69, -0.05]), + "slide cabinet": np.array([0.37]), + "hinge cabinet": np.array([0.0, 1.45]), + "microwave": np.array([-0.75]), + "kettle": np.array([-0.23, 0.75, 1.62, 0.99, 0.0, 0.0, -0.06]), +} +BONUS_THRESH = 0.3 +logger = logging.getLogger() + + +class KitchenBase(KitchenTaskRelaxV1): + # A string of element names. The robot's task is then to modify each of + # these elements appropriately. + TASK_ELEMENTS = [] + ALL_TASKS = [ + "bottom burner", + "top burner", + "light switch", + "slide cabinet", + "hinge cabinet", + "microwave", + "kettle", + ] + REMOVE_TASKS_WHEN_COMPLETE = True + TERMINATE_ON_TASK_COMPLETE = True + TERMINATE_ON_WRONG_COMPLETE = False + COMPLETE_IN_ANY_ORDER = ( + True # This allows for the tasks to be completed in arbitrary order. + ) + + def __init__( + self, dataset_url=None, ref_max_score=None, ref_min_score=None, + use_abs_action=False, + **kwargs + ): + self.tasks_to_complete = list(self.TASK_ELEMENTS) + self.goal_masking = True + super(KitchenBase, self).__init__(use_abs_action=use_abs_action, **kwargs) + + def set_goal_masking(self, goal_masking=True): + """Sets goal masking for goal-conditioned approaches (like RPL).""" + self.goal_masking = goal_masking + + def _get_task_goal(self, task=None, actually_return_goal=False): + if task is None: + task = ["microwave", "kettle", "bottom burner", "light switch"] + new_goal = np.zeros_like(self.goal) + if self.goal_masking and not actually_return_goal: + return new_goal + for element in task: + element_idx = OBS_ELEMENT_INDICES[element] + element_goal = OBS_ELEMENT_GOALS[element] + new_goal[element_idx] = element_goal + + return new_goal + + def reset_model(self): + self.tasks_to_complete = list(self.TASK_ELEMENTS) + return super(KitchenBase, self).reset_model() + + def get_observation(self): + return self._get_obs() + + def _get_reward_n_score(self, obs_dict): + reward_dict, score = super(KitchenBase, self)._get_reward_n_score(obs_dict) + reward = 0.0 + next_q_obs = obs_dict["qp"] + next_obj_obs = obs_dict["obj_qp"] + next_goal = self._get_task_goal( + task=self.TASK_ELEMENTS, actually_return_goal=True + ) # obs_dict['goal'] + idx_offset = len(next_q_obs) + completions = [] + all_completed_so_far = True + for element in self.tasks_to_complete: + element_idx = OBS_ELEMENT_INDICES[element] + distance = np.linalg.norm( + next_obj_obs[..., element_idx - idx_offset] - next_goal[element_idx] + ) + complete = distance < BONUS_THRESH + condition = ( + complete and all_completed_so_far + if not self.COMPLETE_IN_ANY_ORDER + else complete + ) + if condition: # element == self.tasks_to_complete[0]: + print("Task {} completed!".format(element)) + completions.append(element) + all_completed_so_far = all_completed_so_far and complete + if self.REMOVE_TASKS_WHEN_COMPLETE: + [self.tasks_to_complete.remove(element) for element in completions] + bonus = float(len(completions)) + reward_dict["bonus"] = bonus + reward_dict["r_total"] = bonus + score = bonus + return reward_dict, score + + def step(self, a, b=None): + obs, reward, done, env_info = super(KitchenBase, self).step(a, b=b) + if self.TERMINATE_ON_TASK_COMPLETE: + done = not self.tasks_to_complete + if self.TERMINATE_ON_WRONG_COMPLETE: + all_goal = self._get_task_goal(task=self.ALL_TASKS) + for wrong_task in list(set(self.ALL_TASKS) - set(self.TASK_ELEMENTS)): + element_idx = OBS_ELEMENT_INDICES[wrong_task] + distance = np.linalg.norm(obs[..., element_idx] - all_goal[element_idx]) + complete = distance < BONUS_THRESH + if complete: + done = True + break + env_info["completed_tasks"] = set(self.TASK_ELEMENTS) - set( + self.tasks_to_complete + ) + return obs, reward, done, env_info + + def get_goal(self): + """Loads goal state from dataset for goal-conditioned approaches (like RPL).""" + raise NotImplementedError + + def _split_data_into_seqs(self, data): + """Splits dataset object into list of sequence dicts.""" + seq_end_idxs = np.where(data["terminals"])[0] + start = 0 + seqs = [] + for end_idx in seq_end_idxs: + seqs.append( + dict( + states=data["observations"][start : end_idx + 1], + actions=data["actions"][start : end_idx + 1], + ) + ) + start = end_idx + 1 + return seqs diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/kitchen_lowdim_wrapper.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/kitchen_lowdim_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..7c99f752cce4588136f24270bcdada96c3ecb265 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/kitchen_lowdim_wrapper.py @@ -0,0 +1,56 @@ +from typing import List, Dict, Optional, Optional +import numpy as np +import gym +from gym.spaces import Box +from cleandiffuser.env.kitchen.base import KitchenBase + +class KitchenLowdimWrapper(gym.Env): + def __init__(self, + env: KitchenBase, + init_qpos: Optional[np.ndarray]=None, + init_qvel: Optional[np.ndarray]=None, + render_hw = (240,360) + ): + self.env = env + self.all_tasks = env.ALL_TASKS + self.init_qpos = init_qpos + self.init_qvel = init_qvel + self.render_hw = render_hw + + @property + def action_space(self): + return self.env.action_space + + @property + def observation_space(self): + return self.env.observation_space + + @property + def tasks_to_complete(self): + return self.env.tasks_to_complete + + def set_init_q(self, init_qpos, init_qvel): + self.init_qpos = init_qpos + self.init_qvel = init_qvel + + def seed(self, seed=None): + return self.env.seed(seed) + + def reset(self): + if self.init_qpos is not None: + print("Load known initial conditions.") + # reset anyway to be safe, not very expensive + _ = self.env.reset() + # start from known state + self.env.set_state(self.init_qpos, self.init_qvel) + obs = self.env.get_observation() + return obs + else: + return self.env.reset() + + def render(self, mode='rgb_array'): + h, w = self.render_hw + return self.env.render(mode=mode, width=w, height=h) + + def step(self, a): + return self.env.step(a) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/kitchen_util.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/kitchen_util.py new file mode 100644 index 0000000000000000000000000000000000000000..59852c1946c8cc1eea17dbd9bd7bca73df095ebe --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/kitchen_util.py @@ -0,0 +1,51 @@ +import struct +import numpy as np + +def parse_mjl_logs(read_filename, skipamount): + with open(read_filename, mode='rb') as file: + fileContent = file.read() + headers = struct.unpack('iiiiiii', fileContent[:28]) + nq = headers[0] + nv = headers[1] + nu = headers[2] + nmocap = headers[3] + nsensordata = headers[4] + nuserdata = headers[5] + name_len = headers[6] + name = struct.unpack(str(name_len) + 's', fileContent[28:28+name_len])[0] + rem_size = len(fileContent[28 + name_len:]) + num_floats = int(rem_size/4) + dat = np.asarray(struct.unpack(str(num_floats) + 'f', fileContent[28+name_len:])) + recsz = 1 + nq + nv + nu + 7*nmocap + nsensordata + nuserdata + if rem_size % recsz != 0: + print("ERROR") + else: + dat = np.reshape(dat, (int(len(dat)/recsz), recsz)) + dat = dat.T + + time = dat[0,:][::skipamount] - 0*dat[0, 0] + qpos = dat[1:nq + 1, :].T[::skipamount, :] + qvel = dat[nq+1:nq+nv+1,:].T[::skipamount, :] + ctrl = dat[nq+nv+1:nq+nv+nu+1,:].T[::skipamount,:] + mocap_pos = dat[nq+nv+nu+1:nq+nv+nu+3*nmocap+1,:].T[::skipamount, :] + mocap_quat = dat[nq+nv+nu+3*nmocap+1:nq+nv+nu+7*nmocap+1,:].T[::skipamount, :] + sensordata = dat[nq+nv+nu+7*nmocap+1:nq+nv+nu+7*nmocap+nsensordata+1,:].T[::skipamount,:] + userdata = dat[nq+nv+nu+7*nmocap+nsensordata+1:,:].T[::skipamount,:] + + data = dict(nq=nq, + nv=nv, + nu=nu, + nmocap=nmocap, + nsensordata=nsensordata, + name=name, + time=time, + qpos=qpos, + qvel=qvel, + ctrl=ctrl, + mocap_pos=mocap_pos, + mocap_quat=mocap_quat, + sensordata=sensordata, + userdata=userdata, + logName = read_filename + ) + return data diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5e92a7bb514ef0df943ef53f458590861a33a8e2 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/__init__.py @@ -0,0 +1,19 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import adept_envs.franka + +from adept_envs.utils.configurable import global_config diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/base_robot.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/base_robot.py new file mode 100644 index 0000000000000000000000000000000000000000..5c6f30f94f0f87df65f4108eff6b6d849d366cce --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/base_robot.py @@ -0,0 +1,151 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +from collections import deque + +class BaseRobot(object): + """Base class for all robot classes.""" + + def __init__(self, + n_jnt, + n_obj, + pos_bounds=None, + vel_bounds=None, + calibration_path=None, + is_hardware=False, + device_name=None, + overlay=False, + calibration_mode=False, + observation_cache_maxsize=5): + """Create a new robot. + Args: + n_jnt: The number of dofs in the robot. + n_obj: The number of dofs in the object. + pos_bounds: (n_jnt, 2)-shape matrix denoting the min and max joint + position for each joint. + vel_bounds: (n_jnt, 2)-shape matrix denoting the min and max joint + velocity for each joint. + calibration_path: File path to the calibration configuration file to + use. + is_hardware: Whether to run on hardware or not. + device_name: The device path for the robot hardware. Only required + in legacy mode. + overlay: Whether to show a simulation overlay of the hardware. + calibration_mode: Start with motors disengaged. + """ + + assert n_jnt > 0 + assert n_obj >= 0 + + self._n_jnt = n_jnt + self._n_obj = n_obj + self._n_dofs = n_jnt + n_obj + + self._pos_bounds = None + if pos_bounds is not None: + pos_bounds = np.array(pos_bounds, dtype=np.float32) + assert pos_bounds.shape == (self._n_dofs, 2) + for low, high in pos_bounds: + assert low < high + self._pos_bounds = pos_bounds + self._vel_bounds = None + if vel_bounds is not None: + vel_bounds = np.array(vel_bounds, dtype=np.float32) + assert vel_bounds.shape == (self._n_dofs, 2) + for low, high in vel_bounds: + assert low < high + self._vel_bounds = vel_bounds + + self._is_hardware = is_hardware + self._device_name = device_name + self._calibration_path = calibration_path + self._overlay = overlay + self._calibration_mode = calibration_mode + self._observation_cache_maxsize = observation_cache_maxsize + + # Gets updated + self._observation_cache = deque([], maxlen=self._observation_cache_maxsize) + + + @property + def n_jnt(self): + return self._n_jnt + + @property + def n_obj(self): + return self._n_obj + + @property + def n_dofs(self): + return self._n_dofs + + @property + def pos_bounds(self): + return self._pos_bounds + + @property + def vel_bounds(self): + return self._vel_bounds + + @property + def is_hardware(self): + return self._is_hardware + + @property + def device_name(self): + return self._device_name + + @property + def calibration_path(self): + return self._calibration_path + + @property + def overlay(self): + return self._overlay + + @property + def has_obj(self): + return self._n_obj > 0 + + @property + def calibration_mode(self): + return self._calibration_mode + + @property + def observation_cache_maxsize(self): + return self._observation_cache_maxsize + + @property + def observation_cache(self): + return self._observation_cache + + + def clip_positions(self, positions): + """Clips the given joint positions to the position bounds. + + Args: + positions: The joint positions. + + Returns: + The bounded joint positions. + """ + if self.pos_bounds is None: + return positions + assert len(positions) == self.n_jnt or len(positions) == self.n_dofs + pos_bounds = self.pos_bounds[:len(positions)] + return np.clip(positions, pos_bounds[:, 0], pos_bounds[:, 1]) + diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/franka/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/franka/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..528f3447d280872e394915e47bbb9fe8c1dd87a7 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/franka/__init__.py @@ -0,0 +1,24 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from gym.envs.registration import register + +# Relax the robot +register( + id='kitchen_relax-v1', + entry_point='adept_envs.franka.kitchen_multitask_v0:KitchenTaskRelaxV1', + max_episode_steps=280, +) \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/franka/kitchen_multitask_v0.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/franka/kitchen_multitask_v0.py new file mode 100644 index 0000000000000000000000000000000000000000..30bdd331af1db8be963716d59ff1fdc3d3be176e --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/franka/kitchen_multitask_v0.py @@ -0,0 +1,221 @@ +""" Kitchen environment for long horizon manipulation """ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import numpy as np +from adept_envs import robot_env +from adept_envs.utils.configurable import configurable +from gym import spaces +from dm_control.mujoco import engine + +@configurable(pickleable=True) +class KitchenV0(robot_env.RobotEnv): + + CALIBRATION_PATHS = { + 'default': + os.path.join(os.path.dirname(__file__), 'robot/franka_config.xml') + } + # Converted to velocity actuation + ROBOTS = {'robot': 'adept_envs.franka.robot.franka_robot:Robot_VelAct'} + MODEl = os.path.join( + os.path.dirname(__file__), + '../franka/assets/franka_kitchen_jntpos_act_ab.xml') + N_DOF_ROBOT = 9 + N_DOF_OBJECT = 21 + + def __init__(self, + robot_params={}, frame_skip=40, + use_abs_action=False): + self.goal_concat = True + self.obs_dict = {} + self.robot_noise_ratio = 0.1 # 10% as per robot_config specs + self.goal = np.zeros((30,)) + self.use_abs_action = use_abs_action + if use_abs_action: + self.ROBOTS = {'robot': 'adept_envs.franka.robot.franka_robot:Robot_PosAct'} + + super().__init__( + self.MODEl, + robot=self.make_robot( + n_jnt=self.N_DOF_ROBOT, #root+robot_jnts + n_obj=self.N_DOF_OBJECT, + **robot_params), + frame_skip=frame_skip, + camera_settings=dict( + distance=4.5, + azimuth=-66, + elevation=-65, + ), + ) + self.init_qpos = self.sim.model.key_qpos[0].copy() + + # For the microwave kettle slide hinge + self.init_qpos = np.array([ 1.48388023e-01, -1.76848573e+00, 1.84390296e+00, -2.47685760e+00, + 2.60252026e-01, 7.12533105e-01, 1.59515394e+00, 4.79267505e-02, + 3.71350919e-02, -2.66279850e-04, -5.18043486e-05, 3.12877220e-05, + -4.51199853e-05, -3.90842156e-06, -4.22629655e-05, 6.28065475e-05, + 4.04984708e-05, 4.62730939e-04, -2.26906415e-04, -4.65501369e-04, + -6.44129196e-03, -1.77048263e-03, 1.08009684e-03, -2.69397440e-01, + 3.50383255e-01, 1.61944683e+00, 1.00618764e+00, 4.06395120e-03, + -6.62095997e-03, -2.68278933e-04]) + + self.init_qvel = self.sim.model.key_qvel[0].copy() + + self.act_mid = np.zeros(self.N_DOF_ROBOT) + self.act_amp = 2.0 * np.ones(self.N_DOF_ROBOT) + + act_lower = -1*np.ones((self.N_DOF_ROBOT,)) + act_upper = 1*np.ones((self.N_DOF_ROBOT,)) + if use_abs_action: + act_lower = act_lower * 8. + act_upper = act_upper * 8. + self.act_amp = np.ones(self.N_DOF_ROBOT) + + self.action_space = spaces.Box(act_lower, act_upper) + + obs_upper = 8. * np.ones(self.obs_dim) + obs_lower = -obs_upper + self.observation_space = spaces.Box(obs_lower, obs_upper) + + def _get_reward_n_score(self, obs_dict): + raise NotImplementedError() + + def step(self, a, b=None): + if not self.use_abs_action: + a = np.clip(a, -1.0, 1.0) + + if not self.initializing: + a = self.act_mid + a * self.act_amp # mean center and scale + else: + self.goal = self._get_task_goal() # update goal if init + + self.robot.step( + self, a, step_duration=self.skip * self.model.opt.timestep) + + # observations + obs = self._get_obs() + + #rewards + reward_dict, score = self._get_reward_n_score(self.obs_dict) + + # termination + done = False + + # finalize step + env_info = { + 'time': self.obs_dict['t'], + 'obs_dict': self.obs_dict, + 'rewards': reward_dict, + 'score': score, + # don't render every frame + # 'images': np.asarray(self.render(mode='rgb_array')) + } + # self.render() + return obs, reward_dict['r_total'], done, env_info + + def _get_obs(self): + t, qp, qv, obj_qp, obj_qv = self.robot.get_obs( + self, robot_noise_ratio=self.robot_noise_ratio) + + self.obs_dict = {} + self.obs_dict['t'] = t + self.obs_dict['qp'] = qp + self.obs_dict['qv'] = qv + self.obs_dict['obj_qp'] = obj_qp + self.obs_dict['obj_qv'] = obj_qv + self.obs_dict['goal'] = self.goal + if self.goal_concat: + return np.concatenate([self.obs_dict['qp'], self.obs_dict['obj_qp'], self.obs_dict['goal']]) + + def reset_model(self): + reset_pos = self.init_qpos[:].copy() + reset_vel = self.init_qvel[:].copy() + self.robot.reset(self, reset_pos, reset_vel) + self.sim.forward() + self.goal = self._get_task_goal() #sample a new goal on reset + return self._get_obs() + + def evaluate_success(self, paths): + # score + mean_score_per_rollout = np.zeros(shape=len(paths)) + for idx, path in enumerate(paths): + mean_score_per_rollout[idx] = np.mean(path['env_infos']['score']) + mean_score = np.mean(mean_score_per_rollout) + + # success percentage + num_success = 0 + num_paths = len(paths) + for path in paths: + num_success += bool(path['env_infos']['rewards']['bonus'][-1]) + success_percentage = num_success * 100.0 / num_paths + + # fuse results + return np.sign(mean_score) * ( + 1e6 * round(success_percentage, 2) + abs(mean_score)) + + def close_env(self): + self.robot.close() + + def set_goal(self, goal): + self.goal = goal + + def _get_task_goal(self): + return self.goal + + # Only include goal + @property + def goal_space(self): + len_obs = self.observation_space.low.shape[0] + env_lim = np.abs(self.observation_space.low[0]) + return spaces.Box(low=-env_lim, high=env_lim, shape=(len_obs//2,)) + + def convert_to_active_observation(self, observation): + return observation + +class KitchenTaskRelaxV1(KitchenV0): + """Kitchen environment with proper camera and goal setup""" + + def __init__(self, use_abs_action=False): + super(KitchenTaskRelaxV1, self).__init__( + use_abs_action=use_abs_action) + + def _get_reward_n_score(self, obs_dict): + reward_dict = {} + reward_dict['true_reward'] = 0. + reward_dict['bonus'] = 0. + reward_dict['r_total'] = 0. + score = 0. + return reward_dict, score + + def render(self, mode='human', width=1280, height=720, custom=True, **kwargs): + if custom: + camera = engine.MovableCamera(self.sim, height, width) + if 'distance' not in kwargs: + kwargs['distance'] = 2.2 + if 'lookat' not in kwargs: + kwargs['lookat'] = [-0.2, .5, 2.] + if 'azimuth' not in kwargs: + kwargs['azimuth'] = 70 + if 'elevation' not in kwargs: + kwargs['elevation'] = -35 + camera.set_pose(**kwargs) + img = camera.render() + return img + else: + return super(KitchenTaskRelaxV1, self).render( + mode=mode, width=width, height=height, **kwargs) + diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/franka/robot/franka_robot.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/franka/robot/franka_robot.py new file mode 100644 index 0000000000000000000000000000000000000000..3c6cfc3c2669c778c6cb1faf15d94d41fbda8ff6 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/franka/robot/franka_robot.py @@ -0,0 +1,264 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os, getpass +import numpy as np +from termcolor import cprint +import time +import copy +import click + +from adept_envs import base_robot +from adept_envs.utils.config import (get_config_root_node, read_config_from_node) + +# observations structure +from collections import namedtuple +observation = namedtuple('observation', ['time', 'qpos_robot', 'qvel_robot', 'qpos_object', 'qvel_object']) + + + +franka_interface = '' + +class Robot(base_robot.BaseRobot): + + """ + Abstracts away the differences between the robot_simulation and robot_hardware + + """ + + def __init__(self, *args, **kwargs): + super(Robot, self).__init__(*args, **kwargs) + global franka_interface + + # Read robot configurations + self._read_specs_from_config(robot_configs=self.calibration_path) + + + # Robot: Handware + if self.is_hardware: + if franka_interface == '': + raise NotImplementedError() + from handware.franka import franka + + # initialize franka + self.franka_interface = franka() + franka_interface = self.franka_interface + cprint("Initializing %s Hardware (Status:%d)" % (self.robot_name, self.franka.okay(self.robot_hardware_dof)), 'white', 'on_grey') + else: + self.franka_interface = franka_interface + cprint("Reusing previours Franka session", 'white', 'on_grey') + + # Robot: Simulation + else: + self.robot_name = "Franka" + cprint("Initializing %s sim" % self.robot_name, 'white', 'on_grey') + + # Robot's time + self.time_start = time.time() + self.time = time.time()-self.time_start + self.time_render = -1 # time of rendering + + + # read specs from the calibration file + def _read_specs_from_config(self, robot_configs): + root, root_name = get_config_root_node(config_file_name=robot_configs) + self.robot_name = root_name[0] + self.robot_mode = np.zeros(self.n_dofs, dtype=int) + self.robot_mj_dof = np.zeros(self.n_dofs, dtype=int) + self.robot_hardware_dof = np.zeros(self.n_dofs, dtype=int) + self.robot_scale = np.zeros(self.n_dofs, dtype=float) + self.robot_offset = np.zeros(self.n_dofs, dtype=float) + self.robot_pos_bound = np.zeros([self.n_dofs, 2], dtype=float) + self.robot_vel_bound = np.zeros([self.n_dofs, 2], dtype=float) + self.robot_pos_noise_amp = np.zeros(self.n_dofs, dtype=float) + self.robot_vel_noise_amp = np.zeros(self.n_dofs, dtype=float) + + print("Reading configurations for %s" % self.robot_name) + for i in range(self.n_dofs): + self.robot_mode[i] = read_config_from_node(root, "qpos"+str(i), "mode", int) + self.robot_mj_dof[i] = read_config_from_node(root, "qpos"+str(i), "mj_dof", int) + self.robot_hardware_dof[i] = read_config_from_node(root, "qpos"+str(i), "hardware_dof", int) + self.robot_scale[i] = read_config_from_node(root, "qpos"+str(i), "scale", float) + self.robot_offset[i] = read_config_from_node(root, "qpos"+str(i), "offset", float) + self.robot_pos_bound[i] = read_config_from_node(root, "qpos"+str(i), "pos_bound", float) + self.robot_vel_bound[i] = read_config_from_node(root, "qpos"+str(i), "vel_bound", float) + self.robot_pos_noise_amp[i] = read_config_from_node(root, "qpos"+str(i), "pos_noise_amp", float) + self.robot_vel_noise_amp[i] = read_config_from_node(root, "qpos"+str(i), "vel_noise_amp", float) + + + # convert to hardware space + def _de_calib(self, qp_mj, qv_mj=None): + qp_ad = (qp_mj-self.robot_offset)/self.robot_scale + if qv_mj is not None: + qv_ad = qv_mj/self.robot_scale + return qp_ad, qv_ad + else: + return qp_ad + + # convert to mujoco space + def _calib(self, qp_ad, qv_ad): + qp_mj = qp_ad* self.robot_scale + self.robot_offset + qv_mj = qv_ad* self.robot_scale + return qp_mj, qv_mj + + + # refresh the observation cache + def _observation_cache_refresh(self, env): + for _ in range(self.observation_cache_maxsize): + self.get_obs(env, sim_mimic_hardware=False) + + # get past observation + def get_obs_from_cache(self, env, index=-1): + assert (index>=0 and index=-self.observation_cache_maxsize), \ + "cache index out of bound. (cache size is %2d)"%self.observation_cache_maxsize + obs = self.observation_cache[index] + if self.has_obj: + return obs.time, obs.qpos_robot, obs.qvel_robot, obs.qpos_object, obs.qvel_object + else: + return obs.time, obs.qpos_robot, obs.qvel_robot + + + # get observation + def get_obs(self, env, robot_noise_ratio=1, object_noise_ratio=1, sim_mimic_hardware=True): + if self.is_hardware: + raise NotImplementedError() + + else: + #Gather simulated observation + qp = env.sim.data.qpos[:self.n_jnt].copy() + qv = env.sim.data.qvel[:self.n_jnt].copy() + if self.has_obj: + qp_obj = env.sim.data.qpos[-self.n_obj:].copy() + qv_obj = env.sim.data.qvel[-self.n_obj:].copy() + else: + qp_obj = None + qv_obj = None + self.time = env.sim.data.time + + # Simulate observation noise + if not env.initializing: + qp += robot_noise_ratio*self.robot_pos_noise_amp[:self.n_jnt]*env.np_random.uniform(low=-1., high=1., size=self.n_jnt) + qv += robot_noise_ratio*self.robot_vel_noise_amp[:self.n_jnt]*env.np_random.uniform(low=-1., high=1., size=self.n_jnt) + if self.has_obj: + qp_obj += robot_noise_ratio*self.robot_pos_noise_amp[-self.n_obj:]*env.np_random.uniform(low=-1., high=1., size=self.n_obj) + qv_obj += robot_noise_ratio*self.robot_vel_noise_amp[-self.n_obj:]*env.np_random.uniform(low=-1., high=1., size=self.n_obj) + + # cache observations + obs = observation(time=self.time, qpos_robot=qp, qvel_robot=qv, qpos_object=qp_obj, qvel_object=qv_obj) + self.observation_cache.append(obs) + + if self.has_obj: + return obs.time, obs.qpos_robot, obs.qvel_robot, obs.qpos_object, obs.qvel_object + else: + return obs.time, obs.qpos_robot, obs.qvel_robot + + + # enforce position specs. + def ctrl_position_limits(self, ctrl_position): + ctrl_feasible_position = np.clip(ctrl_position, self.robot_pos_bound[:self.n_jnt, 0], self.robot_pos_bound[:self.n_jnt, 1]) + return ctrl_feasible_position + + + # step the robot env + def step(self, env, ctrl_desired, step_duration, sim_override=False): + + # Populate observation cache during startup + if env.initializing: + self._observation_cache_refresh(env) + + # enforce velocity limits + ctrl_feasible = self.ctrl_velocity_limits(ctrl_desired, step_duration) + + # enforce position limits + ctrl_feasible = self.ctrl_position_limits(ctrl_feasible) + + # Send controls to the robot + if self.is_hardware and (not sim_override): + raise NotImplementedError() + else: + env.do_simulation(ctrl_feasible, int(step_duration/env.sim.model.opt.timestep)) # render is folded in here + + # Update current robot state on the overlay + if self.overlay: + env.sim.data.qpos[self.n_jnt:2*self.n_jnt] = env.desired_pose.copy() + env.sim.forward() + + # synchronize time + if self.is_hardware: + time_now = (time.time()-self.time_start) + time_left_in_step = step_duration - (time_now-self.time) + if(time_left_in_step>0.0001): + time.sleep(time_left_in_step) + return 1 + + + def reset(self, env, reset_pose, reset_vel, overlay_mimic_reset_pose=True, sim_override=False): + reset_pose = self.clip_positions(reset_pose) + + if self.is_hardware: + raise NotImplementedError() + else: + env.sim.reset() + env.sim.data.qpos[:self.n_jnt] = reset_pose[:self.n_jnt].copy() + env.sim.data.qvel[:self.n_jnt] = reset_vel[:self.n_jnt].copy() + if self.has_obj: + env.sim.data.qpos[-self.n_obj:] = reset_pose[-self.n_obj:].copy() + env.sim.data.qvel[-self.n_obj:] = reset_vel[-self.n_obj:].copy() + env.sim.forward() + + if self.overlay: + env.sim.data.qpos[self.n_jnt:2*self.n_jnt] = env.desired_pose[:self.n_jnt].copy() + env.sim.forward() + + # refresh observation cache before exit + self._observation_cache_refresh(env) + + + def close(self): + if self.is_hardware: + cprint("Closing Franka hardware... ", 'white', 'on_grey', end='', flush=True) + status = 0 + raise NotImplementedError() + cprint("Closed (Status: {})".format(status), 'white', 'on_grey', flush=True) + else: + cprint("Closing Franka sim", 'white', 'on_grey', flush=True) + + +class Robot_PosAct(Robot): + + # enforce velocity sepcs. + # ALERT: This depends on previous observation. This is not ideal as it breaks MDP addumptions. Be careful + def ctrl_velocity_limits(self, ctrl_position, step_duration): + last_obs = self.observation_cache[-1] + ctrl_desired_vel = (ctrl_position-last_obs.qpos_robot[:self.n_jnt])/step_duration + + ctrl_feasible_vel = np.clip(ctrl_desired_vel, self.robot_vel_bound[:self.n_jnt, 0], self.robot_vel_bound[:self.n_jnt, 1]) + ctrl_feasible_position = last_obs.qpos_robot[:self.n_jnt] + ctrl_feasible_vel*step_duration + return ctrl_feasible_position + + +class Robot_VelAct(Robot): + + # enforce velocity sepcs. + # ALERT: This depends on previous observation. This is not ideal as it breaks MDP addumptions. Be careful + def ctrl_velocity_limits(self, ctrl_velocity, step_duration): + last_obs = self.observation_cache[-1] + + ctrl_feasible_vel = np.clip(ctrl_velocity, self.robot_vel_bound[:self.n_jnt, 0], self.robot_vel_bound[:self.n_jnt, 1]) + ctrl_feasible_position = last_obs.qpos_robot[:self.n_jnt] + ctrl_feasible_vel*step_duration + return ctrl_feasible_position + diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/mujoco_env.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/mujoco_env.py new file mode 100644 index 0000000000000000000000000000000000000000..e8cd9ec8c4b9e34ada17861e6fc8cb78d13618f7 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/mujoco_env.py @@ -0,0 +1,202 @@ +"""Base environment for MuJoCo-based environments.""" + +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import collections +import os +import time +from typing import Dict, Optional + +import gym +from gym import spaces +from gym.utils import seeding +import numpy as np + +from adept_envs.simulation.sim_robot import MujocoSimRobot, RenderMode + +DEFAULT_RENDER_SIZE = 480 + +USE_DM_CONTROL = True + + +class MujocoEnv(gym.Env): + """Superclass for all MuJoCo environments.""" + + def __init__(self, + model_path: str, + frame_skip: int, + camera_settings: Optional[Dict] = None, + use_dm_backend: Optional[bool] = None, + ): + """Initializes a new MuJoCo environment. + + Args: + model_path: The path to the MuJoCo XML file. + frame_skip: The number of simulation steps per environment step. On + hardware this influences the duration of each environment step. + camera_settings: Settings to initialize the simulation camera. This + can contain the keys `distance`, `azimuth`, and `elevation`. + use_dm_backend: A boolean to switch between mujoco-py and dm_control. + """ + self._seed() + if not os.path.isfile(model_path): + raise IOError( + '[MujocoEnv]: Model path does not exist: {}'.format(model_path)) + self.frame_skip = frame_skip + + self.sim_robot = MujocoSimRobot( + model_path, + use_dm_backend=use_dm_backend or USE_DM_CONTROL, + camera_settings=camera_settings) + self.sim = self.sim_robot.sim + self.model = self.sim_robot.model + self.data = self.sim_robot.data + + self.metadata = { + 'render.modes': ['human', 'rgb_array', 'depth_array'], + 'video.frames_per_second': int(np.round(1.0 / self.dt)) + } + self.mujoco_render_frames = False + + self.init_qpos = self.data.qpos.ravel().copy() + self.init_qvel = self.data.qvel.ravel().copy() + observation, _reward, done, _info = self.step(np.zeros(self.model.nu)) + assert not done + + bounds = self.model.actuator_ctrlrange.copy() + act_upper = bounds[:, 1] + act_lower = bounds[:, 0] + + # Define the action and observation spaces. + # HACK: MJRL is still using gym 0.9.x so we can't provide a dtype. + try: + self.action_space = spaces.Box( + act_lower, act_upper, dtype=np.float32) + if isinstance(observation, collections.Mapping): + self.observation_space = spaces.Dict({ + k: spaces.Box(-np.inf, np.inf, shape=v.shape, dtype=np.float32) for k, v in observation.items()}) + else: + self.obs_dim = np.sum([o.size for o in observation]) if type(observation) is tuple else observation.size + self.observation_space = spaces.Box( + -np.inf, np.inf, observation.shape, dtype=np.float32) + + except TypeError: + # Fallback case for gym 0.9.x + self.action_space = spaces.Box(act_lower, act_upper) + assert not isinstance(observation, collections.Mapping), 'gym 0.9.x does not support dictionary observation.' + self.obs_dim = np.sum([o.size for o in observation]) if type(observation) is tuple else observation.size + self.observation_space = spaces.Box( + -np.inf, np.inf, observation.shape) + + def seed(self, seed=None): # Compatibility with new gym + return self._seed(seed) + + def _seed(self, seed=None): + self.np_random, seed = seeding.np_random(seed) + return [seed] + + # methods to override: + # ---------------------------- + + def reset_model(self): + """Reset the robot degrees of freedom (qpos and qvel). + + Implement this in each subclass. + """ + raise NotImplementedError + + # ----------------------------- + + def reset(self): # compatibility with new gym + return self._reset() + + def _reset(self): + self.sim.reset() + self.sim.forward() + ob = self.reset_model() + return ob + + def set_state(self, qpos, qvel): + assert qpos.shape == (self.model.nq,) and qvel.shape == (self.model.nv,) + # we are directly manipulating mujoco state here + data = self.sim.data # MjData + for i in range(self.model.nq): + data.qpos[i] = qpos[i] + for i in range(self.model.nv): + data.qvel[i] = qvel[i] + # state = np.concatenate([self.data.qpos, self.data.qvel, self.data.act]) + # self.sim.set_state(state) + self.sim.forward() + + @property + def dt(self): + return self.model.opt.timestep * self.frame_skip + + def do_simulation(self, ctrl, n_frames): + for i in range(self.model.nu): + self.sim.data.ctrl[i] = ctrl[i] + + for _ in range(n_frames): + self.sim.step() + + # TODO(michaelahn): Remove this; render should be called separately. + if self.mujoco_render_frames is True: + self.mj_render() + + def render(self, + mode='human', + width=DEFAULT_RENDER_SIZE, + height=DEFAULT_RENDER_SIZE, + camera_id=-1): + """Renders the environment. + + Args: + mode: The type of rendering to use. + - 'human': Renders to a graphical window. + - 'rgb_array': Returns the RGB image as an np.ndarray. + - 'depth_array': Returns the depth image as an np.ndarray. + width: The width of the rendered image. This only affects offscreen + rendering. + height: The height of the rendered image. This only affects + offscreen rendering. + camera_id: The ID of the camera to use. By default, this is the free + camera. If specified, only affects offscreen rendering. + """ + if mode == 'human': + self.sim_robot.renderer.render_to_window() + elif mode == 'rgb_array': + assert width and height + return self.sim_robot.renderer.render_offscreen( + width, height, mode=RenderMode.RGB, camera_id=camera_id) + elif mode == 'depth_array': + assert width and height + return self.sim_robot.renderer.render_offscreen( + width, height, mode=RenderMode.DEPTH, camera_id=camera_id) + else: + raise NotImplementedError(mode) + + def close(self): + self.sim_robot.close() + + def mj_render(self): + """Backwards compatibility with MJRL.""" + self.render(mode='human') + + def state_vector(self): + state = self.sim.get_state() + return np.concatenate([state.qpos.flat, state.qvel.flat]) \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/robot_env.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/robot_env.py new file mode 100644 index 0000000000000000000000000000000000000000..c50c2ee5313ce7173056183e79e5000c6724db1b --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/robot_env.py @@ -0,0 +1,166 @@ +"""Base class for robotics environments.""" + +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib +import inspect +import os +from typing import Dict, Optional + +import numpy as np + + +from adept_envs import mujoco_env +from adept_envs.base_robot import BaseRobot +from adept_envs.utils.configurable import import_class_from_path +from adept_envs.utils.constants import MODELS_PATH + + +class RobotEnv(mujoco_env.MujocoEnv): + """Base environment for all adept robots.""" + + # Mapping of robot name to fully qualified class path. + # e.g. 'robot': 'adept_envs.dclaw.robot.Robot' + # Subclasses should override this to specify the Robot classes they support. + ROBOTS = {} + + # Mapping of device path to the calibration file to use. If the device path + # is not found, the 'default' key is used. + # This can be overridden by subclasses. + CALIBRATION_PATHS = {} + + def __init__(self, + model_path: str, + robot: BaseRobot, + frame_skip: int, + camera_settings: Optional[Dict] = None): + """Initializes a robotics environment. + + Args: + model_path: The path to the model to run. Relative paths will be + interpreted as relative to the 'adept_models' folder. + robot: The Robot object to use. + frame_skip: The number of simulation steps per environment step. On + hardware this influences the duration of each environment step. + camera_settings: Settings to initialize the simulation camera. This + can contain the keys `distance`, `azimuth`, and `elevation`. + """ + self._robot = robot + + # Initial pose for first step. + self.desired_pose = np.zeros(self.n_jnt) + + if not model_path.startswith('/'): + model_path = os.path.abspath(os.path.join(MODELS_PATH, model_path)) + + self.remote_viz = None + + try: + from adept_envs.utils.remote_viz import RemoteViz + self.remote_viz = RemoteViz(model_path) + except ImportError: + pass + + + self._initializing = True + super(RobotEnv, self).__init__( + model_path, frame_skip, camera_settings=camera_settings) + self._initializing = False + + + @property + def robot(self): + return self._robot + + @property + def n_jnt(self): + return self._robot.n_jnt + + @property + def n_obj(self): + return self._robot.n_obj + + @property + def skip(self): + """Alias for frame_skip. Needed for MJRL.""" + return self.frame_skip + + @property + def initializing(self): + return self._initializing + + def close_env(self): + if self._robot is not None: + self._robot.close() + + def make_robot(self, + n_jnt, + n_obj=0, + is_hardware=False, + device_name=None, + legacy=False, + **kwargs): + """Creates a new robot for the environment. + + Args: + n_jnt: The number of joints in the robot. + n_obj: The number of object joints in the robot environment. + is_hardware: Whether to run on hardware or not. + device_name: The device path for the robot hardware. + legacy: If true, runs using direct dynamixel communication rather + than DDS. + kwargs: See BaseRobot for other parameters. + + Returns: + A Robot object. + """ + if not self.ROBOTS: + raise NotImplementedError('Subclasses must override ROBOTS.') + + if is_hardware and not device_name: + raise ValueError('Must provide device name if running on hardware.') + + robot_name = 'dds_robot' if not legacy and is_hardware else 'robot' + if robot_name not in self.ROBOTS: + raise KeyError("Unsupported robot '{}', available: {}".format( + robot_name, list(self.ROBOTS.keys()))) + + cls = import_class_from_path(self.ROBOTS[robot_name]) + + calibration_path = None + if self.CALIBRATION_PATHS: + if not device_name: + calibration_name = 'default' + elif device_name not in self.CALIBRATION_PATHS: + print('Device "{}" not in CALIBRATION_PATHS; using default.' + .format(device_name)) + calibration_name = 'default' + else: + calibration_name = device_name + + calibration_path = self.CALIBRATION_PATHS[calibration_name] + if not os.path.isfile(calibration_path): + raise OSError('Could not find calibration file at: {}'.format( + calibration_path)) + + return cls( + n_jnt, + n_obj, + is_hardware=is_hardware, + device_name=device_name, + calibration_path=calibration_path, + **kwargs) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/simulation/module.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/simulation/module.py new file mode 100644 index 0000000000000000000000000000000000000000..a1284c76c346062def8e7eb6d680169db776967c --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/simulation/module.py @@ -0,0 +1,126 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Module for caching Python modules related to simulation.""" + +import sys + +_MUJOCO_PY_MODULE = None + +_DM_MUJOCO_MODULE = None +_DM_VIEWER_MODULE = None +_DM_RENDER_MODULE = None + +_GLFW_MODULE = None + + +def get_mujoco_py(): + """Returns the mujoco_py module.""" + global _MUJOCO_PY_MODULE + if _MUJOCO_PY_MODULE: + return _MUJOCO_PY_MODULE + try: + import mujoco_py + # Override the warning function. + from mujoco_py.builder import cymj + cymj.set_warning_callback(_mj_warning_fn) + except ImportError: + print( + 'Failed to import mujoco_py. Ensure that mujoco_py (using MuJoCo ' + 'v1.50) is installed.', + file=sys.stderr) + sys.exit(1) + _MUJOCO_PY_MODULE = mujoco_py + return mujoco_py + + +def get_mujoco_py_mjlib(): + """Returns the mujoco_py mjlib module.""" + + class MjlibDelegate: + """Wrapper that forwards mjlib calls.""" + + def __init__(self, lib): + self._lib = lib + + def __getattr__(self, name: str): + if name.startswith('mj'): + return getattr(self._lib, '_' + name) + raise AttributeError(name) + + return MjlibDelegate(get_mujoco_py().cymj) + + +def get_dm_mujoco(): + """Returns the DM Control mujoco module.""" + global _DM_MUJOCO_MODULE + if _DM_MUJOCO_MODULE: + return _DM_MUJOCO_MODULE + try: + from dm_control import mujoco + except ImportError: + print( + 'Failed to import dm_control.mujoco. Ensure that dm_control (using ' + 'MuJoCo v2.00) is installed.', + file=sys.stderr) + sys.exit(1) + _DM_MUJOCO_MODULE = mujoco + return mujoco + + +def get_dm_viewer(): + """Returns the DM Control viewer module.""" + global _DM_VIEWER_MODULE + if _DM_VIEWER_MODULE: + return _DM_VIEWER_MODULE + try: + from dm_control import viewer + except ImportError: + print( + 'Failed to import dm_control.viewer. Ensure that dm_control (using ' + 'MuJoCo v2.00) is installed.', + file=sys.stderr) + sys.exit(1) + _DM_VIEWER_MODULE = viewer + return viewer + + +def get_dm_render(): + """Returns the DM Control render module.""" + global _DM_RENDER_MODULE + if _DM_RENDER_MODULE: + return _DM_RENDER_MODULE + try: + try: + from dm_control import _render + render = _render + except ImportError: + print('Warning: DM Control is out of date.') + from dm_control import render + except ImportError: + print( + 'Failed to import dm_control.render. Ensure that dm_control (using ' + 'MuJoCo v2.00) is installed.', + file=sys.stderr) + sys.exit(1) + _DM_RENDER_MODULE = render + return render + + +def _mj_warning_fn(warn_data: bytes): + """Warning function override for mujoco_py.""" + print('WARNING: Mujoco simulation is unstable (has NaNs): {}'.format( + warn_data.decode())) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/simulation/renderer.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/simulation/renderer.py new file mode 100644 index 0000000000000000000000000000000000000000..28b9de294d3d612f15555ba2bbd9ebfab13ea1af --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/simulation/renderer.py @@ -0,0 +1,293 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Module for viewing Physics objects in the DM Control viewer.""" + +import abc +import enum +import sys +from typing import Dict, Optional + +import numpy as np + +from adept_envs.simulation import module + +# Default window dimensions. +DEFAULT_WINDOW_WIDTH = 1024 +DEFAULT_WINDOW_HEIGHT = 768 + +DEFAULT_WINDOW_TITLE = 'MuJoCo Viewer' + +_MAX_RENDERBUFFER_SIZE = 2048 + + +class RenderMode(enum.Enum): + """Rendering modes for offscreen rendering.""" + RGB = 0 + DEPTH = 1 + SEGMENTATION = 2 + + +class Renderer(abc.ABC): + """Base interface for rendering simulations.""" + + def __init__(self, camera_settings: Optional[Dict] = None): + self._camera_settings = camera_settings + + @abc.abstractmethod + def close(self): + """Cleans up any resources being used by the renderer.""" + + @abc.abstractmethod + def render_to_window(self): + """Renders the simulation to a window.""" + + @abc.abstractmethod + def render_offscreen(self, + width: int, + height: int, + mode: RenderMode = RenderMode.RGB, + camera_id: int = -1) -> np.ndarray: + """Renders the camera view as a NumPy array of pixels. + + Args: + width: The viewport width (pixels). + height: The viewport height (pixels). + mode: The rendering mode. + camera_id: The ID of the camera to render from. By default, uses + the free camera. + + Returns: + A NumPy array of the pixels. + """ + + def _update_camera(self, camera): + """Updates the given camera to move to the initial settings.""" + if not self._camera_settings: + return + distance = self._camera_settings.get('distance') + azimuth = self._camera_settings.get('azimuth') + elevation = self._camera_settings.get('elevation') + lookat = self._camera_settings.get('lookat') + + if distance is not None: + camera.distance = distance + if azimuth is not None: + camera.azimuth = azimuth + if elevation is not None: + camera.elevation = elevation + if lookat is not None: + camera.lookat[:] = lookat + + +class MjPyRenderer(Renderer): + """Class for rendering mujoco_py simulations.""" + + def __init__(self, sim, **kwargs): + assert isinstance(sim, module.get_mujoco_py().MjSim), \ + 'MjPyRenderer takes a mujoco_py MjSim object.' + super().__init__(**kwargs) + self._sim = sim + self._onscreen_renderer = None + self._offscreen_renderer = None + + def render_to_window(self): + """Renders the simulation to a window.""" + if not self._onscreen_renderer: + self._onscreen_renderer = module.get_mujoco_py().MjViewer(self._sim) + self._update_camera(self._onscreen_renderer.cam) + + self._onscreen_renderer.render() + + def render_offscreen(self, + width: int, + height: int, + mode: RenderMode = RenderMode.RGB, + camera_id: int = -1) -> np.ndarray: + """Renders the camera view as a NumPy array of pixels. + + Args: + width: The viewport width (pixels). + height: The viewport height (pixels). + mode: The rendering mode. + camera_id: The ID of the camera to render from. By default, uses + the free camera. + + Returns: + A NumPy array of the pixels. + """ + if not self._offscreen_renderer: + self._offscreen_renderer = module.get_mujoco_py() \ + .MjRenderContextOffscreen(self._sim) + + # Update the camera configuration for the free-camera. + if camera_id == -1: + self._update_camera(self._offscreen_renderer.cam) + + self._offscreen_renderer.render(width, height, camera_id) + if mode == RenderMode.RGB: + data = self._offscreen_renderer.read_pixels( + width, height, depth=False) + # Original image is upside-down, so flip it + return data[::-1, :, :] + elif mode == RenderMode.DEPTH: + data = self._offscreen_renderer.read_pixels( + width, height, depth=True)[1] + # Original image is upside-down, so flip it + return data[::-1, :] + else: + raise NotImplementedError(mode) + + def close(self): + """Cleans up any resources being used by the renderer.""" + + +class DMRenderer(Renderer): + """Class for rendering DM Control Physics objects.""" + + def __init__(self, physics, **kwargs): + assert isinstance(physics, module.get_dm_mujoco().Physics), \ + 'DMRenderer takes a DM Control Physics object.' + super().__init__(**kwargs) + self._physics = physics + self._window = None + + # Set the camera to lookat the center of the geoms. (mujoco_py does + # this automatically. + if 'lookat' not in self._camera_settings: + self._camera_settings['lookat'] = [ + np.median(self._physics.data.geom_xpos[:, i]) for i in range(3) + ] + + def render_to_window(self): + """Renders the Physics object to a window. + + The window continuously renders the Physics in a separate thread. + + This function is a no-op if the window was already created. + """ + if not self._window: + self._window = DMRenderWindow() + self._window.load_model(self._physics) + self._update_camera(self._window.camera) + self._window.run_frame() + + def render_offscreen(self, + width: int, + height: int, + mode: RenderMode = RenderMode.RGB, + camera_id: int = -1) -> np.ndarray: + """Renders the camera view as a NumPy array of pixels. + + Args: + width: The viewport width (pixels). + height: The viewport height (pixels). + mode: The rendering mode. + camera_id: The ID of the camera to render from. By default, uses + the free camera. + + Returns: + A NumPy array of the pixels. + """ + mujoco = module.get_dm_mujoco() + # TODO(michaelahn): Consider caching the camera. + camera = mujoco.Camera( + physics=self._physics, + height=height, + width=width, + camera_id=camera_id) + + # Update the camera configuration for the free-camera. + if camera_id == -1: + self._update_camera( + camera._render_camera, # pylint: disable=protected-access + ) + + image = camera.render( + depth=(mode == RenderMode.DEPTH), + segmentation=(mode == RenderMode.SEGMENTATION)) + camera._scene.free() # pylint: disable=protected-access + return image + + def close(self): + """Cleans up any resources being used by the renderer.""" + if self._window: + self._window.close() + self._window = None + + +class DMRenderWindow: + """Class that encapsulates a graphical window.""" + + def __init__(self, + width: int = DEFAULT_WINDOW_WIDTH, + height: int = DEFAULT_WINDOW_HEIGHT, + title: str = DEFAULT_WINDOW_TITLE): + """Creates a graphical render window. + + Args: + width: The width of the window. + height: The height of the window. + title: The title of the window. + """ + dmv = module.get_dm_viewer() + self._viewport = dmv.renderer.Viewport(width, height) + self._window = dmv.gui.RenderWindow(width, height, title) + self._viewer = dmv.viewer.Viewer(self._viewport, self._window.mouse, + self._window.keyboard) + self._draw_surface = None + self._renderer = dmv.renderer.NullRenderer() + + @property + def camera(self): + return self._viewer._camera._camera + + def close(self): + self._viewer.deinitialize() + self._renderer.release() + self._draw_surface.free() + self._window.close() + + def load_model(self, physics): + """Loads the given Physics object to render.""" + self._viewer.deinitialize() + + self._draw_surface = module.get_dm_render().Renderer( + max_width=_MAX_RENDERBUFFER_SIZE, max_height=_MAX_RENDERBUFFER_SIZE) + self._renderer = module.get_dm_viewer().renderer.OffScreenRenderer( + physics.model, self._draw_surface) + + self._viewer.initialize(physics, self._renderer, touchpad=False) + + def run_frame(self): + """Renders one frame of the simulation. + + NOTE: This is extremely slow at the moment. + """ + glfw = module.get_dm_viewer().gui.glfw_gui.glfw + glfw_window = self._window._context.window + if glfw.window_should_close(glfw_window): + sys.exit(0) + + self._viewport.set_size(*self._window.shape) + self._viewer.render() + pixels = self._renderer.pixels + + with self._window._context.make_current() as ctx: + ctx.call(self._window._update_gui_on_render_thread, glfw_window, + pixels) + self._window._mouse.process_events() + self._window._keyboard.process_events() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/simulation/sim_robot.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/simulation/sim_robot.py new file mode 100644 index 0000000000000000000000000000000000000000..195d524b801b95c9ef9ef5d7009e82a18b362659 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/simulation/sim_robot.py @@ -0,0 +1,135 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Module for loading MuJoCo models.""" + +import os +from typing import Dict, Optional + +from adept_envs.simulation import module +from adept_envs.simulation.renderer import DMRenderer, MjPyRenderer, RenderMode + + +class MujocoSimRobot: + """Class that encapsulates a MuJoCo simulation. + + This class exposes methods that are agnostic to the simulation backend. + Two backends are supported: + 1. mujoco_py - MuJoCo v1.50 + 2. dm_control - MuJoCo v2.00 + """ + + def __init__(self, + model_file: str, + use_dm_backend: bool = False, + camera_settings: Optional[Dict] = None): + """Initializes a new simulation. + + Args: + model_file: The MuJoCo XML model file to load. + use_dm_backend: If True, uses DM Control's Physics (MuJoCo v2.0) as + the backend for the simulation. Otherwise, uses mujoco_py (MuJoCo + v1.5) as the backend. + camera_settings: Settings to initialize the renderer's camera. This + can contain the keys `distance`, `azimuth`, and `elevation`. + """ + self._use_dm_backend = use_dm_backend + + if not os.path.isfile(model_file): + raise ValueError( + '[MujocoSimRobot] Invalid model file path: {}'.format( + model_file)) + + if self._use_dm_backend: + dm_mujoco = module.get_dm_mujoco() + if model_file.endswith('.mjb'): + self.sim = dm_mujoco.Physics.from_binary_path(model_file) + else: + self.sim = dm_mujoco.Physics.from_xml_path(model_file) + self.model = self.sim.model + self._patch_mjlib_accessors(self.model, self.sim.data) + self.renderer = DMRenderer( + self.sim, camera_settings=camera_settings) + else: # Use mujoco_py + mujoco_py = module.get_mujoco_py() + self.model = mujoco_py.load_model_from_path(model_file) + self.sim = mujoco_py.MjSim(self.model) + self.renderer = MjPyRenderer( + self.sim, camera_settings=camera_settings) + + self.data = self.sim.data + + def close(self): + """Cleans up any resources being used by the simulation.""" + self.renderer.close() + + def save_binary(self, path: str): + """Saves the loaded model to a binary .mjb file.""" + if os.path.exists(path): + raise ValueError( + '[MujocoSimRobot] Path already exists: {}'.format(path)) + if not path.endswith('.mjb'): + path = path + '.mjb' + if self._use_dm_backend: + self.model.save_binary(path) + else: + with open(path, 'wb') as f: + f.write(self.model.get_mjb()) + + def get_mjlib(self): + """Returns an object that exposes the low-level MuJoCo API.""" + if self._use_dm_backend: + return module.get_dm_mujoco().wrapper.mjbindings.mjlib + else: + return module.get_mujoco_py_mjlib() + + def _patch_mjlib_accessors(self, model, data): + """Adds accessors to the DM Control objects to support mujoco_py API.""" + assert self._use_dm_backend + mjlib = self.get_mjlib() + + def name2id(type_name, name): + obj_id = mjlib.mj_name2id(model.ptr, + mjlib.mju_str2Type(type_name.encode()), + name.encode()) + if obj_id < 0: + raise ValueError('No {} with name "{}" exists.'.format( + type_name, name)) + return obj_id + + if not hasattr(model, 'body_name2id'): + model.body_name2id = lambda name: name2id('body', name) + + if not hasattr(model, 'geom_name2id'): + model.geom_name2id = lambda name: name2id('geom', name) + + if not hasattr(model, 'site_name2id'): + model.site_name2id = lambda name: name2id('site', name) + + if not hasattr(model, 'joint_name2id'): + model.joint_name2id = lambda name: name2id('joint', name) + + if not hasattr(model, 'actuator_name2id'): + model.actuator_name2id = lambda name: name2id('actuator', name) + + if not hasattr(model, 'camera_name2id'): + model.camera_name2id = lambda name: name2id('camera', name) + + if not hasattr(data, 'body_xpos'): + data.body_xpos = data.xpos + + if not hasattr(data, 'body_xquat'): + data.body_xquat = data.xquat diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/config.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/config.py new file mode 100644 index 0000000000000000000000000000000000000000..6444bf7296a09620b28b23f28765b183003f5834 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/config.py @@ -0,0 +1,99 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +try: + import cElementTree as ET +except ImportError: + try: + # Python 2.5 need to import a different module + import xml.etree.cElementTree as ET + except ImportError: + exit_err("Failed to import cElementTree from any known place") + +CONFIG_XML_DATA = """ + + + + + +""" + + +# Read config from root +def read_config_from_node(root_node, parent_name, child_name, dtype=int): + # find parent + parent_node = root_node.find(parent_name) + if parent_node == None: + quit("Parent %s not found" % parent_name) + + # get child data + child_data = parent_node.get(child_name) + if child_data == None: + quit("Child %s not found" % child_name) + + config_val = np.array(child_data.split(), dtype=dtype) + return config_val + + +# get config frlom file or string +def get_config_root_node(config_file_name=None, config_file_data=None): + try: + # get root + if config_file_data is None: + config_file_content = open(config_file_name, "r") + config = ET.parse(config_file_content) + root_node = config.getroot() + else: + root_node = ET.fromstring(config_file_data) + + # get root data + root_data = root_node.get('name') + root_name = np.array(root_data.split(), dtype=str) + except: + quit("ERROR: Unable to process config file %s" % config_file_name) + + return root_node, root_name + + +# Read config from config_file +def read_config_from_xml(config_file_name, parent_name, child_name, dtype=int): + root_node, root_name = get_config_root_node( + config_file_name=config_file_name) + return read_config_from_node(root_node, parent_name, child_name, dtype) + + +# tests +if __name__ == '__main__': + print("Read config and parse -------------------------") + root, root_name = get_config_root_node(config_file_data=CONFIG_XML_DATA) + print("Root:name \t", root_name) + print("limit:low \t", read_config_from_node(root, "limits", "low", float)) + print("limit:high \t", read_config_from_node(root, "limits", "high", float)) + print("scale:joint \t", read_config_from_node(root, "scale", "joint", + float)) + print("data:type \t", read_config_from_node(root, "data", "type", str)) + + # read straight from xml (dumb the XML data as duh.xml for this test) + root, root_name = get_config_root_node(config_file_name="duh.xml") + print("Read from xml --------------------------------") + print("limit:low \t", read_config_from_xml("duh.xml", "limits", "low", + float)) + print("limit:high \t", + read_config_from_xml("duh.xml", "limits", "high", float)) + print("scale:joint \t", + read_config_from_xml("duh.xml", "scale", "joint", float)) + print("data:type \t", read_config_from_xml("duh.xml", "data", "type", str)) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/configurable.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/configurable.py new file mode 100644 index 0000000000000000000000000000000000000000..6685a5007b028c487cba5d0f5be96be297c17fa6 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/configurable.py @@ -0,0 +1,163 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib +import inspect +import os + +from gym.envs.registration import registry as gym_registry + + +def import_class_from_path(class_path): + """Given 'path.to.module:object', imports and returns the object.""" + module_path, class_name = class_path.split(":") + module = importlib.import_module(module_path) + return getattr(module, class_name) + + +class ConfigCache(object): + """Configuration class to store constructor arguments. + + This is used to store parameters to pass to Gym environments at init time. + """ + + def __init__(self): + self._configs = {} + self._default_config = {} + + def set_default_config(self, config): + """Sets the default configuration used for all RobotEnv envs.""" + self._default_config = dict(config) + + def set_config(self, cls_or_env_id, config): + """Sets the configuration for the given environment within a context. + + Args: + cls_or_env_id (Class | str): A class type or Gym environment ID to + configure. + config (dict): The configuration parameters. + """ + config_key = self._get_config_key(cls_or_env_id) + self._configs[config_key] = dict(config) + + def get_config(self, cls_or_env_id): + """Returns the configuration for the given env name. + + Args: + cls_or_env_id (Class | str): A class type or Gym environment ID to + get the configuration of. + """ + config_key = self._get_config_key(cls_or_env_id) + config = dict(self._default_config) + config.update(self._configs.get(config_key, {})) + return config + + def clear_config(self, cls_or_env_id): + """Clears the configuration for the given ID.""" + config_key = self._get_config_key(cls_or_env_id) + if config_key in self._configs: + del self._configs[config_key] + + def _get_config_key(self, cls_or_env_id): + if inspect.isclass(cls_or_env_id): + return cls_or_env_id + env_id = cls_or_env_id + assert isinstance(env_id, str) + if env_id not in gym_registry.env_specs: + raise ValueError("Unregistered environment name {}.".format(env_id)) + entry_point = gym_registry.env_specs[env_id]._entry_point + if callable(entry_point): + return entry_point + else: + return import_class_from_path(entry_point) + + +# Global robot config. +global_config = ConfigCache() + + +def configurable(config_id=None, pickleable=False, config_cache=global_config): + """Class decorator to allow injection of constructor arguments. + + This allows constructor arguments to be passed via ConfigCache. + Example usage: + + @configurable() + class A: + def __init__(b=None, c=2, d='Wow'): + ... + + global_config.set_config(A, {'b': 10, 'c': 20}) + a = A() # b=10, c=20, d='Wow' + a = A(b=30) # b=30, c=20, d='Wow' + + Args: + config_id: ID of the config to use. This defaults to the class type. + pickleable: Whether this class is pickleable. If true, causes the pickle + state to include the config and constructor arguments. + config_cache: The ConfigCache to use to read config data from. Uses + the global ConfigCache by default. + """ + def cls_decorator(cls): + assert inspect.isclass(cls) + + # Overwrite the class constructor to pass arguments from the config. + base_init = cls.__init__ + def __init__(self, *args, **kwargs): + + config = config_cache.get_config(config_id or type(self)) + # Allow kwargs to override the config. + kwargs = {**config, **kwargs} + + # print('Initializing {} with params: {}'.format(type(self).__name__, + # kwargs)) + + if pickleable: + self._pkl_env_args = args + self._pkl_env_kwargs = kwargs + + base_init(self, *args, **kwargs) + cls.__init__ = __init__ + + # If the class is pickleable, overwrite the state methods to save + # the constructor arguments and config. + if pickleable: + # Use same pickle keys as gym.utils.ezpickle for backwards compat. + PKL_ARGS_KEY = '_ezpickle_args' + PKL_KWARGS_KEY = '_ezpickle_kwargs' + + def __getstate__(self): + return { + PKL_ARGS_KEY: self._pkl_env_args, + PKL_KWARGS_KEY: self._pkl_env_kwargs, + } + cls.__getstate__ = __getstate__ + + def __setstate__(self, data): + saved_args = data[PKL_ARGS_KEY] + saved_kwargs = data[PKL_KWARGS_KEY] + + # Override the saved state with the current config. + config = config_cache.get_config(config_id or type(self)) + # Allow kwargs to override the config. + kwargs = {**saved_kwargs, **config} + + inst = type(self)(*saved_args, **kwargs) + self.__dict__.update(inst.__dict__) + cls.__setstate__ = __setstate__ + + return cls + return cls_decorator diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/constants.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..9c63fb7885adf3bf8fc57b441d5b5a3e6ba5fa1d --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/constants.py @@ -0,0 +1,23 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +ENVS_ROOT_PATH = os.path.abspath(os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "../../")) + +MODELS_PATH = os.path.abspath(os.path.join(ENVS_ROOT_PATH, "../adept_models/")) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/parse_demos.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/parse_demos.py new file mode 100644 index 0000000000000000000000000000000000000000..01f9c36ca0131387766915fb417ef188ab2ae2e5 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/parse_demos.py @@ -0,0 +1,221 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import click +import glob +import pickle +import numpy as np +from parse_mjl import parse_mjl_logs, viz_parsed_mjl_logs +from mjrl.utils.gym_env import GymEnv +import adept_envs +import time as timer +import skvideo.io +import gym + +# headless renderer +render_buffer = [] # rendering buffer + + +def viewer(env, + mode='initialize', + filename='video', + frame_size=(640, 480), + camera_id=0, + render=None): + if render == 'onscreen': + env.mj_render() + + elif render == 'offscreen': + + global render_buffer + if mode == 'initialize': + render_buffer = [] + mode = 'render' + + if mode == 'render': + curr_frame = env.render(mode='rgb_array') + render_buffer.append(curr_frame) + + if mode == 'save': + skvideo.io.vwrite(filename, np.asarray(render_buffer)) + print("\noffscreen buffer saved", filename) + + elif render == 'None': + pass + + else: + print("unknown render: ", render) + + +# view demos (physics ignored) +def render_demos(env, data, filename='demo_rendering.mp4', render=None): + FPS = 30 + render_skip = max(1, round(1. / \ + (FPS * env.sim.model.opt.timestep * env.frame_skip))) + t0 = timer.time() + + viewer(env, mode='initialize', render=render) + for i_frame in range(data['ctrl'].shape[0]): + env.sim.data.qpos[:] = data['qpos'][i_frame].copy() + env.sim.data.qvel[:] = data['qvel'][i_frame].copy() + env.sim.forward() + if i_frame % render_skip == 0: + viewer(env, mode='render', render=render) + print(i_frame, end=', ', flush=True) + + viewer(env, mode='save', filename=filename, render=render) + print("time taken = %f" % (timer.time() - t0)) + + +# playback demos and get data(physics respected) +def gather_training_data(env, data, filename='demo_playback.mp4', render=None): + env = env.env + FPS = 30 + render_skip = max(1, round(1. / \ + (FPS * env.sim.model.opt.timestep * env.frame_skip))) + t0 = timer.time() + + # initialize + env.reset() + init_qpos = data['qpos'][0].copy() + init_qvel = data['qvel'][0].copy() + act_mid = env.act_mid + act_rng = env.act_amp + + # prepare env + env.sim.data.qpos[:] = init_qpos + env.sim.data.qvel[:] = init_qvel + env.sim.forward() + viewer(env, mode='initialize', render=render) + + # step the env and gather data + path_obs = None + for i_frame in range(data['ctrl'].shape[0] - 1): + # Reset every time step + # if i_frame % 1 == 0: + # qp = data['qpos'][i_frame].copy() + # qv = data['qvel'][i_frame].copy() + # env.sim.data.qpos[:] = qp + # env.sim.data.qvel[:] = qv + # env.sim.forward() + + obs = env._get_obs() + + # Construct the action + # ctrl = (data['qpos'][i_frame + 1][:9] - obs[:9]) / (env.skip * env.model.opt.timestep) + ctrl = (data['ctrl'][i_frame] - obs[:9])/(env.skip*env.model.opt.timestep) + act = (ctrl - act_mid) / act_rng + act = np.clip(act, -0.999, 0.999) + next_obs, reward, done, env_info = env.step(act) + if path_obs is None: + path_obs = obs + path_act = act + else: + path_obs = np.vstack((path_obs, obs)) + path_act = np.vstack((path_act, act)) + + # render when needed to maintain FPS + if i_frame % render_skip == 0: + viewer(env, mode='render', render=render) + print(i_frame, end=', ', flush=True) + + # finalize + if render: + viewer(env, mode='save', filename=filename, render=render) + + t1 = timer.time() + print("time taken = %f" % (t1 - t0)) + + # note that are one step away from + return path_obs, path_act, init_qpos, init_qvel + + +# MAIN ========================================================= +@click.command(help="parse tele-op demos") +@click.option('--env', '-e', type=str, help='gym env name', required=True) +@click.option( + '--demo_dir', + '-d', + type=str, + help='directory with tele-op logs', + required=True) +@click.option( + '--skip', + '-s', + type=int, + help='number of frames to skip (1:no skip)', + default=1) +@click.option('--graph', '-g', type=bool, help='plot logs', default=False) +@click.option('--save_logs', '-l', type=bool, help='save logs', default=False) +@click.option( + '--view', '-v', type=str, help='render/playback', default='render') +@click.option( + '--render', '-r', type=str, help='onscreen/offscreen', default='onscreen') +def main(env, demo_dir, skip, graph, save_logs, view, render): + + gym_env = gym.make(env) + paths = [] + print("Scanning demo_dir: " + demo_dir + "=========") + for ind, file in enumerate(glob.glob(demo_dir + "*.mjl")): + + # process logs + print("processing: " + file, end=': ') + + data = parse_mjl_logs(file, skip) + + print("log duration %0.2f" % (data['time'][-1] - data['time'][0])) + + # plot logs + if (graph): + print("plotting: " + file) + viz_parsed_mjl_logs(data) + + # save logs + if (save_logs): + pickle.dump(data, open(file[:-4] + ".pkl", 'wb')) + + # render logs to video + if view == 'render': + render_demos( + gym_env, + data, + filename=data['logName'][:-4] + '_demo_render.mp4', + render=render) + + # playback logs and gather data + elif view == 'playback': + try: + obs, act,init_qpos, init_qvel = gather_training_data(gym_env, data,\ + filename=data['logName'][:-4]+'_playback.mp4', render=render) + except Exception as e: + print(e) + continue + path = { + 'observations': obs, + 'actions': act, + 'goals': obs, + 'init_qpos': init_qpos, + 'init_qvel': init_qvel + } + paths.append(path) + # accept = input('accept demo?') + # if accept == 'n': + # continue + pickle.dump(path, open(demo_dir + env + str(ind) + "_path.pkl", 'wb')) + print(demo_dir + env + file + "_path.pkl") + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/quatmath.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/quatmath.py new file mode 100644 index 0000000000000000000000000000000000000000..bae531a41acc0ca5c991aa55a2e452120be32cf1 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/quatmath.py @@ -0,0 +1,180 @@ +#!/usr/bin/python +# +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +# For testing whether a number is close to zero +_FLOAT_EPS = np.finfo(np.float64).eps +_EPS4 = _FLOAT_EPS * 4.0 + + +def mulQuat(qa, qb): + res = np.zeros(4) + res[0] = qa[0]*qb[0] - qa[1]*qb[1] - qa[2]*qb[2] - qa[3]*qb[3] + res[1] = qa[0]*qb[1] + qa[1]*qb[0] + qa[2]*qb[3] - qa[3]*qb[2] + res[2] = qa[0]*qb[2] - qa[1]*qb[3] + qa[2]*qb[0] + qa[3]*qb[1] + res[3] = qa[0]*qb[3] + qa[1]*qb[2] - qa[2]*qb[1] + qa[3]*qb[0] + return res + +def negQuat(quat): + return np.array([quat[0], -quat[1], -quat[2], -quat[3]]) + +def quat2Vel(quat, dt=1): + axis = quat[1:].copy() + sin_a_2 = np.sqrt(np.sum(axis**2)) + axis = axis/(sin_a_2+1e-8) + speed = 2*np.arctan2(sin_a_2, quat[0])/dt + return speed, axis + +def quatDiff2Vel(quat1, quat2, dt): + neg = negQuat(quat1) + diff = mulQuat(quat2, neg) + return quat2Vel(diff, dt) + + +def axis_angle2quat(axis, angle): + c = np.cos(angle/2) + s = np.sin(angle/2) + return np.array([c, s*axis[0], s*axis[1], s*axis[2]]) + +def euler2mat(euler): + """ Convert Euler Angles to Rotation Matrix. See rotation.py for notes """ + euler = np.asarray(euler, dtype=np.float64) + assert euler.shape[-1] == 3, "Invalid shaped euler {}".format(euler) + + ai, aj, ak = -euler[..., 2], -euler[..., 1], -euler[..., 0] + si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak) + ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak) + cc, cs = ci * ck, ci * sk + sc, ss = si * ck, si * sk + + mat = np.empty(euler.shape[:-1] + (3, 3), dtype=np.float64) + mat[..., 2, 2] = cj * ck + mat[..., 2, 1] = sj * sc - cs + mat[..., 2, 0] = sj * cc + ss + mat[..., 1, 2] = cj * sk + mat[..., 1, 1] = sj * ss + cc + mat[..., 1, 0] = sj * cs - sc + mat[..., 0, 2] = -sj + mat[..., 0, 1] = cj * si + mat[..., 0, 0] = cj * ci + return mat + + +def euler2quat(euler): + """ Convert Euler Angles to Quaternions. See rotation.py for notes """ + euler = np.asarray(euler, dtype=np.float64) + assert euler.shape[-1] == 3, "Invalid shape euler {}".format(euler) + + ai, aj, ak = euler[..., 2] / 2, -euler[..., 1] / 2, euler[..., 0] / 2 + si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak) + ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak) + cc, cs = ci * ck, ci * sk + sc, ss = si * ck, si * sk + + quat = np.empty(euler.shape[:-1] + (4,), dtype=np.float64) + quat[..., 0] = cj * cc + sj * ss + quat[..., 3] = cj * sc - sj * cs + quat[..., 2] = -(cj * ss + sj * cc) + quat[..., 1] = cj * cs - sj * sc + return quat + + +def mat2euler(mat): + """ Convert Rotation Matrix to Euler Angles. See rotation.py for notes """ + mat = np.asarray(mat, dtype=np.float64) + assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat) + + cy = np.sqrt(mat[..., 2, 2] * mat[..., 2, 2] + mat[..., 1, 2] * mat[..., 1, 2]) + condition = cy > _EPS4 + euler = np.empty(mat.shape[:-1], dtype=np.float64) + euler[..., 2] = np.where(condition, + -np.arctan2(mat[..., 0, 1], mat[..., 0, 0]), + -np.arctan2(-mat[..., 1, 0], mat[..., 1, 1])) + euler[..., 1] = np.where(condition, + -np.arctan2(-mat[..., 0, 2], cy), + -np.arctan2(-mat[..., 0, 2], cy)) + euler[..., 0] = np.where(condition, + -np.arctan2(mat[..., 1, 2], mat[..., 2, 2]), + 0.0) + return euler + + +def mat2quat(mat): + """ Convert Rotation Matrix to Quaternion. See rotation.py for notes """ + mat = np.asarray(mat, dtype=np.float64) + assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat) + + Qxx, Qyx, Qzx = mat[..., 0, 0], mat[..., 0, 1], mat[..., 0, 2] + Qxy, Qyy, Qzy = mat[..., 1, 0], mat[..., 1, 1], mat[..., 1, 2] + Qxz, Qyz, Qzz = mat[..., 2, 0], mat[..., 2, 1], mat[..., 2, 2] + # Fill only lower half of symmetric matrix + K = np.zeros(mat.shape[:-2] + (4, 4), dtype=np.float64) + K[..., 0, 0] = Qxx - Qyy - Qzz + K[..., 1, 0] = Qyx + Qxy + K[..., 1, 1] = Qyy - Qxx - Qzz + K[..., 2, 0] = Qzx + Qxz + K[..., 2, 1] = Qzy + Qyz + K[..., 2, 2] = Qzz - Qxx - Qyy + K[..., 3, 0] = Qyz - Qzy + K[..., 3, 1] = Qzx - Qxz + K[..., 3, 2] = Qxy - Qyx + K[..., 3, 3] = Qxx + Qyy + Qzz + K /= 3.0 + # TODO: vectorize this -- probably could be made faster + q = np.empty(K.shape[:-2] + (4,)) + it = np.nditer(q[..., 0], flags=['multi_index']) + while not it.finished: + # Use Hermitian eigenvectors, values for speed + vals, vecs = np.linalg.eigh(K[it.multi_index]) + # Select largest eigenvector, reorder to w,x,y,z quaternion + q[it.multi_index] = vecs[[3, 0, 1, 2], np.argmax(vals)] + # Prefer quaternion with positive w + # (q * -1 corresponds to same rotation as q) + if q[it.multi_index][0] < 0: + q[it.multi_index] *= -1 + it.iternext() + return q + + +def quat2euler(quat): + """ Convert Quaternion to Euler Angles. See rotation.py for notes """ + return mat2euler(quat2mat(quat)) + + +def quat2mat(quat): + """ Convert Quaternion to Euler Angles. See rotation.py for notes """ + quat = np.asarray(quat, dtype=np.float64) + assert quat.shape[-1] == 4, "Invalid shape quat {}".format(quat) + + w, x, y, z = quat[..., 0], quat[..., 1], quat[..., 2], quat[..., 3] + Nq = np.sum(quat * quat, axis=-1) + s = 2.0 / Nq + X, Y, Z = x * s, y * s, z * s + wX, wY, wZ = w * X, w * Y, w * Z + xX, xY, xZ = x * X, x * Y, x * Z + yY, yZ, zZ = y * Y, y * Z, z * Z + + mat = np.empty(quat.shape[:-1] + (3, 3), dtype=np.float64) + mat[..., 0, 0] = 1.0 - (yY + zZ) + mat[..., 0, 1] = xY - wZ + mat[..., 0, 2] = xZ + wY + mat[..., 1, 0] = xY + wZ + mat[..., 1, 1] = 1.0 - (xX + zZ) + mat[..., 1, 2] = yZ - wX + mat[..., 2, 0] = xZ - wY + mat[..., 2, 1] = yZ + wX + mat[..., 2, 2] = 1.0 - (xX + yY) + return np.where((Nq > _FLOAT_EPS)[..., np.newaxis, np.newaxis], mat, np.eye(3)) \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_models/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/relay_policy_learning/adept_models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/v0.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/v0.py new file mode 100644 index 0000000000000000000000000000000000000000..1a2c4ed93087eda5900771b1ebd0b612219eae3a --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/kitchen/v0.py @@ -0,0 +1,20 @@ +from cleandiffuser.env.kitchen.base import KitchenBase + + +class KitchenMicrowaveKettleBottomBurnerLightV0(KitchenBase): + TASK_ELEMENTS = ["microwave", "kettle", "bottom burner", "light switch"] + COMPLETE_IN_ANY_ORDER = False + + +class KitchenMicrowaveKettleLightSliderV0(KitchenBase): + TASK_ELEMENTS = ["microwave", "kettle", "light switch", "slide cabinet"] + COMPLETE_IN_ANY_ORDER = False + + +class KitchenKettleMicrowaveLightSliderV0(KitchenBase): + TASK_ELEMENTS = ["kettle", "microwave", "light switch", "slide cabinet"] + COMPLETE_IN_ANY_ORDER = False + + +class KitchenAllV0(KitchenBase): + TASK_ELEMENTS = KitchenBase.ALL_TASKS diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9ff697b63e2bcd1862d592f16e3e5c7b1b48101f --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/__init__.py @@ -0,0 +1,24 @@ +from gym.envs.registration import register + +register( + id='pusht-v0', + entry_point='cleandiffuser.env.pusht.pusht_env:PushTEnv', + max_episode_steps=300, + reward_threshold=1.0 +) + +register( + id='pusht-keypoints-v0', + entry_point='cleandiffuser.env.pusht.pusht_keypoints_env:PushTKeypointsEnv', + max_episode_steps=300, + reward_threshold=1.0 +) + +register( + id='pusht-image-v0', + entry_point='cleandiffuser.env.pusht.pusht_image_env:PushTImageEnv', + max_episode_steps=300, + reward_threshold=1.0 +) + + diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pusht_env.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pusht_env.py new file mode 100644 index 0000000000000000000000000000000000000000..a3c925f1ee1403acfefe69c0b868b433347f4b63 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pusht_env.py @@ -0,0 +1,367 @@ +import gym +from gym import spaces + +import collections +import numpy as np +import pygame +import pymunk +import pymunk.pygame_util +from pymunk.vec2d import Vec2d +import shapely.geometry as sg +import cv2 +import skimage.transform as st +from cleandiffuser.env.pusht.pymunk_override import DrawOptions + + +def pymunk_to_shapely(body, shapes): + geoms = list() + for shape in shapes: + if isinstance(shape, pymunk.shapes.Poly): + verts = [body.local_to_world(v) for v in shape.get_vertices()] + verts += [verts[0]] + geoms.append(sg.Polygon(verts)) + else: + raise RuntimeError(f'Unsupported shape type {type(shape)}') + geom = sg.MultiPolygon(geoms) + return geom + +class PushTEnv(gym.Env): + metadata = {"render.modes": ["human", "rgb_array"], "video.frames_per_second": 10} + reward_range = (0., 1.) + + def __init__(self, + legacy=False, + block_cog=None, damping=None, + render_action=True, + render_size=96, + reset_to_state=None + ): + self._seed = None + self.seed() + self.window_size = ws = 512 # The size of the PyGame window + self.render_size = render_size + self.sim_hz = 100 + # Local controller params. + self.k_p, self.k_v = 100, 20 # PD control.z + self.control_hz = self.metadata['video.frames_per_second'] + # legcay set_state for data compatibility + self.legacy = legacy + + # agent_pos, block_pos, block_angle + self.observation_space = spaces.Box( + low=np.array([0,0,0,0,0], dtype=np.float64), + high=np.array([ws,ws,ws,ws,np.pi*2], dtype=np.float64), + shape=(5,), + dtype=np.float64 + ) + + # positional goal for agent + self.action_space = spaces.Box( + low=np.array([0,0], dtype=np.float64), + high=np.array([ws,ws], dtype=np.float64), + shape=(2,), + dtype=np.float64 + ) + + self.block_cog = block_cog + self.damping = damping + self.render_action = render_action + + """ + If human-rendering is used, `self.window` will be a reference + to the window that we draw to. `self.clock` will be a clock that is used + to ensure that the environment is rendered at the correct framerate in + human-mode. They will remain `None` until human-mode is used for the + first time. + """ + self.window = None + self.clock = None + self.screen = None + + self.space = None + self.teleop = None + self.render_buffer = None + self.latest_action = None + self.reset_to_state = reset_to_state + + def reset(self): + seed = self._seed + self._setup() + if self.block_cog is not None: + self.block.center_of_gravity = self.block_cog + if self.damping is not None: + self.space.damping = self.damping + + # use legacy RandomState for compatibility + state = self.reset_to_state + if state is None: + rs = np.random.RandomState(seed=seed) + state = np.array([ + rs.randint(50, 450), rs.randint(50, 450), + rs.randint(100, 400), rs.randint(100, 400), + rs.randn() * 2 * np.pi - np.pi + ]) + self._set_state(state) + + observation = self._get_obs() + return observation + + def step(self, action): + dt = 1.0 / self.sim_hz + self.n_contact_points = 0 + n_steps = self.sim_hz // self.control_hz + if action is not None: + self.latest_action = action + for i in range(n_steps): + # Step PD control. + # self.agent.velocity = self.k_p * (act - self.agent.position) # P control works too. + acceleration = self.k_p * (action - self.agent.position) + self.k_v * (Vec2d(0, 0) - self.agent.velocity) + self.agent.velocity += acceleration * dt + + # Step physics. + self.space.step(dt) + + # compute reward + goal_body = self._get_goal_pose_body(self.goal_pose) + goal_geom = pymunk_to_shapely(goal_body, self.block.shapes) + block_geom = pymunk_to_shapely(self.block, self.block.shapes) + + intersection_area = goal_geom.intersection(block_geom).area + goal_area = goal_geom.area + coverage = intersection_area / goal_area + reward = np.clip(coverage / self.success_threshold, 0, 1) + done = coverage > self.success_threshold + + observation = self._get_obs() + info = self._get_info() + + return observation, reward, done, info + + def render(self, mode): + return self._render_frame(mode) + + def teleop_agent(self): + TeleopAgent = collections.namedtuple('TeleopAgent', ['act']) + def act(obs): + act = None + mouse_position = pymunk.pygame_util.from_pygame(Vec2d(*pygame.mouse.get_pos()), self.screen) + if self.teleop or (mouse_position - self.agent.position).length < 30: + self.teleop = True + act = mouse_position + return act + return TeleopAgent(act) + + def _get_obs(self): + obs = np.array( + tuple(self.agent.position) \ + + tuple(self.block.position) \ + + (self.block.angle % (2 * np.pi),)) + return obs + + def _get_goal_pose_body(self, pose): + mass = 1 + inertia = pymunk.moment_for_box(mass, (50, 100)) + body = pymunk.Body(mass, inertia) + # preserving the legacy assignment order for compatibility + # the order here doesn't matter somehow, maybe because CoM is aligned with body origin + body.position = pose[:2].tolist() + body.angle = pose[2] + return body + + def _get_info(self): + n_steps = self.sim_hz // self.control_hz + n_contact_points_per_step = int(np.ceil(self.n_contact_points / n_steps)) + info = { + 'pos_agent': np.array(self.agent.position), + 'vel_agent': np.array(self.agent.velocity), + 'block_pose': np.array(list(self.block.position) + [self.block.angle]), + 'goal_pose': self.goal_pose, + 'n_contacts': n_contact_points_per_step} + return info + + def _render_frame(self, mode): + + if self.window is None and mode == "human": + pygame.init() + pygame.display.init() + self.window = pygame.display.set_mode((self.window_size, self.window_size)) + if self.clock is None and mode == "human": + self.clock = pygame.time.Clock() + + canvas = pygame.Surface((self.window_size, self.window_size)) + canvas.fill((255, 255, 255)) + self.screen = canvas + + draw_options = DrawOptions(canvas) + + # Draw goal pose. + goal_body = self._get_goal_pose_body(self.goal_pose) + for shape in self.block.shapes: + goal_points = [pymunk.pygame_util.to_pygame(goal_body.local_to_world(v), draw_options.surface) for v in shape.get_vertices()] + goal_points += [goal_points[0]] + pygame.draw.polygon(canvas, self.goal_color, goal_points) + + # Draw agent and block. + self.space.debug_draw(draw_options) + + if mode == "human": + # The following line copies our drawings from `canvas` to the visible window + self.window.blit(canvas, canvas.get_rect()) + pygame.event.pump() + pygame.display.update() + + # the clock is already ticked during in step for "human" + + + img = np.transpose( + np.array(pygame.surfarray.pixels3d(canvas)), axes=(1, 0, 2) + ) + img = cv2.resize(img, (self.render_size, self.render_size)) + if self.render_action: + if self.render_action and (self.latest_action is not None): + action = np.array(self.latest_action) + coord = (action / 512 * 96).astype(np.int32) + marker_size = int(8/96*self.render_size) + thickness = int(1/96*self.render_size) + cv2.drawMarker(img, coord, + color=(255,0,0), markerType=cv2.MARKER_CROSS, + markerSize=marker_size, thickness=thickness) + return img + + + def close(self): + if self.window is not None: + pygame.display.quit() + pygame.quit() + + def seed(self, seed=None): + if seed is None: + seed = np.random.randint(0,25536) + self._seed = seed + self.np_random = np.random.default_rng(seed) + + def _handle_collision(self, arbiter, space, data): + self.n_contact_points += len(arbiter.contact_point_set.points) + + def _set_state(self, state): + if isinstance(state, np.ndarray): + state = state.tolist() + pos_agent = state[:2] + pos_block = state[2:4] + rot_block = state[4] + self.agent.position = pos_agent + # setting angle rotates with respect to center of mass + # therefore will modify the geometric position + # if not the same as CoM + # therefore should be modified first. + if self.legacy: + # for compatibility with legacy data + self.block.position = pos_block + self.block.angle = rot_block + else: + self.block.angle = rot_block + self.block.position = pos_block + + # Run physics to take effect + self.space.step(1.0 / self.sim_hz) + + def _set_state_local(self, state_local): + agent_pos_local = state_local[:2] + block_pose_local = state_local[2:] + tf_img_obj = st.AffineTransform( + translation=self.goal_pose[:2], + rotation=self.goal_pose[2]) + tf_obj_new = st.AffineTransform( + translation=block_pose_local[:2], + rotation=block_pose_local[2] + ) + tf_img_new = st.AffineTransform( + matrix=tf_img_obj.params @ tf_obj_new.params + ) + agent_pos_new = tf_img_new(agent_pos_local) + new_state = np.array( + list(agent_pos_new[0]) + list(tf_img_new.translation) \ + + [tf_img_new.rotation]) + self._set_state(new_state) + return new_state + + def _setup(self): + self.space = pymunk.Space() + self.space.gravity = 0, 0 + self.space.damping = 0 + self.teleop = False + self.render_buffer = list() + + # Add walls. + walls = [ + self._add_segment((5, 506), (5, 5), 2), + self._add_segment((5, 5), (506, 5), 2), + self._add_segment((506, 5), (506, 506), 2), + self._add_segment((5, 506), (506, 506), 2) + ] + self.space.add(*walls) + + # Add agent, block, and goal zone. + self.agent = self.add_circle((256, 400), 15) + self.block = self.add_tee((256, 300), 0) + self.goal_color = pygame.Color('LightGreen') + self.goal_pose = np.array([256,256,np.pi/4]) # x, y, theta (in radians) + + # Add collision handling + self.collision_handeler = self.space.add_collision_handler(0, 0) + self.collision_handeler.post_solve = self._handle_collision + self.n_contact_points = 0 + + self.max_score = 50 * 100 + self.success_threshold = 0.95 # 95% coverage. + + def _add_segment(self, a, b, radius): + shape = pymunk.Segment(self.space.static_body, a, b, radius) + shape.color = pygame.Color('LightGray') # https://htmlcolorcodes.com/color-names + return shape + + def add_circle(self, position, radius): + body = pymunk.Body(body_type=pymunk.Body.KINEMATIC) + body.position = position + body.friction = 1 + shape = pymunk.Circle(body, radius) + shape.color = pygame.Color('RoyalBlue') + self.space.add(body, shape) + return body + + def add_box(self, position, height, width): + mass = 1 + inertia = pymunk.moment_for_box(mass, (height, width)) + body = pymunk.Body(mass, inertia) + body.position = position + shape = pymunk.Poly.create_box(body, (height, width)) + shape.color = pygame.Color('LightSlateGray') + self.space.add(body, shape) + return body + + def add_tee(self, position, angle, scale=30, color='LightSlateGray', mask=pymunk.ShapeFilter.ALL_MASKS()): + mass = 1 + length = 4 + vertices1 = [(-length*scale/2, scale), + ( length*scale/2, scale), + ( length*scale/2, 0), + (-length*scale/2, 0)] + inertia1 = pymunk.moment_for_poly(mass, vertices=vertices1) + vertices2 = [(-scale/2, scale), + (-scale/2, length*scale), + ( scale/2, length*scale), + ( scale/2, scale)] + inertia2 = pymunk.moment_for_poly(mass, vertices=vertices1) + body = pymunk.Body(mass, inertia1 + inertia2) + shape1 = pymunk.Poly(body, vertices1) + shape2 = pymunk.Poly(body, vertices2) + shape1.color = pygame.Color(color) + shape2.color = pygame.Color(color) + shape1.filter = pymunk.ShapeFilter(mask=mask) + shape2.filter = pymunk.ShapeFilter(mask=mask) + body.center_of_gravity = (shape1.center_of_gravity + shape2.center_of_gravity) / 2 + body.position = position + body.angle = angle + body.friction = 1 + self.space.add(body, shape1, shape2) + return body diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pusht_image_env.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pusht_image_env.py new file mode 100644 index 0000000000000000000000000000000000000000..7871365d031e2ca8b172733d20bfb4d3b770ee51 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pusht_image_env.py @@ -0,0 +1,66 @@ +from gym import spaces +from cleandiffuser.env.pusht.pusht_env import PushTEnv +import numpy as np +import cv2 + +class PushTImageEnv(PushTEnv): + metadata = {"render.modes": ["rgb_array"], "video.frames_per_second": 10} + + def __init__(self, + legacy=False, + block_cog=None, + damping=None, + render_size=96): + super().__init__( + legacy=legacy, + block_cog=block_cog, + damping=damping, + render_size=render_size, + render_action=False) + ws = self.window_size + self.observation_space = spaces.Dict({ + 'image': spaces.Box( + low=0, + high=1, + shape=(3,render_size,render_size), + dtype=np.float32 + ), + 'agent_pos': spaces.Box( + low=0, + high=ws, + shape=(2,), + dtype=np.float32 + ) + }) + self.render_cache = None + + def _get_obs(self): + img = super()._render_frame(mode='rgb_array') + + agent_pos = np.array(self.agent.position) + img_obs = np.moveaxis(img.astype(np.float32) / 255, -1, 0) + obs = { + 'image': img_obs, + 'agent_pos': agent_pos + } + + # draw action + if self.latest_action is not None: + action = np.array(self.latest_action) + coord = (action / 512 * 96).astype(np.int32) + marker_size = int(8/96*self.render_size) + thickness = int(1/96*self.render_size) + cv2.drawMarker(img, coord, + color=(255,0,0), markerType=cv2.MARKER_CROSS, + markerSize=marker_size, thickness=thickness) + self.render_cache = img + + return obs + + def render(self, mode): + assert mode == 'rgb_array' + + if self.render_cache is None: + self._get_obs() + + return self.render_cache diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pusht_keypoints_env.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pusht_keypoints_env.py new file mode 100644 index 0000000000000000000000000000000000000000..74de409c9e2ef598a951351f73cef6c6a9a296fa --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pusht_keypoints_env.py @@ -0,0 +1,132 @@ +from typing import Dict, Sequence, Union, Optional +from gym import spaces +from cleandiffuser.env.pusht.pusht_env import PushTEnv +from cleandiffuser.env.pusht.pymunk_keypoint_manager import PymunkKeypointManager +import numpy as np + + +class PushTKeypointsEnv(PushTEnv): + def __init__(self, + legacy=False, + block_cog=None, + damping=None, + render_size=96, + keypoint_visible_rate=1.0, + agent_keypoints=False, + draw_keypoints=False, + reset_to_state=None, + render_action=True, + local_keypoint_map: Dict[str, np.ndarray]=None, + color_map: Optional[Dict[str, np.ndarray]]=None): + super().__init__( + legacy=legacy, + block_cog=block_cog, + damping=damping, + render_size=render_size, + reset_to_state=reset_to_state, + render_action=render_action) + ws = self.window_size + + if local_keypoint_map is None: + # create default keypoint definition + kp_kwargs = self.genenerate_keypoint_manager_params() + local_keypoint_map = kp_kwargs['local_keypoint_map'] + color_map = kp_kwargs['color_map'] + + # create observation spaces + Dblockkps = np.prod(local_keypoint_map['block'].shape) + Dagentkps = np.prod(local_keypoint_map['agent'].shape) + Dagentpos = 2 + + Do = Dblockkps + if agent_keypoints: + # blockkp + agnet_pos + Do += Dagentkps + else: + # blockkp + agnet_kp + Do += Dagentpos + # obs + obs_mask + Dobs = Do * 2 + + low = np.zeros((Dobs,), dtype=np.float64) + high = np.full_like(low, ws) + # mask range 0-1 + high[Do:] = 1. + + # (block_kps+agent_kps, xy+confidence) + self.observation_space = spaces.Box( + low=low, + high=high, + shape=low.shape, + dtype=np.float64 + ) + + self.keypoint_visible_rate = keypoint_visible_rate + self.agent_keypoints = agent_keypoints + self.draw_keypoints = draw_keypoints + self.kp_manager = PymunkKeypointManager( + local_keypoint_map=local_keypoint_map, + color_map=color_map) + self.draw_kp_map = None + + @classmethod + def genenerate_keypoint_manager_params(cls): + env = PushTEnv() + kp_manager = PymunkKeypointManager.create_from_pusht_env(env) + kp_kwargs = kp_manager.kwargs + return kp_kwargs + + def _get_obs(self): + # get keypoints + obj_map = { + 'block': self.block + } + if self.agent_keypoints: + obj_map['agent'] = self.agent + + kp_map = self.kp_manager.get_keypoints_global( + pose_map=obj_map, is_obj=True) + # python dict guerentee order of keys and values + kps = np.concatenate(list(kp_map.values()), axis=0) + + # select keypoints to drop + n_kps = kps.shape[0] + visible_kps = self.np_random.random(size=(n_kps,)) < self.keypoint_visible_rate + kps_mask = np.repeat(visible_kps[:,None], 2, axis=1) + + # save keypoints for rendering + vis_kps = kps.copy() + vis_kps[~visible_kps] = 0 + draw_kp_map = { + 'block': vis_kps[:len(kp_map['block'])] + } + if self.agent_keypoints: + draw_kp_map['agent'] = vis_kps[len(kp_map['block']):] + self.draw_kp_map = draw_kp_map + + # construct obs + obs = kps.flatten() + obs_mask = kps_mask.flatten() + if not self.agent_keypoints: + # passing agent position when keypoints are not available + agent_pos = np.array(self.agent.position) + obs = np.concatenate([ + obs, agent_pos + ]) + obs_mask = np.concatenate([ + obs_mask, np.ones((2,), dtype=bool) + ]) + + # obs, obs_mask + obs = np.concatenate([ + obs, obs_mask.astype(obs.dtype) + ], axis=0) + return obs + + + def _render_frame(self, mode): + img = super()._render_frame(mode) + if self.draw_keypoints: + self.kp_manager.draw_keypoints( + img, self.draw_kp_map, radius=int(img.shape[0]/96)) + return img diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pymunk_keypoint_manager.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pymunk_keypoint_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..f4f2db2ee64bb7e90e0d6657d10b0e573877d29d --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pymunk_keypoint_manager.py @@ -0,0 +1,146 @@ +from typing import Dict, Sequence, Union, Optional +import numpy as np +import skimage.transform as st +import pymunk +import pygame +from matplotlib import cm +import cv2 +from cleandiffuser.env.pusht.pymunk_override import DrawOptions + + +def farthest_point_sampling(points: np.ndarray, n_points: int, init_idx: int): + """ + Naive O(N^2) + """ + assert(n_points >= 1) + chosen_points = [points[init_idx]] + for _ in range(n_points-1): + cpoints = np.array(chosen_points) + all_dists = np.linalg.norm(points[:,None,:] - cpoints[None,:,:], axis=-1) + min_dists = all_dists.min(axis=1) + next_idx = np.argmax(min_dists) + next_pt = points[next_idx] + chosen_points.append(next_pt) + result = np.array(chosen_points) + return result + + +class PymunkKeypointManager: + def __init__(self, + local_keypoint_map: Dict[str, np.ndarray], + color_map: Optional[Dict[str, np.ndarray]]=None): + """ + local_keypoint_map: + "": (N,2) floats in object local coordinate + """ + if color_map is None: + cmap = cm.get_cmap('tab10') + color_map = dict() + for i, key in enumerate(local_keypoint_map.keys()): + color_map[key] = (np.array(cmap.colors[i]) * 255).astype(np.uint8) + + self.local_keypoint_map = local_keypoint_map + self.color_map = color_map + + @property + def kwargs(self): + return { + 'local_keypoint_map': self.local_keypoint_map, + 'color_map': self.color_map + } + + @classmethod + def create_from_pusht_env(cls, env, n_block_kps=9, n_agent_kps=3, seed=0, **kwargs): + rng = np.random.default_rng(seed=seed) + local_keypoint_map = dict() + for name in ['block','agent']: + self = env + self.space = pymunk.Space() + if name == 'agent': + self.agent = obj = self.add_circle((256, 400), 15) + n_kps = n_agent_kps + else: + self.block = obj = self.add_tee((256, 300), 0) + n_kps = n_block_kps + + self.screen = pygame.Surface((512,512)) + self.screen.fill(pygame.Color("white")) + draw_options = DrawOptions(self.screen) + self.space.debug_draw(draw_options) + # pygame.display.flip() + img = np.uint8(pygame.surfarray.array3d(self.screen).transpose(1, 0, 2)) + obj_mask = (img != np.array([255,255,255],dtype=np.uint8)).any(axis=-1) + + tf_img_obj = cls.get_tf_img_obj(obj) + xy_img = np.moveaxis(np.array(np.indices((512,512))), 0, -1)[:,:,::-1] + local_coord_img = tf_img_obj.inverse(xy_img.reshape(-1,2)).reshape(xy_img.shape) + obj_local_coords = local_coord_img[obj_mask] + + # furthest point sampling + init_idx = rng.choice(len(obj_local_coords)) + obj_local_kps = farthest_point_sampling(obj_local_coords, n_kps, init_idx) + small_shift = rng.uniform(0, 1, size=obj_local_kps.shape) + obj_local_kps += small_shift + + local_keypoint_map[name] = obj_local_kps + + return cls(local_keypoint_map=local_keypoint_map, **kwargs) + + @staticmethod + def get_tf_img(pose: Sequence): + pos = pose[:2] + rot = pose[2] + tf_img_obj = st.AffineTransform( + translation=pos, rotation=rot) + return tf_img_obj + + @classmethod + def get_tf_img_obj(cls, obj: pymunk.Body): + pose = tuple(obj.position) + (obj.angle,) + return cls.get_tf_img(pose) + + def get_keypoints_global(self, + pose_map: Dict[set, Union[Sequence, pymunk.Body]], + is_obj=False): + kp_map = dict() + for key, value in pose_map.items(): + if is_obj: + tf_img_obj = self.get_tf_img_obj(value) + else: + tf_img_obj = self.get_tf_img(value) + kp_local = self.local_keypoint_map[key] + kp_global = tf_img_obj(kp_local) + kp_map[key] = kp_global + return kp_map + + def draw_keypoints(self, img, kps_map, radius=1): + scale = np.array(img.shape[:2]) / np.array([512,512]) + for key, value in kps_map.items(): + color = self.color_map[key].tolist() + coords = (value * scale).astype(np.int32) + for coord in coords: + cv2.circle(img, coord, radius=radius, color=color, thickness=-1) + return img + + def draw_keypoints_pose(self, img, pose_map, is_obj=False, **kwargs): + kp_map = self.get_keypoints_global(pose_map, is_obj=is_obj) + return self.draw_keypoints(img, kps_map=kp_map, **kwargs) + + +def test(): + from cleandiffuser.environment.push_t_env import PushTEnv + from matplotlib import pyplot as plt + + env = PushTEnv(headless=True, obs_state=False, draw_action=False) + kp_manager = PymunkKeypointManager.create_from_pusht_env(env=env) + env.reset() + obj_map = { + 'block': env.block, + 'agent': env.agent + } + + obs = env.render() + img = obs.astype(np.uint8) + kp_manager.draw_keypoints_pose(img=img, pose_map=obj_map, is_obj=True) + + plt.imshow(img) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pymunk_override.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pymunk_override.py new file mode 100644 index 0000000000000000000000000000000000000000..2439020a13e01ad48f3677919157bb7e49e50569 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/pusht/pymunk_override.py @@ -0,0 +1,248 @@ +# ---------------------------------------------------------------------------- +# pymunk +# Copyright (c) 2007-2016 Victor Blomqvist +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# ---------------------------------------------------------------------------- + +"""This submodule contains helper functions to help with quick prototyping +using pymunk together with pygame. + +Intended to help with debugging and prototyping, not for actual production use +in a full application. The methods contained in this module is opinionated +about your coordinate system and not in any way optimized. +""" + +__docformat__ = "reStructuredText" + +__all__ = [ + "DrawOptions", + "get_mouse_pos", + "to_pygame", + "from_pygame", + "lighten", + "positive_y_is_up", +] + +from typing import List, Sequence, Tuple + +import pygame + +import numpy as np + +import pymunk +from pymunk.space_debug_draw_options import SpaceDebugColor +from pymunk.vec2d import Vec2d + +positive_y_is_up: bool = False +"""Make increasing values of y point upwards. + +When True:: + + y + ^ + | . (3, 3) + | + | . (2, 2) + | + +------ > x + +When False:: + + +------ > x + | + | . (2, 2) + | + | . (3, 3) + v + y + +""" + + +class DrawOptions(pymunk.SpaceDebugDrawOptions): + def __init__(self, surface: pygame.Surface) -> None: + """Draw a pymunk.Space on a pygame.Surface object. + + Typical usage:: + + >>> import pymunk + >>> surface = pygame.Surface((10,10)) + >>> space = pymunk.Space() + >>> options = pymunk.pygame_util.DrawOptions(surface) + >>> space.debug_draw(options) + + You can control the color of a shape by setting shape.color to the color + you want it drawn in:: + + >>> c = pymunk.Circle(None, 10) + >>> c.color = pygame.Color("pink") + + See pygame_util.demo.py for a full example + + Since pygame uses a coordinate system where y points down (in contrast + to many other cases), you either have to make the physics simulation + with Pymunk also behave in that way, or flip everything when you draw. + + The easiest is probably to just make the simulation behave the same + way as Pygame does. In that way all coordinates used are in the same + orientation and easy to reason about:: + + >>> space = pymunk.Space() + >>> space.gravity = (0, -1000) + >>> body = pymunk.Body() + >>> body.position = (0, 0) # will be positioned in the top left corner + >>> space.debug_draw(options) + + To flip the drawing its possible to set the module property + :py:data:`positive_y_is_up` to True. Then the pygame drawing will flip + the simulation upside down before drawing:: + + >>> positive_y_is_up = True + >>> body = pymunk.Body() + >>> body.position = (0, 0) + >>> # Body will be position in bottom left corner + + :Parameters: + surface : pygame.Surface + Surface that the objects will be drawn on + """ + self.surface = surface + super(DrawOptions, self).__init__() + + def draw_circle( + self, + pos: Vec2d, + angle: float, + radius: float, + outline_color: SpaceDebugColor, + fill_color: SpaceDebugColor, + ) -> None: + p = to_pygame(pos, self.surface) + + pygame.draw.circle(self.surface, fill_color.as_int(), p, round(radius), 0) + pygame.draw.circle(self.surface, light_color(fill_color).as_int(), p, round(radius-4), 0) + + circle_edge = pos + Vec2d(radius, 0).rotated(angle) + p2 = to_pygame(circle_edge, self.surface) + line_r = 2 if radius > 20 else 1 + # pygame.draw.lines(self.surface, outline_color.as_int(), False, [p, p2], line_r) + + def draw_segment(self, a: Vec2d, b: Vec2d, color: SpaceDebugColor) -> None: + p1 = to_pygame(a, self.surface) + p2 = to_pygame(b, self.surface) + + pygame.draw.aalines(self.surface, color.as_int(), False, [p1, p2]) + + def draw_fat_segment( + self, + a: Tuple[float, float], + b: Tuple[float, float], + radius: float, + outline_color: SpaceDebugColor, + fill_color: SpaceDebugColor, + ) -> None: + p1 = to_pygame(a, self.surface) + p2 = to_pygame(b, self.surface) + + r = round(max(1, radius * 2)) + pygame.draw.lines(self.surface, fill_color.as_int(), False, [p1, p2], r) + if r > 2: + orthog = [abs(p2[1] - p1[1]), abs(p2[0] - p1[0])] + if orthog[0] == 0 and orthog[1] == 0: + return + scale = radius / (orthog[0] * orthog[0] + orthog[1] * orthog[1]) ** 0.5 + orthog[0] = round(orthog[0] * scale) + orthog[1] = round(orthog[1] * scale) + points = [ + (p1[0] - orthog[0], p1[1] - orthog[1]), + (p1[0] + orthog[0], p1[1] + orthog[1]), + (p2[0] + orthog[0], p2[1] + orthog[1]), + (p2[0] - orthog[0], p2[1] - orthog[1]), + ] + pygame.draw.polygon(self.surface, fill_color.as_int(), points) + pygame.draw.circle( + self.surface, + fill_color.as_int(), + (round(p1[0]), round(p1[1])), + round(radius), + ) + pygame.draw.circle( + self.surface, + fill_color.as_int(), + (round(p2[0]), round(p2[1])), + round(radius), + ) + + def draw_polygon( + self, + verts: Sequence[Tuple[float, float]], + radius: float, + outline_color: SpaceDebugColor, + fill_color: SpaceDebugColor, + ) -> None: + ps = [to_pygame(v, self.surface) for v in verts] + ps += [ps[0]] + + radius = 2 + pygame.draw.polygon(self.surface, light_color(fill_color).as_int(), ps) + + if radius > 0: + for i in range(len(verts)): + a = verts[i] + b = verts[(i + 1) % len(verts)] + self.draw_fat_segment(a, b, radius, fill_color, fill_color) + + def draw_dot( + self, size: float, pos: Tuple[float, float], color: SpaceDebugColor + ) -> None: + p = to_pygame(pos, self.surface) + pygame.draw.circle(self.surface, color.as_int(), p, round(size), 0) + + +def get_mouse_pos(surface: pygame.Surface) -> Tuple[int, int]: + """Get position of the mouse pointer in pymunk coordinates.""" + p = pygame.mouse.get_pos() + return from_pygame(p, surface) + + +def to_pygame(p: Tuple[float, float], surface: pygame.Surface) -> Tuple[int, int]: + """Convenience method to convert pymunk coordinates to pygame surface + local coordinates. + + Note that in case positive_y_is_up is False, this function won't actually do + anything except converting the point to integers. + """ + if positive_y_is_up: + return round(p[0]), surface.get_height() - round(p[1]) + else: + return round(p[0]), round(p[1]) + + +def from_pygame(p: Tuple[float, float], surface: pygame.Surface) -> Tuple[int, int]: + """Convenience method to convert pygame surface local coordinates to + pymunk coordinates + """ + return to_pygame(p, surface) + + +def light_color(color: SpaceDebugColor): + color = np.minimum(1.2 * np.float32([color.r, color.g, color.b, color.a]), np.float32([255])) + color = SpaceDebugColor(r=color[0], g=color[1], b=color[2], a=color[3]) + return color diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/robomimic/robomimic_image_wrapper.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/robomimic/robomimic_image_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..7ee209b8d1eea642d08ce2c675ed3df9f8ba45a8 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/robomimic/robomimic_image_wrapper.py @@ -0,0 +1,165 @@ +from typing import List, Optional +from matplotlib.pyplot import fill +import numpy as np +import gym +from gym import spaces +from omegaconf import OmegaConf +from robomimic.envs.env_robosuite import EnvRobosuite + +class RobomimicImageWrapper(gym.Env): + def __init__(self, + env: EnvRobosuite, + shape_meta: dict, + init_state: Optional[np.ndarray]=None, + render_obs_key='agentview_image', + ): + + self.env = env + self.render_obs_key = render_obs_key + self.init_state = init_state + self.seed_state_map = dict() + self._seed = None + self.shape_meta = shape_meta + self.render_cache = None + self.has_reset_before = False + + # setup spaces + action_shape = shape_meta['action']['shape'] + action_space = spaces.Box( + low=-1, + high=1, + shape=action_shape, + dtype=np.float32 + ) + self.action_space = action_space + + observation_space = spaces.Dict() + for key, value in shape_meta['obs'].items(): + shape = value['shape'] + min_value, max_value = -1, 1 + if key.endswith('image'): + min_value, max_value = 0, 1 + elif key.endswith('quat'): + min_value, max_value = -1, 1 + elif key.endswith('qpos'): + min_value, max_value = -1, 1 + elif key.endswith('pos'): + # better range? + min_value, max_value = -1, 1 + else: + raise RuntimeError(f"Unsupported type {key}") + + this_space = spaces.Box( + low=min_value, + high=max_value, + shape=shape, + dtype=np.float32 + ) + observation_space[key] = this_space + self.observation_space = observation_space + + + def get_observation(self, raw_obs=None): + if raw_obs is None: + raw_obs = self.env.get_observation() + + self.render_cache = raw_obs[self.render_obs_key] + + obs = dict() + # print(raw_obs.keys()) + for key in self.observation_space.keys(): + obs[key] = raw_obs[key] + return obs + + def seed(self, seed=None): + np.random.seed(seed=seed) + self._seed = seed + + def reset(self): + if self.init_state is not None: + if not self.has_reset_before: + # the env must be fully reset at least once to ensure correct rendering + self.env.reset() + self.has_reset_before = True + + # always reset to the same state + # to be compatible with gym + raw_obs = self.env.reset_to({'states': self.init_state}) + elif self._seed is not None: + # reset to a specific seed + seed = self._seed + if seed in self.seed_state_map: + # env.reset is expensive, use cache + raw_obs = self.env.reset_to({'states': self.seed_state_map[seed]}) + else: + # robosuite's initializes all use numpy global random state + np.random.seed(seed=seed) + raw_obs = self.env.reset() + state = self.env.get_state()['states'] + self.seed_state_map[seed] = state + self._seed = None + else: + # random reset + raw_obs = self.env.reset() + + # return obs + obs = self.get_observation(raw_obs) + return obs + + def step(self, action): + raw_obs, reward, done, info = self.env.step(action) + obs = self.get_observation(raw_obs) + return obs, reward, done, info + + def render(self, mode='rgb_array'): + if self.render_cache is None: + raise RuntimeError('Must run reset or step before render.') + img = np.moveaxis(self.render_cache, 0, -1) + img = (img * 255).astype(np.uint8) + return img + + +def test(): + import os + from omegaconf import OmegaConf + cfg_path = os.path.expanduser('~/dev/cleandiffuser/cleandiffuser/config/task/lift_image.yaml') + cfg = OmegaConf.load(cfg_path) + shape_meta = cfg['shape_meta'] + + + import robomimic.utils.file_utils as FileUtils + import robomimic.utils.env_utils as EnvUtils + from matplotlib import pyplot as plt + + dataset_path = os.path.expanduser('~/dev/cleandiffuser/data/robomimic/datasets/square/ph/image.hdf5') + env_meta = FileUtils.get_env_metadata_from_dataset( + dataset_path) + + env = EnvUtils.create_env_from_metadata( + env_meta=env_meta, + render=False, + render_offscreen=False, + use_image_obs=True, + ) + + wrapper = RobomimicImageWrapper( + env=env, + shape_meta=shape_meta + ) + wrapper.seed(0) + obs = wrapper.reset() + img = wrapper.render() + plt.imshow(img) + + + # states = list() + # for _ in range(2): + # wrapper.seed(0) + # wrapper.reset() + # states.append(wrapper.env.get_state()['states']) + # assert np.allclose(states[0], states[1]) + + # img = wrapper.render() + # plt.imshow(img) + # wrapper.seed() + # states.append(wrapper.env.get_state()['states']) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/robomimic/robomimic_lowdim_wrapper.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/robomimic/robomimic_lowdim_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..3130ee91d821d63629b28273f67cb214be7833aa --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/robomimic/robomimic_lowdim_wrapper.py @@ -0,0 +1,97 @@ +from typing import List, Dict, Optional +import numpy as np +import gym +from gym.spaces import Box +from robomimic.envs.env_robosuite import EnvRobosuite + +class RobomimicLowdimWrapper(gym.Env): + def __init__(self, + env: EnvRobosuite, + obs_keys: List[str]=[ + 'object', + 'robot0_eef_pos', + 'robot0_eef_quat', + 'robot0_gripper_qpos'], + init_state: Optional[np.ndarray]=None, + render_hw=(256,256), + render_camera_name='agentview' + ): + + self.env = env + self.obs_keys = obs_keys + self.init_state = init_state + self.render_hw = render_hw + self.render_camera_name = render_camera_name + self.seed_state_map = dict() + self._seed = None + + # setup spaces + low = np.full(env.action_dimension, fill_value=-1) + high = np.full(env.action_dimension, fill_value=1) + self.action_space = Box( + low=low, + high=high, + shape=low.shape, + dtype=low.dtype + ) + obs_example = self.get_observation() + low = np.full_like(obs_example, fill_value=-1) + high = np.full_like(obs_example, fill_value=1) + self.observation_space = Box( + low=low, + high=high, + shape=low.shape, + dtype=low.dtype + ) + + def get_observation(self): + raw_obs = self.env.get_observation() + obs = np.concatenate([ + raw_obs[key] for key in self.obs_keys + ], axis=0) + return obs + + def seed(self, seed=None): + np.random.seed(seed=seed) + self._seed = seed + + def reset(self): + if self.init_state is not None: + # always reset to the same state + # to be compatible with gym + self.env.reset_to({'states': self.init_state}) + elif self._seed is not None: + # reset to a specific seed + seed = self._seed + if seed in self.seed_state_map: + # env.reset is expensive, use cache + self.env.reset_to({'states': self.seed_state_map[seed]}) + else: + # robosuite's initializes all use numpy global random state + np.random.seed(seed=seed) + self.env.reset() + state = self.env.get_state()['states'] + self.seed_state_map[seed] = state + self._seed = None + else: + # random reset + self.env.reset() + + # return obs + obs = self.get_observation() + return obs + + def step(self, action): + raw_obs, reward, done, info = self.env.step(action) + obs = np.concatenate([ + raw_obs[key] for key in self.obs_keys + ], axis=0) + return obs, reward, done, info + + def render(self, mode='rgb_array'): + h, w = self.render_hw + return self.env.render(mode=mode, + height=h, width=w, + camera_name=self.render_camera_name) + + diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/utils.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..59f12491c1fdcc7527497e217cddd7e35e54d8f7 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/utils.py @@ -0,0 +1,158 @@ +from typing import List, Tuple, Optional, Dict +import math +import numpy as np +import av + + +def get_accumulate_timestamp_idxs( + timestamps: List[float], + start_time: float, + dt: float, + eps:float=1e-5, + next_global_idx: Optional[int]=0, + allow_negative=False + ) -> Tuple[List[int], List[int], int]: + """ + For each dt window, choose the first timestamp in the window. + Assumes timestamps sorted. One timestamp might be chosen multiple times due to dropped frames. + next_global_idx should start at 0 normally, and then use the returned next_global_idx. + However, when overwiting previous values are desired, set last_global_idx to None. + + Returns: + local_idxs: which index in the given timestamps array to chose from + global_idxs: the global index of each chosen timestamp + next_global_idx: used for next call. + """ + local_idxs = list() + global_idxs = list() + for local_idx, ts in enumerate(timestamps): + # add eps * dt to timestamps so that when ts == start_time + k * dt + # is always recorded as kth element (avoiding floating point errors) + global_idx = math.floor((ts - start_time) / dt + eps) + if (not allow_negative) and (global_idx < 0): + continue + if next_global_idx is None: + next_global_idx = global_idx + + n_repeats = max(0, global_idx - next_global_idx + 1) + for i in range(n_repeats): + local_idxs.append(local_idx) + global_idxs.append(next_global_idx + i) + next_global_idx += n_repeats + return local_idxs, global_idxs, next_global_idx + + +class VideoRecorder: + def __init__(self, + fps, + codec, + input_pix_fmt, + # options for codec + **kwargs + ): + """ + input_pix_fmt: rgb24, bgr24 see https://github.com/PyAV-Org/PyAV/blob/bc4eedd5fc474e0f25b22102b2771fe5a42bb1c7/av/video/frame.pyx#L352 + """ + + self.fps = fps + self.codec = codec + self.input_pix_fmt = input_pix_fmt + self.kwargs = kwargs + # runtime set + self._reset_state() + + def _reset_state(self): + self.container = None + self.stream = None + self.shape = None + self.dtype = None + self.start_time = None + self.next_global_idx = 0 + + @classmethod + def create_h264(cls, + fps, + codec='h264', + input_pix_fmt='rgb24', + output_pix_fmt='yuv420p', + crf=18, + profile='high', + **kwargs + ): + obj = cls( + fps=fps, + codec=codec, + input_pix_fmt=input_pix_fmt, + pix_fmt=output_pix_fmt, + options={ + 'crf': str(crf), + 'profile': profile + }, + **kwargs + ) + return obj + + + def __del__(self): + self.stop() + + def is_ready(self): + return self.stream is not None + + def start(self, file_path, start_time=None): + if self.is_ready(): + # if still recording, stop first and start anew. + self.stop() + + self.container = av.open(file_path, mode='w') + self.stream = self.container.add_stream(self.codec, rate=self.fps) + codec_context = self.stream.codec_context + for k, v in self.kwargs.items(): + setattr(codec_context, k, v) + self.start_time = start_time + + def write_frame(self, img: np.ndarray, frame_time=None): + if not self.is_ready(): + raise RuntimeError('Must run start() before writing!') + + n_repeats = 1 + if self.start_time is not None: + local_idxs, global_idxs, self.next_global_idx \ + = get_accumulate_timestamp_idxs( + # only one timestamp + timestamps=[frame_time], + start_time=self.start_time, + dt=1/self.fps, + next_global_idx=self.next_global_idx + ) + # number of appearance means repeats + n_repeats = len(local_idxs) + + if self.shape is None: + self.shape = img.shape + self.dtype = img.dtype + h,w,c = img.shape + self.stream.width = w + self.stream.height = h + assert img.shape == self.shape + assert img.dtype == self.dtype + + frame = av.VideoFrame.from_ndarray( + img, format=self.input_pix_fmt) + for i in range(n_repeats): + for packet in self.stream.encode(frame): + self.container.mux(packet) + + def stop(self): + if not self.is_ready(): + return + + # Flush stream + for packet in self.stream.encode(): + self.container.mux(packet) + + # Close the file + self.container.close() + + # reset runtime parameters + self._reset_state() \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/wrapper.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..d27ea8e1585399c545aff44ea78407d792fbbbcc --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/env/wrapper.py @@ -0,0 +1,257 @@ +import gym +from gym import spaces +import numpy as np +from collections import defaultdict, deque +import dill +from cleandiffuser.env.utils import VideoRecorder + +# ------------------ MultiStepWrapper ------------------------ + +def stack_repeated(x, n): + return np.repeat(np.expand_dims(x,axis=0),n,axis=0) + +def repeated_box(box_space, n): + return spaces.Box( + low=stack_repeated(box_space.low, n), + high=stack_repeated(box_space.high, n), + shape=(n,) + box_space.shape, + dtype=box_space.dtype + ) + +def repeated_space(space, n): + if isinstance(space, spaces.Box): + return repeated_box(space, n) + elif isinstance(space, spaces.Dict): + result_space = spaces.Dict() + for key, value in space.items(): + result_space[key] = repeated_space(value, n) + return result_space + else: + raise RuntimeError(f'Unsupported space type {type(space)}') + +def take_last_n(x, n): + x = list(x) + n = min(len(x), n) + return np.array(x[-n:]) + +def dict_take_last_n(x, n): + result = dict() + for key, value in x.items(): + result[key] = take_last_n(value, n) + return result + +def aggregate(data, method='max'): + if method == 'max': + # equivalent to any + return np.max(data) + elif method == 'min': + # equivalent to all + return np.min(data) + elif method == 'mean': + return np.mean(data) + elif method == 'sum': + return np.sum(data) + else: + raise NotImplementedError() + +def stack_last_n_obs(all_obs, n_steps): + assert(len(all_obs) > 0) + all_obs = list(all_obs) + result = np.zeros((n_steps,) + all_obs[-1].shape, + dtype=all_obs[-1].dtype) + start_idx = -min(n_steps, len(all_obs)) + result[start_idx:] = np.array(all_obs[start_idx:]) + if n_steps > len(all_obs): + # pad + result[:start_idx] = result[start_idx] + return result + +class MultiStepWrapper(gym.Wrapper): + def __init__(self, + env, + n_obs_steps, + n_action_steps, + max_episode_steps=None, + reward_agg_method='max' + ): + super().__init__(env) + self._action_space = repeated_space(env.action_space, n_action_steps) + self._observation_space = repeated_space(env.observation_space, n_obs_steps) + self.max_episode_steps = max_episode_steps + self.n_obs_steps = n_obs_steps + self.n_action_steps = n_action_steps + self.reward_agg_method = reward_agg_method + + self.obs = deque(maxlen=n_obs_steps+1) + self.reward = list() + self.done = list() + self.info = defaultdict(lambda : deque(maxlen=n_obs_steps+1)) + + def reset(self): + """Resets the environment using kwargs.""" + obs = super().reset() + + self.obs = deque([obs], maxlen=self.n_obs_steps+1) + self.reward = list() + self.done = list() + self.info = defaultdict(lambda : deque(maxlen=self.n_obs_steps+1)) + + obs = self._get_obs(self.n_obs_steps) + return obs + + def step(self, action): + """ + actions: (n_action_steps,) + action_shape + """ + for act in action: + if len(self.done) > 0 and self.done[-1]: + # termination + break + observation, reward, done, info = super().step(act) + + self.obs.append(observation) + self.reward.append(reward) + if (self.max_episode_steps is not None) \ + and (len(self.reward) >= self.max_episode_steps): + # truncation + done = True + self.done.append(done) + self._add_info(info) + + observation = self._get_obs(self.n_obs_steps) + reward = aggregate(self.reward, self.reward_agg_method) + done = aggregate(self.done, 'max') + info = dict_take_last_n(self.info, self.n_obs_steps) + return observation, reward, done, info + + def _get_obs(self, n_steps=1): + """ + Output (n_steps,) + obs_shape + """ + assert(len(self.obs) > 0) + if isinstance(self.observation_space, spaces.Box): + return stack_last_n_obs(self.obs, n_steps) + elif isinstance(self.observation_space, spaces.Dict): + result = dict() + for key in self.observation_space.keys(): + result[key] = stack_last_n_obs( + [obs[key] for obs in self.obs], + n_steps + ) + return result + else: + raise RuntimeError('Unsupported space type') + + def _add_info(self, info): + for key, value in info.items(): + self.info[key].append(value) + + def get_rewards(self): + return self.reward + + def get_attr(self, name): + return getattr(self, name) + + def run_dill_function(self, dill_fn): + fn = dill.loads(dill_fn) + return fn(self) + + def get_infos(self): + result = dict() + for k, v in self.info.items(): + result[k] = list(v) + return result + +# ------------------ VideoWrapper ------------------------ + +class VideoWrapper(gym.Wrapper): + def __init__(self, + env, + mode='rgb_array', + enabled=True, + steps_per_render=1, + **kwargs + ): + super().__init__(env) + + self.mode = mode + self.enabled = enabled + self.render_kwargs = kwargs + self.steps_per_render = steps_per_render + + self.frames = list() + self.step_count = 0 + + def reset(self, **kwargs): + obs = super().reset(**kwargs) + self.frames = list() + self.step_count = 1 + if self.enabled: + frame = self.env.render( + mode=self.mode, **self.render_kwargs) + assert frame.dtype == np.uint8 + self.frames.append(frame) + return obs + + def step(self, action): + result = super().step(action) + self.step_count += 1 + if self.enabled and ((self.step_count % self.steps_per_render) == 0): + frame = self.env.render( + mode=self.mode, **self.render_kwargs) + assert frame.dtype == np.uint8 + self.frames.append(frame) + return result + + def render(self, mode='rgb_array', **kwargs): + return self.frames + +# ------------------ VideoRecordingWrapper ------------------------ + +class VideoRecordingWrapper(gym.Wrapper): + def __init__(self, + env, + video_recoder: VideoRecorder, + mode='rgb_array', + file_path=None, + steps_per_render=1, + **kwargs + ): + """ + When file_path is None, don't record. + """ + super().__init__(env) + + self.mode = mode + self.render_kwargs = kwargs + self.steps_per_render = steps_per_render + self.file_path = file_path + self.video_recoder = video_recoder + + self.step_count = 0 + + def reset(self, **kwargs): + obs = super().reset(**kwargs) + self.frames = list() + self.step_count = 1 + self.video_recoder.stop() + return obs + + def step(self, action): + result = super().step(action) + self.step_count += 1 + if self.file_path is not None \ + and ((self.step_count % self.steps_per_render) == 0): + if not self.video_recoder.is_ready(): + self.video_recoder.start(self.file_path) + + frame = self.env.render( + mode=self.mode, **self.render_kwargs) + assert frame.dtype == np.uint8 + self.video_recoder.write_frame(frame) + return result + + def render(self, mode='rgb_array', **kwargs): + if self.video_recoder.is_ready(): + self.video_recoder.stop() + return self.file_path \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/invdynamic/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/invdynamic/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e75a6dc4cff0c87e0ae70f2deb8575f5cfe0888f --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/invdynamic/__init__.py @@ -0,0 +1,2 @@ +from .common import BasicInvDynamic +from .mlp import MlpInvDynamic, FancyMlpInvDynamic diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/invdynamic/common.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/invdynamic/common.py new file mode 100644 index 0000000000000000000000000000000000000000..e471e5f7d0c8bec8d04dff1396088a6a73c215e9 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/invdynamic/common.py @@ -0,0 +1,7 @@ +class BasicInvDynamic: + + def predict(self, **kwargs): + raise NotImplementedError + + def __call__(self, **kwargs): + return self.predict(**kwargs) \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/invdynamic/mlp.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/invdynamic/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..8350ac0862c2c2d6f986193415df41bab239c437 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/invdynamic/mlp.py @@ -0,0 +1,293 @@ +import torch +import torch.nn as nn + +from cleandiffuser.utils import Mlp + + +class MlpInvDynamic: + """ Simple MLP-based inverse dynamics model. The model is a 3-layer MLP with ReLU activation. + + Args: + o_dim: int, + Dimension of observation space. + a_dim: int, + Dimension of action space. + hidden_dim: int, + Dimension of hidden layers. Default: 512. + out_activation: nn.Module, + Activation function for output layer. Default: nn.Tanh(). + optim_params: dict, + Optimizer parameters. Default: {}. + device: str, + Device for the model. Default: "cpu". + + Examples: + >>> invdyn = MlpInvDynamic(3, 2) + >>> invdyn.train() + >>> batch = ... + >>> obs, act, obs_next = batch + >>> loss = invdyn.update(obs, act, obs_next) + >>> invdyn.eval() + >>> pred_act = invdyn.predict(obs, obs_next) + """ + def __init__( + self, + o_dim: int, + a_dim: int, + hidden_dim: int = 512, + out_activation: nn.Module = nn.Tanh(), + optim_params: dict = {}, + device: str = "cpu", + ): + self.device = device + self.o_dim, self.a_dim, self.hidden_dim = o_dim, a_dim, hidden_dim + self.out_activation = out_activation + self.optim_params = optim_params + params = {"lr": 5e-4} + params.update(optim_params) + self.mlp = Mlp( + 2 * o_dim, [hidden_dim, hidden_dim], a_dim, + nn.ReLU(), out_activation).to(device) + self.optim = torch.optim.Adam(self.mlp.parameters(), **optim_params) + self._init_weights() + + def _init_weights(self): + for m in self.mlp.modules(): + if isinstance(m, nn.Linear): + nn.init.orthogonal_(m.weight) + nn.init.zeros_(m.bias) + + def forward(self, o, o_next): + return self.mlp(torch.cat([o, o_next], dim=-1)) + + def update(self, o, a, o_next): + self.optim.zero_grad() + a_pred = self.forward(o, o_next) + loss = ((a_pred - a) ** 2).mean() + loss.backward() + self.optim.step() + return {"loss": loss.item()} + + @torch.no_grad() + def predict(self, o, o_next): + return self.forward(o, o_next) + + def __call__(self, o, o_next): + return self.predict(o, o_next) + + def train(self): + self.mlp.train() + + def eval(self): + self.mlp.eval() + + def save(self, path): + torch.save(self.mlp.state_dict(), path) + + def load(self, path): + self.mlp.load_state_dict(torch.load(path, self.device)) + + +class FancyMlpInvDynamic: + """ Fancy MLP-based inverse dynamics model. The model is a 3-layer MLP with GELU activation. It also includes + optional LayerNorm and Dropout. We suggest using 0.1 Dropout and LayerNorm for better performance. + + Args: + o_dim: int, + Dimension of observation space. + a_dim: int, + Dimension of action space. + hidden_dim: int, + Dimension of hidden layers. Default: 256. + out_activation: nn.Module, + Activation function for output layer. Default: nn.Tanh(). + add_norm: bool, + Whether to add LayerNorm. Default: False. + add_dropout: bool, + Whether to add Dropout. Default: False. + optim_params: dict, + Optimizer parameters. Default: {}. + device: str, + Device for the model. Default: "cpu". + + Examples: + >>> invdyn = FancyMlpInvDynamic(3, 2, add_norm=True, add_dropout=True) + >>> invdyn.train() + >>> batch = ... + >>> obs, act, obs_next = batch + >>> loss = invdyn.update(obs, act, obs_next) + >>> invdyn.eval() + >>> pred_act = invdyn.predict(obs, obs_next) + """ + def __init__( + self, o_dim: int, a_dim: int, hidden_dim: int = 256, + out_activation: nn.Module = nn.Tanh(), + add_norm: bool = False, add_dropout: bool = False, + optim_params: dict = {}, device: str = "cpu", + ): + self.device = device + self.o_dim, self.a_dim, self.hidden_dim = o_dim, a_dim, hidden_dim + self.out_activation = out_activation + self.optim_params = optim_params + params = {"lr": 3e-4} + params.update(optim_params) + + self.model = nn.Sequential( + nn.Linear(2 * o_dim, hidden_dim), nn.GELU(), + nn.LayerNorm(hidden_dim) if add_norm else nn.Identity(), + nn.Dropout(0.1) if add_dropout else nn.Identity(), + nn.Linear(hidden_dim, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, a_dim), out_activation).to(device) + + self.optim = torch.optim.Adam(self.model.parameters(), **optim_params) + + def forward(self, o, o_next): + return self.model(torch.cat([o, o_next], dim=-1)) + + def update(self, o, a, o_next): + self.optim.zero_grad() + a_pred = self.forward(o, o_next) + loss = ((a_pred - a) ** 2).mean() + loss.backward() + self.optim.step() + return {"loss": loss.item()} + + @torch.no_grad() + def predict(self, o, o_next): + return self.forward(o, o_next) + + def __call__(self, o, o_next): + return self.predict(o, o_next) + + def train(self): + self.model.train() + + def eval(self): + self.model.eval() + + def save(self, path): + torch.save(self.model.state_dict(), path) + + def load(self, path): + self.model.load_state_dict(torch.load(path, self.device)) + + +class EnsembleMlpInvDynamic(MlpInvDynamic): + def __init__( + self, + o_dim: int, + a_dim: int, + hidden_dim: int = 512, + out_activation: nn.Module = nn.Identity(), + optim_params: dict = {}, + n_models=5, + mlp_type="standard", + device: str = "cpu", + ): + assert mlp_type in ["standard", "fancy"] + super().__init__(o_dim, a_dim, hidden_dim, out_activation, optim_params, device) + self.n_models = n_models + if mlp_type == "standard": + self.mlp = nn.ModuleList([Mlp( + 2 * self.o_dim, [self.hidden_dim, self.hidden_dim], self.a_dim, + nn.ReLU(), self.out_activation) for _ in range(n_models)]).to(device) + else: + self.mlp = nn.ModuleList([nn.Sequential( + nn.Linear(2 * self.o_dim, self.hidden_dim), nn.LayerNorm(self.hidden_dim), nn.Mish(), + nn.Dropout(0.1), + nn.Linear(self.hidden_dim, self.hidden_dim), nn.LayerNorm(self.hidden_dim), nn.Mish(), + nn.Linear(self.hidden_dim, self.hidden_dim), self.out_activation) for _ in range(n_models)]).to(device) + self.optim = torch.optim.Adam(self.mlp.parameters(), **self.optim_params) + self._init_weights() + + def forward(self, o, o_next, idx=None): + if idx is None: + return sum([m(torch.cat([o, o_next], dim=-1)) for m in self.mlp]) / self.n_models + else: + return self.mlp[idx](torch.cat([o, o_next], dim=-1)) + + def update_idx(self, idx, o, a, o_next): + self.optim.zero_grad() + a_pred = self.forward(o, o_next, idx) + loss = ((a_pred - a) ** 2).mean() + loss.backward() + self.optim.step() + return loss.item() + + +# =============================== Development ================================= + +class ResidualBlock(nn.Module): + def __init__(self, hidden_dim: int = 256, add_norm: bool = False, add_dropout: bool = False): + super().__init__() + self.norm = nn.LayerNorm(hidden_dim) if add_norm else nn.Identity() + self.mlp = nn.Sequential( + nn.Linear(hidden_dim, hidden_dim), nn.GELU(), + nn.Dropout(0.1) if add_dropout else nn.Identity(), + nn.Linear(hidden_dim, hidden_dim)) + + def forward(self, x): + x = self.norm(x) + return x + self.mlp(x) + + +class ResInvDynamic: + def __init__( + self, o_dim: int, a_dim: int, hidden_dim: int = 256, + out_activation: nn.Module = nn.Tanh(), + add_norm: bool = False, add_dropout: bool = False, + n_blocks: int = 1, + optim_params: dict = {}, device: str = "cpu", + ): + self.device = device + self.n_blocks = n_blocks + self.o_dim, self.a_dim, self.hidden_dim = o_dim, a_dim, hidden_dim + self.out_activation = out_activation + self.optim_params = optim_params + params = {"lr": 3e-4} + params.update(optim_params) + + self.model = nn.ModuleDict({ + "pre_linear": nn.Sequential( + nn.Linear(2 * o_dim, hidden_dim), nn.GELU()).to(device), + "post_linear": nn.Sequential( + nn.Linear(hidden_dim, a_dim), out_activation).to(device)}) + for i in range(n_blocks): + self.model[f"res_block{i}"] = ResidualBlock(hidden_dim, add_norm, add_dropout).to(device) + + # Adam for pre_linear+res_blocks+post_linear + self.optim = torch.optim.Adam(self.model.parameters(), **optim_params) + + def forward(self, o, o_next): + feature = self.model["pre_linear"](torch.cat([o, o_next], dim=-1)) + for i in range(self.n_blocks): + feature = self.model[f"res_block{i}"](feature) + return self.model["post_linear"](feature) + + def update(self, o, a, o_next): + self.optim.zero_grad() + a_pred = self.forward(o, o_next) + loss = ((a_pred - a) ** 2).mean() + loss.backward() + self.optim.step() + return {"loss": loss.item()} + + @torch.no_grad() + def predict(self, o, o_next): + return self.forward(o, o_next) + + def __call__(self, o, o_next): + return self.predict(o, o_next) + + def train(self): + self.model.train() + + def eval(self): + self.model.eval() + + def save(self, path): + torch.save(self.model.state_dict(), path) + + def load(self, path): + self.model.load_state_dict(torch.load(path, self.device)) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dad808cc87ed1031e88a787472b9104e1b47009d --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/__init__.py @@ -0,0 +1,4 @@ +from .base_nn_classifier import BaseNNClassifier +from .half_dit import HalfDiT1d +from .half_jannerunet import HalfJannerUNet1d +from .mlp import MLPNNClassifier, QGPONNClassifier diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/base_nn_classifier.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/base_nn_classifier.py new file mode 100644 index 0000000000000000000000000000000000000000..b2fff6a8e5aa9966d35623c02b8631997286a4b6 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/base_nn_classifier.py @@ -0,0 +1,44 @@ +from typing import Optional + +import torch +import torch.nn as nn + +from cleandiffuser.utils import SUPPORTED_TIMESTEP_EMBEDDING + + +class BaseNNClassifier(nn.Module): + """ + Base Neural Network (NN) for Classifiers. + + This NN is supposed to implement the mapping: (x, t, y) -> log p(y | x, t) + C, where C is a constant. + From a coding perspective, the output of the NN should be a real number with dimension=1. + + Parameters: + - emb_dim: int + Dimension of the embedding for the time variable t. + - timestep_emb_type: str + Type of embedding for the time variable t. Options are: "positional" or "fourier". + """ + + def __init__( + self, emb_dim: int, + timestep_emb_type: str = "positional", + timestep_emb_params: Optional[dict] = None + ): + assert timestep_emb_type in SUPPORTED_TIMESTEP_EMBEDDING.keys() + super().__init__() + timestep_emb_params = timestep_emb_params or {} + self.map_noise = SUPPORTED_TIMESTEP_EMBEDDING[timestep_emb_type](emb_dim, **timestep_emb_params) + + +def forward(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor): + """ + Input: + x: (b, *x_shape) + t: (b, ) + y: (b, *y_shape) + + Output: + logp: (b, 1) + """ + raise NotImplementedError diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/half_dit.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/half_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..502fa53086f6114925ba2d66128d85479b0da9d2 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/half_dit.py @@ -0,0 +1,76 @@ +from typing import Optional + +import torch +import torch.nn as nn + +from cleandiffuser.nn_diffusion.dit import DiT1d, FinalLayer1d + + +class HalfDiT1d(DiT1d): + """ Half DiT1d for diffusion classifier. + + Args: + in_dim: int, + Dimension of the input sequence. + out_dim: int, + Dimension of the output tensor. + emb_dim: int, + Dimension of the condition and time embedding. + d_model: int, + Dimension of the transformer. Default: 384. + n_heads: int, + Number of heads in the transformer. Default: 6. + depth: int, + Number of transformer layers. Default: 12. + dropout: float, + Dropout rate. Default: 0.0. + timestep_emb_type: str, + Type of the timestep embedding. Default: "positional". + + Examples: + >>> nn_classifier = HalfDiT1d(in_dim=10, out_dim=1, emb_dim=64) + >>> x = torch.randn(2, 32, 10) + >>> t = torch.randint(1000, (2,)) + >>> condition = torch.randn(2, 64) + >>> nn_classifier(x, t).shape + torch.Size([2, 1]) + >>> nn_classifier(x, t, condition).shape + torch.Size([2, 1]) + """ + def __init__(self, + in_dim: int, + out_dim: int, + emb_dim: int, + d_model: int = 384, + n_heads: int = 6, + depth: int = 12, + dropout: float = 0.0, + timestep_emb_type: str = "positional", + ): + super().__init__(in_dim, emb_dim, d_model, n_heads, depth, dropout, timestep_emb_type) + self.final_layer = FinalLayer1d(d_model, d_model // 2) + nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0) + nn.init.constant_(self.final_layer.linear.weight, 0) + nn.init.constant_(self.final_layer.linear.bias, 0) + + self.proj = nn.Sequential( + nn.LayerNorm(d_model // 2), nn.SiLU(), + nn.Linear(d_model // 2, d_model // 4), + nn.LayerNorm(d_model // 4), nn.SiLU(), + nn.Linear(d_model // 4, out_dim)) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, horizon, in_dim) + noise: (b, ) + condition: (b, emb_dim) + + Output: + logp(x | noise, condition): (b, 1) + """ + feat = super().forward(x, noise, condition).mean(1) + return self.proj(feat) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/half_jannerunet.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/half_jannerunet.py new file mode 100644 index 0000000000000000000000000000000000000000..d75a0122bd63ec2ac98b081700ed8a21a2a40751 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/half_jannerunet.py @@ -0,0 +1,125 @@ +from typing import Optional, Tuple + +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from cleandiffuser.nn_diffusion.jannerunet import ResidualBlock, Downsample1d + + +class HalfJannerUNet1d(BaseNNDiffusion): + """ Half JannerUNet1d for diffusion classifier. Adapted from https://github.com/jannerm/diffuser. + + Args: + horizon: int, + Length of the input sequence. + in_dim: int, + Number of input dimensions. + out_dim: int, + Number of output dimensions. Default is 1. + kernel_size: int, + Size of the convolution kernel. Default is 3. + model_dim: int, + Initial CNN model dimension. Default is 32. + emb_dim: int, + Number of dimensions in the embedding. Default is 32. + dim_mult: Tuple[int], + UNet dimension multiplier. Default is (1, 2, 2, 2). + timestep_emb_type: str, + Type of the timestep embedding. Default is "positional". + norm_type: str, + Type of the normalization layer. Default is "groupnorm". + + Examples: + >>> nn_classifier = HalfJannerUNet1d(horizon=32, in_dim=10, out_dim=1, emb_dim=64) + >>> x = torch.randn(2, 32, 10) + >>> t = torch.randint(1000, (2,)) + >>> condition = torch.randn(2, 64) + >>> nn_classifier(x, t).shape + torch.Size([2, 1]) + >>> nn_classifier(x, t, condition).shape + torch.Size([2, 1]) + """ + def __init__( + self, + horizon: int, + in_dim: int, + out_dim: int = 1, + kernel_size: int = 3, + model_dim: int = 32, + emb_dim: int = 32, + dim_mult: Tuple[int] = (1, 2, 2, 2), + timestep_emb_type: str = "positional", + norm_type: str = "groupnorm", + ): + super().__init__(emb_dim, timestep_emb_type) + + dims = [in_dim] + [model_dim * m for m in np.cumprod(dim_mult)] + in_out = list(zip(dims[:-1], dims[1:])) + + self.map_emb = nn.Sequential( + nn.Linear(emb_dim, model_dim * 4), nn.Mish(), + nn.Linear(model_dim * 4, model_dim)) + + self.downs = nn.ModuleList([]) + self.ups = nn.ModuleList([]) + num_resolutions = len(in_out) + + for ind, (dim_in, dim_out) in enumerate(in_out): + is_last = ind >= (num_resolutions - 1) + + self.downs.append(nn.ModuleList([ + ResidualBlock(dim_in, dim_out, model_dim, kernel_size, norm_type), + ResidualBlock(dim_out, dim_out, model_dim, kernel_size, norm_type), + Downsample1d(dim_out) if not is_last else nn.Identity() + ])) + + if not is_last: + horizon = horizon // 2 + + mid_dim = dims[-1] + mid_dim_2 = mid_dim // 2 + mid_dim_3 = mid_dim // 4 + + self.mid_block1 = nn.ModuleList([ + ResidualBlock(mid_dim, mid_dim_2, model_dim, kernel_size=5, norm_type=norm_type), + Downsample1d(mid_dim_2)]) + horizon = horizon // 2 + + self.mid_block2 = nn.ModuleList([ + ResidualBlock(mid_dim_2, mid_dim_3, model_dim, kernel_size=5, norm_type=norm_type), + Downsample1d(mid_dim_3)]) + horizon = horizon // 2 + + fc_dim = mid_dim_3 * max(horizon, 1) + + self.final_block = nn.Sequential( + nn.Linear(fc_dim + model_dim, fc_dim // 2), + nn.Mish(), + nn.Linear(fc_dim // 2, out_dim)) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + + x = x.permute(0, 2, 1) + + emb = self.map_noise(noise) + if condition is not None: + emb = emb + condition + emb = self.map_emb(emb) + + for resnet1, resnet2, downsample in self.downs: + x = resnet1(x, emb) + x = resnet2(x, emb) + x = downsample(x) + + x = self.mid_block1[0](x, emb) + x = self.mid_block1[1](x) + x = self.mid_block2[0](x, emb) + x = self.mid_block2[1](x) + + x = x.flatten(1) + out = self.final_block(torch.cat([x, emb], dim=-1)) + return out diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/mlp.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..f93c6e0523b7d65ef58bfc68d1fb97abac7297aa --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_classifier/mlp.py @@ -0,0 +1,55 @@ +from typing import List + +import torch +import torch.nn as nn + +from cleandiffuser.utils import Mlp +from .base_nn_classifier import BaseNNClassifier + + +class MLPNNClassifier(BaseNNClassifier): + def __init__( + self, x_dim: int, out_dim: int, emb_dim: int, + hidden_dims: List[int], activation: nn.Module = nn.ReLU(), out_activation: nn.Module = nn.Identity(), + timestep_emb_type: str = "positional"): + super().__init__(emb_dim, timestep_emb_type) + self.mlp = Mlp( + in_dim=x_dim + emb_dim, hidden_dims=hidden_dims, out_dim=out_dim, + activation=activation, out_activation=out_activation) + + def forward(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor = None): + out = self.mlp(torch.cat([x, self.map_noise(t)], dim=-1)) + return out + + +class QGPONNClassifier(BaseNNClassifier): + """ NN of energy guidance model f_phi in QGPO """ + def __init__( + self, obs_dim: int, act_dim: int, emb_dim: int, + hidden_dims: List[int], + timestep_emb_type: str = "positional"): + super().__init__(emb_dim, timestep_emb_type) + self.obs_proj = nn.Linear(obs_dim, emb_dim) + self.act_proj = nn.Linear(act_dim, emb_dim) + self.mlp = Mlp( + in_dim=3 * emb_dim, hidden_dims=hidden_dims, out_dim=1, activation=nn.SiLU()) + + def forward(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor): + """ + + Args: + x: torch.Tensor, + Noisy actions. Shape (batch_size, act_dim) + t: torch.Tensor, + Timestep. Shape (batch_size, ) + y: torch.Tensor, + Observations. Shape (batch_size, obs_dim) + + Returns: + f: torch.Tensor, + Energy prediction. Shape (batch_size, 1) + """ + y = self.obs_proj(y) + x = self.act_proj(x) + out = self.mlp(torch.cat([y, x, self.map_noise(t)], dim=-1)) + return torch.tanh(out / 10) * 10 diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d22869e18df8accd87a63db3498a6650f8c0eb8b --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/__init__.py @@ -0,0 +1,7 @@ +from .base_nn_condition import * +from .mlp import MLPCondition, MLPSieveObsCondition, LinearCondition +from .positional import FourierCondition, PositionalCondition +from .pearce_obs_condition import PearceObsCondition +from .multi_image_condition import MultiImageObsCondition +from .early_conv_vit import EarlyConvViTMultiViewImageCondition +from .resnets import ResNet18ImageCondition, ResNet18MultiViewImageCondition diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/base_nn_condition.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/base_nn_condition.py new file mode 100644 index 0000000000000000000000000000000000000000..128d7571cefa9f92fac4ba4419cca3ef4e36e466 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/base_nn_condition.py @@ -0,0 +1,57 @@ +import torch +import torch.nn as nn + +from cleandiffuser.utils import at_least_ndim + + +def get_mask(mask: torch.Tensor, mask_shape: tuple, dropout: float, train: bool, device: torch.device): + if train: + mask = (torch.rand(mask_shape, device=device) > dropout).float() + else: + mask = 1. if mask is None else mask + return mask + + +class BaseNNCondition(nn.Module): + """ + In decision-making tasks, generating condition selections can be very diverse, + including cumulative rewards, languages, images, demonstrations, and so on. + It can even be a combination of these conditions. Therefore, we aim for + nn_condition to handle diverse condition selections flexibly and + ultimately output a tensor of shape (b, *cond_out_shape). + + Input: + - condition: Any + - mask : Optional, (b, ) or None, None means no mask + + Output: + - condition: (b, *cond_out_shape) + """ + + def __init__(self, ): + super().__init__() + + def forward(self, condition: torch.Tensor, mask: torch.Tensor = None): + raise NotImplementedError + + +class IdentityCondition(BaseNNCondition): + """ + Identity condition does not change the input condition. + + Input: + - condition: (b, *cond_in_shape) + - mask : (b, ) or None, None means no mask + + Output: + - condition: (b, *cond_in_shape) + """ + + def __init__(self, dropout: float = 0.25): + super().__init__() + self.dropout = dropout + + def forward(self, condition: torch.Tensor, mask: torch.Tensor = None): + mask = at_least_ndim(get_mask( + mask, (condition.shape[0],), self.dropout, self.training, condition.device), condition.dim()) + return condition * mask diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/early_conv_vit.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/early_conv_vit.py new file mode 100644 index 0000000000000000000000000000000000000000..43b8ad8feabc7227bb4d93b0a6dd7372b942f932 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/early_conv_vit.py @@ -0,0 +1,197 @@ +from typing import Tuple, Optional, Dict + +import einops +import torch +import torch.nn as nn + +from cleandiffuser.nn_condition import BaseNNCondition +from cleandiffuser.utils import SinusoidalEmbedding, Transformer + + +class SmallStem(nn.Module): + def __init__( + self, patch_size: int = 16, + in_channels: int = 3, + channels_per_group: int = 16, + kernel_sizes: tuple = (3, 3, 3, 3), + strides: tuple = (2, 2, 2, 2), + features: tuple = (32, 64, 128, 256), + padding: tuple = (1, 1, 1, 1), + num_features: int = 256, + ): + super().__init__() + + self.patch_size = patch_size + + cnn = [] + for n, (k, s, f, p) in enumerate(zip(kernel_sizes, strides, features, padding)): + cnn.append( + nn.Sequential( + nn.Conv2d( + in_channels if n == 0 else features[n - 1], f, k, s, p), + nn.GroupNorm(f // channels_per_group, f), + nn.ReLU())) + self.cnn = nn.Sequential(*cnn) + + self.patchify = nn.Conv2d( + features[-1], num_features, + kernel_size=patch_size // 16, + stride=patch_size // 16, padding=0) + + def forward(self, x): + x = self.cnn(x) + x = self.patchify(x) + return einops.rearrange(x, "b c h w -> b (h w) c") + + +class EarlyConvViTMultiViewImageCondition(BaseNNCondition): + """ Early-CNN Vision Transformer (ViT) for multi-view image condition. + + A ViT model that uses a shallow CNN instead of a patchify layer to extract image tokens. + This architecture is proposed in https://arxiv.org/pdf/2106.14881 and demonstrated to be + effective for CV tasks. The vision encoder in Octo (https://arxiv.org/pdf/2405.12213) is + mainly based on this architecture. + Each view is processed by a separate CNN and the resulting tokens are concatenated along + the token dimension, which are then processed by a transformer. The output of the + learnable 'readout' token is returned as the final representation. + + Args: + image_sz: Tuple[int], + The size of the input image for each view. Assumes square images. + in_channels: Tuple[int], + The number of input channels for each view. + lowdim_sz: Optional[int], + The size of the low-dimensional condition. If None, no low-dimensional condition is used. + To: int, + The number of frames for each view. + + # Transformer arguments + d_model: int: + The dimension of the transformer token. + nhead: int: + The number of heads in the transformer. + num_layers: int: + The number of transformer layers. + attn_dropout: float: + The dropout rate for the attention layer. + ffn_dropout: float: + The dropout rate for the feedforward layer. + + # CNN arguments + patch_size: Tuple[int]: + The size of the patch for each view. + channels_per_group: Tuple[int]: + The number of channels per group in the CNN. + kernel_sizes: Tuple[Tuple[int]]: + The kernel sizes for each CNN layer. + strides: Tuple[Tuple[int]]: + The strides for each CNN layer. + features: Tuple[Tuple[int]]: + The number of features for each CNN layer. + padding: Tuple[Tuple[int]]: + The padding for each CNN layer. + + Examples: + >>> d_model = 384 + >>> batch, view, To, C, H, W, D = 4, 2, 1, 3, 64, 64, 9 + >>> nn_condition = EarlyConvViTMultiViewImageCondition(d_model=d_model, ...) + >>> condition = { + ... "image": torch.randn((batch, view, To, C, H, W)), + ... "lowdim": torch.randn((batch, To, D)),} + >>> nn_condition(condition).shape + torch.Size([batch, d_model]) + """ + def __init__( + self, + image_sz: Tuple[int] = (64, 64), + in_channels: Tuple[int] = (3, 3), + lowdim_sz: Optional[int] = None, + To: int = 1, + + # Transformer parameters + d_model: int = 384, + nhead: int = 6, + num_layers: int = 2, + attn_dropout: float = 0., + ffn_dropout: float = 0., + + # CNN parameters + patch_size: Tuple[int] = (16, 16), + channels_per_group: Tuple[int] = (16, 16), + kernel_sizes: Tuple[Tuple[int]] = ((3, 3, 3, 3), (3, 3, 3, 3)), + strides: Tuple[Tuple[int]] = ((2, 2, 2, 2), (2, 2, 2, 2)), + features: Tuple[Tuple[int]] = ((32, 64, 128, 256), (32, 64, 128, 256)), + padding: Tuple[Tuple[int]] = ((1, 1, 1, 1), (1, 1, 1, 1)), + ): + super().__init__() + + self.image_sz, self.in_channels = image_sz, in_channels + self.n_views = len(image_sz) + + self.patchifies = nn.ModuleList([ + SmallStem( + patch_size=patch_size[i], + in_channels=in_channels[i], + channels_per_group=channels_per_group[i], + kernel_sizes=kernel_sizes[i], + strides=strides[i], + features=features[i], + padding=padding[i], + num_features=d_model) for i in range(self.n_views)]) + + self.pos_emb = nn.ParameterList([ + nn.Parameter( + SinusoidalEmbedding(d_model)(torch.arange(To * self.image_token_lens[i]))[None, :], + requires_grad=False) + for i in range(self.n_views)]) + + self.view_emb = nn.ParameterList([ + nn.Parameter( + torch.zeros((1, 1, d_model)), requires_grad=True) + for _ in range(self.n_views)]) + + self.lowdim_proj = nn.Linear(lowdim_sz, d_model) if lowdim_sz is not None else None + self.lowdim_emb = nn.Parameter( + torch.zeros((1, 1, d_model)), requires_grad=True) if lowdim_sz is not None else None + + self.readout_emb = nn.Parameter( + torch.zeros((1, 1, d_model)), requires_grad=True) + + self.tfm = Transformer( + d_model, nhead, num_layers, 4, attn_dropout, ffn_dropout) + + self.mask_cache = None + + @property + def image_token_lens(self): + examples = [ + torch.randn((1, self.in_channels[i], self.image_sz[i], self.image_sz[i])) + for i in range(self.n_views)] + return [self.patchifies[i](examples[i]).shape[1] for i in range(self.n_views)] + + def forward(self, condition: Dict[str, torch.Tensor], mask: Optional[torch.Tensor] = None): + + b, v, t, c, h, w = condition["image"].shape + + tokens = [] + + if self.lowdim_proj is not None: + tokens.append( + self.lowdim_proj(condition["lowdim"]) + self.lowdim_emb) + + for i in range(v): + view_tokens = self.patchifies[i]( + einops.rearrange(condition["image"][:, i], "b t c h w -> (b t) c h w")) + view_tokens = (einops.rearrange(view_tokens, "(b t) n d -> b (t n) d", b=b) + + self.view_emb[i] + self.pos_emb[i]) + tokens.append(view_tokens) + + tokens.append(self.readout_emb.repeat(b, 1, 1)) + + tokens = torch.cat(tokens, dim=1) + + if self.mask_cache is None or tokens.shape[1] != self.mask_cache.shape[1]: + self.mask_cache = torch.tril( + torch.ones(tokens.shape[1], tokens.shape[1], device=condition["image"].device), diagonal=0) + + return self.tfm(tokens, self.mask_cache)[0][:, -1] diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/mlp.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..0eba46cf28d41f9a180d49cc16aab3bdc0913a54 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/mlp.py @@ -0,0 +1,92 @@ +from typing import List + +import torch +import torch.nn as nn + +from cleandiffuser.utils import at_least_ndim, Mlp +from cleandiffuser.nn_condition import IdentityCondition, get_mask + + +class LinearCondition(IdentityCondition): + """ A sample affine condition + + Use a linear layer to project the input condition to the desired dimension. + + Args: + in_dim: int, + The input dimension of the condition + out_dim: int, + The output dimension of the condition + dropout: float, + The label dropout rate, Default: 0.25 + + Examples: + >>> nn_condition = LinearCondition(in_dim=5, out_dim=10) + >>> condition = torch.randn(2, 5) + >>> nn_condition(condition).shape + torch.Size([2, 10]) + >>> condition = torch.randn(2, 20, 5) + >>> nn_condition(condition).shape + torch.Size([2, 20, 10]) + """ + def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.25): + super().__init__(dropout) + self.affine = nn.Linear(in_dim, out_dim) + + def forward(self, condition: torch.Tensor, mask: torch.Tensor = None): + mask = at_least_ndim(get_mask( + mask, (condition.shape[0],), self.dropout, self.training, condition.device), condition.dim()) + return self.affine(condition) * mask + + +class MLPCondition(IdentityCondition): + """ A simple MLP condition + + Use a simple MLP to project the input condition to the desired dimension. + + Args: + in_dim: int, + The input dimension of the condition + out_dim: int, + The output dimension of the condition + hidden_dims: List[int], + The hidden dimensions of the MLP + act: nn.Module, + The activation function of the MLP + dropout: float, + The label dropout rate + + Examples: + >>> nn_condition = MLPCondition(in_dim=5, out_dim=10) + >>> condition = torch.randn(2, 5) + >>> nn_condition(condition).shape + torch.Size([2, 10]) + >>> condition = torch.randn(2, 20, 5) + >>> nn_condition(condition).shape + torch.Size([2, 20, 10]) + """ + def __init__(self, in_dim: int, out_dim: int, hidden_dims: List[int], + act=nn.LeakyReLU(), dropout: float = 0.25): + super().__init__(dropout) + hidden_dims = [hidden_dims, ] if isinstance(hidden_dims, int) else hidden_dims + self.mlp = Mlp( + in_dim, hidden_dims, out_dim, act) + + def forward(self, condition: torch.Tensor, mask: torch.Tensor = None): + mask = at_least_ndim(get_mask( + mask, (condition.shape[0],), self.dropout, self.training, condition.device), condition.dim()) + return self.mlp(condition) * mask + + +class MLPSieveObsCondition(IdentityCondition): + def __init__(self, o_dim: int, emb_dim: int = 128, hidden_dim: int = 512, dropout: float = 0.25): + super().__init__(dropout) + hidden_dims = [hidden_dim, ] + self.mlp = Mlp( + o_dim, hidden_dims, emb_dim, nn.LeakyReLU()) + + def forward(self, obs: torch.Tensor, mask: torch.Tensor = None): + mask = at_least_ndim(get_mask( + mask, (obs.shape[0],), self.dropout, self.training, obs.device), 2) + embs = self.mlp(obs) # (b, history_len, emb_dim) + return torch.flatten(embs, 1) * mask diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/multi_image_condition.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/multi_image_condition.py new file mode 100644 index 0000000000000000000000000000000000000000..ea479b4021765e0584d79ff01101e237559e5a44 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/multi_image_condition.py @@ -0,0 +1,327 @@ +from typing import Dict, Tuple, Union, Callable +import copy +import torch +import torch.nn as nn +import torchvision + +from cleandiffuser.utils.crop_randomizer import CropRandomizer +from cleandiffuser.nn_condition import BaseNNCondition + + +def replace_submodules( + root_module: nn.Module, + predicate: Callable[[nn.Module], bool], + func: Callable[[nn.Module], nn.Module]) -> nn.Module: + """ + predicate: Return true if the module is to be replaced. + func: Return new module to use. + """ + if predicate(root_module): + return func(root_module) + + bn_list = [k.split('.') for k, m + in root_module.named_modules(remove_duplicate=True) + if predicate(m)] + for *parent, k in bn_list: + parent_module = root_module + if len(parent) > 0: + parent_module = root_module.get_submodule('.'.join(parent)) + if isinstance(parent_module, nn.Sequential): + src_module = parent_module[int(k)] + else: + src_module = getattr(parent_module, k) + tgt_module = func(src_module) + if isinstance(parent_module, nn.Sequential): + parent_module[int(k)] = tgt_module + else: + setattr(parent_module, k, tgt_module) + # verify that all BN are replaced + bn_list = [k.split('.') for k, m + in root_module.named_modules(remove_duplicate=True) + if predicate(m)] + assert len(bn_list) == 0 + return root_module + + +def get_resnet(name, weights=None, **kwargs): + """ + name: resnet18, resnet34, resnet50 + weights: "IMAGENET1K_V1", "r3m" + """ + func = getattr(torchvision.models, name) + resnet = func(weights=weights, **kwargs) + resnet.fc = torch.nn.Identity() + return resnet + + +class MultiImageObsCondition(BaseNNCondition): + + """ + Input: + - condition: {"cond1": (b, *cond1_shape), "cond2": (b, *cond2_shape), ...} or (b, *cond_in_shape) + - mask : (b, *mask_shape) or None, None means no mask + + Output: + - condition: (b, *cond_out_shape) + + Assumes rgb input: B, C, H, W or B, seq_len, C,H,W + Assumes low_dim input: B, D or B, seq_len, D + """ + def __init__(self, + shape_meta: dict, + rgb_model_name: str, + emb_dim: int = 256, + resize_shape: Union[Tuple[int,int], Dict[str,tuple], None]=None, + crop_shape: Union[Tuple[int,int], Dict[str,tuple], None]=None, + random_crop: bool=True, + # replace BatchNorm with GroupNorm + use_group_norm: bool=False, + # use single rgb model for all rgb inputs + share_rgb_model: bool=False, + # renormalize rgb input with imagenet normalization + # assuming input in [0,1] + imagenet_norm: bool=False, + # use_seq: B, seq_len, C, H, W or B, C, H, W + use_seq=False, + # if True: (bs, seq_len, embed_dim) + keep_horizon_dims=False + ): + super().__init__() + rgb_keys = list() + low_dim_keys = list() + key_model_map = nn.ModuleDict() + key_transform_map = nn.ModuleDict() + key_shape_map = dict() + + # rgb_model + if 'resnet' in rgb_model_name: + rgb_model = get_resnet(rgb_model_name) + else: + raise ValueError("Fatal rgb_model") + + # handle sharing vision backbone + if share_rgb_model: + assert isinstance(rgb_model, nn.Module) + key_model_map['rgb'] = rgb_model + + obs_shape_meta = shape_meta['obs'] + for key, attr in obs_shape_meta.items(): + # print(key, attr) + shape = tuple(attr['shape']) + type = attr.get('type', 'low_dim') + key_shape_map[key] = shape + if type == 'rgb': + rgb_keys.append(key) + # configure model for this key + this_model = None + if not share_rgb_model: + if isinstance(rgb_model, dict): + # have provided model for each key + this_model = rgb_model[key] + else: + assert isinstance(rgb_model, nn.Module) + # have a copy of the rgb model + this_model = copy.deepcopy(rgb_model) + + if this_model is not None: + if use_group_norm: + this_model = replace_submodules( + root_module=this_model, + predicate=lambda x: isinstance(x, nn.BatchNorm2d), + func=lambda x: nn.GroupNorm( + num_groups=x.num_features//16, + num_channels=x.num_features) + ) + key_model_map[key] = this_model + + # configure resize + input_shape = shape + this_resizer = nn.Identity() + if resize_shape is not None: + if isinstance(resize_shape, dict): + h, w = resize_shape[key] + else: + h, w = resize_shape + this_resizer = torchvision.transforms.Resize( + size=(h, w) + ) + input_shape = (shape[0], h, w) + + # configure randomizer + this_randomizer = nn.Identity() + if crop_shape is not None: + if isinstance(crop_shape, dict): + h, w = crop_shape[key] + else: + h, w = crop_shape + if random_crop: + this_randomizer = CropRandomizer( + input_shape=input_shape, + crop_height=h, + crop_width=w, + num_crops=1, + pos_enc=False + ) + else: + this_randomizer = torchvision.transforms.CenterCrop( + size=(h,w) + ) + # configure normalizer + this_normalizer = nn.Identity() + if imagenet_norm: + this_normalizer = torchvision.transforms.Normalize( + mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + + this_transform = nn.Sequential(this_resizer, this_randomizer, this_normalizer) + key_transform_map[key] = this_transform + elif type == 'low_dim': + low_dim_keys.append(key) + else: + raise RuntimeError(f"Unsupported obs type: {type}") + rgb_keys = sorted(rgb_keys) + low_dim_keys = sorted(low_dim_keys) + + self.shape_meta = shape_meta + self.key_model_map = key_model_map + self.key_transform_map = key_transform_map + self.share_rgb_model = share_rgb_model + self.rgb_keys = rgb_keys + self.low_dim_keys = low_dim_keys + self.key_shape_map = key_shape_map + + self.use_seq = use_seq + self.keep_horizon_dims = keep_horizon_dims + self.mlp = nn.Sequential( + nn.Linear(self.output_shape(), emb_dim), nn.LeakyReLU(), nn.Linear(emb_dim, emb_dim)) + + def multi_image_forward(self, obs_dict): + batch_size = None + features = list() + + if self.use_seq: + # input: (bs, horizon, c, h, w) + for k in obs_dict.keys(): + obs_dict[k] = obs_dict[k].flatten(end_dim=1) + + # process rgb input + if self.share_rgb_model: + # pass all rgb obs to rgb model + imgs = list() + for key in self.rgb_keys: + img = obs_dict[key] + if batch_size is None: + batch_size = img.shape[0] + else: + assert batch_size == img.shape[0] + assert img.shape[1:] == self.key_shape_map[key] + img = self.key_transform_map[key](img) + imgs.append(img) + # (N*B,C,H,W) + imgs = torch.cat(imgs, dim=0) + # (N*B,D) + feature = self.key_model_map['rgb'](imgs) + # (N,B,D) + feature = feature.reshape(-1,batch_size,*feature.shape[1:]) + # (B,N,D) + feature = torch.moveaxis(feature,0,1) + # (B,N*D) + feature = feature.reshape(batch_size,-1) + features.append(feature) + else: + # run each rgb obs to independent models + for key in self.rgb_keys: + img = obs_dict[key] + if batch_size is None: + batch_size = img.shape[0] + else: + assert batch_size == img.shape[0] + assert img.shape[1:] == self.key_shape_map[key] + img = self.key_transform_map[key](img) + feature = self.key_model_map[key](img) + features.append(feature) + + # process lowdim input + for key in self.low_dim_keys: + data = obs_dict[key] + if batch_size is None: + batch_size = data.shape[0] + else: + assert batch_size == data.shape[0] + assert data.shape[1:] == self.key_shape_map[key] + features.append(data) + + # concatenate all features + features = torch.cat(features, dim=-1) + return features + + def forward(self, obs_dict, mask=None): + ori_batch_size, ori_seq_len = self.get_batch_size(obs_dict) + features = self.multi_image_forward(obs_dict) + # linear embedding + result = self.mlp(features) + if self.use_seq: + if self.keep_horizon_dims: + result = result.reshape(ori_batch_size, ori_seq_len, -1) + else: + result = result.reshape(ori_batch_size, -1) + return result + + @torch.no_grad() + def output_shape(self): + example_obs_dict = dict() + obs_shape_meta = self.shape_meta['obs'] + batch_size = 1 + for key, attr in obs_shape_meta.items(): + shape = tuple(attr['shape']) + if self.use_seq: + prefix = (batch_size, 1) + else: + prefix = (batch_size,) + this_obs = torch.zeros( + prefix + shape, + dtype=self.dtype, + device=self.device) + example_obs_dict[key] = this_obs + example_output = self.multi_image_forward(example_obs_dict) + output_shape = example_output.shape[1:] + return output_shape[0] + + def get_batch_size(self, obs_dict): + any_key = next(iter(obs_dict)) + any_tensor = obs_dict[any_key] + return any_tensor.size(0), any_tensor.size(1) + + @property + def device(self): + return next(iter(self.parameters())).device + + @property + def dtype(self): + return next(iter(self.parameters())).dtype + + +# if __name__ == "__main__": +# shape_meta = { +# 'obs': { +# 'image': { +# 'shape': (3, 96, 96), +# 'type': 'rgb' +# }, +# 'agent_pos': { +# 'shape': (2, ), +# 'type': 'low_dim' +# } +# } +# } +# rgb_model = "resnet18" + +# resize_shape=None +# crop_shape=(84, 84) +# random_crop=True +# use_group_norm=True +# share_rgb_model=False +# imagenet_norm=False +# im = MultiImageObsCondition(shape_meta, rgb_model=rgb_model, resize_shape=resize_shape, crop_shape=crop_shape, +# random_crop=random_crop, use_group_norm=use_group_norm, share_rgb_model=share_rgb_model, +# imagenet_norm=imagenet_norm) +# im.output_shape() \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/pearce_obs_condition.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/pearce_obs_condition.py new file mode 100644 index 0000000000000000000000000000000000000000..0e53893efb63a5ee8853dca0aa1f1d95a3248a13 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/pearce_obs_condition.py @@ -0,0 +1,50 @@ +from typing import Optional + +import torch +import torch.nn as nn + +from cleandiffuser.nn_condition import IdentityCondition, get_mask +from cleandiffuser.utils import at_least_ndim + + +class PearceObsCondition(IdentityCondition): + """ Observation condition from DiffusionBC: https://arxiv.org/abs/2301.10677 + + The model suggests using multi-frame observations as conditions. + It encodes each frame of observation using the same MLP, + then flattens them to create a condition embedding. + + Args: + obs_dim: int, + The dimension of the observation. Suppose the observation has shape (b, To, obs_dim), + where b is the batch size, To is the number of frames, and obs_dim is the dimension of each frame. + emb_dim: int, + The dimension of the condition embedding. Default: 128 + flatten: bool, + Whether to flatten the condition embedding. Default: False + dropout: float, + The label dropout rate. Default: 0.25 + + Examples: + >>> nn_condition = PearceObsCondition(obs_dim=3, emb_dim=128, flatten=False) + >>> obs = torch.randn(2, 10, 3) + >>> nn_condition(obs).shape + torch.Size([2, 10, 128]) + >>> nn_condition = PearceObsCondition(obs_dim=3, emb_dim=128, flatten=True) + >>> obs = torch.randn(2, 10, 3) + >>> nn_condition(obs).shape + torch.Size([2, 1280]) + """ + def __init__(self, obs_dim: int, emb_dim: int = 128, flatten: bool = False, dropout: float = 0.25): + super().__init__(dropout) + self.mlp = nn.Sequential( + nn.Linear(obs_dim, emb_dim), nn.LeakyReLU(), nn.Linear(emb_dim, emb_dim)) + self.flatten = flatten + + def forward(self, obs: torch.Tensor, mask: Optional[torch.Tensor] = None): + mask = at_least_ndim(get_mask( + mask, (obs.shape[0],), self.dropout, self.training, obs.device), + 2 if self.flatten else 3) + embs = self.mlp(obs) # (b, To, emb_dim) + embs = torch.flatten(embs, 1) if self.flatten else embs + return embs * mask diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/positional.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/positional.py new file mode 100644 index 0000000000000000000000000000000000000000..a4ce35680d8e75eb56d226e54f6649bf7d740bb3 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/positional.py @@ -0,0 +1,54 @@ +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.nn_condition import MLPCondition + + +class FourierCondition(MLPCondition): + """ + Fourier condition is a simple fourier basis to process the input condition. + + Input: + - condition: (b, 1) + - mask : (b, ) or None, None means no mask + + Output: + - condition: (b, out_dim) + """ + + def __init__(self, out_dim, hidden_dim, scale=16, dropout=0.25): + super().__init__(hidden_dim, out_dim, hidden_dim, nn.Mish(), dropout) + self.register_buffer('freqs', torch.randn(hidden_dim // 2) * scale) + + def forward(self, condition: torch.Tensor, mask: torch.Tensor = None): + emb = condition.squeeze(-1).ger((2 * np.pi * self.freqs).to(condition.dtype)) + emb = torch.cat([emb.cos(), emb.sin()], -1) + return super().forward(emb, mask) + + +class PositionalCondition(MLPCondition): + """ + Positional condition is a simple positional encoding to process the input condition. + + Input: + - condition: (b, 1) + - mask : (b, ) or None, None means no mask + + Output: + - condition: (b, out_dim) + """ + + def __init__(self, out_dim, hidden_dim, dropout=0.25, max_positions: int = 10000, endpoint: bool = False): + super().__init__(hidden_dim, out_dim, hidden_dim, nn.Mish(), dropout) + self.max_positions = max_positions + self.endpoint = endpoint + self.dim = out_dim + + def forward(self, condition: torch.Tensor, mask: torch.Tensor = None): + freqs = torch.arange(start=0, end=self.dim // 2, dtype=torch.float32, device=condition.device) + freqs = freqs / (self.dim // 2 - (1 if self.endpoint else 0)) + freqs = (1 / self.max_positions) ** freqs + x = condition.ger(freqs.to(condition.dtype)) + x = torch.cat([x.cos(), x.sin()], dim=1) + return super().forward(x, mask) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/resnets.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/resnets.py new file mode 100644 index 0000000000000000000000000000000000000000..9bc445e77fa9b11bb506bcd9c2f519a34ad8f160 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_condition/resnets.py @@ -0,0 +1,295 @@ +import einops +import numpy as np +import torch +import torch.nn as nn + +from .base_nn_condition import BaseNNCondition, get_mask, at_least_ndim + + +def get_norm(channel: int, use_group_norm: bool = True, group_channels: int = 16): + if use_group_norm: + return nn.GroupNorm(channel // group_channels, channel, affine=True) + else: + return nn.BatchNorm2d(channel, affine=True) + + +def get_image_coordinates(h, w, normalise): + x_range = torch.arange(w, dtype=torch.float32) + y_range = torch.arange(h, dtype=torch.float32) + if normalise: + x_range = (x_range / (w - 1)) * 2 - 1 + y_range = (y_range / (h - 1)) * 2 - 1 + image_x = x_range.unsqueeze(0).repeat_interleave(h, 0) + image_y = y_range.unsqueeze(0).repeat_interleave(w, 0).t() + return image_x, image_y + + +class ResidualBlock(nn.Module): + def __init__( + self, in_channel: int, out_channel: int, downsample: bool = False, + use_group_norm: bool = True, group_channels: int = 16, + activation: nn.Module = nn.ReLU(), + ): + super().__init__() + + stride = 2 if downsample else 1 + + self.cnn = nn.Sequential( + nn.Conv2d(in_channel, out_channel, 3, stride, 1, bias=False), + get_norm(out_channel, use_group_norm, group_channels), activation, + nn.Conv2d(out_channel, out_channel, 3, 1, 1, bias=False), + get_norm(out_channel, use_group_norm, group_channels)) + + self.skip = nn.Sequential( + nn.Conv2d(in_channel, out_channel, 1, stride, bias=False), + get_norm(out_channel, use_group_norm, group_channels)) \ + if downsample else nn.Identity() + + def forward(self, x): + return self.cnn(x) + self.skip(x) + + +class SpatialSoftmax(nn.Module): + """ + Spatial Softmax layer as described in [1]. + [1]: "Deep Spatial Autoencoders for Visuomotor Learning" + Chelsea Finn, Xin Yu Tan, Yan Duan, Trevor Darrell, Sergey Levine, Pieter Abbeel + + Adapted from https://github.com/gorosgobe/dsae-torch/blob/master/dsae.py + """ + + def __init__(self, temperature=None, normalise=True): + super().__init__() + self.temperature = nn.Parameter(torch.ones(1)) if temperature is None else torch.tensor([temperature]) + self.normalise = normalise + + def forward(self, x): + n, c, h, w = x.size() + spatial_softmax_per_map = nn.functional.softmax(x.view(n * c, h * w) / self.temperature, dim=1) + spatial_softmax = spatial_softmax_per_map.view(n, c, h, w) + + # calculate image coordinate maps + image_x, image_y = get_image_coordinates(h, w, normalise=self.normalise) + # size (H, W, 2) + image_coordinates = torch.cat((image_x.unsqueeze(-1), image_y.unsqueeze(-1)), dim=-1) + # send to device + image_coordinates = image_coordinates.to(device=x.device) + + # multiply coordinates by the softmax and sum over height and width, like in [2] + expanded_spatial_softmax = spatial_softmax.unsqueeze(-1) + image_coordinates = image_coordinates.unsqueeze(0) + out = torch.sum(expanded_spatial_softmax * image_coordinates, dim=[2, 3]) + # (N, C, 2) + return out + + +class ResNet18(nn.Module): + def __init__( + self, image_sz: int, in_channel: int, emb_dim: int, act_fn=lambda: nn.ReLU(), + use_group_norm: bool = True, group_channels: int = 16, + use_spatial_softmax: bool = True): + super().__init__() + + self.image_sz, self.in_channel = image_sz, in_channel + + self.cnn = nn.Sequential( + + # Initial convolution + nn.Conv2d(in_channel, 64, 7, 2, 3, bias=False), + get_norm(64, use_group_norm, group_channels), + act_fn(), nn.MaxPool2d(3, 2, 1), + + # Residual blocks + ResidualBlock(64, 64, False, + use_group_norm, group_channels, act_fn()), + ResidualBlock(64, 64, False, + use_group_norm, group_channels, act_fn()), + + ResidualBlock(64, 128, True, + use_group_norm, group_channels, act_fn()), + ResidualBlock(128, 128, False, + use_group_norm, group_channels, act_fn()), + + ResidualBlock(128, 256, True, + use_group_norm, group_channels, act_fn()), + ResidualBlock(256, 256, False, + use_group_norm, group_channels, act_fn()), + + ResidualBlock(256, 512, True, + use_group_norm, group_channels, act_fn()), + ResidualBlock(512, 512, False, + use_group_norm, group_channels, act_fn()), + + # Final pooling + nn.AvgPool2d(7, 1, 0) if not use_spatial_softmax else + SpatialSoftmax(None, True) + ) + + self.mlp = nn.Sequential( + nn.Linear(np.prod(self.cnn_output_shape), emb_dim), nn.SiLU(), + nn.Linear(emb_dim, emb_dim)) + + @property + def device(self): + return next(iter(self.parameters())).device + + @property + def dtype(self): + return next(iter(self.parameters())).dtype + + @property + def cnn_output_shape(self): + example = torch.zeros((1, self.in_channel, self.image_sz, self.image_sz), + device=self.device, dtype=self.dtype) + return self.cnn(example).shape + + def forward(self, x): + feat = self.cnn(x) + return self.mlp(torch.flatten(feat, 1)) + + +class ResNet18ImageCondition(BaseNNCondition): + """ ResNet18 for image condition. + + ResNet18Condition encodes the input image into a fixed-size embedding. + The implementation is adapted from `DiffusionPolicy`. + Compared to the original implementation, we replace `BatchNorm2d` with `GroupNorm`, + and use a SpatialSoftmax instead of an average pooling layer. + + Args: + image_sz: int, + Size of the input image. The image is assumed to be square. + in_channel: int, + Number of input channels. 3 for RGB images. + emb_dim: int, + Dimension of the output embedding. + + act_fn: callable, + Activation function to use in the network. Default is ReLU. + use_group_norm: bool, + Whether to use GroupNorm instead of BatchNorm. Default is True. + group_channels: int, + Number of channels per group in GroupNorm. Default is 16. + use_spatial_softmax: bool, + Whether to use SpatialSoftmax instead of average pooling. Default is True. + + dropout: float, + Condition Dropout rate. Default is 0.0. + + Examples: + >>> nn_condition = ResNet18ImageCondition(image_sz=64, in_channel=3, emb_dim=256) + >>> condition = torch.randn(32, 3, 64, 64) + >>> nn_condition(condition).shape + torch.Size([32, 256]) + >>> condition = torch.randn(32, 4, 3, 64, 64) + >>> nn_condition(condition).shape + torch.Size([32, 4, 256]) + """ + def __init__( + self, image_sz: int, in_channel: int, emb_dim: int, act_fn=lambda: nn.ReLU(), + use_group_norm: bool = True, group_channels: int = 16, + use_spatial_softmax: bool = True, dropout: float = 0.0, + ): + super().__init__() + self.dropout = dropout + + self.resnet18 = ResNet18( + image_sz, in_channel, emb_dim, act_fn, use_group_norm, group_channels, use_spatial_softmax) + + def forward(self, condition: torch.Tensor, mask: torch.Tensor = None): + + condition_dim = condition.dim() + assert condition_dim == 4 or condition_dim == 5, \ + "Input condition must be 4D or 5D tensor, got shape {}".format(condition.shape) + + b = condition.shape[0] + if condition_dim == 5: + condition = einops.rearrange(condition, 'b n c h w -> (b n) c h w') + + mask = get_mask( + mask, (b,), self.dropout, self.training, condition.device) + + emb = self.resnet18(condition) + + if condition_dim == 5: + emb = einops.rearrange(emb, '(b n) d -> b n d', b=b) * mask[:, None, None] + + return emb + + +class ResNet18MultiViewImageCondition(BaseNNCondition): + """ ResNet18 for multi-view image condition. + + ResNet18Condition encodes the input image into a fixed-size embedding. + The implementation is adapted from `DiffusionPolicy`. + Compared to the original implementation, we replace `BatchNorm2d` with `GroupNorm`, + and use a SpatialSoftmax instead of an average pooling layer. + The Multi-view version uses different ResNet18 networks for each view. + + Args: + image_sz: int, + Size of the input image. The image is assumed to be square. + in_channel: int, + Number of input channels. 3 for RGB images. + emb_dim: int, + Dimension of the output embedding. + n_views: int, + Number of views. + + act_fn: callable, + Activation function to use in the network. Default is ReLU. + use_group_norm: bool, + Whether to use GroupNorm instead of BatchNorm. Default is True. + group_channels: int, + Number of channels per group in GroupNorm. Default is 16. + use_spatial_softmax: bool, + Whether to use SpatialSoftmax instead of average pooling. Default is True. + + dropout: float, + Condition Dropout rate. Default is 0.0. + + Examples: + >>> nn_condition = ResNet18MultiViewImageCondition(image_sz=64, in_channel=3, emb_dim=256, n_views=2) + >>> condition = torch.randn(32, 2, 3, 64, 64) + >>> nn_condition(condition).shape + torch.Size([32, 2, 256]) + >>> condition = torch.randn(32, 2, 4, 3, 64, 64) + >>> nn_condition(condition).shape + torch.Size([32, 2, 4, 256]) + """ + def __init__( + self, image_sz: int, in_channel: int, emb_dim: int, n_views: int, act_fn=lambda: nn.ReLU(), + use_group_norm: bool = True, group_channels: int = 16, + use_spatial_softmax: bool = True, dropout: float = 0.0, + ): + super().__init__() + self.dropout = dropout + self.n_views = n_views + + self.resnet18 = nn.ModuleList([ + ResNet18( + image_sz, in_channel, emb_dim, act_fn, use_group_norm, group_channels, use_spatial_softmax) + for _ in range(n_views) + ]) + + def forward(self, condition: torch.Tensor, mask: torch.Tensor = None): + + condition_dim = condition.dim() + assert condition_dim == 5 or condition_dim == 6, \ + "Input condition must be 5D or 6D tensor, got shape {}".format(condition.shape) + + b = condition.shape[0] + if condition.dim() == 6: + condition = einops.rearrange(condition, 'b v n c h w -> (b n) v c h w') + + mask = get_mask( + mask, (b,), self.dropout, self.training, condition.device) + + emb = [self.resnet18[i](condition[:, i]) for i in range(self.n_views)] + + emb = torch.stack(emb, dim=1) + + if condition_dim == 6: + emb = einops.rearrange(emb, '(b n) v d -> b v n d', b=b) * at_least_ndim(mask, 4) + + return emb diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3366e4176f12782e8cbbb078b6c9f9e3a9bb4a4f --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/__init__.py @@ -0,0 +1,12 @@ +from .base_nn_diffusion import * +from .chitransformer import ChiTransformer +from .chiunet import ChiUNet1d +from .dit import DiT1d +from .dqlmlp import DQLMlp +from .idqlmlp import IDQLMlp +from .dvinvmlp import DVInvMlp +from .jannerunet import JannerUNet1d +from .pearcemlp import PearceMlp +from .pearcetransformer import PearceTransformer +from .sfbc_unet import SfBCUNet +from .mlps import MlpNNDiffusion diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/base_nn_diffusion.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/base_nn_diffusion.py new file mode 100644 index 0000000000000000000000000000000000000000..54c7d5cee60decc858305568c5397d764d147bbe --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/base_nn_diffusion.py @@ -0,0 +1,44 @@ +from typing import Optional + +import torch +import torch.nn as nn + +from cleandiffuser.utils import SUPPORTED_TIMESTEP_EMBEDDING + + +class BaseNNDiffusion(nn.Module): + """ + The neural network backbone for the Diffusion model used for score matching + (or training a noise predictor) should take in three inputs. + The first input is the noisy data. + The second input is the denoising time step, which can be either as a discrete variable + or a continuous variable, specified by the parameter `discrete_t`. + The third input is the condition embedding that has been processed through the `nn_condition`. + In the general case, we assume that there may be multiple conditions, + which are inputted as a tensor dictionary, or a single condition, directly inputted as a tensor. + """ + + def __init__( + self, emb_dim: int, + timestep_emb_type: str = "positional", + timestep_emb_params: Optional[dict] = None + ): + assert timestep_emb_type in SUPPORTED_TIMESTEP_EMBEDDING.keys() + super().__init__() + timestep_emb_params = timestep_emb_params or {} + self.map_noise = SUPPORTED_TIMESTEP_EMBEDDING[timestep_emb_type](emb_dim, **timestep_emb_params) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, horizon, in_dim) + noise: (b, ) + condition: (b, emb_dim) or None / No condition indicates zeros((b, emb_dim)) + + Output: + y: (b, horizon, in_dim) + """ + raise NotImplementedError + diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/chitransformer.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/chitransformer.py new file mode 100644 index 0000000000000000000000000000000000000000..b79f981230e59ec208c50b4e22798bb19cb91f20 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/chitransformer.py @@ -0,0 +1,158 @@ +from typing import Optional + +import torch +import torch.nn as nn + +from cleandiffuser.utils import (FourierEmbedding, PositionalEmbedding, + SinusoidalEmbedding) + +from .base_nn_diffusion import BaseNNDiffusion + + +def init_weight(module): + ignore_types = ( + nn.Dropout, + SinusoidalEmbedding, + FourierEmbedding, + PositionalEmbedding, + nn.TransformerEncoderLayer, + nn.TransformerDecoderLayer, + nn.TransformerEncoder, + nn.TransformerDecoder, + nn.ModuleList, + nn.Mish, + nn.Sequential) + + if isinstance(module, (nn.Linear, nn.Embedding)): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + if isinstance(module, nn.Linear) and module.bias is not None: + torch.nn.init.zeros_(module.bias) + + elif isinstance(module, nn.MultiheadAttention): + weight_names = [ + 'in_proj_weight', 'q_proj_weight', 'k_proj_weight', 'v_proj_weight'] + for name in weight_names: + weight = getattr(module, name) + if weight is not None: + torch.nn.init.normal_(weight, mean=0.0, std=0.02) + + bias_names = ['in_proj_bias', 'bias_k', 'bias_v'] + for name in bias_names: + bias = getattr(module, name) + if bias is not None: + torch.nn.init.zeros_(bias) + + elif isinstance(module, nn.LayerNorm): + torch.nn.init.zeros_(module.bias) + torch.nn.init.ones_(module.weight) + + elif isinstance(module, ChiTransformer): + torch.nn.init.normal_(module.pos_emb, mean=0.0, std=0.02) + if module.obs_emb is not None: + torch.nn.init.normal_(module.cond_pos_emb, mean=0.0, std=0.02) + + elif isinstance(module, ignore_types): + # no param + pass + else: + raise RuntimeError("Unaccounted module {}".format(module)) + + +class ChiTransformer(BaseNNDiffusion): + """ condition: (1 + To) | x: (Ta) """ + + def __init__( + self, + act_dim: int, obs_dim: int, Ta: int, To: int, + d_model: int = 256, nhead: int = 4, num_layers: int = 8, + p_drop_emb: float = 0.0, p_drop_attn: float = 0.3, + n_cond_layers: int = 0, + timestep_emb_type: str = "positional", + timestep_emb_params: Optional[dict] = None + ): + super().__init__(d_model, timestep_emb_type, timestep_emb_params) + + T = Ta + T_cond = 1 + To + self.To = To + self.obs_dim = obs_dim + + self.act_emb = nn.Linear(act_dim, d_model) + self.pos_emb = nn.Parameter(torch.zeros(1, Ta, d_model)) + + self.obs_emb = nn.Linear(obs_dim, d_model) + self.cond_pos_emb = nn.Parameter(torch.zeros(1, 1 + To, d_model)) + + self.drop = nn.Dropout(p_drop_emb) + self.cond_encoder = nn.Sequential( + nn.Linear(d_model, 4 * d_model), nn.Mish(), nn.Linear(4 * d_model, d_model)) + + # encoder + if n_cond_layers > 0: + encoder_layer = nn.TransformerEncoderLayer( + d_model, nhead, 4 * d_model, p_drop_attn, activation='gelu', batch_first=True, norm_first=True) + self.encoder = nn.TransformerEncoder( + encoder_layer=encoder_layer, num_layers=n_cond_layers) + else: + self.encoder = nn.Sequential( + nn.Linear(d_model, 4 * d_model), nn.Mish(), nn.Linear(4 * d_model, d_model)) + + # decoder + decoder_layer = nn.TransformerDecoderLayer( + d_model, nhead, 4 * d_model, p_drop_attn, activation='gelu', batch_first=True, norm_first=True) + self.decoder = nn.TransformerDecoder( + decoder_layer=decoder_layer, num_layers=num_layers) + + # attention mask + mask = (torch.triu(torch.ones(Ta, Ta)) == 1).transpose(0, 1) + mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) + self.mask = nn.Parameter(mask, requires_grad=False) + + t, s = torch.meshgrid( + torch.arange(Ta), torch.arange(To + 1), indexing='ij') + mask = (t >= (s - 1)) + mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) + self.memory_mask = nn.Parameter(mask, requires_grad=False) + + # decoder head + self.ln_f = nn.LayerNorm(d_model) + self.head = nn.Linear(d_model, act_dim) + + # constant + self.T = T + self.T_cond = T_cond + + self.apply(init_weight) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, Ta, act_dim) + noise: (b, ) + condition: (b, To, obs_dim) + + Output: + y: (b, Ta, act_dim) + """ + if condition is None: + condition = torch.zeros((x.shape[0], self.To, self.obs_dim)).to(x.device) # (b, To, obs_dim) + t_emb = self.map_noise(noise).unsqueeze(1) # (b, 1, d_model) + + act_emb = self.act_emb(x) + obs_emb = self.obs_emb(condition) + + cond_emb = torch.cat([t_emb, obs_emb], dim=1) # (b, 1+To, d_model) + cond_pos_emb = self.cond_pos_emb[:, :cond_emb.shape[1], :] + memory = self.drop(cond_emb + cond_pos_emb) + memory = self.encoder(memory) # (b, 1+To, d_model) + + act_pos_emb = self.pos_emb[:, :act_emb.shape[1], :] + x = self.drop(act_emb + act_pos_emb) # (b, Ta, d_model) + x = self.decoder(tgt=x, memory=memory, tgt_mask=self.mask, memory_mask=self.memory_mask) + + x = self.ln_f(x) + x = self.head(x) + + return x diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/chiunet.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/chiunet.py new file mode 100644 index 0000000000000000000000000000000000000000..ec46345923745b95bc05583872908ab5046d51b8 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/chiunet.py @@ -0,0 +1,193 @@ +from typing import List, Optional + +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.utils import GroupNorm1d + +from .base_nn_diffusion import BaseNNDiffusion +from .jannerunet import Downsample1d, Upsample1d + + +class ChiResidualBlock(nn.Module): + def __init__(self, in_dim: int, out_dim: int, emb_dim: int, kernel_size: int = 3, cond_predict_scale: bool = False): + super().__init__() + + self.conv1 = nn.Sequential( + nn.Conv1d(in_dim, out_dim, kernel_size, padding=kernel_size // 2), + GroupNorm1d(out_dim, 8, 4), nn.Mish()) + self.conv2 = nn.Sequential( + nn.Conv1d(out_dim, out_dim, kernel_size, padding=kernel_size // 2), + GroupNorm1d(out_dim, 8, 4), nn.Mish()) + + cond_dim = 2 * out_dim if cond_predict_scale else out_dim + self.cond_predict_scale = cond_predict_scale + self.out_dim = out_dim + self.cond_encoder = nn.Sequential( + nn.Mish(), nn.Linear(emb_dim, cond_dim)) + + self.residual_conv = nn.Conv1d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + + def forward(self, x, emb): + out = self.conv1(x) + embed = self.cond_encoder(emb) + if self.cond_predict_scale: + embed = embed.reshape( + embed.shape[0], 2, self.out_dim, 1) + scale = embed[:, 0, ...] + bias = embed[:, 1, ...] + out = scale * out + bias + else: + out = out + embed.unsqueeze(-1) + out = self.conv2(out) + out = out + self.residual_conv(x) + return out + + +class ChiUNet1d(BaseNNDiffusion): + def __init__( + self, + act_dim: int, + obs_dim: int, + To: int, + model_dim: int = 256, + emb_dim: int = 256, + kernel_size: int = 5, + cond_predict_scale: bool = True, + obs_as_global_cond: bool = True, + dim_mult: List[int] = [1, 2, 2], + timestep_emb_type: str = "positional", + timestep_emb_params: Optional[dict] = None + ): + super().__init__(emb_dim, timestep_emb_type, timestep_emb_params) + + self.obs_as_global_cond = obs_as_global_cond + self.model_dim = model_dim + self.emb_dim = emb_dim + + dims = [act_dim] + [model_dim * m for m in np.cumprod(dim_mult)] + + self.map_emb = nn.Sequential( + nn.Linear(emb_dim, emb_dim * 4), nn.Mish(), + nn.Linear(emb_dim * 4, emb_dim)) + + if obs_as_global_cond: + self.global_cond_encoder = nn.Linear(To * obs_dim, emb_dim) + emb_dim = emb_dim * 2 # cat obs and emb + self.local_cond_encoder = None + else: + self.global_cond_encoder = None + emb_dim = emb_dim + self.local_cond_encoder = nn.ModuleList([ + ChiResidualBlock( + obs_dim, model_dim, emb_dim, kernel_size, cond_predict_scale), + ChiResidualBlock( + obs_dim, model_dim, emb_dim, kernel_size, cond_predict_scale), + Downsample1d(model_dim)]) + + in_out = list(zip(dims[:-1], dims[1:])) + mid_dim = dims[-1] + + self.downs = nn.ModuleList([]) + self.ups = nn.ModuleList([]) + num_resolutions = len(in_out) + + for ind, (dim_in, dim_out) in enumerate(in_out): + is_last = ind >= (num_resolutions - 1) + self.downs.append(nn.ModuleList([ + ChiResidualBlock( + dim_in, dim_out, emb_dim, kernel_size, cond_predict_scale), + ChiResidualBlock( + dim_out, dim_out, emb_dim, kernel_size, cond_predict_scale), + Downsample1d(dim_out) if not is_last else nn.Identity() + ])) + + self.mids = nn.ModuleList([ + ChiResidualBlock( + mid_dim, mid_dim, emb_dim, kernel_size, cond_predict_scale), + ChiResidualBlock( + mid_dim, mid_dim, emb_dim, kernel_size, cond_predict_scale)]) + + for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): + is_last = ind >= (num_resolutions - 1) + self.ups.append(nn.ModuleList([ + ChiResidualBlock( + dim_out * 2, dim_in, emb_dim, kernel_size, cond_predict_scale), + ChiResidualBlock( + dim_in, dim_in, emb_dim, kernel_size, cond_predict_scale), + Upsample1d(dim_in) if not is_last else nn.Identity() + ])) + + self.final_conv = nn.Sequential( + nn.Conv1d(model_dim, model_dim, kernel_size, padding=kernel_size // 2), + GroupNorm1d(model_dim, 8, 4), nn.Mish(), + nn.Conv1d(model_dim, act_dim, 1)) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, Ta, act_dim) + noise: (b, ) + condition: (b, To, obs_dim) + + Output: + y: (b, Ta, act_dim) + """ + # check Ta dimension + assert x.shape[1] & (x.shape[1] - 1) == 0, "Ta dimension must be 2^n" + + x = x.permute(0, 2, 1) + + emb = self.map_noise(noise) + emb = self.map_emb(emb) + + # If obs_as_global_cond, concatenate obs and emb + if self.obs_as_global_cond: + condition = self.global_cond_encoder(torch.flatten(condition, 1)) + if condition is not None: + emb = torch.cat([emb, condition], dim=-1) + else: + emb = torch.cat([emb, torch.zeros_like(emb)], dim=-1) + h_local = None + else: + condition = condition.permute(0, 2, 1) + assert x.shape[-1] == condition.shape[-1] + if condition is not None: + resnet1, resnet2, dowmsample = self.local_cond_encoder + h_local = [ + resnet1(condition, emb), dowmsample(resnet2(condition, emb))] + else: + zero_cond = torch.zeros((x.shape[0], self.emb_dim, x.shape[1])) + resnet1, resnet2, dowmsample = self.local_cond_encoder + h_local = [ + resnet1(condition, emb), dowmsample(resnet2(zero_cond, emb))] + + h = [] + + for idx, (resnet1, resnet2, downsample) in enumerate(self.downs): + x = resnet1(x, emb) + if idx == 0 and h_local is not None: + x = x + h_local[0] + x = resnet2(x, emb) + h.append(x) + x = downsample(x) + + for mid_module in self.mids: + x = mid_module(x, emb) + + for idx, (resnet1, resnet2, upsample) in enumerate(self.ups): + x = torch.cat((x, h.pop()), dim=1) + x = resnet1(x, emb) + if idx == (len(self.ups) - 1) and h_local is not None: + x = x + h_local[1] + x = resnet2(x, emb) + x = upsample(x) + + x = self.final_conv(x) + + x = x.permute(0, 2, 1) + return x + diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/dit.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/dit.py new file mode 100644 index 0000000000000000000000000000000000000000..e07a5672556818c20ab395ca3f2faa22737c1ec6 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/dit.py @@ -0,0 +1,180 @@ +from typing import Optional + +import torch +import torch.nn as nn + +from cleandiffuser.utils import SinusoidalEmbedding +from cleandiffuser.nn_diffusion import BaseNNDiffusion + + +def modulate(x, shift, scale): + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +class DiTBlock(nn.Module): + """ A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning. """ + + def __init__(self, hidden_size: int, n_heads: int, dropout: float = 0.0): + super().__init__() + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.attn = nn.MultiheadAttention(hidden_size, n_heads, dropout, batch_first=True) + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + + def approx_gelu(): return nn.GELU(approximate="tanh") + + self.mlp = nn.Sequential( + nn.Linear(hidden_size, hidden_size * 4), approx_gelu(), nn.Dropout(dropout), + nn.Linear(hidden_size * 4, hidden_size)) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), nn.Linear(hidden_size, hidden_size * 6)) + + def forward(self, x: torch.Tensor, t: torch.Tensor): + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(t).chunk(6, dim=1) + x = modulate(self.norm1(x), shift_msa, scale_msa) + x = x + gate_msa.unsqueeze(1) * self.attn(x, x, x)[0] + x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp)) + return x + + +class FinalLayer1d(nn.Module): + def __init__(self, hidden_size: int, out_dim: int): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, out_dim) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size)) + + def forward(self, x: torch.Tensor, t: torch.Tensor): + shift, scale = self.adaLN_modulation(t).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift, scale) + return self.linear(x) + + +class DiT1d(BaseNNDiffusion): + def __init__( + self, + in_dim: int, + emb_dim: int, + d_model: int = 384, + n_heads: int = 6, + depth: int = 12, + dropout: float = 0.0, + timestep_emb_type: str = "positional", + timestep_emb_params: Optional[dict] = None + ): + super().__init__(emb_dim, timestep_emb_type, timestep_emb_params) + self.in_dim, self.emb_dim = in_dim, emb_dim + self.d_model = d_model + + self.x_proj = nn.Linear(in_dim, d_model) + self.map_emb = nn.Sequential( + nn.Linear(emb_dim, d_model), nn.Mish(), nn.Linear(d_model, d_model), nn.Mish()) + + self.pos_emb = SinusoidalEmbedding(d_model) + self.pos_emb_cache = None + + self.blocks = nn.ModuleList([ + DiTBlock(d_model, n_heads, dropout) for _ in range(depth)]) + self.final_layer = FinalLayer1d(d_model, in_dim) + self.initialize_weights() + + def initialize_weights(self): + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + self.apply(_basic_init) + + # Initialize time step embedding MLP: + nn.init.normal_(self.map_emb[0].weight, std=0.02) + nn.init.normal_(self.map_emb[2].weight, std=0.02) + + # Zero-out adaLN modulation layers in DiT blocks: + for block in self.blocks: + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + + # Zero-out output layers: + nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0) + nn.init.constant_(self.final_layer.linear.weight, 0) + nn.init.constant_(self.final_layer.linear.bias, 0) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, horizon, in_dim) + noise: (b, ) + condition: (b, emb_dim) or None / No condition indicates zeros((b, emb_dim)) + + Output: + y: (b, horizon, in_dim) + """ + if self.pos_emb_cache is None or self.pos_emb_cache.shape[0] != x.shape[1]: + self.pos_emb_cache = self.pos_emb(torch.arange(x.shape[1], device=x.device)) + + x = self.x_proj(x) + self.pos_emb_cache[None,] + emb = self.map_noise(noise) + if condition is not None: + emb = emb + condition + else: + emb = emb + torch.zeros_like(emb) + emb = self.map_emb(emb) + + for block in self.blocks: + x = block(x, emb) + x = self.final_layer(x, emb) + return x + + +class DiT1Ref(DiT1d): + def __init__( + self, + in_dim: int, + emb_dim: int, + d_model: int = 384, + n_heads: int = 6, + depth: int = 12, + dropout: float = 0.0, + timestep_emb_type: str = "positional", + ): + super().__init__(in_dim, emb_dim, d_model, n_heads, depth, dropout, timestep_emb_type) + self.cross_attns = nn.ModuleList([ + nn.MultiheadAttention(d_model, n_heads, batch_first=True) for _ in range(depth)]) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, horizon, in_dim * 2), where the first half is the reference signal + noise: (b, ) + condition: (b, emb_dim) or None / No condition indicates zeros((b, emb_dim)) + + Output: + y: (b, horizon, in_dim) + """ + if self.pos_emb_cache is None or self.pos_emb_cache.shape[0] != x.shape[1]: + self.pos_emb_cache = self.pos_emb(torch.arange(x.shape[1], device=x.device)) + + x_ref, x = torch.chunk(x, 2, -1) + x_ref_bkp = x_ref.clone() + + x_ref = self.x_proj(x_ref) + self.pos_emb_cache[None, ] + x = self.x_proj(x) + self.pos_emb_cache[None, ] + emb = self.map_noise(noise) + + if condition is not None: + emb = emb + condition + emb = self.map_emb(emb) + + for cross_attn, block in zip(self.cross_attns, self.blocks): + x, _ = cross_attn(x, x_ref, x_ref) + x = block(x, emb) + x = self.final_layer(x, emb) + return torch.cat([x_ref_bkp, x], -1) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/dqlmlp.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/dqlmlp.py new file mode 100644 index 0000000000000000000000000000000000000000..27950411e5dd4d0aa3858472c447e919f61c356e --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/dqlmlp.py @@ -0,0 +1,52 @@ +from typing import Optional + +import torch +import torch.nn as nn + +from cleandiffuser.nn_diffusion import BaseNNDiffusion + + +class DQLMlp(BaseNNDiffusion): + def __init__( + self, + obs_dim: int, + act_dim: int, + emb_dim: int = 16, + timestep_emb_type: str = "positional", + timestep_emb_params: Optional[dict] = None + ): + super().__init__(emb_dim, timestep_emb_type, timestep_emb_params) + + self.obs_dim = obs_dim + self.time_mlp = nn.Sequential( + nn.Linear(emb_dim, emb_dim * 2), nn.Mish(), nn.Linear(emb_dim * 2, emb_dim)) + + self.mid_layer = nn.Sequential( + nn.Linear(obs_dim + act_dim + emb_dim, 256), nn.Mish(), + nn.Linear(256, 256), nn.Mish(), + nn.Linear(256, 256), nn.Mish()) + + self.final_layer = nn.Linear(256, act_dim) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, act_dim) + noise: (b, ) + condition: (b, obs_dim) + + Output: + y: (b, act_dim) + """ + if condition is None: + condition = torch.zeros(x.shape[0], self.obs_dim).to(x.device) + t = self.time_mlp(self.map_noise(noise)) + x = torch.cat([x, t, condition], -1) + x = self.mid_layer(x) + + return self.final_layer(x) + + + diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/dvinvmlp.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/dvinvmlp.py new file mode 100644 index 0000000000000000000000000000000000000000..e0926d9ffeef42c159fad40625b53f3306231463 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/dvinvmlp.py @@ -0,0 +1,47 @@ +from typing import Optional + +import torch +import torch.nn as nn + +from cleandiffuser.nn_diffusion import BaseNNDiffusion + + +class DVInvMlp(BaseNNDiffusion): + def __init__( + self, + obs_dim: int, + act_dim: int, + emb_dim: int = 16, + hidden_dim: int = 256, + timestep_emb_type: str = "positional", + timestep_emb_params: Optional[dict] = None + ): + super().__init__(emb_dim, timestep_emb_type, timestep_emb_params) + + self.time_mlp = nn.Sequential( + nn.Linear(emb_dim, emb_dim * 2), nn.Mish(), nn.Linear(emb_dim * 2, emb_dim)) + + self.mid_layer = nn.Sequential( + nn.Linear(obs_dim * 2 + act_dim + emb_dim, hidden_dim), nn.Mish(), + nn.Linear(hidden_dim, hidden_dim), nn.Mish(), + nn.Linear(hidden_dim, hidden_dim), nn.Mish()) + + self.final_layer = nn.Linear(hidden_dim, act_dim) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: torch.Tensor = None): + """ + Input: + x: (b, act_dim) + noise: (b, ) + condition: (b, obs_dim * 2) + + Output: + y: (b, act_dim) + """ + t = self.time_mlp(self.map_noise(noise)) + x = torch.cat([x, t, condition], -1) + x = self.mid_layer(x) + + return self.final_layer(x) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/idqlmlp.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/idqlmlp.py new file mode 100644 index 0000000000000000000000000000000000000000..5efb1ba338fa348970e6962b8341b3de065a9435 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/idqlmlp.py @@ -0,0 +1,113 @@ +from typing import Optional + +import torch +import torch.nn as nn + +from cleandiffuser.nn_diffusion import BaseNNDiffusion + + +class ResidualBlock(nn.Module): + def __init__(self, hidden_dim: int, dropout: float = 0.1): + super().__init__() + self.net = nn.Sequential( + nn.Dropout(dropout), nn.LayerNorm(hidden_dim), + nn.Linear(hidden_dim, hidden_dim * 4), nn.Mish(), + nn.Linear(hidden_dim * 4, hidden_dim)) + + def forward(self, x): + return x + self.net(x) + + +class IDQLMlp(BaseNNDiffusion): + def __init__( + self, + obs_dim: int, + act_dim: int, + emb_dim: int = 64, + hidden_dim: int = 256, + n_blocks: int = 3, + dropout: float = 0.1, + timestep_emb_type: str = "positional", + timestep_emb_params: Optional[dict] = None + ): + super().__init__(emb_dim, timestep_emb_type, timestep_emb_params) + + self.obs_dim = obs_dim + self.time_mlp = nn.Sequential( + nn.Linear(emb_dim, emb_dim * 2), nn.Mish(), nn.Linear(emb_dim * 2, emb_dim)) + + self.affine_in = nn.Linear(obs_dim + act_dim + emb_dim, hidden_dim) + + self.ln_resnet = nn.Sequential(*[ResidualBlock(hidden_dim, dropout) for _ in range(n_blocks)]) + + self.affine_out = nn.Linear(hidden_dim, act_dim) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, act_dim) + noise: (b, ) + condition: (b, obs_dim) + + Output: + y: (b, act_dim) + """ + if condition is None: + condition = torch.zeros(x.shape[0], self.obs_dim).to(x.device) + + t = self.time_mlp(self.map_noise(noise)) + x = torch.cat([x, t, condition], -1) + x = self.affine_in(x) + x = self.ln_resnet(x) + + return self.affine_out(x) + + +class NewIDQLMlp(BaseNNDiffusion): + def __init__( + self, + obs_dim: int, + act_dim: int, + emb_dim: int = 64, + hidden_dim: int = 256, + n_blocks: int = 3, + dropout: float = 0.1, + timestep_emb_type: str = "positional", + timestep_emb_params: Optional[dict] = None + ): + super().__init__(emb_dim, timestep_emb_type, timestep_emb_params) + + self.obs_dim = obs_dim + self.time_mlp = nn.Sequential( + nn.Linear(emb_dim, emb_dim * 2), nn.Mish(), nn.Linear(emb_dim * 2, emb_dim)) + + self.affine_in = nn.Linear(obs_dim + act_dim + emb_dim, hidden_dim) + + self.ln_resnet = nn.Sequential(*[ResidualBlock(hidden_dim, dropout) for _ in range(n_blocks)]) + + self.affine_out = nn.Sequential(nn.Mish(), nn.Linear(hidden_dim, act_dim)) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, act_dim) + noise: (b, ) + condition: (b, obs_dim) + + Output: + y: (b, act_dim) + """ + if condition is None: + condition = torch.zeros(x.shape[0], self.obs_dim).to(x.device) + + t = self.time_mlp(self.map_noise(noise)) + x = torch.cat([x, t, condition], -1) + x = self.affine_in(x) + x = self.ln_resnet(x) + + return self.affine_out(x) + \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/jannerunet.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/jannerunet.py new file mode 100644 index 0000000000000000000000000000000000000000..caa1d6b9d55911874bfc09ee6dba72b9341dcc35 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/jannerunet.py @@ -0,0 +1,201 @@ +from typing import List, Optional + +import einops +import numpy as np +import torch +import torch.nn as nn + +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from cleandiffuser.utils import GroupNorm1d + + +def get_norm(dim: int, norm_type: str = "groupnorm"): + if norm_type == "groupnorm": + return GroupNorm1d(dim, 8, 4) + elif norm_type == "layernorm": + return LayerNorm(dim) + else: + return nn.Identity() + + +class Downsample1d(nn.Module): + def __init__(self, dim): + super().__init__() + self.conv = nn.Conv1d(dim, dim, 3, 2, 1) + + def forward(self, x): + return self.conv(x) + + +class Upsample1d(nn.Module): + def __init__(self, dim): + super().__init__() + self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1) + + def forward(self, x): + return self.conv(x) + + +class LayerNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.eps = eps + self.g = nn.Parameter(torch.ones(1, dim, 1)) + self.b = nn.Parameter(torch.zeros(1, dim, 1)) + + def forward(self, x): + var = torch.var(x, dim=1, unbiased=False, keepdim=True) + mean = torch.mean(x, dim=1, keepdim=True) + return (x - mean) / (var + self.eps).sqrt() * self.g + self.b + + +class ResidualBlock(nn.Module): + def __init__(self, in_dim: int, out_dim: int, emb_dim: int, kernel_size: int = 3, norm_type: str = "groupnorm"): + super().__init__() + + self.conv1 = nn.Sequential( + nn.Conv1d(in_dim, out_dim, kernel_size, padding=kernel_size // 2), + get_norm(out_dim, norm_type), nn.Mish()) + self.conv2 = nn.Sequential( + nn.Conv1d(out_dim, out_dim, kernel_size, padding=kernel_size // 2), + get_norm(out_dim, norm_type), nn.Mish()) + self.emb_mlp = nn.Sequential( + nn.Mish(), nn.Linear(emb_dim, out_dim)) + self.residual_conv = nn.Conv1d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + + def forward(self, x, emb): + out = self.conv1(x) + self.emb_mlp(emb).unsqueeze(-1) + out = self.conv2(out) + return out + self.residual_conv(x) + + +class LinearAttention(nn.Module): + def __init__(self, dim, heads=4, dim_head=32): + super().__init__() + self.norm = LayerNorm(dim) + self.scale = dim_head ** -0.5 + self.heads = heads + hidden_dim = dim_head * heads + self.to_qkv = nn.Conv1d(dim, hidden_dim * 3, 1, bias=False) + self.to_out = nn.Conv1d(hidden_dim, dim, 1) + + def forward(self, x): + x = self.norm(x) + + qkv = self.to_qkv(x).chunk(3, dim=1) + q, k, v = map(lambda t: einops.rearrange(t, 'b (h c) d -> b h c d', h=self.heads), qkv) + q = q * self.scale + + k = k.softmax(dim=-1) + context = torch.einsum('b h d n, b h e n -> b h d e', k, v) + + out = torch.einsum('b h d e, b h d n -> b h e n', context, q) + out = einops.rearrange(out, 'b h c d -> b (h c) d') + out = self.to_out(out) + return out + x + + +class JannerUNet1d(BaseNNDiffusion): + def __init__( + self, + in_dim: int, + model_dim: int = 32, + emb_dim: int = 32, + kernel_size: int = 3, + dim_mult: List[int] = [1, 2, 2, 2], + norm_type: str = "groupnorm", + attention: bool = False, + timestep_emb_type: str = "positional", + timestep_emb_params: Optional[dict] = None + ): + super().__init__(emb_dim, timestep_emb_type, timestep_emb_params) + + dims = [in_dim] + [model_dim * m for m in np.cumprod(dim_mult)] + in_out = list(zip(dims[:-1], dims[1:])) + + self.map_emb = nn.Sequential( + nn.Linear(emb_dim, model_dim * 4), nn.Mish(), + nn.Linear(model_dim * 4, model_dim)) + + self.downs = nn.ModuleList([]) + self.ups = nn.ModuleList([]) + num_resolutions = len(in_out) + + for ind, (dim_in, dim_out) in enumerate(in_out): + is_last = ind >= (num_resolutions - 1) + + self.downs.append(nn.ModuleList([ + ResidualBlock(dim_in, dim_out, model_dim, kernel_size, norm_type), + ResidualBlock(dim_out, dim_out, model_dim, kernel_size, norm_type), + LinearAttention(dim_out) if attention else nn.Identity(), + Downsample1d(dim_out) if not is_last else nn.Identity() + ])) + + mid_dim = dims[-1] + self.mid_block1 = ResidualBlock(mid_dim, mid_dim, model_dim, kernel_size, norm_type) + self.mid_attn = LinearAttention(mid_dim) if attention else nn.Identity() + self.mid_block2 = ResidualBlock(mid_dim, mid_dim, model_dim, kernel_size, norm_type) + + for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): + is_last = ind >= (num_resolutions - 1) + + self.ups.append(nn.ModuleList([ + ResidualBlock(dim_out * 2, dim_in, model_dim, kernel_size, norm_type), + ResidualBlock(dim_in, dim_in, model_dim, kernel_size, norm_type), + LinearAttention(dim_in) if attention else nn.Identity(), + Upsample1d(dim_in) if not is_last else nn.Identity() + ])) + + self.final_conv = nn.Sequential( + nn.Conv1d(model_dim, model_dim, 5, padding=2), + get_norm(model_dim, norm_type), nn.Mish(), + nn.Conv1d(model_dim, in_dim, 1)) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, horizon, in_dim) + noise: (b, ) + condition: (b, emb_dim) or None / No condition indicates zeros((b, emb_dim)) + + Output: + y: (b, horizon, in_dim) + """ + # check horizon dimension + assert x.shape[1] & (x.shape[1] - 1) == 0, "Ta dimension must be 2^n" + + x = x.permute(0, 2, 1) + + emb = self.map_noise(noise) + if condition is not None: + emb = emb + condition + else: + emb = emb + torch.zeros_like(emb) + emb = self.map_emb(emb) + + h = [] + + for resnet1, resnet2, attn, downsample in self.downs: + x = resnet1(x, emb) + x = resnet2(x, emb) + x = attn(x) + h.append(x) + x = downsample(x) + + x = self.mid_block1(x, emb) + x = self.mid_attn(x) + x = self.mid_block2(x, emb) + + for resnet1, resnet2, attn, upsample in self.ups: + x = torch.cat([x, h.pop()], dim=1) + x = resnet1(x, emb) + x = resnet2(x, emb) + x = attn(x) + x = upsample(x) + + x = self.final_conv(x) + + x = x.permute(0, 2, 1) + return x diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/mlps.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/mlps.py new file mode 100644 index 0000000000000000000000000000000000000000..2ebcef0d45c44b5f654eef2157f0697d8cc46c7c --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/mlps.py @@ -0,0 +1,44 @@ +from typing import List, Optional + +import torch +import torch.nn as nn + +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from cleandiffuser.utils import Mlp + + +class MlpNNDiffusion(BaseNNDiffusion): + def __init__( + self, + x_dim: int, + emb_dim: int = 16, + hidden_dims: List[int] = (256, 256), + activation: nn.Module = nn.ReLU(), + timestep_emb_type: str = "positional", + timestep_emb_params: Optional[dict] = None + ): + super().__init__(emb_dim, timestep_emb_type, timestep_emb_params) + self.mlp = Mlp( + x_dim + emb_dim, hidden_dims, x_dim, activation) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, x_dim) + noise: (b, ) + condition: (b, emd_dim) + + Output: + y: (b, x_dim) + """ + t = self.map_noise(noise) + if condition is not None: + t += condition + else: + t += torch.zeros_like(t) + return self.mlp(torch.cat([x, t], -1)) + + + \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/pearcemlp.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/pearcemlp.py new file mode 100644 index 0000000000000000000000000000000000000000..9339fe3fcfd0ae8bd670dcba6008c6ea0f7cb1e4 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/pearcemlp.py @@ -0,0 +1,79 @@ +from typing import Optional + +import torch +import torch.nn as nn + +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from cleandiffuser.utils import GroupNorm1d + + +class TimeSiren(nn.Module): + def __init__(self, input_dim, emb_dim): + super(TimeSiren, self).__init__() + # just a fully connected NN with sin activations + self.lin1 = nn.Linear(input_dim, emb_dim, bias=False) + self.lin2 = nn.Linear(emb_dim, emb_dim) + + def forward(self, x): + x = torch.sin(self.lin1(x)) + x = self.lin2(x) + return x + + +class FCBlock(nn.Module): + def __init__(self, in_feats, out_feats): + super().__init__() + # one layer of non-linearities (just a useful building block to use below) + self.model = nn.Sequential( + nn.Linear(in_feats, out_feats), + GroupNorm1d(out_feats, 8, 4), + nn.GELU()) + + def forward(self, x): + return self.model(x) + + +class PearceMlp(BaseNNDiffusion): + def __init__( + self, act_dim: int, To: int = 1, timestep_emb_type: str = "positional", + emb_dim: int = 128, hidden_dim: int = 512, timestep_emb_params: Optional[dict] = None, + ): + super().__init__(emb_dim, timestep_emb_type, timestep_emb_params) + + self.act_emb = nn.Sequential( + nn.Linear(act_dim, emb_dim), nn.LeakyReLU(), nn.Linear(emb_dim, emb_dim)) + + self.fcs = nn.ModuleList([ + FCBlock(emb_dim * (2 + To), hidden_dim), + FCBlock(hidden_dim + act_dim + 1, hidden_dim), + FCBlock(hidden_dim + act_dim + 1, hidden_dim), + nn.Linear(hidden_dim + act_dim + 1, act_dim)]) + + self.To = To + self.emb_dim = emb_dim + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, act_dim) + noise: (b, ) + condition: (b, To, emb_dim) + + Output: + y: (b, act_dim) + """ + x_e, t_e = self.act_emb(x), self.map_noise(noise) + t = noise.unsqueeze(-1) + + if condition is not None: + nn1 = self.fcs[0](torch.cat([x_e, t_e, torch.flatten(condition, 1)], -1)) + else: + condition = torch.zeros(x.shape[0], self.To, self.emb_dim).to(x.device) + nn1 = self.fcs[0](torch.cat([x_e, t_e, torch.flatten(condition, 1)], -1)) + nn2 = self.fcs[1](torch.cat([nn1 / 1.414, x, t], -1)) + nn1 / 1.414 + nn3 = self.fcs[2](torch.cat([nn2 / 1.414, x, t], -1)) + nn2 / 1.414 + out = self.fcs[3](torch.cat([nn3, x, t], -1)) + + return out diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/pearcetransformer.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/pearcetransformer.py new file mode 100644 index 0000000000000000000000000000000000000000..ee13efda86e9d9233b8477b0db8d5770e98efeff --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/pearcetransformer.py @@ -0,0 +1,151 @@ +import torch +import torch.nn as nn + +from cleandiffuser.nn_diffusion import BaseNNDiffusion +from typing import Optional + + +class TimeSiren(nn.Module): + def __init__(self, input_dim, emb_dim): + super(TimeSiren, self).__init__() + # just a fully connected NN with sin activations + self.lin1 = nn.Linear(input_dim, emb_dim, bias=False) + self.lin2 = nn.Linear(emb_dim, emb_dim) + + def forward(self, x): + x = torch.sin(self.lin1(x)) + x = self.lin2(x) + return x + + +class TransformerEncoderBlock(nn.Module): + def __init__(self, trans_emb_dim, transformer_dim, nheads): + super(TransformerEncoderBlock, self).__init__() + # mainly going off of https://jalammar.github.io/illustrated-transformer/ + + self.trans_emb_dim = trans_emb_dim + self.transformer_dim = transformer_dim + self.nheads = nheads + + self.input_to_qkv1 = nn.Linear(self.trans_emb_dim, self.transformer_dim * 3) + self.multihead_attn1 = nn.MultiheadAttention(self.transformer_dim, num_heads=self.nheads) + self.attn1_to_fcn = nn.Linear(self.transformer_dim, self.trans_emb_dim) + self.attn1_fcn = nn.Sequential( + nn.Linear(self.trans_emb_dim, self.trans_emb_dim * 4), + nn.GELU(), + nn.Linear(self.trans_emb_dim * 4, self.trans_emb_dim), + ) + self.norm1a = nn.BatchNorm1d(self.trans_emb_dim) + self.norm1b = nn.BatchNorm1d(self.trans_emb_dim) + + def split_qkv(self, qkv): + assert qkv.shape[-1] == self.transformer_dim * 3 + q = qkv[:, :, :self.transformer_dim] + k = qkv[:, :, self.transformer_dim: 2 * self.transformer_dim] + v = qkv[:, :, 2 * self.transformer_dim:] + return q, k, v + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + qkvs1 = self.input_to_qkv1(inputs) + # shape out = [3, batchsize, transformer_dim*3] + + qs1, ks1, vs1 = self.split_qkv(qkvs1) + # shape out = [3, batchsize, transformer_dim] + + attn1_a = self.multihead_attn1(qs1, ks1, vs1, need_weights=False) + attn1_a = attn1_a[0] + # shape out = [3, batchsize, transformer_dim = trans_emb_dim x nheads] + + attn1_b = self.attn1_to_fcn(attn1_a) + attn1_b = attn1_b / 1.414 + inputs / 1.414 # add residual + # shape out = [3, batchsize, trans_emb_dim] + + # normalise + attn1_b = self.norm1a(attn1_b.transpose(0, 2).transpose(0, 1)) + attn1_b = attn1_b.transpose(0, 1).transpose(0, 2) + # batchnorm likes shape = [batchsize, trans_emb_dim, 3] + # so have to shape like this, then return + + # fully connected layer + attn1_c = self.attn1_fcn(attn1_b) / 1.414 + attn1_b / 1.414 + # shape out = [3, batchsize, trans_emb_dim] + + # normalise + # attn1_c = self.norm1b(attn1_c) + attn1_c = self.norm1b(attn1_c.transpose(0, 2).transpose(0, 1)) + attn1_c = attn1_c.transpose(0, 1).transpose(0, 2) + return attn1_c + + +class EmbeddingBlock(nn.Module): + def __init__(self, in_dim: int, emb_dim: int): + super().__init__() + self.model = nn.Sequential( + nn.Linear(in_dim, emb_dim), nn.LeakyReLU(), + nn.Linear(emb_dim, emb_dim)) + + def forward(self, x): + return self.model(x) + + +class PearceTransformer(BaseNNDiffusion): + def __init__( + self, act_dim: int, To: int = 1, + emb_dim: int = 128, trans_emb_dim: int = 64, nhead: int = 16, + timestep_emb_type: str = "positional", + timestep_emb_params: Optional[dict] = None + ): + super().__init__(emb_dim, timestep_emb_type, timestep_emb_params) + + self.To = To + self.emb_dim = emb_dim + self.act_emb = nn.Sequential( + nn.Linear(act_dim, emb_dim), nn.LeakyReLU(), nn.Linear(emb_dim, emb_dim)) + + transformer_dim = trans_emb_dim * nhead + + self.act_to_input = nn.Linear(emb_dim, trans_emb_dim) + self.t_to_input = nn.Linear(emb_dim, trans_emb_dim) + self.cond_to_input = nn.Linear(emb_dim, trans_emb_dim) + + self.pos_embed = TimeSiren(1, trans_emb_dim) + + self.transformer_blocks = nn.Sequential( + TransformerEncoderBlock(trans_emb_dim, transformer_dim, nhead), + TransformerEncoderBlock(trans_emb_dim, transformer_dim, nhead), + TransformerEncoderBlock(trans_emb_dim, transformer_dim, nhead), + TransformerEncoderBlock(trans_emb_dim, transformer_dim, nhead)) + + self.final = nn.Linear(trans_emb_dim * (2 + To), act_dim) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, act_dim) + noise: (b, ) + condition: (b, To, emb_dim) + + Output: + y: (b, act_dim) + """ + if condition is None: + condition = torch.zero((x.shape[0], self.To, self.emb_dim)).to(x.device) + + x_e, t_e = self.act_emb(x), self.map_noise(noise) + + x_input, t_input, c_input = self.act_to_input(x_e), self.t_to_input(t_e), self.cond_to_input(condition) + + x_input += self.pos_embed(torch.zeros(1, 1, device=x.device) + 1.0) + t_input += self.pos_embed(torch.zeros(1, 1, device=x.device) + 2.0) + c_input += self.pos_embed( + torch.arange(3, 3 + condition.shape[1], device=x.device, dtype=torch.float32)[None, :, None]) + + f = torch.cat([x_input.unsqueeze(1), t_input.unsqueeze(1), c_input], dim=1) + f = self.transformer_blocks(f.permute(1, 0, 2)).permute(1, 0, 2) + + flat = torch.flatten(f, start_dim=1) + + out = self.final(flat) + return out diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/sfbc_unet.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/sfbc_unet.py new file mode 100644 index 0000000000000000000000000000000000000000..5484597d519034ca80928e9433e7c9ad7b818c8a --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/nn_diffusion/sfbc_unet.py @@ -0,0 +1,82 @@ +from typing import Optional, List + +import torch +import torch.nn as nn + +from .base_nn_diffusion import BaseNNDiffusion + + +class ResidualBlock(nn.Module): + def __init__(self, in_dim: int, out_dim: int, emb_dim: int): + super().__init__() + self.linear1 = nn.Sequential(nn.Linear(in_dim, out_dim), nn.SiLU()) + self.linear2 = nn.Sequential(nn.Linear(out_dim, out_dim), nn.SiLU()) + self.linearc = nn.Linear(emb_dim, out_dim) + + self.skip = nn.Linear(in_dim, out_dim) if in_dim != out_dim else nn.Identity() + + def forward(self, x: torch.Tensor, c: torch.Tensor): + return self.linear2(self.linear1(x) + self.linearc(c)) + self.skip(x) + + +class SfBCUNet(BaseNNDiffusion): + def __init__( + self, + act_dim: int, + emb_dim: int = 64, + hidden_dims: List[int] = (512, 256, 128), + timestep_emb_type: str = "untrainable_fourier", + timestep_emb_params: Optional[dict] = None + ): + super().__init__(emb_dim, timestep_emb_type, timestep_emb_params) + n_layers = len(hidden_dims) + + self.t_layer = nn.Sequential( + nn.Linear(emb_dim, emb_dim), nn.SiLU(), nn.Linear(emb_dim, emb_dim)) + + self.down_blocks = nn.ModuleList() + self.up_blocks = nn.ModuleList() + + in_dim = act_dim + for i in range(n_layers): + self.down_blocks.append(ResidualBlock(in_dim, hidden_dims[i], emb_dim)) + in_dim = hidden_dims[i] + + self.mid_block = ResidualBlock(in_dim, in_dim, emb_dim) + + for i in range(n_layers - 1): + self.up_blocks.append(ResidualBlock(in_dim + hidden_dims[-1 - i], hidden_dims[-2 - i], emb_dim)) + in_dim = hidden_dims[-2 - i] + + self.out_layer = nn.Linear(in_dim, act_dim) + + def forward(self, + x: torch.Tensor, noise: torch.Tensor, + condition: Optional[torch.Tensor] = None): + """ + Input: + x: (b, horizon, act_dim) + noise: (b, ) + condition: (b, emb_dim) or None / No condition indicates zeros((b, emb_dim)) + + Output: + y: (b, horizon, act_dim) + """ + c = self.t_layer(self.map_noise(noise)) + if condition is not None: + c += condition + else: + c += torch.zeros_like(c) + + buffer = [] + for block in self.down_blocks: + x = block(x, c) + buffer.append(x) + + x = self.mid_block(x, c) + + for block in self.up_blocks: + x = torch.cat([x, buffer.pop()], dim=-1) + x = block(x, c) + + return self.out_layer(x) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dee9803286412da38454e6c83e475d3774f78c10 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/__init__.py @@ -0,0 +1,5 @@ +from .utils import * +from .iql import IQL +from .building_blocks import * +from .normalizers import * +from .typings import * diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/building_blocks.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/building_blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..85806da44caa9fd7a9b4ed04af7fb8ae6617d824 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/building_blocks.py @@ -0,0 +1,373 @@ +import torch +import torch.nn as nn +import numpy as np +import einops +from typing import List +from .iql import TwinQ, V +from cleandiffuser.utils import SinusoidalEmbedding + +IDQLQNet = TwinQ +IDQLVNet = V + + +class Mlp(nn.Module): + """ **Multilayer perceptron.** A simple pytorch MLP module. + + Args: + in_dim: int, + The dimension of the input tensor. + hidden_dims: List[int], + A list of integers, each element is the dimension of the hidden layer. + out_dim: int, + The dimension of the output tensor. + activation: nn.Module, + The activation function used in the hidden layers. + out_activation: nn.Module, + The activation function used in the output layer. + """ + + def __init__( + self, + in_dim: int, + hidden_dims: List[int], + out_dim: int, + activation: nn.Module = nn.ReLU(), + out_activation: nn.Module = nn.Identity(), + ): + super().__init__() + self.mlp = nn.Sequential( + *[ + nn.Sequential( + nn.Linear(in_dim if i == 0 else hidden_dims[i - 1], hidden_dims[i]), + activation, + ) + for i in range(len(hidden_dims)) + ], + nn.Linear(hidden_dims[-1], out_dim), + out_activation + ) + + def _init_weights(self): + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + nn.init.zeros_(m.bias) + + def forward(self, x): + return self.mlp(x) + + +class GroupNorm1d(nn.Module): + def __init__(self, dim, num_groups=32, min_channels_per_group=4, eps=1e-5): + super().__init__() + self.num_groups = min(num_groups, dim // min_channels_per_group) + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + self.bias = nn.Parameter(torch.zeros(dim)) + + def forward(self, x): + x = torch.nn.functional.group_norm( + x.unsqueeze(2), + num_groups=self.num_groups, + weight=self.weight.to(x.dtype), + bias=self.bias.to(x.dtype), + eps=self.eps, + ) + return x.squeeze(2) + + +class SoftLowerBound(nn.Module): + """ Soft lower bound. + + Args: + lower_bound: float, + The lower bound of the output. + """ + + def __init__(self, lower_bound: float): + super().__init__() + self.lower_bound = lower_bound + + def forward(self, x): + return self.lower_bound + torch.nn.functional.softplus(x - self.lower_bound) + + +class SoftUpperBound(nn.Module): + """ Soft upper bound. + + Args: + upper_bound: float, + The upper bound of the output. + """ + + def __init__(self, upper_bound: float): + super().__init__() + self.upper_bound = upper_bound + + def forward(self, x): + return self.upper_bound - torch.nn.functional.softplus(self.upper_bound - x) + + +class DQLCritic(nn.Module): + """ **Deep Q-Learning Critic.** A pytorch critic module for DQL. The module incorporates double Q trick. + + Args: + obs_dim: int, + The dimension of the observation space. + act_dim: int, + The dimension of the action space. + hidden_dim: int, + The dimension of the hidden layers. + """ + + def __init__(self, obs_dim: int, act_dim: int, hidden_dim: int = 256): + super().__init__() + self.q1_model = nn.Sequential( + nn.Linear(obs_dim + act_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.Tanh(), + nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.Mish(), + nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.Mish(), + nn.Linear(hidden_dim, 1)) + + self.q2_model = nn.Sequential( + nn.Linear(obs_dim + act_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.Tanh(), + nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.Mish(), + nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.Mish(), + nn.Linear(hidden_dim, 1)) + + def forward(self, obs, act): + x = torch.cat([obs, act], dim=-1) + return self.q1_model(x), self.q2_model(x) + + def q1(self, obs, act): + x = torch.cat([obs, act], dim=-1) + return self.q1_model(x) + + def q_min(self, obs, act): + q1, q2 = self.forward(obs, act) + return torch.min(q1, q2) + +class DVTransformerBlock(nn.Module): + def __init__(self, hidden_size: int, n_heads: int, dropout: float = 0.0, norm_type="post"): + super().__init__() + self.norm_type = norm_type + + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.attn = nn.MultiheadAttention(hidden_size, n_heads, dropout, batch_first=True) + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + + def approx_gelu(): return nn.GELU(approximate="tanh") + + self.mlp = nn.Sequential( + nn.Linear(hidden_size, hidden_size * 4), approx_gelu(), nn.Dropout(dropout), + nn.Linear(hidden_size * 4, hidden_size)) + + def forward(self, x: torch.Tensor): + if self.norm_type == "post": + x = self.norm1(x + self.attn(x, x, x)[0]) + x = self.norm2(x + self.mlp(x)) + elif self.norm_type == "pre": + x = self.norm1(x) + x = x + self.attn(x, x, x)[0] + x = x + self.mlp(self.norm2(x)) + else: + raise NotImplementedError + return x + +class DVHorizonCritic(nn.Module): + def __init__( + self, + in_dim: int, + emb_dim: int, + d_model: int = 384, + n_heads: int = 6, + depth: int = 12, + dropout: float = 0.0, + norm_type: str = "post" + ): + super().__init__() + self.in_dim, self.emb_dim = in_dim, emb_dim + self.d_model = d_model + + self.x_proj = nn.Linear(in_dim, d_model) + + self.pos_emb = SinusoidalEmbedding(d_model) + self.pos_emb_cache = None + + self.blocks = nn.ModuleList([DVTransformerBlock(d_model, n_heads, dropout, norm_type) for _ in range(depth)]) + self.final_layer = nn.Linear(d_model, 1) + self.initialize_weights() + + def initialize_weights(self): + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + self.apply(_basic_init) + + def forward(self, x: torch.Tensor): + """ + Input: + x: (b, horizon, in_dim) + + Output: + y: (b, horizon, in_dim) + """ + if self.pos_emb_cache is None or self.pos_emb_cache.shape[0] != x.shape[1]: + self.pos_emb_cache = self.pos_emb(torch.arange(x.shape[1], device=x.device)) + + x = self.x_proj(x) + self.pos_emb_cache[None,] + + for block in self.blocks: + x = block(x) + x = self.final_layer(x) + + x = x[:, 0, :] + + return x + +class PreNorm(nn.Module): + """Layer normalization before the function. + + output = fn(LayerNorm(x), **kwargs) + """ + + def __init__(self, dim, fn): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.fn = fn + + def forward(self, x, **kwargs): + return self.fn(self.norm(x), **kwargs) + + +class Residual(nn.Module): + """Residual connection. + + output = fn(x, **kwargs) + x + """ + + def __init__(self, fn): + super().__init__() + self.fn = fn + + def forward(self, x, **kwargs): + return self.fn(x, **kwargs) + x + + +class FeedForward(nn.Module): + def __init__(self, d_model: int, hidden_scale: int = 4, dropout: float = 0.0): + super().__init__() + hidden_dim = int(d_model * hidden_scale) + self.net = nn.Sequential( + nn.Linear(d_model, hidden_dim), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(hidden_dim, d_model), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +class MultiHeadAttention(nn.Module): + def __init__( + self, d_model: int, nhead: int, dropout: float = 0.1, bias: bool = False + ): + super().__init__() + assert d_model % nhead == 0, "`d_model` must be divisible by `nhead`." + + self.nhead, self.d_k = nhead, d_model // nhead + self.scale = 1 / np.sqrt(self.d_k) + + self.q_layer = nn.Linear(d_model, d_model, bias=bias) + self.k_layer = nn.Linear(d_model, d_model, bias=bias) + self.v_layer = nn.Linear(d_model, d_model, bias=True) + + self.dropout = nn.Dropout(dropout) + + def forward(self, q, k, v, mask=None): + + if mask is not None: + if mask.dim() == 2: + assert mask.shape[0] == q.shape[1] and mask.shape[1] == k.shape[1] + mask = mask.unsqueeze(0) + elif mask.dim() == 3: + assert ( + mask.shape[0] == q.shape[0] + and mask.shape[1] == q.shape[1] + and mask.shape[2] == k.shape[1] + ) + else: + raise ValueError("`mask` shape should be either (i, j) or (b, i, j)") + mask = mask.unsqueeze(-1) + + q = einops.rearrange(self.q_layer(q), "b i (h d) -> b i h d", h=self.nhead) + k = einops.rearrange(self.k_layer(k), "b j (h d) -> b j h d", h=self.nhead) + v = einops.rearrange(self.v_layer(v), "b j (h d) -> b j h d", h=self.nhead) + + scores = torch.einsum("b i h d, b j h d -> b i j h", q, k) * self.scale + if mask is not None: + scores.masked_fill_(mask == 0, float("-inf")) + + attn = torch.softmax(scores, dim=2) + + attn = self.dropout(attn) + + out = torch.einsum("b i j h, b j h d -> b i h d", attn, v) + + out = einops.rearrange(out, "b i h d -> b i (h d)") + + return out, attn.detach() + + +def generate_causal_mask(length: int, device: torch.device = "cpu"): + mask = torch.tril(torch.ones(length, length, device=device), diagonal=0) + return mask + + +class Transformer(nn.Module): + def __init__( + self, + d_model: int, + nhead: int, + num_layers: int, + hidden_scale: int = 4, + attn_dropout: float = 0.0, + ffn_dropout: float = 0.0, + bias: bool = False, + ): + super().__init__() + self.layers = nn.ModuleList( + [ + nn.ModuleList( + [ + nn.LayerNorm(d_model), + MultiHeadAttention(d_model, nhead, attn_dropout, bias), + nn.LayerNorm(d_model), + FeedForward(d_model, hidden_scale, ffn_dropout), + ] + ) + for _ in range(num_layers) + ] + ) + + def forward(self, x, mask=None): + attn_maps = [] + for norm1, attn, norm2, ffn in self.layers: + _x = norm1(x) + + _x, attn_map = attn(_x, _x, _x, mask=mask) + + attn_maps.append(attn_map) + + x = _x + x + + _x = norm2(x) + + x = ffn(_x) + x + + return x, attn_maps diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/crop_randomizer.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/crop_randomizer.py new file mode 100644 index 0000000000000000000000000000000000000000..471cc95b70d526a19320ab7b047b0fca06e36671 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/crop_randomizer.py @@ -0,0 +1,291 @@ +import torch +import torch.nn as nn +import torchvision.transforms.functional as ttf +import cleandiffuser.utils.tensor_utils as tu + + +class CropRandomizer(nn.Module): + """ + Randomly sample crops at input, and then average across crop features at output. + """ + + def __init__( + self, + input_shape, + crop_height, + crop_width, + num_crops=1, + pos_enc=False, + ): + """ + Args: + input_shape (tuple, list): shape of input (not including batch dimension) + crop_height (int): crop height + crop_width (int): crop width + num_crops (int): number of random crops to take + pos_enc (bool): if True, add 2 channels to the output to encode the spatial + location of the cropped pixels in the source image + """ + super().__init__() + + assert len(input_shape) == 3 # (C, H, W) + assert crop_height < input_shape[1] + assert crop_width < input_shape[2] + + self.input_shape = input_shape + self.crop_height = crop_height + self.crop_width = crop_width + self.num_crops = num_crops + self.pos_enc = pos_enc + + def output_shape_in(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. Corresponds to + the @forward_in operation, where raw inputs (usually observation modalities) + are passed in. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + + # outputs are shape (C, CH, CW), or maybe C + 2 if using position encoding, because + # the number of crops are reshaped into the batch dimension, increasing the batch + # size from B to B * N + out_c = self.input_shape[0] + 2 if self.pos_enc else self.input_shape[0] + return [out_c, self.crop_height, self.crop_width] + + def output_shape_out(self, input_shape=None): + """ + Function to compute output shape from inputs to this module. Corresponds to + the @forward_out operation, where processed inputs (usually encoded observation + modalities) are passed in. + + Args: + input_shape (iterable of int): shape of input. Does not include batch dimension. + Some modules may not need this argument, if their output does not depend + on the size of the input, or if they assume fixed size input. + + Returns: + out_shape ([int]): list of integers corresponding to output shape + """ + + # since the forward_out operation splits [B * N, ...] -> [B, N, ...] + # and then pools to result in [B, ...], only the batch dimension changes, + # and so the other dimensions retain their shape. + return list(input_shape) + + def forward_in(self, inputs): + """ + Samples N random crops for each input in the batch, and then reshapes + inputs to [B * N, ...]. + """ + assert len(inputs.shape) >= 3 # must have at least (C, H, W) dimensions + if self.training: + # generate random crops + out, _ = sample_random_image_crops( + images=inputs, + crop_height=self.crop_height, + crop_width=self.crop_width, + num_crops=self.num_crops, + pos_enc=self.pos_enc, + ) + # [B, N, ...] -> [B * N, ...] + return tu.join_dimensions(out, 0, 1) + else: + # take center crop during eval + out = ttf.center_crop(img=inputs, output_size=( + self.crop_height, self.crop_width)) + if self.num_crops > 1: + B, C, H, W = out.shape + out = out.unsqueeze(1).expand(B, self.num_crops, C, H, W).reshape(-1, C, H, W) + # [B * N, ...] + return out + + def forward_out(self, inputs): + """ + Splits the outputs from shape [B * N, ...] -> [B, N, ...] and then average across N + to result in shape [B, ...] to make sure the network output is consistent with + what would have happened if there were no randomization. + """ + if self.num_crops <= 1: + return inputs + else: + batch_size = (inputs.shape[0] // self.num_crops) + out = tu.reshape_dimensions(inputs, begin_axis=0, end_axis=0, + target_dims=(batch_size, self.num_crops)) + return out.mean(dim=1) + + def forward(self, inputs): + return self.forward_in(inputs) + + def __repr__(self): + """Pretty print network.""" + header = '{}'.format(str(self.__class__.__name__)) + msg = header + "(input_shape={}, crop_size=[{}, {}], num_crops={})".format( + self.input_shape, self.crop_height, self.crop_width, self.num_crops) + return msg + + +def crop_image_from_indices(images, crop_indices, crop_height, crop_width): + """ + Crops images at the locations specified by @crop_indices. Crops will be + taken across all channels. + + Args: + images (torch.Tensor): batch of images of shape [..., C, H, W] + + crop_indices (torch.Tensor): batch of indices of shape [..., N, 2] where + N is the number of crops to take per image and each entry corresponds + to the pixel height and width of where to take the crop. Note that + the indices can also be of shape [..., 2] if only 1 crop should + be taken per image. Leading dimensions must be consistent with + @images argument. Each index specifies the top left of the crop. + Values must be in range [0, H - CH - 1] x [0, W - CW - 1] where + H and W are the height and width of @images and CH and CW are + @crop_height and @crop_width. + + crop_height (int): height of crop to take + + crop_width (int): width of crop to take + + Returns: + crops (torch.Tesnor): cropped images of shape [..., C, @crop_height, @crop_width] + """ + + # make sure length of input shapes is consistent + assert crop_indices.shape[-1] == 2 + ndim_im_shape = len(images.shape) + ndim_indices_shape = len(crop_indices.shape) + assert (ndim_im_shape == ndim_indices_shape + 1) or (ndim_im_shape == ndim_indices_shape + 2) + + # maybe pad so that @crop_indices is shape [..., N, 2] + is_padded = False + if ndim_im_shape == ndim_indices_shape + 2: + crop_indices = crop_indices.unsqueeze(-2) + is_padded = True + + # make sure leading dimensions between images and indices are consistent + assert images.shape[:-3] == crop_indices.shape[:-2] + + device = images.device + image_c, image_h, image_w = images.shape[-3:] + num_crops = crop_indices.shape[-2] + + # make sure @crop_indices are in valid range + assert (crop_indices[..., 0] >= 0).all().item() + assert (crop_indices[..., 0] < (image_h - crop_height)).all().item() + assert (crop_indices[..., 1] >= 0).all().item() + assert (crop_indices[..., 1] < (image_w - crop_width)).all().item() + + # convert each crop index (ch, cw) into a list of pixel indices that correspond to the entire window. + + # 2D index array with columns [0, 1, ..., CH - 1] and shape [CH, CW] + crop_ind_grid_h = torch.arange(crop_height).to(device) + crop_ind_grid_h = tu.unsqueeze_expand_at(crop_ind_grid_h, size=crop_width, dim=-1) + # 2D index array with rows [0, 1, ..., CW - 1] and shape [CH, CW] + crop_ind_grid_w = torch.arange(crop_width).to(device) + crop_ind_grid_w = tu.unsqueeze_expand_at(crop_ind_grid_w, size=crop_height, dim=0) + # combine into shape [CH, CW, 2] + crop_in_grid = torch.cat((crop_ind_grid_h.unsqueeze(-1), crop_ind_grid_w.unsqueeze(-1)), dim=-1) + + # Add above grid with the offset index of each sampled crop to get 2d indices for each crop. + # After broadcasting, this will be shape [..., N, CH, CW, 2] and each crop has a [CH, CW, 2] + # shape array that tells us which pixels from the corresponding source image to grab. + grid_reshape = [1] * len(crop_indices.shape[:-1]) + [crop_height, crop_width, 2] + all_crop_inds = crop_indices.unsqueeze(-2).unsqueeze(-2) + crop_in_grid.reshape(grid_reshape) + + # For using @torch.gather, convert to flat indices from 2D indices, and also + # repeat across the channel dimension. To get flat index of each pixel to grab for + # each sampled crop, we just use the mapping: ind = h_ind * @image_w + w_ind + all_crop_inds = all_crop_inds[..., 0] * image_w + all_crop_inds[..., 1] # shape [..., N, CH, CW] + all_crop_inds = tu.unsqueeze_expand_at(all_crop_inds, size=image_c, dim=-3) # shape [..., N, C, CH, CW] + all_crop_inds = tu.flatten(all_crop_inds, begin_axis=-2) # shape [..., N, C, CH * CW] + + # Repeat and flatten the source images -> [..., N, C, H * W] and then use gather to index with crop pixel inds + images_to_crop = tu.unsqueeze_expand_at(images, size=num_crops, dim=-4) + images_to_crop = tu.flatten(images_to_crop, begin_axis=-2) + crops = torch.gather(images_to_crop, dim=-1, index=all_crop_inds) + # [..., N, C, CH * CW] -> [..., N, C, CH, CW] + reshape_axis = len(crops.shape) - 1 + crops = tu.reshape_dimensions(crops, begin_axis=reshape_axis, end_axis=reshape_axis, + target_dims=(crop_height, crop_width)) + + if is_padded: + # undo padding -> [..., C, CH, CW] + crops = crops.squeeze(-4) + return crops + + +def sample_random_image_crops(images, crop_height, crop_width, num_crops, pos_enc=False): + """ + For each image, randomly sample @num_crops crops of size (@crop_height, @crop_width), from + @images. + + Args: + images (torch.Tensor): batch of images of shape [..., C, H, W] + + crop_height (int): height of crop to take + + crop_width (int): width of crop to take + + num_crops (n): number of crops to sample + + pos_enc (bool): if True, also add 2 channels to the outputs that gives a spatial + encoding of the original source pixel locations. This means that the + output crops will contain information about where in the source image + it was sampled from. + + Returns: + crops (torch.Tensor): crops of shape (..., @num_crops, C, @crop_height, @crop_width) + if @pos_enc is False, otherwise (..., @num_crops, C + 2, @crop_height, @crop_width) + + crop_inds (torch.Tensor): sampled crop indices of shape (..., N, 2) + """ + device = images.device + + # maybe add 2 channels of spatial encoding to the source image + source_im = images + if pos_enc: + # spatial encoding [y, x] in [0, 1] + h, w = source_im.shape[-2:] + pos_y, pos_x = torch.meshgrid(torch.arange(h), torch.arange(w)) + pos_y = pos_y.float().to(device) / float(h) + pos_x = pos_x.float().to(device) / float(w) + position_enc = torch.stack((pos_y, pos_x)) # shape [C, H, W] + + # unsqueeze and expand to match leading dimensions -> shape [..., C, H, W] + leading_shape = source_im.shape[:-3] + position_enc = position_enc[(None,) * len(leading_shape)] + position_enc = position_enc.expand(*leading_shape, -1, -1, -1) + + # concat across channel dimension with input + source_im = torch.cat((source_im, position_enc), dim=-3) + + # make sure sample boundaries ensure crops are fully within the images + image_c, image_h, image_w = source_im.shape[-3:] + max_sample_h = image_h - crop_height + max_sample_w = image_w - crop_width + + # Sample crop locations for all tensor dimensions up to the last 3, which are [C, H, W]. + # Each gets @num_crops samples - typically this will just be the batch dimension (B), so + # we will sample [B, N] indices, but this supports having more than one leading dimension, + # or possibly no leading dimension. + # + # Trick: sample in [0, 1) with rand, then re-scale to [0, M) and convert to long to get sampled ints + crop_inds_h = (max_sample_h * torch.rand(*source_im.shape[:-3], num_crops).to(device)).long() + crop_inds_w = (max_sample_w * torch.rand(*source_im.shape[:-3], num_crops).to(device)).long() + crop_inds = torch.cat((crop_inds_h.unsqueeze(-1), crop_inds_w.unsqueeze(-1)), dim=-1) # shape [..., N, 2] + + crops = crop_image_from_indices( + images=source_im, + crop_indices=crop_inds, + crop_height=crop_height, + crop_width=crop_width, + ) + + return crops, crop_inds diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/eval_utils.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/eval_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1a6844bd29f7f04bbf41df89904e929c369e69c5 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/eval_utils.py @@ -0,0 +1,17 @@ +"""Evaluation compatibility without probing (and killing) async Gym workers.""" + + +def reset_seeded_env(env, seeds=None): + if seeds is None: + result = env.reset() + elif callable(getattr(env, "seed", None)): + # D4RL Kitchen exposes seed() but reset() does not accept seed=. With + # AsyncVectorEnv, trying reset(seed=...) first kills the workers, so a + # try/except fallback is too late. + env.seed(seeds) + result = env.reset() + else: + result = env.reset(seed=seeds) + if isinstance(result, tuple) and len(result) == 2: + return result[0] + return result diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/iql.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/iql.py new file mode 100644 index 0000000000000000000000000000000000000000..90ae4d7a1126dc62455317934ed28a159b660ade --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/iql.py @@ -0,0 +1,95 @@ +from copy import deepcopy + +import torch +import torch.nn as nn + + +class TwinQ(nn.Module): + def __init__(self, obs_dim, act_dim, hidden_dim: int = 256): + super().__init__() + self.Q1 = nn.Sequential( + nn.Linear(obs_dim + act_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.Mish(), + nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.Mish(), + nn.Linear(hidden_dim, 1)) + self.Q2 = nn.Sequential( + nn.Linear(obs_dim + act_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.Mish(), + nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.Mish(), + nn.Linear(hidden_dim, 1)) + + def both(self, obs, act): + q1, q2 = self.Q1(torch.cat([obs, act], -1)), self.Q2(torch.cat([obs, act], -1)) + return q1, q2 + + def forward(self, obs, act): + return torch.min(*self.both(obs, act)) + + +class V(nn.Module): + def __init__(self, obs_dim, hidden_dim: int = 256): + super().__init__() + self.V = nn.Sequential( + nn.Linear(obs_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.Mish(), + nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.Mish(), + nn.Linear(hidden_dim, 1)) + + def forward(self, obs): + v = self.V(obs) + return v + + +class IQL(nn.Module): + """ Simple Implicit Q-Learning (IQL) pytorch implementation. + + Args: + obs_dim: int, observation space dimension. + act_dim: int, action space dimension. + tau: float, quantile level. Default is 0.7. + discount: float, discount factor. Default is 0.99. + hidden_dim: int, hidden dimension. Default is 256. + + Example: + >>> iql = IQL(...) + >>> batch = ... + >>> obs, act, rew, obs_next, done = batch + >>> loss_v = iql.update_V(obs, act) + >>> loss_q = iql.update_Q(obs, act, rew, obs_next, done) + >>> iql.update_target() + """ + def __init__(self, obs_dim: int, act_dim: int, tau: float = 0.7, discount: float = 0.99, hidden_dim: int = 256): + super().__init__() + self.iql_tau, self.discount = tau, discount + self.Q = TwinQ(obs_dim, act_dim, hidden_dim) + self.Q_targ = deepcopy(self.Q).requires_grad_(False).eval() + self.V = V(obs_dim, hidden_dim) + self.optimV = torch.optim.Adam(self.V.parameters(), lr=3e-4) + self.optimQ = torch.optim.Adam(self.Q.parameters(), lr=3e-4) + + def update_target(self, mu=0.995): + for p, p_targ in zip(self.Q.parameters(), self.Q_targ.parameters()): + p_targ.data = mu * p_targ.data + (1 - mu) * p.data + + def update_V(self, obs, act): + q = self.Q_targ(obs, act) + v = self.V(obs) + loss = (torch.abs(self.iql_tau - ((q - v) < 0).float()) * (q - v) ** 2).mean() + self.optimV.zero_grad() + loss.backward() + self.optimV.step() + return loss.item() + + def update_Q(self, obs, act, rew, obs_next, done): + with torch.no_grad(): + td_target = rew + self.discount * (1 - done) * self.V(obs_next) + q1, q2 = self.Q.both(obs, act) + loss = ((q1 - td_target) ** 2 + (q2 - td_target) ** 2).mean() + self.optimQ.zero_grad() + loss.backward() + self.optimQ.step() + self.update_target() + return loss.item() + + def save(self, path): + torch.save(self.state_dict(), path) + + def load(self, path, device): + self.load_state_dict(torch.load(path, map_location=device)) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/normalizers.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/normalizers.py new file mode 100644 index 0000000000000000000000000000000000000000..50987d941cbb5fef0fb7fdc24449352a63fad9d3 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/normalizers.py @@ -0,0 +1,127 @@ +from typing import Optional + +import numpy as np + +from .utils import at_least_ndim + + +class EmptyNormalizer: + """ Empty Normalizer + + Does nothing to the input data. + """ + + def normalize(self, x: np.ndarray): + return x + + def unnormalize(self, x: np.ndarray): + return x + + +class GaussianNormalizer(EmptyNormalizer): + """ Gaussian Normalizer + + Normalizes data to have zero mean and unit variance. + For those dimensions with zero variance, the normalized value will be zero. + + Args: + X: np.ndarray, + dataset with shape (..., *x_shape) + start_dim: int, + the dimension to start normalization from, Default: -1 + + Examples: + >>> x_dataset = np.random.randn(100000, 3, 10) + + >>> normalizer = GaussianNormalizer(x_dataset, 1) + >>> x = np.random.randn(1, 3, 10) + >>> norm_x = normalizer.normalize(x) + >>> unnorm_x = normalizer.unnormalize(norm_x) + + >>> normalizer = GaussianNormalizer(x_dataset, 2) + >>> x = np.random.randn(1, 10) + >>> norm_x = normalizer.normalize(x) + >>> unnorm_x = normalizer.unnormalize(norm_x) + """ + + def __init__(self, X: np.ndarray, start_dim: int = -1): + total_dims = X.ndim + if start_dim < 0: + start_dim = total_dims + start_dim + + axes = tuple(range(start_dim)) + + self.mean = np.mean(X, axis=axes) + self.std = np.std(X, axis=axes) + self.std[self.std == 0] = 1. + + def normalize(self, x: np.ndarray): + ndim = x.ndim + return (x - at_least_ndim(self.mean, ndim, 1)) / at_least_ndim(self.std, ndim, 1) + + def unnormalize(self, x: np.ndarray): + ndim = x.ndim + return x * at_least_ndim(self.std, ndim, 1) + at_least_ndim(self.mean, ndim, 1) + + +class MinMaxNormalizer(EmptyNormalizer): + """ MinMax Normalizer + + Normalizes data from range [min, max] to [-1, 1]. + For those dimensions with zero range, the normalized value will be zero. + + Args: + X: np.ndarray, + dataset with shape (..., *x_shape) + start_dim: int, + the dimension to start normalization from, Default: -1 + X_max: Optional[np.ndarray], + Maximum value for each dimension. If None, it will be calculated from X. Default: None + X_min: Optional[np.ndarray], + Minimum value for each dimension. If None, it will be calculated from X. Default: None + + Examples: + >>> x_dataset = np.random.randn(100000, 3, 10) + + >>> x_min = np.random.randn(3, 10) + >>> normalizer = MinMaxNormalizer(x_dataset, 1, X_min=x_min) + >>> x = np.random.randn(1, 3, 10) + >>> norm_x = normalizer.normalize(x) + >>> unnorm_x = normalizer.unnormalize(norm_x) + + >>> x_max = np.random.randn(10) + >>> normalizer = MinMaxNormalizer(x_dataset, 2, X_max=x_max) + >>> x = np.random.randn(1, 10) + >>> norm_x = normalizer.normalize(x) + >>> unnorm_x = normalizer.unnormalize(norm_x) + """ + + def __init__( + self, X: np.ndarray, start_dim: int = -1, + X_max: Optional[np.ndarray] = None, X_min: Optional[np.ndarray] = None): + total_dims = X.ndim + if start_dim < 0: + start_dim = total_dims + start_dim + + axes = tuple(range(start_dim)) + + self.max = np.max(X, axis=axes) if X_max is None else X_max + self.min = np.min(X, axis=axes) if X_min is None else X_min + self.mask = np.ones_like(self.max) + self.range = self.max - self.min + self.mask[self.max == self.min] = 0. + self.range[self.range == 0] = 1. + + def normalize(self, x: np.ndarray): + ndim = x.ndim + x = (x - at_least_ndim(self.min, ndim, 1)) / at_least_ndim(self.range, ndim, 1) + x = x * 2 - 1 + x = x * at_least_ndim(self.mask, ndim, 1) + return x + + def unnormalize(self, x: np.ndarray): + ndim = x.ndim + x = (x + 1) / 2 + x = x * at_least_ndim(self.mask, ndim, 1) + x = x * at_least_ndim(self.range, ndim, 1) + at_least_ndim(self.min, ndim, 1) + return x diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/selector_normalizer.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/selector_normalizer.py new file mode 100644 index 0000000000000000000000000000000000000000..13dad14fdd264a71fd20ec84642574bd11781b86 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/selector_normalizer.py @@ -0,0 +1,37 @@ +"""Serializable Gaussian observation statistics for frozen IQL selectors.""" +import hashlib + +import numpy as np +import torch + +from cleandiffuser.utils.normalizers import GaussianNormalizer + + +def normalizer_state(normalizer): + return {"normalizer_mean": torch.as_tensor(normalizer.mean).cpu().clone(), + "normalizer_std": torch.as_tensor(normalizer.std).cpu().clone()} + + +def checkpoint_normalizer(checkpoint, obs_dim): + normalizer = GaussianNormalizer.__new__(GaussianNormalizer) + for field in ("mean", "std"): + value = checkpoint["normalizer_" + field] + if torch.is_tensor(value): + value = value.detach().cpu().numpy() + value = np.asarray(value, dtype=np.float32) + if value.shape != (obs_dim,) or not np.isfinite(value).all(): + raise ValueError(f"invalid IQL normalizer {field}: expected finite ({obs_dim},) values") + if field == "std" and not (value > 0).all(): + raise ValueError("IQL normalizer std must be strictly positive") + setattr(normalizer, field, value.copy()) + return normalizer + + +def normalizer_sha256(normalizer): + digest = hashlib.sha256() + for field in ("mean", "std"): + value = np.asarray(getattr(normalizer, field), dtype=np.float32) + digest.update(field.encode()) + digest.update(str(value.shape).encode()) + digest.update(value.tobytes()) + return digest.hexdigest() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/tensor_utils.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/tensor_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2900dfd05afab400c8b512c32bf2723917bb25b8 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/tensor_utils.py @@ -0,0 +1,960 @@ +""" +A collection of utilities for working with nested tensor structures consisting +of numpy arrays and torch tensors. +""" +import collections +import numpy as np +import torch + + +def recursive_dict_list_tuple_apply(x, type_func_dict): + """ + Recursively apply functions to a nested dictionary or list or tuple, given a dictionary of + {data_type: function_to_apply}. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + type_func_dict (dict): a mapping from data types to the functions to be + applied for each data type. + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + assert(list not in type_func_dict) + assert(tuple not in type_func_dict) + assert(dict not in type_func_dict) + + if isinstance(x, (dict, collections.OrderedDict)): + new_x = collections.OrderedDict() if isinstance(x, collections.OrderedDict) else dict() + for k, v in x.items(): + new_x[k] = recursive_dict_list_tuple_apply(v, type_func_dict) + return new_x + elif isinstance(x, (list, tuple)): + ret = [recursive_dict_list_tuple_apply(v, type_func_dict) for v in x] + if isinstance(x, tuple): + ret = tuple(ret) + return ret + else: + for t, f in type_func_dict.items(): + if isinstance(x, t): + return f(x) + else: + raise NotImplementedError( + 'Cannot handle data type %s' % str(type(x))) + + +def map_tensor(x, func): + """ + Apply function @func to torch.Tensor objects in a nested dictionary or + list or tuple. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + func (function): function to apply to each tensor + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: func, + type(None): lambda x: x, + } + ) + + +def map_ndarray(x, func): + """ + Apply function @func to np.ndarray objects in a nested dictionary or + list or tuple. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + func (function): function to apply to each array + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + np.ndarray: func, + type(None): lambda x: x, + } + ) + + +def map_tensor_ndarray(x, tensor_func, ndarray_func): + """ + Apply function @tensor_func to torch.Tensor objects and @ndarray_func to + np.ndarray objects in a nested dictionary or list or tuple. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + tensor_func (function): function to apply to each tensor + ndarray_Func (function): function to apply to each array + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: tensor_func, + np.ndarray: ndarray_func, + type(None): lambda x: x, + } + ) + + +def clone(x): + """ + Clones all torch tensors and numpy arrays in nested dictionary or list + or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.clone(), + np.ndarray: lambda x: x.copy(), + type(None): lambda x: x, + } + ) + + +def detach(x): + """ + Detaches all torch tensors in nested dictionary or list + or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.detach(), + } + ) + + +def to_batch(x): + """ + Introduces a leading batch dimension of 1 for all torch tensors and numpy + arrays in nested dictionary or list or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x[None, ...], + np.ndarray: lambda x: x[None, ...], + type(None): lambda x: x, + } + ) + + +def to_sequence(x): + """ + Introduces a time dimension of 1 at dimension 1 for all torch tensors and numpy + arrays in nested dictionary or list or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x[:, None, ...], + np.ndarray: lambda x: x[:, None, ...], + type(None): lambda x: x, + } + ) + + +def index_at_time(x, ind): + """ + Indexes all torch tensors and numpy arrays in dimension 1 with index @ind in + nested dictionary or list or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + ind (int): index + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x[:, ind, ...], + np.ndarray: lambda x: x[:, ind, ...], + type(None): lambda x: x, + } + ) + + +def unsqueeze(x, dim): + """ + Adds dimension of size 1 at dimension @dim in all torch tensors and numpy arrays + in nested dictionary or list or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + dim (int): dimension + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.unsqueeze(dim=dim), + np.ndarray: lambda x: np.expand_dims(x, axis=dim), + type(None): lambda x: x, + } + ) + + +def contiguous(x): + """ + Makes all torch tensors and numpy arrays contiguous in nested dictionary or + list or tuple and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.contiguous(), + np.ndarray: lambda x: np.ascontiguousarray(x), + type(None): lambda x: x, + } + ) + + +def to_device(x, device): + """ + Sends all torch tensors in nested dictionary or list or tuple to device + @device, and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + device (torch.Device): device to send tensors to + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x, d=device: x.to(d), + type(None): lambda x: x, + } + ) + + +def to_tensor(x): + """ + Converts all numpy arrays in nested dictionary or list or tuple to + torch tensors (and leaves existing torch Tensors as-is), and returns + a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x, + np.ndarray: lambda x: torch.tensor(x), + type(None): lambda x: x, + } + ) + + +def to_numpy(x): + """ + Converts all torch tensors in nested dictionary or list or tuple to + numpy (and leaves existing numpy arrays as-is), and returns + a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + def f(tensor): + if tensor.is_cuda: + return tensor.detach().cpu().numpy() + else: + return tensor.detach().numpy() + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: f, + np.ndarray: lambda x: x, + type(None): lambda x: x, + } + ) + + +def to_list(x): + """ + Converts all torch tensors and numpy arrays in nested dictionary or list + or tuple to a list, and returns a new nested structure. Useful for + json encoding. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + def f(tensor): + if tensor.is_cuda: + return tensor.detach().cpu().numpy().tolist() + else: + return tensor.detach().numpy().tolist() + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: f, + np.ndarray: lambda x: x.tolist(), + type(None): lambda x: x, + } + ) + + +def to_float(x): + """ + Converts all torch tensors and numpy arrays in nested dictionary or list + or tuple to float type entries, and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.float(), + np.ndarray: lambda x: x.astype(np.float32), + type(None): lambda x: x, + } + ) + + +def to_uint8(x): + """ + Converts all torch tensors and numpy arrays in nested dictionary or list + or tuple to uint8 type entries, and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.byte(), + np.ndarray: lambda x: x.astype(np.uint8), + type(None): lambda x: x, + } + ) + + +def to_torch(x, device): + """ + Converts all numpy arrays and torch tensors in nested dictionary or list or tuple to + torch tensors on device @device and returns a new nested structure. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + device (torch.Device): device to send tensors to + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return to_device(to_float(to_tensor(x)), device) + + +def to_one_hot_single(tensor, num_class): + """ + Convert tensor to one-hot representation, assuming a certain number of total class labels. + + Args: + tensor (torch.Tensor): tensor containing integer labels + num_class (int): number of classes + + Returns: + x (torch.Tensor): tensor containing one-hot representation of labels + """ + x = torch.zeros(tensor.size() + (num_class,)).to(tensor.device) + x.scatter_(-1, tensor.unsqueeze(-1), 1) + return x + + +def to_one_hot(tensor, num_class): + """ + Convert all tensors in nested dictionary or list or tuple to one-hot representation, + assuming a certain number of total class labels. + + Args: + tensor (dict or list or tuple): a possibly nested dictionary or list or tuple + num_class (int): number of classes + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return map_tensor(tensor, func=lambda x, nc=num_class: to_one_hot_single(x, nc)) + + +def flatten_single(x, begin_axis=1): + """ + Flatten a tensor in all dimensions from @begin_axis onwards. + + Args: + x (torch.Tensor): tensor to flatten + begin_axis (int): which axis to flatten from + + Returns: + y (torch.Tensor): flattened tensor + """ + fixed_size = x.size()[:begin_axis] + _s = list(fixed_size) + [-1] + return x.reshape(*_s) + + +def flatten(x, begin_axis=1): + """ + Flatten all tensors in nested dictionary or list or tuple, from @begin_axis onwards. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + begin_axis (int): which axis to flatten from + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x, b=begin_axis: flatten_single(x, begin_axis=b), + } + ) + + +def reshape_dimensions_single(x, begin_axis, end_axis, target_dims): + """ + Reshape selected dimensions in a tensor to a target dimension. + + Args: + x (torch.Tensor): tensor to reshape + begin_axis (int): begin dimension + end_axis (int): end dimension + target_dims (tuple or list): target shape for the range of dimensions + (@begin_axis, @end_axis) + + Returns: + y (torch.Tensor): reshaped tensor + """ + assert(begin_axis <= end_axis) + assert(begin_axis >= 0) + assert(end_axis < len(x.shape)) + assert(isinstance(target_dims, (tuple, list))) + s = x.shape + final_s = [] + for i in range(len(s)): + if i == begin_axis: + final_s.extend(target_dims) + elif i < begin_axis or i > end_axis: + final_s.append(s[i]) + return x.reshape(*final_s) + + +def reshape_dimensions(x, begin_axis, end_axis, target_dims): + """ + Reshape selected dimensions for all tensors in nested dictionary or list or tuple + to a target dimension. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + begin_axis (int): begin dimension + end_axis (int): end dimension + target_dims (tuple or list): target shape for the range of dimensions + (@begin_axis, @end_axis) + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x, b=begin_axis, e=end_axis, t=target_dims: reshape_dimensions_single( + x, begin_axis=b, end_axis=e, target_dims=t), + np.ndarray: lambda x, b=begin_axis, e=end_axis, t=target_dims: reshape_dimensions_single( + x, begin_axis=b, end_axis=e, target_dims=t), + type(None): lambda x: x, + } + ) + + +def join_dimensions(x, begin_axis, end_axis): + """ + Joins all dimensions between dimensions (@begin_axis, @end_axis) into a flat dimension, for + all tensors in nested dictionary or list or tuple. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + begin_axis (int): begin dimension + end_axis (int): end dimension + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x, b=begin_axis, e=end_axis: reshape_dimensions_single( + x, begin_axis=b, end_axis=e, target_dims=[-1]), + np.ndarray: lambda x, b=begin_axis, e=end_axis: reshape_dimensions_single( + x, begin_axis=b, end_axis=e, target_dims=[-1]), + type(None): lambda x: x, + } + ) + + +def expand_at_single(x, size, dim): + """ + Expand a tensor at a single dimension @dim by @size + + Args: + x (torch.Tensor): input tensor + size (int): size to expand + dim (int): dimension to expand + + Returns: + y (torch.Tensor): expanded tensor + """ + assert dim < x.ndimension() + assert x.shape[dim] == 1 + expand_dims = [-1] * x.ndimension() + expand_dims[dim] = size + return x.expand(*expand_dims) + + +def expand_at(x, size, dim): + """ + Expand all tensors in nested dictionary or list or tuple at a single + dimension @dim by @size. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + size (int): size to expand + dim (int): dimension to expand + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return map_tensor(x, lambda t, s=size, d=dim: expand_at_single(t, s, d)) + + +def unsqueeze_expand_at(x, size, dim): + """ + Unsqueeze and expand a tensor at a dimension @dim by @size. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + size (int): size to expand + dim (int): dimension to unsqueeze and expand + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + x = unsqueeze(x, dim) + return expand_at(x, size, dim) + + +def repeat_by_expand_at(x, repeats, dim): + """ + Repeat a dimension by combining expand and reshape operations. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + repeats (int): number of times to repeat the target dimension + dim (int): dimension to repeat on + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + x = unsqueeze_expand_at(x, repeats, dim + 1) + return join_dimensions(x, dim, dim + 1) + + +def named_reduce_single(x, reduction, dim): + """ + Reduce tensor at a dimension by named reduction functions. + + Args: + x (torch.Tensor): tensor to be reduced + reduction (str): one of ["sum", "max", "mean", "flatten"] + dim (int): dimension to be reduced (or begin axis for flatten) + + Returns: + y (torch.Tensor): reduced tensor + """ + assert x.ndimension() > dim + assert reduction in ["sum", "max", "mean", "flatten"] + if reduction == "flatten": + x = flatten(x, begin_axis=dim) + elif reduction == "max": + x = torch.max(x, dim=dim)[0] # [B, D] + elif reduction == "sum": + x = torch.sum(x, dim=dim) + else: + x = torch.mean(x, dim=dim) + return x + + +def named_reduce(x, reduction, dim): + """ + Reduces all tensors in nested dictionary or list or tuple at a dimension + using a named reduction function. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + reduction (str): one of ["sum", "max", "mean", "flatten"] + dim (int): dimension to be reduced (or begin axis for flatten) + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return map_tensor(x, func=lambda t, r=reduction, d=dim: named_reduce_single(t, r, d)) + + +def gather_along_dim_with_dim_single(x, target_dim, source_dim, indices): + """ + This function indexes out a target dimension of a tensor in a structured way, + by allowing a different value to be selected for each member of a flat index + tensor (@indices) corresponding to a source dimension. This can be interpreted + as moving along the source dimension, using the corresponding index value + in @indices to select values for all other dimensions outside of the + source and target dimensions. A common use case is to gather values + in target dimension 1 for each batch member (target dimension 0). + + Args: + x (torch.Tensor): tensor to gather values for + target_dim (int): dimension to gather values along + source_dim (int): dimension to hold constant and use for gathering values + from the other dimensions + indices (torch.Tensor): flat index tensor with same shape as tensor @x along + @source_dim + + Returns: + y (torch.Tensor): gathered tensor, with dimension @target_dim indexed out + """ + assert len(indices.shape) == 1 + assert x.shape[source_dim] == indices.shape[0] + + # unsqueeze in all dimensions except the source dimension + new_shape = [1] * x.ndimension() + new_shape[source_dim] = -1 + indices = indices.reshape(*new_shape) + + # repeat in all dimensions - but preserve shape of source dimension, + # and make sure target_dimension has singleton dimension + expand_shape = list(x.shape) + expand_shape[source_dim] = -1 + expand_shape[target_dim] = 1 + indices = indices.expand(*expand_shape) + + out = x.gather(dim=target_dim, index=indices) + return out.squeeze(target_dim) + + +def gather_along_dim_with_dim(x, target_dim, source_dim, indices): + """ + Apply @gather_along_dim_with_dim_single to all tensors in a nested + dictionary or list or tuple. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + target_dim (int): dimension to gather values along + source_dim (int): dimension to hold constant and use for gathering values + from the other dimensions + indices (torch.Tensor): flat index tensor with same shape as tensor @x along + @source_dim + + Returns: + y (dict or list or tuple): new nested dict-list-tuple + """ + return map_tensor(x, + lambda y, t=target_dim, s=source_dim, i=indices: gather_along_dim_with_dim_single(y, t, s, i)) + + +def gather_sequence_single(seq, indices): + """ + Given a tensor with leading dimensions [B, T, ...], gather an element from each sequence in + the batch given an index for each sequence. + + Args: + seq (torch.Tensor): tensor with leading dimensions [B, T, ...] + indices (torch.Tensor): tensor indices of shape [B] + + Return: + y (torch.Tensor): indexed tensor of shape [B, ....] + """ + return gather_along_dim_with_dim_single(seq, target_dim=1, source_dim=0, indices=indices) + + +def gather_sequence(seq, indices): + """ + Given a nested dictionary or list or tuple, gathers an element from each sequence of the batch + for tensors with leading dimensions [B, T, ...]. + + Args: + seq (dict or list or tuple): a possibly nested dictionary or list or tuple with tensors + of leading dimensions [B, T, ...] + indices (torch.Tensor): tensor indices of shape [B] + + Returns: + y (dict or list or tuple): new nested dict-list-tuple with tensors of shape [B, ...] + """ + return gather_along_dim_with_dim(seq, target_dim=1, source_dim=0, indices=indices) + + +def pad_sequence_single(seq, padding, batched=False, pad_same=True, pad_values=None): + """ + Pad input tensor or array @seq in the time dimension (dimension 1). + + Args: + seq (np.ndarray or torch.Tensor): sequence to be padded + padding (tuple): begin and end padding, e.g. [1, 1] pads both begin and end of the sequence by 1 + batched (bool): if sequence has the batch dimension + pad_same (bool): if pad by duplicating + pad_values (scalar or (ndarray, Tensor)): values to be padded if not pad_same + + Returns: + padded sequence (np.ndarray or torch.Tensor) + """ + assert isinstance(seq, (np.ndarray, torch.Tensor)) + assert pad_same or pad_values is not None + if pad_values is not None: + assert isinstance(pad_values, float) + repeat_func = np.repeat if isinstance(seq, np.ndarray) else torch.repeat_interleave + concat_func = np.concatenate if isinstance(seq, np.ndarray) else torch.cat + ones_like_func = np.ones_like if isinstance(seq, np.ndarray) else torch.ones_like + seq_dim = 1 if batched else 0 + + begin_pad = [] + end_pad = [] + + if padding[0] > 0: + pad = seq[[0]] if pad_same else ones_like_func(seq[[0]]) * pad_values + begin_pad.append(repeat_func(pad, padding[0], seq_dim)) + if padding[1] > 0: + pad = seq[[-1]] if pad_same else ones_like_func(seq[[-1]]) * pad_values + end_pad.append(repeat_func(pad, padding[1], seq_dim)) + + return concat_func(begin_pad + [seq] + end_pad, seq_dim) + + +def pad_sequence(seq, padding, batched=False, pad_same=True, pad_values=None): + """ + Pad a nested dictionary or list or tuple of sequence tensors in the time dimension (dimension 1). + + Args: + seq (dict or list or tuple): a possibly nested dictionary or list or tuple with tensors + of leading dimensions [B, T, ...] + padding (tuple): begin and end padding, e.g. [1, 1] pads both begin and end of the sequence by 1 + batched (bool): if sequence has the batch dimension + pad_same (bool): if pad by duplicating + pad_values (scalar or (ndarray, Tensor)): values to be padded if not pad_same + + Returns: + padded sequence (dict or list or tuple) + """ + return recursive_dict_list_tuple_apply( + seq, + { + torch.Tensor: lambda x, p=padding, b=batched, ps=pad_same, pv=pad_values: + pad_sequence_single(x, p, b, ps, pv), + np.ndarray: lambda x, p=padding, b=batched, ps=pad_same, pv=pad_values: + pad_sequence_single(x, p, b, ps, pv), + type(None): lambda x: x, + } + ) + + +def assert_size_at_dim_single(x, size, dim, msg): + """ + Ensure that array or tensor @x has size @size in dim @dim. + + Args: + x (np.ndarray or torch.Tensor): input array or tensor + size (int): size that tensors should have at @dim + dim (int): dimension to check + msg (str): text to display if assertion fails + """ + assert x.shape[dim] == size, msg + + +def assert_size_at_dim(x, size, dim, msg): + """ + Ensure that arrays and tensors in nested dictionary or list or tuple have + size @size in dim @dim. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + size (int): size that tensors should have at @dim + dim (int): dimension to check + """ + map_tensor(x, lambda t, s=size, d=dim, m=msg: assert_size_at_dim_single(t, s, d, m)) + + +def get_shape(x): + """ + Get all shapes of arrays and tensors in nested dictionary or list or tuple. + + Args: + x (dict or list or tuple): a possibly nested dictionary or list or tuple + + Returns: + y (dict or list or tuple): new nested dict-list-tuple that contains each array or + tensor's shape + """ + return recursive_dict_list_tuple_apply( + x, + { + torch.Tensor: lambda x: x.shape, + np.ndarray: lambda x: x.shape, + type(None): lambda x: x, + } + ) + + +def list_of_flat_dict_to_dict_of_list(list_of_dict): + """ + Helper function to go from a list of flat dictionaries to a dictionary of lists. + By "flat" we mean that none of the values are dictionaries, but are numpy arrays, + floats, etc. + + Args: + list_of_dict (list): list of flat dictionaries + + Returns: + dict_of_list (dict): dictionary of lists + """ + assert isinstance(list_of_dict, list) + dic = collections.OrderedDict() + for i in range(len(list_of_dict)): + for k in list_of_dict[i]: + if k not in dic: + dic[k] = [] + dic[k].append(list_of_dict[i][k]) + return dic + + +def flatten_nested_dict_list(d, parent_key='', sep='_', item_key=''): + """ + Flatten a nested dict or list to a list. + + For example, given a dict + { + a: 1 + b: { + c: 2 + } + c: 3 + } + + the function would return [(a, 1), (b_c, 2), (c, 3)] + + Args: + d (dict, list): a nested dict or list to be flattened + parent_key (str): recursion helper + sep (str): separator for nesting keys + item_key (str): recursion helper + Returns: + list: a list of (key, value) tuples + """ + items = [] + if isinstance(d, (tuple, list)): + new_key = parent_key + sep + item_key if len(parent_key) > 0 else item_key + for i, v in enumerate(d): + items.extend(flatten_nested_dict_list(v, new_key, sep=sep, item_key=str(i))) + return items + elif isinstance(d, dict): + new_key = parent_key + sep + item_key if len(parent_key) > 0 else item_key + for k, v in d.items(): + assert isinstance(k, str) + items.extend(flatten_nested_dict_list(v, new_key, sep=sep, item_key=k)) + return items + else: + new_key = parent_key + sep + item_key if len(parent_key) > 0 else item_key + return [(new_key, d)] + + +def time_distributed(inputs, op, activation=None, inputs_as_kwargs=False, inputs_as_args=False, **kwargs): + """ + Apply function @op to all tensors in nested dictionary or list or tuple @inputs in both the + batch (B) and time (T) dimension, where the tensors are expected to have shape [B, T, ...]. + Will do this by reshaping tensors to [B * T, ...], passing through the op, and then reshaping + outputs to [B, T, ...]. + + Args: + inputs (list or tuple or dict): a possibly nested dictionary or list or tuple with tensors + of leading dimensions [B, T, ...] + op: a layer op that accepts inputs + activation: activation to apply at the output + inputs_as_kwargs (bool): whether to feed input as a kwargs dict to the op + inputs_as_args (bool) whether to feed input as a args list to the op + kwargs (dict): other kwargs to supply to the op + + Returns: + outputs (dict or list or tuple): new nested dict-list-tuple with tensors of leading dimension [B, T]. + """ + batch_size, seq_len = flatten_nested_dict_list(inputs)[0][1].shape[:2] + inputs = join_dimensions(inputs, 0, 1) + if inputs_as_kwargs: + outputs = op(**inputs, **kwargs) + elif inputs_as_args: + outputs = op(*inputs, **kwargs) + else: + outputs = op(inputs, **kwargs) + + if activation is not None: + outputs = map_tensor(outputs, activation) + outputs = reshape_dimensions(outputs, begin_axis=0, end_axis=0, target_dims=(batch_size, seq_len)) + return outputs \ No newline at end of file diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/typings.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/typings.py new file mode 100644 index 0000000000000000000000000000000000000000..baf93c993ce1c07108e36dbbdbcd4d5d36c69af3 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/typings.py @@ -0,0 +1,4 @@ +from typing import Dict, Union +import torch + +TensorDict = Dict[str, Union['TensorDict', torch.Tensor]] diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/utils.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6cd9a804cebed3c2d4f6d987fb54a171519f3398 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/cleandiffuser/utils/utils.py @@ -0,0 +1,483 @@ +import math +import os +import random +from typing import Union, Dict, Callable + +import numpy as np +import torch +import torch.nn as nn + + +def set_seed(seed: int): + """ Set seed for reproducibility """ + random.seed(seed) + os.environ["PYTHONHASHSEED"] = str(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def at_least_ndim(x: Union[np.ndarray, torch.Tensor, int, float], ndim: int, pad: int = 0): + """ Add dimensions to the input tensor to make it at least ndim-dimensional. + + Args: + x: Union[np.ndarray, torch.Tensor, int, float], input tensor + ndim: int, minimum number of dimensions + pad: int, padding direction. `0`: pad in the last dimension, `1`: pad in the first dimension + + Returns: + Any of these 2 options + + - np.ndarray or torch.Tensor: reshaped tensor + - int or float: input value + + Examples: + >>> x = np.random.rand(3, 4) + >>> at_least_ndim(x, 3, 0).shape + (3, 4, 1) + >>> x = torch.randn(3, 4) + >>> at_least_ndim(x, 4, 1).shape + (1, 1, 3, 4) + >>> x = 1 + >>> at_least_ndim(x, 3) + 1 + """ + if isinstance(x, np.ndarray): + if ndim > x.ndim: + if pad == 0: + return np.reshape(x, x.shape + (1,) * (ndim - x.ndim)) + else: + return np.reshape(x, (1,) * (ndim - x.ndim) + x.shape) + else: + return x + elif isinstance(x, torch.Tensor): + if ndim > x.ndim: + if pad == 0: + return torch.reshape(x, x.shape + (1,) * (ndim - x.ndim)) + else: + return torch.reshape(x, (1,) * (ndim - x.ndim) + x.shape) + else: + return x + elif isinstance(x, (int, float)): + return x + else: + raise ValueError(f"Unsupported type {type(x)}") + + +def to_tensor(x, device=None): + if isinstance(x, torch.Tensor): + return x.to(device) + elif isinstance(x, (np.ndarray, list, tuple, int, float)): + return torch.tensor(x, device=device) + else: + raise ValueError(f"Unsupported type {type(x)}") + + +def linear_beta_schedule(beta_min: float = 1e-4, beta_max: float = 0.02, T: int = 1000): + return np.linspace(beta_min, beta_max, T) + + +def cosine_beta_schedule(s: float = 0.008, T: int = 1000): + f = np.cos((np.arange(T + 1) / T + s) / (1 + s) * np.pi / 2.0) ** 2 + alpha_bar = f / f[0] + beta = 1 - alpha_bar[1:] / alpha_bar[:-1] + return beta.clip(None, 0.999) + + +# ================ Discretization ================== +def uniform_discretization(T: int = 1000, eps: float = 1e-3): + return torch.linspace(eps, 1.0, T) + + +SUPPORTED_DISCRETIZATIONS = { + "uniform": uniform_discretization, +} + + +# ================= Noise schedules ================= +def linear_noise_schedule( + t_diffusion: torch.Tensor, beta0: float = 0.1, beta1: float = 20.0 +): + log_alpha = -(beta1 - beta0) / 4.0 * (t_diffusion**2) - beta0 / 2.0 * t_diffusion + alpha = log_alpha.exp() + sigma = (1.0 - alpha**2).sqrt() + return alpha, sigma + + +def inverse_linear_noise_schedule( + alpha: torch.Tensor = None, + sigma: torch.Tensor = None, + logSNR: torch.Tensor = None, + beta0: float = 0.1, + beta1: float = 20.0, +): + assert (logSNR is not None) or (alpha is not None and sigma is not None) + lmbda = (alpha / sigma).log() if logSNR is None else logSNR + t_diffusion = (2 * (1 + (-2 * lmbda).exp()).log() / + (beta0 + (beta0**2 + 2 * (beta1 - beta0) * (1 + (-2 * lmbda).exp()).log()))) + return t_diffusion + + +def cosine_noise_schedule(t_diffusion: torch.Tensor, s: float = 0.008): + alpha = (np.pi / 2.0 * ((t_diffusion).clip(0., 0.9946) + s) / (1 + s)).cos() / np.cos( + np.pi / 2.0 * s / (1 + s)) + sigma = (1.0 - alpha**2).sqrt() + return alpha, sigma + + +def inverse_cosine_noise_schedule( + alpha: torch.Tensor = None, + sigma: torch.Tensor = None, + logSNR: torch.Tensor = None, + s: float = 0.008, +): + assert (logSNR is not None) or (alpha is not None and sigma is not None) + lmbda = (alpha / sigma).log() if logSNR is None else logSNR + t_diffusion = ( + 2 * (1 + s) / np.pi * torch.arccos(( + -0.5 * (1 + (-2 * lmbda).exp()).log() + + np.log(np.cos(np.pi * s / 2 / (s + 1)))).exp()) - s) + return t_diffusion + + +SUPPORTED_NOISE_SCHEDULES = { + "linear": { + "forward": linear_noise_schedule, + "reverse": inverse_linear_noise_schedule, + }, + "cosine": { + "forward": cosine_noise_schedule, + "reverse": inverse_cosine_noise_schedule, + }, +} + + +# ================= Sampling step schedule =============== +def uniform_sampling_step_schedule(T: int = 1000, sampling_steps: int = 10): + return torch.linspace(0, T - 1, sampling_steps + 1, dtype=torch.long) + + +def uniform_sampling_step_schedule_continuous(trange=None, sampling_steps: int = 10): + if trange is None: + trange = [1e-3, 1.0] + return torch.linspace(trange[0], trange[1], sampling_steps + 1, dtype=torch.float32) + + +def quad_sampling_step_schedule(T: int = 1000, sampling_steps: int = 10, n: int = 1.5): + schedule = (T - 1) * ( + torch.linspace(0, 1, sampling_steps + 1, dtype=torch.float32) ** n) + return schedule.to(torch.long) + + +def quad_sampling_step_schedule_continuous( + trange=None, sampling_steps: int = 10, n: int = 1.5 +): + if trange is None: + trange = [1e-3, 1.0] + schedule = (trange[1] - trange[0]) * ( + torch.linspace(0, 1, sampling_steps + 1, dtype=torch.float32) ** n + ) + trange[0] + return schedule + + +def cat_cos_sampling_step_schedule( + T: int = 1000, sampling_steps: int = 10, n: int = 2.0 +): + idx = torch.linspace(0, 1, sampling_steps + 1, dtype=torch.float32) + idx = (0.5 * (2 * (idx > 0.5) - 1) * torch.sin(np.pi * torch.abs(idx - 0.5)) ** (1 / n) + 0.5) + schedule = (T - 1) * idx + return schedule.to(torch.long) + + +def cat_cos_sampling_step_schedule_continuous( + trange=None, sampling_steps: int = 10, n: int = 2.0 +): + if trange is None: + trange = [1e-3, 1.0] + idx = torch.linspace(0, 1, sampling_steps + 1, dtype=torch.float32) + idx = (0.5 * (2 * (idx > 0.5) - 1) * torch.sin(np.pi * torch.abs(idx - 0.5)) ** (1 / n) + 0.5) + schedule = (trange[1] - trange[0]) * idx + trange[0] + return schedule + + +def quad_cos_sampling_step_schedule( + T: int = 1000, sampling_steps: int = 10, n: int = 2.0 +): + idx = torch.linspace(0, 1, sampling_steps + 1, dtype=torch.float32) + idx = ((torch.sin(np.pi * (idx - 0.5)) + 1) / 2) ** n + schedule = (T - 1) * idx + return schedule.to(torch.long) + + +def quad_cos_sampling_step_schedule_continuous( + trange=None, sampling_steps: int = 10, n: int = 2.0 +): + if trange is None: + trange = [1e-3, 1.0] + idx = torch.linspace(0, 1, sampling_steps + 1, dtype=torch.float32) + idx = ((torch.sin(np.pi * (idx - 0.5)) + 1) / 2) ** n + schedule = (trange[1] - trange[0]) * idx + trange[0] + return schedule + + +SUPPORTED_SAMPLING_STEP_SCHEDULE = { + "uniform": uniform_sampling_step_schedule, + "uniform_continuous": uniform_sampling_step_schedule_continuous, + "quad": quad_sampling_step_schedule, + "quad_continuous": quad_sampling_step_schedule_continuous, + "cat_cos": cat_cos_sampling_step_schedule, + "cat_cos_continuous": cat_cos_sampling_step_schedule_continuous, + "quad_cos": quad_cos_sampling_step_schedule, + "quad_cos_continuous": quad_cos_sampling_step_schedule_continuous, +} + + +def count_parameters(model: nn.Module): + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + +def ema_update(model: nn.Module, model_ema: nn.Module, ema_rate: float): + for param, param_ema in zip(model.parameters(), model_ema.parameters()): + param_ema.data.mul_(ema_rate).add_(param.data, alpha=1 - ema_rate) + + +# ----------------------------------------------------------- +# Timestep embedding used in the DDPM++ and ADM architectures, +# from https://github.com/NVlabs/edm/blob/main/training/networks.py#L269 +class PositionalEmbedding(nn.Module): + def __init__(self, dim: int, max_positions: int = 10000, endpoint: bool = False): + super().__init__() + self.dim = dim + self.max_positions = max_positions + self.endpoint = endpoint + + def forward(self, x): + freqs = torch.arange( + start=0, end=self.dim // 2, dtype=torch.float32, device=x.device + ) + freqs = freqs / (self.dim // 2 - (1 if self.endpoint else 0)) + freqs = (1 / self.max_positions) ** freqs + x = x.ger(freqs.to(x.dtype)) + x = torch.cat([x.cos(), x.sin()], dim=1) + return x + + +class UntrainablePositionalEmbedding(nn.Module): + def __init__(self, dim: int, max_positions: int = 10000, endpoint: bool = False): + super().__init__() + self.dim = dim + self.max_positions = max_positions + self.endpoint = endpoint + + def forward(self, x): + freqs = torch.arange( + start=0, end=self.dim // 2, dtype=torch.float32, device=x.device) + freqs = freqs / (self.dim // 2 - (1 if self.endpoint else 0)) + freqs = (1 / self.max_positions) ** freqs + x = torch.einsum('...i,j->...ij', x, freqs.to(x.dtype)) + # x = x.ger(freqs.to(x.dtype)) + x = torch.cat([x.cos(), x.sin()], dim=1) + return x + + +# ----------------------------------------------------------- +# Timestep embedding used in Transformer +class SinusoidalEmbedding(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.dim = dim + + def forward(self, x): + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device) * -emb) + emb = torch.einsum('...i,j->...ij', x, emb.to(x.dtype)) + # emb = x[:, None] * emb[None, :] + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +# ----------------------------------------------------------- +# Timestep embedding used in the DDPM++ and ADM architectures +class FourierEmbedding(nn.Module): + def __init__(self, dim: int, scale=16): + super().__init__() + self.freqs = nn.Parameter(torch.randn(dim // 8) * scale, requires_grad=False) + self.mlp = nn.Sequential( + nn.Linear(dim // 4, dim), nn.Mish(), nn.Linear(dim, dim) + ) + + def forward(self, x: torch.Tensor): + emb = torch.einsum('...i,j->...ij', x, (2 * np.pi * self.freqs).to(x.dtype)) + # emb = x.ger((2 * np.pi * self.freqs).to(x.dtype)) + emb = torch.cat([emb.cos(), emb.sin()], -1) + return self.mlp(emb) + + +class UntrainableFourierEmbedding(nn.Module): + def __init__(self, dim: int, scale=16): + super().__init__() + self.freqs = nn.Parameter(torch.randn(dim // 2) * scale, requires_grad=False) + + def forward(self, x: torch.Tensor): + emb = torch.einsum('...i,j->...ij', x, (2 * np.pi * self.freqs).to(x.dtype)) + # emb = x.ger((2 * np.pi * self.freqs).to(x.dtype)) + emb = torch.cat([emb.cos(), emb.sin()], -1) + return emb + + +SUPPORTED_TIMESTEP_EMBEDDING = { + "positional": PositionalEmbedding, + "fourier": FourierEmbedding, + "untrainable_fourier": UntrainableFourierEmbedding, + "untrainable_positional": UntrainablePositionalEmbedding, +} + + +# ----------------------------------------------------------- +# Beautiful model size visualization from https://github.com/jannerm/diffuser/tree/main + + +def _to_str(num): + if num >= 1e6: + return f"{(num / 1e6):.2f} M" + else: + return f"{(num / 1e3):.2f} k" + + +def param_to_module(param): + module_name = param[::-1].split(".", maxsplit=1)[-1][::-1] + return module_name + + +def report_parameters(model, topk=10): + counts = {k: p.numel() for k, p in model.named_parameters() if p.requires_grad} + n_parameters = sum(counts.values()) + print(f"Total parameters: {_to_str(n_parameters)}") + + modules = dict(model.named_modules()) + sorted_keys = sorted(counts, key=lambda x: -counts[x]) + for i in range(topk): + key = sorted_keys[i] + count = counts[key] + module = param_to_module(key) + print(" " * 8, f"{key:10}: {_to_str(count)} | {modules[module]}") + + remaining_parameters = sum([counts[k] for k in sorted_keys[topk:]]) + print( + " " * 8, + f"... and {len(counts) - topk} others accounting for {_to_str(remaining_parameters)} parameters", + ) + return n_parameters + + +# ----------------------- DD return scale + +# discount = 0.997 +DD_RETURN_SCALE = { + "halfcheetah-medium-expert-v2": 3600, + "halfcheetah-medium-replay-v2": 1600, + "halfcheetah-medium-v2": 1700, + "hopper-medium-expert-v2": 1200, + "hopper-medium-replay-v2": 1000, + "hopper-medium-v2": 1000, + "walker2d-medium-expert-v2": 1600, + "walker2d-medium-replay-v2": 1300, + "walker2d-medium-v2": 1300, + "kitchen-partial-v0": 470, + "kitchen-mixed-v0": 400, + "antmaze-medium-play-v2": 100, + "antmaze-medium-diverse-v2": 100, + "antmaze-large-play-v2": 100, + "antmaze-large-diverse-v2": 100, +} + + +# ---------------------- Freeze and Unfreeze + + +class FreezeModules: + def __init__(self, modules): + self.modules = modules + self.original_grad_status = {} + + def __enter__(self): + for module in self.modules: + for param in module.parameters(): + self.original_grad_status[id(param)] = param.requires_grad + param.requires_grad = False + + def __exit__(self, type, value, traceback): + for module in self.modules: + for param in module.parameters(): + param.requires_grad = self.original_grad_status[id(param)] + + +class UnfreezeModules: + def __init__(self, modules): + self.modules = modules + self.original_grad_status = {} + + def __enter__(self): + for module in self.modules: + for param in module.parameters(): + self.original_grad_status[id(param)] = param.requires_grad + param.requires_grad = True + + def __exit__(self, type, value, traceback): + for module in self.modules: + for param in module.parameters(): + param.requires_grad = self.original_grad_status[id(param)] + + +class EvalModules: + def __init__(self, modules): + self.modules = modules + self.original_status = {} + + def __enter__(self): + for module in self.modules: + self.original_status[id(module)] = module.training + module.eval() + + def __exit__(self, type, value, traceback): + for module in self.modules: + module.train(self.original_status[id(module)]) + + +class TrainModules: + def __init__(self, modules): + self.modules = modules + self.original_status = {} + + def __enter__(self): + for module in self.modules: + self.original_status[id(module)] = module.training + module.train() + + def __exit__(self, type, value, traceback): + for module in self.modules: + module.train(self.original_status[id(module)]) + + +def dict_apply( + x: Dict[str, torch.Tensor], + func: Callable[[torch.Tensor], torch.Tensor] +) -> Dict[str, torch.Tensor]: + result = dict() + for key, value in x.items(): + if isinstance(value, dict): + result[key] = dict_apply(value, func) + elif value is None: + result[key] = None + else: + result[key] = func(value) + return result + + +def loop_dataloader(dl): + while True: + for b in dl: + yield b diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_manipulation/common.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_manipulation/common.py new file mode 100644 index 0000000000000000000000000000000000000000..d18d8e239819f0d218742693bc3c258747130df8 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_manipulation/common.py @@ -0,0 +1,74 @@ +"""Read-only task/config adapters for the twelve matched imitation tasks.""" +import copy +from pathlib import Path + +from omegaconf import OmegaConf + +ROOT = Path(__file__).resolve().parents[2] +TASKS = { + "pusht": ("pusht", "pusht", False), + "pusht_keypoints": ("pusht", "pusht_keypoint", False), + "pusht_image": ("pusht", "pusht_image", True), + "relay_kitchen": ("kitchen", "kitchen_abs", False), + **{f"{name}_{mode}": ("robomimic", f"{name}_abs", mode == "image") + for mode in ("lowdim", "image") for name in ("lift", "can", "square", "tool_hang")}, +} + + +def config_for(task, seed=0, device="cpu"): + family, name, image = TASKS[task] + folder = "robomimic_multi_modal" if family == "robomimic" and image else family + config = OmegaConf.load(ROOT / f"configs/dp/{folder}/chi_unet/{name}.yaml") + config.seed, config.device, config.save_video = int(seed), str(device), False + config.dataset_path = str((ROOT / config.dataset_path).resolve()) + for key in ("hydra", "work_dir"): + if key in config: + del config[key] + return config + + +def shape_meta_for_task(task): + if TASKS[task][0] == "pusht": + from experiments.consistency_pusht.data import shape_meta_for_task as pusht_shape + return pusht_shape(task) + config = config_for(task) + if TASKS[task][2]: + result = copy.deepcopy(OmegaConf.to_container(config.shape_meta, resolve=True)) + else: + result = {"obs": {"state": {"shape": [int(config.obs_dim)], "type": "low_dim"}}} + result["action"] = {"shape": [int(config.action_dim)]} + return result + + +def build_dataset(task, config): + """Original CleanDiffuser dataset, without mutating the frozen task registry. + +The CP image adapter does NOT call this full image builder; it implements a +pixel-equivalent lazy HDF5 reader instead. This helper is for parity tests and +the small Kitchen / low-dimensional datasets. + """ + from experiments.table2 import common as original + if task in original.TASKS: + return original.build_dataset(task, config) + from cleandiffuser.dataset.robomimic_dataset import RobomimicDataset, RobomimicImageDataset + arguments = {"horizon": config.horizon, "pad_before": config.obs_steps - 1, + "pad_after": config.action_steps - 1, "abs_action": config.abs_action} + if TASKS[task][2]: + return RobomimicImageDataset(config.dataset_path, shape_meta=config.shape_meta, + n_obs_steps=config.obs_steps, **arguments) + return RobomimicDataset(config.dataset_path, obs_keys=list(config.obs_keys), **arguments) + + +def dataset_manifest(task): + if task in ("can_lowdim", "tool_hang_lowdim", "can_image", "tool_hang_image"): + # The addon verifies the exact task member, not just a family archive. + from experiments.table2_addon.common import dataset_manifest as addon_manifest + return addon_manifest(task) + else: + family, _, image = TASKS[task] + name = ("kitchen" if family == "kitchen" else "pusht" if family == "pusht" + else "robomimic_image" if image else "robomimic_lowdim") + path = ROOT / f"outputs/table2/downloads/{name}_manifest.json" + if not path.is_file(): + raise FileNotFoundError(f"Dataset has not passed archive verification: {path}") + return path diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_manipulation/data.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_manipulation/data.py new file mode 100644 index 0000000000000000000000000000000000000000..faa4a38eb9994ff731112f610b15ce3938fe911e --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_manipulation/data.py @@ -0,0 +1,221 @@ +"""Raw-dictionary CP datasets with exactly the local Table-2 preprocessing. + +Kitchen retains all sixty coordinates, including the thirty masked goal zeros +and deterministic observation noise in KitchenMjlDataset. Robomimic images are +read lazily, two history frames per item, and pass through the same JPEG2000 +level-50 codec as the original image replay buffer. Image RAM is bounded. +""" +import copy +import json +import os +from collections import OrderedDict +from pathlib import Path + +import numpy as np +import torch + +from diffusion_policy.dataset.base_dataset import BaseImageDataset +from diffusion_policy.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer +from diffusion_policy.common.normalize_util import get_image_range_normalizer + +from experiments.consistency_manipulation.common import TASKS, build_dataset, config_for, shape_meta_for_task + + +def affine_from_local(normalizer, values, *, action=False): + """Export local affine constants, including Kitchen's masked channels. + +Zero-scale observation coordinates are intentional and never inverted by the +official policy. Constant Kitchen action channels would require a custom inverse +and are rejected; the actual Kitchen actions have no constant channels. + """ + values = np.asarray(values, dtype=np.float32).reshape(-1, values.shape[-1]) + minimum = np.asarray(normalizer.min, dtype=np.float32) + span = np.asarray(normalizer.range, dtype=np.float32) + mask = np.asarray(getattr(normalizer, "mask", np.ones_like(span)), dtype=np.float32) + if action and np.any(mask == 0): + raise ValueError("Masked constant actions cannot use official affine inverse normalization") + scale = (np.float32(2) / span) * mask + offset = (-np.float32(1) - minimum * (np.float32(2) / span)) * mask + field = SingleFieldLinearNormalizer.create_manual( + scale=scale.astype(np.float32), offset=offset.astype(np.float32), + input_stats_dict={"min": minimum, "max": np.asarray(normalizer.max, dtype=np.float32), + "mean": values.mean(axis=0), "std": values.std(axis=0)}) + field.requires_grad_(False) + return field + + +def transformed_actions(raw_actions, transformer): + raw = np.asarray(raw_actions, dtype=np.float32) + if raw.ndim != 2 or raw.shape[-1] != 7: + raise ValueError("These Robomimic tasks require seven-dimensional absolute source actions") + return np.concatenate((raw[:, :3], transformer.forward(raw[:, 3:6]), raw[:, 6:]), axis=-1).astype(np.float32) + + +class ManipulationDictDataset(BaseImageDataset): + def __init__(self, task, dataset_path, horizon=16, pad_before=1, pad_after=7, + n_obs_steps=2, seed=42, val_ratio=0.0, image_cache_frames=64): + super().__init__() + if task not in TASKS or TASKS[task][0] == "pusht": + raise ValueError("PushT must use its unchanged original PushTDictDataset adapter") + if val_ratio != 0: + raise ValueError("Matched Table-2 protocol trains on every demonstration; validation is in-sample") + if (horizon, pad_before, pad_after, n_obs_steps) != (16, 1, 7, 2): + raise ValueError("Matched protocol requires horizon16, padding1/7 and observation history2") + if image_cache_frames < 0: + raise ValueError("Image cache size cannot be negative") + self.task, self.seed, self.val_ratio = task, int(seed), float(val_ratio) + self.dataset_path = str(Path(dataset_path).resolve(strict=True)) + self.horizon, self.pad_before, self.pad_after, self.n_obs_steps = horizon, pad_before, pad_after, n_obs_steps + self.shape_meta = shape_meta_for_task(task) + self.is_image = TASKS[task][2] + self.validation_in_sample, self.role = True, "train" + self._h5, self._h5_pid = None, None + self._image_cache = OrderedDict() + self.image_cache_frames = int(image_cache_frames) + config = config_for(task) + config.dataset_path = self.dataset_path + self.config = config + self.parse_audit = None + if self.is_image: + self._prepare_lazy_images() + else: + self.local_dataset = build_dataset(task, config) + self.replay_buffer, self.sampler = self.local_dataset.replay_buffer, self.local_dataset.sampler + self.local_normalizer = self.local_dataset.normalizer + self.parse_audit = getattr(self.local_dataset, "table2_parse_audit", None) + if len(self.sampler) == 0: + raise ValueError("Empty sequence dataset") + self._normalizer = self._make_normalizer() + print("[CP_DATASET] " + json.dumps(self.metadata, allow_nan=False), flush=True) + + def _prepare_lazy_images(self): + import h5py + from cleandiffuser.dataset.dataset_utils import SequenceSampler, MinMaxNormalizer, ImageNormalizer, RotationTransformer + from cleandiffuser.dataset.replay_buffer import ReplayBuffer + from cleandiffuser.dataset.imagecodecs import Jpeg2k + self.rotation_transformer = RotationTransformer(from_rep="axis_angle", to_rep="rotation_6d") + self._codec = Jpeg2k(level=50) + self.rgb_keys = [key for key, value in self.shape_meta["obs"].items() if value.get("type") == "rgb"] + self.lowdim_keys = [key for key in self.shape_meta["obs"] if key not in self.rgb_keys] + self.replay_buffer = ReplayBuffer.create_empty_numpy() + self._demo_names = [] + with h5py.File(self.dataset_path, "r") as file: + demos = file["data"] + for idx in range(len(demos)): + name = f"demo_{idx}" + demo = demos[name] + actions = transformed_actions(demo["actions"][:], self.rotation_transformer) + episode = {"action": actions} + for key in self.lowdim_keys: + value = np.asarray(demo["obs"][key][:], dtype=np.float32) + if value.shape != (len(actions), *self.shape_meta["obs"][key]["shape"]): + raise ValueError(f"Wrong {name}/{key} shape: {value.shape}") + episode[key] = value + for key in self.rgb_keys: + channels, height, width = self.shape_meta["obs"][key]["shape"] + image = demo["obs"][key] + if image.shape != (len(actions), height, width, channels) or image.dtype != np.uint8: + raise ValueError(f"Unexpected raw image storage {name}/{key}: {image.shape}, {image.dtype}") + if any(not np.isfinite(value).all() for value in episode.values()): + raise FloatingPointError("Nonfinite low-dimensional HDF5 values") + self.replay_buffer.add_episode(episode) + self._demo_names.append(name) + self.sampler = SequenceSampler(self.replay_buffer, sequence_length=self.horizon, + pad_before=self.pad_before, pad_after=self.pad_after) + self.local_normalizer = {"obs": {key: MinMaxNormalizer(self.replay_buffer[key][:]) for key in self.lowdim_keys}, + "action": MinMaxNormalizer(self.replay_buffer["action"][:])} + self.local_normalizer["obs"].update({key: ImageNormalizer() for key in self.rgb_keys}) + + def _make_normalizer(self): + result = LinearNormalizer() + result["action"] = affine_from_local(self.local_normalizer["action"], self.replay_buffer["action"][:], action=True) + for key, local in self.local_normalizer["obs"].items(): + if key in getattr(self, "rgb_keys", ()): + result[key] = get_image_range_normalizer() + else: + source_key = "obs" if key == "state" and TASKS[self.task][0] == "robomimic" else key + result[key] = affine_from_local(local, self.replay_buffer[source_key][:]) + result.requires_grad_(False) + return result + + @property + def metadata(self): + return {"task": self.task, "dataset_path": self.dataset_path, + "episodes_total": int(self.replay_buffer.n_episodes), + "frames_total": int(self.replay_buffer.n_steps), "samples": len(self), + "validation_in_sample": True, "validation_use": "diagnostic_only_not_checkpoint_selection", + "role": self.role, "seed": self.seed, "val_ratio": self.val_ratio, + "horizon": self.horizon, "pad_before": self.pad_before, "pad_after": self.pad_after, + "normalizer_rule": "exact local Table-2 affine; Kitchen constant observation channels map to zero", + "image_storage": "lazy HDF5 with local JPEG2000 level50 round-trip" if self.is_image else None, + "image_cache_frames": self.image_cache_frames if self.is_image else 0, + "shape_meta": self.shape_meta, "kitchen_parse_audit": self.parse_audit} + + def get_normalizer(self, **kwargs): + if kwargs: + raise ValueError("Normalization is fixed to the local Table-2 preprocessing") + return copy.deepcopy(self._normalizer) + + def get_validation_dataset(self): + view = copy.copy(self) + view.role = "in_sample_diagnostic" + view._h5, view._h5_pid, view._image_cache = None, None, OrderedDict() + return view + + def get_all_actions(self): + return torch.from_numpy(np.asarray(self.replay_buffer["action"][:], dtype=np.float32).copy()) + + def __len__(self): + return len(self.sampler) + + def _image_frame(self, demo_idx, key, frame): + cache_key = (demo_idx, key, frame) + if cache_key in self._image_cache: + self._image_cache.move_to_end(cache_key) + return self._image_cache[cache_key] + import h5py + if self._h5 is None or self._h5_pid != os.getpid(): + self.close() + self._h5 = h5py.File(self.dataset_path, "r") + self._h5_pid = os.getpid() + raw = self._h5["data"][self._demo_names[demo_idx]]["obs"][key][frame] + value = self._codec.decode(self._codec.encode(raw)) + if self.image_cache_frames: + self._image_cache[cache_key] = value + while len(self._image_cache) > self.image_cache_frames: + self._image_cache.popitem(last=False) + return value + + def __getitem__(self, index): + sample = self.sampler.sample_sequence(index) + if self.is_image: + obs = {key: np.asarray(sample[key][:self.n_obs_steps], dtype=np.float32) for key in self.lowdim_keys} + start, end, sample_start, _ = self.sampler.indices[index] + episode_ends = np.asarray(self.replay_buffer.episode_ends) + demo_idx = int(np.searchsorted(episode_ends, start, side="right")) + episode_start = 0 if demo_idx == 0 else int(episode_ends[demo_idx - 1]) + frames = np.clip(start + np.arange(self.n_obs_steps) - sample_start, start, end - 1) - episode_start + for key in self.rgb_keys: + pixels = np.stack([self._image_frame(demo_idx, key, int(frame)) for frame in frames]) + obs[key] = np.moveaxis(pixels, -1, 1).astype(np.float32) / np.float32(255) + else: + key = "state" if self.task == "relay_kitchen" else "obs" + obs = {"state": np.asarray(sample[key], dtype=np.float32)} + return {"obs": {key: torch.from_numpy(np.ascontiguousarray(value)) for key, value in obs.items()}, + "action": torch.from_numpy(np.array(sample["action"], dtype=np.float32, copy=True))} + + def close(self): + if self._h5 is not None: + self._h5.close() + self._h5, self._h5_pid = None, None + + def __getstate__(self): + state = self.__dict__.copy() + state.update(_h5=None, _h5_pid=None, _image_cache=OrderedDict()) + return state + + def __del__(self): + try: + self.close() + except (AttributeError, TypeError): + pass diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_manipulation/evaluate.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_manipulation/evaluate.py new file mode 100644 index 0000000000000000000000000000000000000000..c0803a28eb0659b226bdf02f8ff1cb27656b5a3e --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_manipulation/evaluate.py @@ -0,0 +1,249 @@ +"""Official EDM/CP on the twelve-task, one-candidate manipulation protocol. + +The original PushT evaluator is reused unchanged. Kitchen and Robomimic use +the same local environments as Table 2, raw observation dictionaries and the +checkpoint's own normalizer. These concurrent timings are diagnostic only. +""" +import argparse +import copy +import json +import math +import random +import sys +from pathlib import Path + +import numpy as np +import torch + +from experiments.consistency_pusht import evaluate as pusht_eval +from experiments.consistency_manipulation.common import TASKS, config_for, shape_meta_for_task +from experiments.table2.common import ROOT +from experiments.table2.evaluate import make_env as table2_make_env +from tools.consistency_policy.evaluate import ( + DEFAULT_REPO, TimedPolicy, checkpoint_hash, load_policy, parse_steps, + set_sampling, verify_source, +) + + +def observation_dict(task, observations, device): + if task in pusht_eval.TASKS: + return pusht_eval.observation_dict(task, observations, device) + if task not in TASKS or not observations: + raise ValueError("Unknown task or empty observation batch") + shape = shape_meta_for_task(task)["obs"] + values = ({key: np.stack([obs[key] for obs in observations]) for key in shape} + if any(spec.get("type") == "rgb" for spec in shape.values()) + else {"state": np.stack(observations)}) + result = {} + for key, spec in shape.items(): + value = np.asarray(values[key]) + if value.shape != (len(observations), 2, *spec["shape"]): + raise ValueError(f"Incorrect {key} observation shape: {value.shape}") + if not np.isfinite(value).all(): + raise FloatingPointError("Nonfinite observation: " + key) + if spec.get("type") == "rgb" and (value.min() < 0 or value.max() > 1): + raise ValueError("Raw image observations must be CHW float pixels in [0,1]") + result[key] = torch.as_tensor(value, dtype=torch.float32, device=device) + return result + + +def validate_policy(policy, task, shape_meta=None): + if task not in TASKS: + raise ValueError("Unsupported task: " + task) + expected = shape_meta_for_task(task) + for attr, value in (("horizon", 16), ("n_obs_steps", 2), ("n_action_steps", 8), + ("action_dim", expected["action"]["shape"][0])): + if int(getattr(policy, attr)) != value: + raise ValueError(f"Checkpoint {attr} must equal {value}") + if getattr(policy, "use_kde", False): + raise ValueError("KDE violates the one-candidate protocol") + if shape_meta is not None: + def canonical(meta): + return {"action": tuple(meta["action"]["shape"]), + "obs": {k: (tuple(v["shape"]), v.get("type", "low_dim")) + for k, v in meta["obs"].items()}} + if canonical(shape_meta) != canonical(expected): + raise ValueError("Checkpoint observation/action shapes do not match --task") + + +def make_env(task, config, max_steps): + from experiments.table2_addon.common import register_tasks + register_tasks() + local = copy.deepcopy(config) + local.max_episode_steps = max_steps + return table2_make_env(task, local) + + +def robomimic_actions(actions): + """Exactly the local dataset's inverse 10D rotation-6D -> 7D pose map.""" + from cleandiffuser.dataset.dataset_utils import RotationTransformer + if actions.shape[-1] != 10: + raise ValueError("Single-arm absolute Robomimic actions must be 10D") + transformer = RotationTransformer(from_rep="axis_angle", to_rep="rotation_6d") + rotation = transformer.inverse(actions[..., 3:9]) + return np.concatenate([actions[..., :3], rotation, actions[..., -1:]], axis=-1) + + +def evaluate(policy, task, episodes=50, envs=10, seed_base=100000, + eval_seed=0, max_steps=None, warmup=3): + local = config_for(task, seed=eval_seed, device=str(policy.device)) + full_limit = int(local.max_episode_steps) + limit = full_limit if max_steps is None else max_steps + if task in pusht_eval.TASKS: + return pusht_eval.evaluate(policy, task, episodes, envs, seed_base, + eval_seed, limit, warmup) + validate_policy(policy, task) + if episodes < 1 or envs < 1 or not 1 <= limit <= full_limit or warmup < 0: + raise ValueError("Invalid episode/environment count, limit or warmup") + device = torch.device(policy.device) + devices = [device.index or 0] if device.type == "cuda" else [] + numpy_state, python_state = np.random.get_state(), random.getstate() + was_training = policy.training + environments, rows, batch_sizes = [], [], [] + timed = TimedPolicy(policy, warmup) + action_dim = shape_meta_for_task(task)["action"]["shape"][0] + kitchen = task == "relay_kitchen" + try: + with torch.random.fork_rng(devices=devices), torch.no_grad(): + random.seed(eval_seed) + np.random.seed(eval_seed) + torch.manual_seed(eval_seed) + policy.eval() + for _ in range(min(envs, episodes)): + environments.append(make_env(task, local, limit)) + for start in range(0, episodes, len(environments)): + count = min(len(environments), episodes - start) + observations = [] + for idx, env in enumerate(environments[:count]): + env.seed(seed_base + start + idx) + observations.append(env.reset()) + policy.reset() + finished, completed = [False] * count, [0] * count + decision_rewards = [[] for _ in range(count)] + for _ in range(math.ceil(limit / 8)): + if all(finished): + break + output = timed.predict_action(observation_dict(task, observations, device)) + if tuple(output["action"].shape) != (count, 8, action_dim): + raise ValueError("Policy must return an already-sliced B x 8 action chunk") + actions = output["action"].detach().cpu().numpy() + if not np.isfinite(actions).all(): + raise FloatingPointError("Nonfinite policy action") + # predict_action already unnormalizes and slices [1:9]. + if not kitchen: + actions = robomimic_actions(actions) + if not np.isfinite(actions).all(): + raise FloatingPointError("Nonfinite converted environment action") + batch_sizes.append(count) + for idx, env in enumerate(environments[:count]): + if finished[idx]: + continue + obs, reward, done, info = env.step(actions[idx]) + if not math.isfinite(float(reward)): + raise FloatingPointError("Nonfinite environment reward") + observations[idx] = obs + decision_rewards[idx].append(float(reward)) + finished[idx] = bool(done) + if kitchen and "completed_tasks" in info and len(info["completed_tasks"]): + completed[idx] = max(completed[idx], len(info["completed_tasks"][-1])) + if not all(finished): + raise RuntimeError("Environment did not terminate at its primitive-step limit") + for idx, env in enumerate(environments[:count]): + rewards = [float(r) for r in env.get_rewards()] + if not rewards or len(rewards) > limit or not np.isfinite(rewards).all(): + raise ValueError("Invalid primitive reward trace") + score = min(completed[idx], 4) / 4 if kitchen else float(any(r > 0 for r in rewards)) + rows.append({"episode_seed": seed_base + start + idx, "score": score, + "primitive_steps": len(rewards), "rewards": rewards, + "decision_rewards": decision_rewards[idx], "done": finished[idx], + "completed_tasks": completed[idx] if kitchen else None}) + print(f"[CP-EVAL] {task} {len(rows)}/{episodes} " + f"score={np.mean([r['score'] for r in rows]):.6f}", flush=True) + finally: + timed.close() + for env in environments: + env.close() + policy.train(was_training) + np.random.set_state(numpy_state) + random.setstate(python_state) + if not timed.nfes or min(timed.nfes) < 1: + raise ValueError("Denoiser hook recorded no calls") + return {"task": task, "score": float(np.mean([r["score"] for r in rows])), + "score_metric": "mean_capped_subtask_completion" if kitchen else "success_rate", + "episodes": rows, "num_episodes": episodes, "seed_base": seed_base, "eval_seed": eval_seed, + "num_candidates": 1, "horizon": 16, "obs_steps": 2, "action_steps": 8, + "max_episode_steps": limit, "pusht_legacy": None, + "env_implementation": "frozen local Table-2 Kitchen/Robomimic environments", + "truncated_smoke": limit < full_limit, "paper_ready": False, + "latency": {"scope": "predict_action_device_resident; excludes transfers, action rotation conversion and simulator", + "concurrent_gpu_diagnostic": True, "paper_ready": False, + "mean_ms": float(np.mean(timed.durations)), + "p50_ms": float(np.percentile(timed.durations, 50)), + "p95_ms": float(np.percentile(timed.durations, 95)), + "raw_ms": timed.durations, "warmup_calls": warmup, + "batch_sizes": batch_sizes, "n_envs": min(envs, episodes)}, + "nfe": {"mean": float(np.mean(timed.nfes)), "min": min(timed.nfes), + "max": max(timed.nfes), "raw": timed.nfes, + "unit": "denoiser forwards per batched predict_action, not per episode"}} + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--task", choices=TASKS, required=True) + parser.add_argument("--method", choices=("edm", "cp"), required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--repo", type=Path, default=DEFAULT_REPO) + parser.add_argument("--steps") + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--episodes", type=int, default=50) + parser.add_argument("--envs", type=int, default=10) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--smoke", action="store_true") + args = parser.parse_args(argv) + if not args.smoke and (args.episodes != 50 or args.envs != 10): + parser.error("Formal evaluations require 50 episodes and 10 environments") + source = verify_source(args.repo) + sys.path.insert(0, source["repo"]) + from omegaconf import OmegaConf + if not OmegaConf.has_resolver("eval"): + OmegaConf.register_new_resolver("eval", eval) + checkpoint = args.checkpoint.resolve(strict=True) + policy, cfg, weights = load_policy(checkpoint, args.method, "model", args.device) + validate_policy(policy, args.task, cfg.policy.shape_meta) + if cfg.experiment.task != args.task or bool(cfg.experiment.smoke) != args.smoke: + raise ValueError("Checkpoint task/smoke metadata does not match evaluation") + steps_list = parse_steps(args.steps or ("1,3" if args.method == "cp" else "80")) + paths = [args.output / f"{args.method}_{n}_{weights}.json" for n in steps_list] + for steps, path in zip(steps_list, paths): + set_sampling(policy, args.method, steps, "official") + if path.exists(): + raise FileExistsError("Refusing to overwrite result: " + str(path)) + args.output.mkdir(parents=True, exist_ok=True) + runtime = pusht_eval.runtime_provenance(args.device) + runtime["local_source_sha256"][str(Path(__file__).relative_to(ROOT))] = checkpoint_hash(Path(__file__)) + for steps, path in zip(steps_list, paths): + times = set_sampling(policy, args.method, steps, "official") + result = evaluate(policy, args.task, episodes=2 if args.smoke else args.episodes, + envs=2 if args.smoke else args.envs, eval_seed=int(cfg.training.seed), + max_steps=8 if args.smoke else None, warmup=args.warmup) + expected_nfe = steps if args.method == "cp" else 2 * steps - 2 + if result["nfe"]["min"] != expected_nfe or result["nfe"]["max"] != expected_nfe: + raise ValueError("Measured denoiser NFE differs from the declared sampling protocol") + result.update(schema_version=1, method=args.method, weights=weights, smoke=args.smoke, + training_seed=int(cfg.training.seed), checkpoint=str(checkpoint), + checkpoint_sha256=checkpoint_hash(checkpoint), source=source, runtime=runtime, + teacher_reference="raw_model_actually_distilled" if args.method == "edm" else None, + sampling={"steps": steps, "steps_unit": "scheduler_bins" if args.method == "edm" else "denoiser_calls", + "chain_mode": "official" if args.method == "cp" else None, + "chaining_times": times, "solver": policy.noise_scheduler.solver}, + architecture=OmegaConf.to_container(cfg.policy, resolve=False), + training=OmegaConf.to_container(cfg.training, resolve=False)) + with path.open("x") as stream: + json.dump(result, stream, indent=2, allow_nan=False) + stream.write("\n") + print(f"[CP-EVAL] {args.method} {steps}: {result['score']:.6f} -> {path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_manipulation/prepare.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_manipulation/prepare.py new file mode 100644 index 0000000000000000000000000000000000000000..14948b5eff009f395c77a01fd911790459a722b4 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_manipulation/prepare.py @@ -0,0 +1,114 @@ +"""Prepare the official EDM-to-CP route for twelve matched imitation tasks. + +Only configs/manifests are written; this command does not train or evaluate. +PushT model/data settings delegate to the unchanged existing PushT recipe. +""" +import argparse +import hashlib +import json +from pathlib import Path + +from omegaconf import OmegaConf + +from experiments.consistency_manipulation.common import ROOT, TASKS, config_for, dataset_manifest, shape_meta_for_task +from experiments.consistency_pusht.prepare import make_config as original_pusht_config +from experiments.table2.common import atomic_json +from tools.consistency_policy.evaluate import DEFAULT_REPO, verify_source + + +def make_config(task, method, directory, seed=0, smoke=False): + if task not in TASKS or method not in ("edm", "ctmp") or seed < 0: + raise ValueError("Choose a supported task/method and nonnegative training seed") + directory = Path(directory).resolve() + if TASKS[task][0] == "pusht": + return original_pusht_config(task, method, directory, seed, smoke) + cfg = original_pusht_config("pusht", method, directory, seed, smoke) + local = config_for(task, seed, "cuda:0") + shape = shape_meta_for_task(task) + cfg.task = {"name": task, "shape_meta": shape, "dataset": { + "_target_": "experiments.consistency_manipulation.data.ManipulationDictDataset", + "task": task, "dataset_path": str(local.dataset_path), + "horizon": 16, "pad_before": 1, "pad_after": 7, "n_obs_steps": 2, + "seed": 42, "val_ratio": 0.0, "image_cache_frames": 64, + }} + cfg.task_name = task + cfg.name = f"official_{method}_{task}" + cfg.exp_name = "manipulation_table2_protocol_adaptation" + cfg.shape_meta = cfg.policy.shape_meta = shape + cfg.policy.crop_shape = list(local.crop_shape) if TASKS[task][2] else None + cfg.logging.name = cfg.name + cfg.multi_run.wandb_name_base = cfg.name + cfg.experiment.task = task + cfg.experiment.max_episode_steps = int(local.max_episode_steps) + cfg.experiment.legacy = None + cfg.experiment.local_task_config = OmegaConf.to_container(local, resolve=True) + cfg.experiment.adaptations = [ + "Unchanged official EDM/CTMPP models, warm-start, losses and schedulers", + "400 EDM / 500 CP epochs; actual update counts depend on task dataset size", + "All local Table-2 demonstrations, horizon16 / history2 / action8 and exact padding", + "Raw observation/action dictionaries; checkpoint policy owns normalization", + "Exact local affine normalization, including Kitchen's thirty zero-masked goal coordinates", + "Robomimic action rotation axis-angle to six-dimensional representation; same local transform", + "Robomimic image reads are lazy, including the local JPEG2000 level50 codec round-trip", + "Separate EDM teacher, not the existing DDPM teacher; raw model is the distilled reference", + "In-sample diagnostic validation only; final-epoch checkpoint, not Last/Max selection", + "Modern pinned Torch 2.8/CUDA 12.8 runtime; deferred simulator/GPU preflight required", + ] + OmegaConf.resolve(cfg) + if OmegaConf.missing_keys(cfg): + raise ValueError(f"Unresolved configuration: {OmegaConf.missing_keys(cfg)}") + return cfg + + +def prepare(output, tasks=tuple(TASKS), seed=0, smoke=False): + output = Path(output).resolve() + if output.exists(): + raise FileExistsError(f"Choose a fresh prepared-run directory: {output}") + if not tasks or len(set(tasks)) != len(tasks) or any(task not in TASKS for task in tasks) or seed < 0: + raise ValueError("Choose distinct supported tasks and a nonnegative training seed") + source = verify_source(DEFAULT_REPO) + configs, data_manifests = {}, {} + for task in tasks: + data_manifests[task] = dataset_manifest(task) + local = config_for(task, seed) + if not Path(local.dataset_path).exists(): + raise FileNotFoundError(f"Missing verified task dataset: {local.dataset_path}") + directory = output / task / f"seed{seed}" + configs[task] = {method: make_config(task, method, directory, seed, smoke) + for method in ("edm", "ctmp")} + output.mkdir(parents=True) + for task, methods in configs.items(): + directory = output / task / f"seed{seed}" + (directory / "configs").mkdir(parents=True) + paths = {} + for method, config in methods.items(): + path = directory / "configs" / f"{method}.yaml" + OmegaConf.save(config, path, resolve=True) + paths[method] = {"path": str(path), "sha256": hashlib.sha256(path.read_bytes()).hexdigest()} + reference = ROOT / "outputs/table2/chi_unet_20260909_v1/runs" / task / f"seed{seed}/model_1000000.pt" + manifest = data_manifests[task] + atomic_json(directory / "manifest.json", { + "upstream": source, "task": task, "seed": seed, "smoke": smoke, + "configs": paths, "dataset_manifest": str(manifest), + "dataset_manifest_sha256": hashlib.sha256(manifest.read_bytes()).hexdigest(), + "existing_ddpm_reference": str(reference) if reference.is_file() else None, + "ddpm_reference_used_for_distillation": False, + "paper_ready": False, "training_started": False, + }) + atomic_json(output / "preparation.json", {"tasks": list(tasks), "seed": seed, + "smoke": smoke, "upstream": source, "paper_ready": False}) + print(json.dumps({"prepared": str(output), "tasks": list(tasks), "seed": seed, "smoke": smoke})) + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", required=True) + parser.add_argument("--tasks", default=",".join(TASKS)) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--smoke", action="store_true") + args = parser.parse_args(argv) + prepare(args.output, args.tasks.split(","), args.seed, args.smoke) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/data.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/data.py new file mode 100644 index 0000000000000000000000000000000000000000..3990556d6e2906dc720b4b38fd8dcda3d6fce51b --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/data.py @@ -0,0 +1,203 @@ +"""PushT representations for the pinned official EDM and CTMPP policies. + +Samples are RAW observations/actions: the official policy owns normalization. +With the default ``val_ratio=0``, all demonstrations train both methods, matching +Table 2. ``get_validation_dataset`` then returns an explicitly IN-SAMPLE +diagnostic view; its losses must not select checkpoints or tune hyperparameters. +An optional positive ratio instead holds out whole episodes and fits statistics +only on the training episodes. Such a split differs from the Table-2 experiment. + +State is five physical coordinates. Keypoints are nine XY pairs, flattened to +18 coordinates, plus the two-coordinate pusher position. Keypoint normalization +shares XY bounds across all nine points, exactly as CleanDiffuser does. Image +observations are RGB CHW floats in [0,1], plus pusher position. Every variant +predicts the same sequence of absolute two-coordinate pusher targets. +""" +from __future__ import annotations + +import copy +from pathlib import Path + +import numpy as np +import torch + +from diffusion_policy.common.normalize_util import get_image_range_normalizer +from diffusion_policy.common.replay_buffer import ReplayBuffer +from diffusion_policy.common.sampler import SequenceSampler, get_val_mask +from diffusion_policy.dataset.base_dataset import BaseImageDataset +from diffusion_policy.model.common.normalizer import LinearNormalizer, SingleFieldLinearNormalizer + + +TASKS = ("pusht", "pusht_keypoints", "pusht_image") + + +def shape_meta_for_task(task): + if task == "pusht_state": + task = "pusht" + observations = { + "pusht": {"state": {"shape": [5], "type": "low_dim"}}, + "pusht_keypoints": {"keypoint": {"shape": [18], "type": "low_dim"}, + "agent_pos": {"shape": [2], "type": "low_dim"}}, + "pusht_image": {"image": {"shape": [3, 96, 96], "type": "rgb"}, + "agent_pos": {"shape": [2], "type": "low_dim"}}, + } + if task not in observations: + raise ValueError(f"Unsupported PushT task: {task}") + return {"action": {"shape": [2]}, "obs": observations[task]} + + +def _minmax_field(data, repetitions=1): + """Export CleanDiffuser's exact affine rule, including constant channels.""" + values = np.asarray(data, dtype=np.float32).reshape(-1, data.shape[-1]) + minimum, maximum = values.min(axis=0), values.max(axis=0) + span = maximum - minimum + span = np.where(span == 0, np.float32(1), span) + scale = np.float32(2) / span + offset = -np.float32(1) - minimum * scale + stats = {"min": minimum, "max": maximum, "mean": values.mean(axis=0), + "std": values.std(axis=0)} + if repetitions != 1: + scale, offset = np.tile(scale, repetitions), np.tile(offset, repetitions) + stats = {key: np.tile(value, repetitions) for key, value in stats.items()} + result = SingleFieldLinearNormalizer.create_manual( + scale=scale.astype(np.float32), offset=offset.astype(np.float32), + input_stats_dict={key: value.astype(np.float32) for key, value in stats.items()}) + result.requires_grad_(False) + return result + + +class PushTDictDataset(BaseImageDataset): + """Dictionary-observation adapter; also supports low-dimensional policies. + +``BaseImageDataset`` is the upstream workspace's interface name, not a promise +that the dictionary contains images. Hybrid policies accept low-dim-only +``shape_meta`` without a vision encoder. + """ + + def __init__(self, task, dataset_path, horizon=16, pad_before=1, pad_after=7, + seed=42, val_ratio=0.0): + super().__init__() + self.task = "pusht" if task == "pusht_state" else task + self.shape_meta = shape_meta_for_task(self.task) + if not 0 <= float(val_ratio) < 1: + raise ValueError("val_ratio must be in [0, 1)") + if horizon < 1 or not 0 <= pad_before < horizon or not 0 <= pad_after < horizon: + raise ValueError("Require positive horizon and padding smaller than horizon") + self.dataset_path = str(Path(dataset_path).expanduser().resolve(strict=True)) + self.horizon, self.pad_before, self.pad_after = horizon, pad_before, pad_after + self.seed, self.val_ratio = int(seed), float(val_ratio) + keys = ["state", "action"] + if self.task == "pusht_keypoints": + keys.append("keypoint") + elif self.task == "pusht_image": + keys.append("img") + self.replay_buffer = ReplayBuffer.copy_from_path(self.dataset_path, keys=keys) + if self.replay_buffer.n_episodes < 1: + raise ValueError("PushT dataset has no episodes") + if self.val_ratio > 0 and self.replay_buffer.n_episodes < 2: + raise ValueError("A disjoint split needs at least two episodes") + expected_shapes = {"state": (5,), "action": (2,), "keypoint": (9, 2), + "img": (96, 96, 3)} + for key in keys: + value = self.replay_buffer[key] + if value.shape != (self.replay_buffer.n_steps,) + expected_shapes[key]: + raise ValueError(f"Unexpected PushT {key} shape: {value.shape}") + if not np.isfinite(value[:]).all(): + raise ValueError(f"Nonfinite values in PushT {key}") + if self.task == "pusht_image": + pixels = self.replay_buffer["img"] + # The released Zarr stores 0..255 pixels as float32; some copies + # use uint8. Neither storage dtype means they are already [0,1]. + if pixels.min() < 0 or pixels.max() > 255: + raise ValueError("PushT source images must be in [0,255]") + if np.issubdtype(pixels.dtype, np.floating) and pixels.max() <= 1: + raise ValueError("Ambiguous [0,1] source pixels; expected the original 0..255 archive") + + self.val_mask = get_val_mask(self.replay_buffer.n_episodes, self.val_ratio, self.seed) + self.train_mask = ~self.val_mask + self.episode_mask = self.train_mask.copy() + self.validation_in_sample = self.val_ratio == 0 + self.role = "train" + self.sampler = self._make_sampler(self.episode_mask) + if len(self.sampler) == 0: + raise ValueError("No training sequences fit this horizon/padding") + self._normalizer = self._make_normalizer() + + def _make_sampler(self, mask): + return SequenceSampler(self.replay_buffer, sequence_length=self.horizon, + pad_before=self.pad_before, pad_after=self.pad_after, + episode_mask=mask) + + def _training_frames(self): + ends = self.replay_buffer.episode_ends[:] + lengths = np.diff(np.concatenate(([0], ends))) + return np.repeat(self.train_mask, lengths) + + def _make_normalizer(self): + mask = self._training_frames() + state = self.replay_buffer["state"][mask] + result = LinearNormalizer() + result["action"] = _minmax_field(self.replay_buffer["action"][mask]) + if self.task == "pusht": + result["state"] = _minmax_field(state) + else: + result["agent_pos"] = _minmax_field(state[:, :2]) + if self.task == "pusht_keypoints": + result["keypoint"] = _minmax_field( + self.replay_buffer["keypoint"][mask], repetitions=9) + else: + result["image"] = get_image_range_normalizer() + result.requires_grad_(False) + return result + + @property + def metadata(self): + evaluation_mask = self.train_mask if self.validation_in_sample else self.val_mask + return {"task": self.task, "dataset_path": self.dataset_path, + "episodes_total": int(self.replay_buffer.n_episodes), + "frames_total": int(self.replay_buffer.n_steps), "samples": len(self), + "train_episode_indices": np.flatnonzero(self.train_mask).tolist(), + "validation_episode_indices": np.flatnonzero(evaluation_mask).tolist(), + "view_episode_indices": np.flatnonzero(self.episode_mask).tolist(), + "role": self.role, "validation_in_sample": self.validation_in_sample, + "validation_use": "diagnostic_only_not_checkpoint_selection", + "normalizer_fit": "training_frames_only", + "normalizer_rule": "CleanDiffuser_MinMax_constant_range_1; image_2x_minus_1", + "keypoint_scaling": "shared_xy_across_nine_keypoints", + "seed": self.seed, "val_ratio": self.val_ratio, + "horizon": self.horizon, "pad_before": self.pad_before, + "pad_after": self.pad_after, "shape_meta": copy.deepcopy(self.shape_meta)} + + def get_validation_dataset(self): + view = copy.copy(self) + view.episode_mask = (self.train_mask if self.validation_in_sample else self.val_mask).copy() + view.role = "in_sample_diagnostic" if self.validation_in_sample else "heldout_diagnostic" + view.sampler = self._make_sampler(view.episode_mask) + if len(view.sampler) == 0: + raise ValueError("No diagnostic sequences fit this horizon/padding") + return view + + def get_normalizer(self, **kwargs): + if kwargs: + raise ValueError("Normalizer settings are fixed to the matched Table-2 representation") + return copy.deepcopy(self._normalizer) + + def get_all_actions(self): + return torch.from_numpy(self.replay_buffer["action"][self._training_frames()].copy()) + + def __len__(self): + return len(self.sampler) + + def __getitem__(self, index): + sample = self.sampler.sample_sequence(index) + state = np.asarray(sample["state"], dtype=np.float32) + if self.task == "pusht": + obs = {"state": state} + else: + obs = {"agent_pos": state[:, :2]} + if self.task == "pusht_keypoints": + obs["keypoint"] = np.asarray(sample["keypoint"], dtype=np.float32).reshape(self.horizon, 18) + else: + obs["image"] = np.moveaxis(sample["img"], -1, 1).astype(np.float32) / np.float32(255) + return {"obs": {key: torch.from_numpy(np.ascontiguousarray(value)) for key, value in obs.items()}, + "action": torch.from_numpy(np.array(sample["action"], dtype=np.float32, copy=True))} diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/evaluate.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/evaluate.py new file mode 100644 index 0000000000000000000000000000000000000000..29ffbca0da7dd8f75efb478424471d0a363dc980 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/evaluate.py @@ -0,0 +1,276 @@ +"""Official EDM / CP policies on the frozen local Table-2 PushT protocol. + +Raw observations enter the checkpoint's own normalizer. In particular, these +are not already-normalized CleanDiffuser dataset observations. The default EDM +reference uses raw ``model`` weights, which official CP actually distills. +Concurrent-GPU timings are diagnostic, never paper-ready latency measurements. +""" +import argparse +import importlib.metadata +import json +import math +import platform +import random +import sys +from pathlib import Path + +import numpy as np +import torch + +ROOT = Path(__file__).resolve().parents[2] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from tools.consistency_policy.evaluate import ( + DEFAULT_REPO, TimedPolicy, checkpoint_hash, load_policy, parse_steps, + set_sampling, verify_source, +) + +TASKS = { + "pusht": {"state": (5,)}, + "pusht_keypoints": {"keypoint": (18,), "agent_pos": (2,)}, + "pusht_image": {"image": (3, 96, 96), "agent_pos": (2,)}, +} +ENV_NAMES = {"pusht": "pusht-v0", "pusht_keypoints": "pusht-keypoints-v0", + "pusht_image": "pusht-image-v0"} + + +def observation_dict(task, observations, device): + """Convert a batch of two-frame *raw* local env observations to DP keys.""" + if task not in TASKS or not observations: + raise ValueError("Unknown task or empty observation batch") + if task == "pusht_image": + values = {key: np.stack([obs[key] for obs in observations]) for key in TASKS[task]} + else: + values = np.stack(observations) + if task == "pusht": + values = {"state": values} + else: + if values.shape[1:] != (2, 40): + raise ValueError("Expected 2 x 40 keypoint observation including visibility masks") + if not np.all(values[..., 20:] == 1): + raise ValueError("This full-visibility protocol cannot discard missing-keypoint masks") + values = {"keypoint": values[..., :18], "agent_pos": values[..., 18:20]} + result = {} + for key, shape in TASKS[task].items(): + value = np.asarray(values[key]) + if value.shape != (len(observations), 2, *shape): + raise ValueError(f"Incorrect {key} observation shape: {value.shape}") + if not np.isfinite(value).all(): + raise FloatingPointError("Nonfinite observation: " + key) + if key == "image" and (value.min() < 0 or value.max() > 1): + raise ValueError("Raw image observations must be CHW float pixels in [0, 1]") + result[key] = torch.as_tensor(value, device=device, dtype=torch.float32) + return result + + +def validate_policy(policy, task, shape_meta=None): + if task not in TASKS: + raise ValueError("Unsupported PushT task: " + task) + for attribute, expected in (("horizon", 16), ("n_obs_steps", 2), + ("n_action_steps", 8), ("action_dim", 2)): + if int(getattr(policy, attribute)) != expected: + raise ValueError(f"Checkpoint {attribute} must equal {expected}") + if getattr(policy, "use_kde", False): + raise ValueError("KDE selection violates the one-candidate protocol") + if shape_meta is not None: + expected = TASKS[task] + actual = {key: tuple(spec["shape"]) for key, spec in shape_meta["obs"].items()} + if actual != expected or tuple(shape_meta["action"]["shape"]) != (2,): + raise ValueError("Checkpoint observation/action shapes do not match --task") + for key, spec in shape_meta["obs"].items(): + if spec.get("type", "low_dim") != ("rgb" if key == "image" else "low_dim"): + raise ValueError("Incorrect observation modality: " + key) + + +def make_env(task, max_steps=300): + import gym + from cleandiffuser.env import pusht # noqa: F401: registers local environments + from cleandiffuser.env.wrapper import MultiStepWrapper + return MultiStepWrapper(gym.make(ENV_NAMES[task], legacy=False), + n_obs_steps=2, n_action_steps=8, + max_episode_steps=max_steps, reward_agg_method="max") + + +def evaluate(policy, task, episodes=50, envs=10, seed_base=100000, + eval_seed=0, max_steps=300, warmup=3): + """RNG-preserving batched rollouts; terminated environments are never stepped.""" + validate_policy(policy, task) + if episodes < 1 or envs < 1 or not 1 <= max_steps <= 300 or warmup < 0: + raise ValueError("Invalid episode/environment count, limit or warmup") + device = torch.device(policy.device) + devices = [device.index or 0] if device.type == "cuda" else [] + numpy_state, python_state = np.random.get_state(), random.getstate() + was_training = policy.training + environments, rows, batch_sizes = [], [], [] + timed = TimedPolicy(policy, warmup) + try: + with torch.random.fork_rng(devices=devices), torch.no_grad(): + random.seed(eval_seed) + np.random.seed(eval_seed) + torch.manual_seed(eval_seed) + policy.eval() + for _ in range(min(envs, episodes)): + environments.append(make_env(task, max_steps)) + for start in range(0, episodes, len(environments)): + count = min(len(environments), episodes - start) + observations = [] + for idx, env in enumerate(environments[:count]): + env.seed(seed_base + start + idx) + observations.append(env.reset()) + policy.reset() + finished = [False] * count + decision_rewards = [[] for _ in range(count)] + for _ in range(math.ceil(max_steps / 8)): + if all(finished): + break + batch = observation_dict(task, observations, device) + output = timed.predict_action(batch) + if tuple(output["action"].shape) != (count, 8, 2): + raise ValueError("Policy must return the already-sliced B x 8 x 2 action chunk") + actions = output["action"].detach().cpu().numpy() + if not np.isfinite(actions).all(): + raise FloatingPointError("Nonfinite policy action") + # Official predict_action already unnormalizes and slices + # action_pred[:, 1:9]. Do not clip or slice a second time. + batch_sizes.append(count) + for idx, env in enumerate(environments[:count]): + if finished[idx]: + continue + obs, reward, done, _ = env.step(actions[idx]) + if not math.isfinite(float(reward)): + raise FloatingPointError("Nonfinite environment reward") + observations[idx] = obs + decision_rewards[idx].append(float(reward)) + finished[idx] = bool(done) + if not all(finished): + raise RuntimeError("Rollout failed to terminate at the configured primitive-step limit") + for idx, env in enumerate(environments[:count]): + rewards = [float(value) for value in env.get_rewards()] + if not rewards or not np.isfinite(rewards).all() or len(rewards) > max_steps: + raise ValueError("Invalid primitive reward trace") + score = max(rewards) + if not 0 <= score <= 1: + raise ValueError("PushT normalized coverage must be in [0, 1]") + rows.append({"episode_seed": seed_base + start + idx, "score": score, + "primitive_steps": len(rewards), "rewards": rewards, + "decision_rewards": decision_rewards[idx], "done": finished[idx]}) + print(f"[CP-EVAL] {task} {len(rows)}/{episodes} " + f"score={np.mean([row['score'] for row in rows]):.6f}", flush=True) + finally: + timed.close() + for env in environments: + env.close() + policy.train(was_training) + np.random.set_state(numpy_state) + random.setstate(python_state) + if not timed.nfes or min(timed.nfes) < 1: + raise ValueError("Denoiser forward hook recorded no calls") + return { + "task": task, "score": float(np.mean([row["score"] for row in rows])), + "score_metric": "mean_episode_max_normalized_coverage", "episodes": rows, + "num_episodes": episodes, "seed_base": seed_base, "eval_seed": eval_seed, + "num_candidates": 1, "horizon": 16, "obs_steps": 2, "action_steps": 8, + "max_episode_steps": max_steps, "pusht_legacy": False, + "env_implementation": "cleandiffuser.env.pusht (same as local Table 2)", + "truncated_smoke": max_steps < 300, "paper_ready": False, + "latency": {"scope": "predict_action_device_resident; includes normalization and observation encoder; excludes observation H2D, action D2H and simulator", + "concurrent_gpu_diagnostic": True, "paper_ready": False, + "mean_ms": float(np.mean(timed.durations)), + "p50_ms": float(np.percentile(timed.durations, 50)), + "p95_ms": float(np.percentile(timed.durations, 95)), + "raw_ms": timed.durations, "warmup_calls": warmup, + "batch_sizes": batch_sizes, "n_envs": min(envs, episodes)}, + "nfe": {"mean": float(np.mean(timed.nfes)), "min": min(timed.nfes), + "max": max(timed.nfes), "raw": timed.nfes, + "unit": "denoiser forwards per batched predict_action, not per episode"}, + } + + +def runtime_provenance(device): + device = torch.device(device) + packages = {} + for name in ("torch", "numpy", "gym", "pymunk", "pygame", "robomimic", "hydra-core", "diffusers"): + try: + packages[name] = importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + packages[name] = None + files = ("experiments/consistency_pusht/evaluate.py", "tools/consistency_policy/evaluate.py", + "cleandiffuser/env/wrapper.py", "cleandiffuser/env/pusht/pusht_env.py", + "cleandiffuser/env/pusht/pusht_keypoints_env.py", "cleandiffuser/env/pusht/pusht_image_env.py") + return {"python": platform.python_version(), "packages": packages, + "device": str(device), "cuda": torch.version.cuda, + "device_name": torch.cuda.get_device_name(device) if device.type == "cuda" else platform.machine(), + "dtype": "float32", "torch_threads": torch.get_num_threads(), + "cudnn_benchmark": torch.backends.cudnn.benchmark, + "tf32_matmul": torch.backends.cuda.matmul.allow_tf32, + "tf32_cudnn": torch.backends.cudnn.allow_tf32, + "local_source_sha256": {name: checkpoint_hash(ROOT / name) for name in files}} + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--task", choices=TASKS, required=True) + parser.add_argument("--method", choices=("edm", "cp"), required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True, help="New result directory") + parser.add_argument("--repo", type=Path, default=DEFAULT_REPO) + parser.add_argument("--steps", help="Default CP: 1,3; EDM: 80 scheduler bins") + parser.add_argument("--weights", choices=("model", "ema_model"), default="model") + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--episodes", type=int, default=50) + parser.add_argument("--envs", type=int, default=10) + parser.add_argument("--seed-base", type=int, default=100000) + parser.add_argument("--eval-seed", type=int, help="Defaults to checkpoint training seed") + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--smoke", action="store_true", help="Two episodes, eight primitive steps; not a result") + args = parser.parse_args(argv) + if args.episodes < 1 or args.envs < 1 or args.warmup < 0: + parser.error("episodes/envs must be positive, warmup nonnegative") + source = verify_source(args.repo) + sys.path.insert(0, source["repo"]) + from omegaconf import OmegaConf + # Official configs use this resolver; unlike upstream eval, never evaluate + # untrusted input. Only locally trained checkpoint pickles are supported. + if not OmegaConf.has_resolver("eval"): + OmegaConf.register_new_resolver("eval", eval) + checkpoint = args.checkpoint.resolve(strict=True) + digest = checkpoint_hash(checkpoint) + steps_list = parse_steps(args.steps or ("1,3" if args.method == "cp" else "80")) + policy, cfg, weights = load_policy(checkpoint, args.method, args.weights, args.device) + validate_policy(policy, args.task, cfg.policy.shape_meta) + paths = [args.output / f"{args.method}_{steps}_{weights}.json" for steps in steps_list] + for steps, path in zip(steps_list, paths): + set_sampling(policy, args.method, steps, "official") + if path.exists(): + raise FileExistsError("Refusing to overwrite result: " + str(path)) + args.output.mkdir(parents=True, exist_ok=True) + training_seed = int(cfg.training.seed) + eval_seed = training_seed if args.eval_seed is None else args.eval_seed + runtime = runtime_provenance(args.device) + for steps, path in zip(steps_list, paths): + times = set_sampling(policy, args.method, steps, "official") + result = evaluate(policy, args.task, episodes=2 if args.smoke else args.episodes, + envs=args.envs, seed_base=args.seed_base, eval_seed=eval_seed, + max_steps=8 if args.smoke else 300, warmup=args.warmup) + if args.method == "cp" and (result["nfe"]["min"] != steps or result["nfe"]["max"] != steps): + raise ValueError("Measured CP denoiser NFE differs from requested 1/3-step mode") + result.update(schema_version=1, method=args.method, weights=weights, + training_seed=training_seed, smoke=args.smoke, + checkpoint=str(checkpoint), checkpoint_sha256=digest, + teacher_reference="raw_model_actually_distilled" if args.method == "edm" and weights == "model" else None, + source=source, runtime=runtime, + sampling={"steps": steps, "steps_unit": "scheduler_bins" if args.method == "edm" else "denoiser_calls", + "chain_mode": "official" if args.method == "cp" else None, + "chaining_times": times, "solver": policy.noise_scheduler.solver}, + architecture=OmegaConf.to_container(cfg.policy, resolve=False), + training=OmegaConf.to_container(cfg.training, resolve=False)) + with path.open("x") as file: + json.dump(result, file, indent=2, allow_nan=False) + file.write("\n") + print(f"[CP-EVAL] {args.method} {steps}: score={result['score']:.6f}, " + f"NFE={result['nfe']['mean']:.1f} -> {path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/launch.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/launch.py new file mode 100644 index 0000000000000000000000000000000000000000..39040ddc3b19866c9c1e352216302b63982442e0 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/launch.py @@ -0,0 +1,243 @@ +"""Detached, fail-fast EDM -> official CP queue, independent of Table-2 jobs. + +Run prepare first. Each task trains one EDM teacher, evaluates it, distils CTMP, +then evaluates CP at one and three steps. No existing run is resumed or replaced. +""" +from __future__ import annotations + +import argparse +import fcntl +import hashlib +import json +import os +from pathlib import Path +import subprocess +import time + + +ROOT = Path(__file__).resolve().parents[2] +TASKS = ("pusht", "pusht_keypoints", "pusht_image") +UPSTREAM_COMMIT = "eed0c465d3c1aa6dcf22f10e784df07dec06beac" + + +def atomic_json(path, value): + path = Path(path) + temporary = path.with_name(f".{path.name}.{os.getpid()}.{time.time_ns()}.tmp") + try: + with temporary.open("x") as stream: + json.dump(value, stream, indent=2, allow_nan=False) + stream.write("\n") + temporary.replace(path) + finally: + temporary.unlink(missing_ok=True) + + +def file_hash(path): + return hashlib.sha256(Path(path).read_bytes()).hexdigest() + + +def source_manifest(): + """Reject upstream edits and fingerprint the small shared execution surface.""" + upstream = ROOT / "third_party/consistency-policy" + revision = subprocess.check_output( + ["git", "-C", str(upstream), "rev-parse", "HEAD"], text=True).strip() + dirty = subprocess.check_output( + ["git", "-C", str(upstream), "status", "--porcelain", "--untracked-files=no"], + text=True).strip() + if revision != UPSTREAM_COMMIT or dirty: + raise RuntimeError("Official Consistency Policy must be clean at the pinned revision") + source_dir = ROOT / "experiments/consistency_pusht" + paths = set(source_dir.glob("*.py")) | set(source_dir.glob("*.sh")) + paths.update(ROOT / name for name in ( + "experiments/table2/common.py", "experiments/table2/evaluate.py", + "experiments/table2/run.sh", "tools/consistency_policy/evaluate.py", + "cleandiffuser/dataset/pusht_dataset.py", "cleandiffuser/utils/normalizers.py", + "cleandiffuser/env/wrapper.py", "cleandiffuser/env/pusht/pusht_env.py", + "cleandiffuser/env/pusht/pusht_keypoints_env.py", + "cleandiffuser/env/pusht/pusht_image_env.py")) + files = {str(path.relative_to(ROOT)): file_hash(path) for path in sorted(paths)} + payload = {"files": files, "upstream_commit": revision} + payload["sha256"] = hashlib.sha256( + json.dumps(payload, sort_keys=True).encode()).hexdigest() + return payload + + +def phase_specs(directory, tasks, seed=0, smoke=False): + launcher = ["bash", str(ROOT / "experiments/consistency_pusht/run.sh"), "-m"] + phases = [] + for task in tasks: + run = directory / task / f"seed{seed}" + for method in ("edm", "ctmp"): + checkpoint = run / method / "checkpoints/final.ckpt" + command = launcher + ["experiments.consistency_pusht.train", "--config", + str(run / "configs" / f"{method}.yaml")] + if smoke: + command.append("--smoke") + phases.append({"task": task, "seed": seed, "phase": f"train_{method}", + "command": command, "checkpoint": str(checkpoint)}) + label = "edm" if method == "edm" else "cp" + command = launcher + ["experiments.consistency_pusht.evaluate", "--task", task, + "--method", label, "--checkpoint", str(checkpoint), + "--output", str(run / f"eval_{label}"), + "--episodes", "2" if smoke else "50", "--envs", "2" if smoke else "10"] + if label == "cp": + command += ["--steps", "1,3"] + if smoke: + command.append("--smoke") + steps = (80,) if label == "edm" else (1, 3) + phases.append({"task": task, "seed": seed, "phase": f"eval_{label}", + "command": command, + "results": [str(run / f"eval_{label}" / f"{label}_{n}_model.json") + for n in steps]}) + return phases + + +def prepare_plan(directory, tasks, seed=0, smoke=False): + directory = Path(directory).resolve() + base = ROOT / "outputs/consistency_pusht" + if not directory.is_relative_to(base) or directory == base: + raise ValueError("Use a prepared run below outputs/consistency_pusht") + if not tasks or len(tasks) != len(set(tasks)) or any(task not in TASKS for task in tasks): + raise ValueError("Choose unique supported PushT tasks") + if seed < 0: + raise ValueError("Seed must be nonnegative") + configs = {} + for task in tasks: + run = directory / task / f"seed{seed}" + for name in ("edm", "ctmp", "eval_edm", "eval_cp"): + if (run / name).exists(): + raise FileExistsError(f"Refusing to reuse an existing run: {run / name}") + for name in ("edm.yaml", "ctmp.yaml"): + path = run / "configs" / name + configs[str(path.relative_to(directory))] = file_hash(path) + for path in run.glob("*manifest*.json"): + configs[str(path.relative_to(directory))] = file_hash(path) + manifest = json.loads(path.read_text()) + if manifest.get("smoke", smoke) != smoke: + raise ValueError(f"--smoke differs from prepared manifest: {path}") + for name in ("plan.json", "status.json", "launcher.json", "queue.log", "logs", ".worker.lock"): + if (directory / name).exists(): + raise FileExistsError(f"Refusing to restart an existing queue: {directory / name}") + plan = {"tasks": list(tasks), "seed": seed, "smoke": smoke, + "created_unix": time.time(), "source": source_manifest(), "configs": configs, + "phases": phase_specs(directory, tasks, seed, smoke), + "serial_within_queue": True, "independent_of_table2_queue": True, + "concurrent_gpu_timing_is_diagnostic_only": True, "paper_ready": False, + "note": "Official EDM-to-CP route; one seed; does not reproduce Table-2 Last/Max"} + # Exclusive creation is the launch claim: two simultaneous launchers cannot + # both start workers, even if both passed the preceding read-only checks. + with (directory / "plan.json").open("x") as stream: + json.dump(plan, stream, indent=2, allow_nan=False) + stream.write("\n") + (directory / "logs").mkdir() + atomic_json(directory / "status.json", {"state": "launching", "completed": [], + "phase_count": len(plan["phases"]), "created_unix": time.time()}) + return plan + + +def verify_frozen_inputs(directory, plan): + if source_manifest()["sha256"] != plan["source"]["sha256"]: + raise RuntimeError("Source files changed; refusing to start the next phase") + for relative, digest in plan["configs"].items(): + if file_hash(directory / relative) != digest: + raise RuntimeError(f"Prepared config/manifest changed: {relative}") + + +def worker(directory): + directory = Path(directory).resolve() + plan = json.loads((directory / "plan.json").read_text()) + with (directory / ".worker.lock").open("a") as lock: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + status = json.loads((directory / "status.json").read_text()) + if status["state"] != "launching": + raise RuntimeError("Queue already started; automatic resume is intentionally disabled") + os.nice(10) # CPU priority only; this does not imply GPU scheduling priority. + completed = [] + status = {"state": "running", "pid": os.getpid(), "completed": completed, + "phase_count": len(plan["phases"]), "smoke": plan["smoke"], + "started_unix": time.time(), "concurrent_gpu_timing_is_diagnostic_only": True} + atomic_json(directory / "status.json", status) + try: + for phase in plan["phases"]: + verify_frozen_inputs(directory, plan) + log = directory / "logs" / f"{phase['task']}_seed{phase['seed']}_{phase['phase']}.log" + status.update(state="running", task=phase["task"], seed=phase["seed"], + phase=phase["phase"], command=phase["command"], log=str(log), + phase_started_unix=time.time(), child_pid=None) + atomic_json(directory / "status.json", status) + with log.open("x") as stream: + process = subprocess.Popen(phase["command"], cwd=ROOT, + stdin=subprocess.DEVNULL, stdout=stream, stderr=subprocess.STDOUT, + env={**os.environ, "PYTHONUNBUFFERED": "1"}) + status["child_pid"] = process.pid + atomic_json(directory / "status.json", status) + print(json.dumps(status), flush=True) + code = process.wait() + status["child_pid"] = None + status["last_exit_code"] = code + if code: + raise RuntimeError(f"{phase['task']} {phase['phase']} exited {code}; see {log}") + if "checkpoint" in phase: + checkpoint = Path(phase["checkpoint"]) + if not checkpoint.is_file() or checkpoint.stat().st_size == 0: + raise RuntimeError(f"Training exited without a final checkpoint: {checkpoint}") + for result_path in phase.get("results", []): + result_path = Path(result_path) + if not result_path.is_file(): + raise RuntimeError(f"Evaluation exited without its result: {result_path}") + result = json.loads(result_path.read_text()) + if result.get("smoke") is not plan["smoke"]: + raise RuntimeError(f"Evaluation smoke mode differs from plan: {result_path}") + completed.append({"task": phase["task"], "seed": phase["seed"], + "phase": phase["phase"], "finished_unix": time.time()}) + atomic_json(directory / "status.json", status) + status.update(state="complete", finished_unix=time.time(), child_pid=None) + atomic_json(directory / "status.json", status) + except BaseException as error: + status.update(state="failed", error=repr(error), finished_unix=time.time()) + atomic_json(directory / "status.json", status) + raise + + +def launch(directory, tasks=TASKS, seed=0, smoke=False, foreground=False): + directory = Path(directory).resolve() + plan = prepare_plan(directory, tasks, seed, smoke) + if foreground: + worker(directory) + return {"output": str(directory), "foreground": True} + command = ["bash", str(ROOT / "experiments/consistency_pusht/run.sh"), "-m", + "experiments.consistency_pusht.launch", "--output", str(directory), "--worker"] + try: + with (directory / "queue.log").open("x") as stream: + process = subprocess.Popen(command, cwd=ROOT, stdin=subprocess.DEVNULL, + stdout=stream, stderr=subprocess.STDOUT, start_new_session=True, + env={**os.environ, "PYTHONUNBUFFERED": "1"}) + receipt = {"queue_pid": process.pid, "output": str(directory), "command": command, + "created_unix": time.time(), "phase_count": len(plan["phases"]), + "smoke": smoke, "log": str(directory / "queue.log")} + atomic_json(directory / "launcher.json", receipt) + return receipt + except Exception as error: + atomic_json(directory / "status.json", {"state": "failed", "error": repr(error), + "completed": [], "finished_unix": time.time()}) + raise + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--tasks", default=",".join(TASKS)) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--smoke", action="store_true") + parser.add_argument("--foreground", action="store_true") + parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS) + args = parser.parse_args() + if args.worker: + worker(args.output) + else: + print(json.dumps(launch(args.output, args.tasks.split(","), args.seed, + args.smoke, args.foreground), indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/prepare.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/prepare.py new file mode 100644 index 0000000000000000000000000000000000000000..b845bba7b6f63dc2a14e6beba1b6bbc096276359 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/prepare.py @@ -0,0 +1,149 @@ +"""Prepare official EDM/CTMP with the Table-2 PushT representations.""" +import argparse +import hashlib +import json +from pathlib import Path + +from omegaconf import OmegaConf + +from experiments.consistency_pusht.data import shape_meta_for_task +from experiments.table2.common import ROOT, atomic_json, config_for, dataset_manifest +from tools.consistency_policy.evaluate import DEFAULT_REPO, verify_source + +TASKS = ("pusht", "pusht_keypoints", "pusht_image") + + +def make_config(task, method, directory, seed=0, smoke=False): + cfg = OmegaConf.load(DEFAULT_REPO / "configs" / f"{method}_square.yaml") + local = config_for(task, seed) + shape = shape_meta_for_task(task) + cfg.task = { + "name": task, "shape_meta": shape, + "dataset": { + "_target_": "experiments.consistency_pusht.data.PushTDictDataset", + "task": task, "dataset_path": str(local.dataset_path), + "horizon": 16, "pad_before": 1, "pad_after": 7, + "seed": 42, "val_ratio": 0.0, + }, + } + cfg.task_name = task + cfg.name = f"official_{method}_{task}" + cfg.exp_name = "pusht_table2_protocol_adaptation" + cfg.shape_meta = cfg.policy.shape_meta = shape + cfg.policy.crop_shape = [84, 84] if task == "pusht_image" else None + cfg.policy.down_dims = [512, 1024, 2048] + cfg.policy.diffusion_step_embed_dim = 128 + cfg.horizon = cfg.policy.horizon = 16 + cfg.n_obs_steps = cfg.dataset_obs_steps = cfg.policy.n_obs_steps = 2 + cfg.n_action_steps = cfg.policy.n_action_steps = 8 + cfg.n_latency_steps = 0 + cfg.past_action_visible = False + cfg.training.seed = seed + cfg.training.device = "cuda:0" + cfg.training.output_dir = str(directory / method) + cfg.training.resume = False + cfg.training.resume_path = "None" + cfg.training.inference_mode = cfg.policy.inference_mode = False + cfg.training.online_rollouts = False + # Preserve upstream epochs/optimizer/losses. Validation is a small in-sample + # diagnostic, never a selection criterion; evaluation uses a final checkpoint. + cfg.training.max_val_steps = 2 + cfg.training.val_every = 50 + cfg.training.val_sample_every = 50 + cfg.training.sample_every = 50 + cfg.training.tqdm_interval_sec = 30 + cfg.training.gradient_accumulate_every = 1 + cfg.training.checkpoint_every = 50 + for loader in (cfg.dataloader, cfg.val_dataloader): + loader.batch_size = 64 + loader.num_workers = 2 + loader.persistent_workers = False + cfg.val_dataloader.shuffle = False + if method == "ctmp": + teacher = directory / "edm/checkpoints/final.ckpt" + cfg.policy.teacher_path = cfg.policy.edm = str(teacher) + cfg.logging.mode = "disabled" + cfg.logging.name = cfg.name + cfg.logging.resume = False + cfg.logging.id = None + cfg.checkpoint.topk.k = 0 + cfg.checkpoint.save_last_snapshot = False + cfg.checkpoint.save_last_ckpt = True + cfg.multi_run.run_dir = str(directory / method) + cfg.multi_run.wandb_name_base = cfg.name + if smoke: + # Real task architecture and batch size; only the budget is truncated. + cfg.training.num_epochs = 1 + cfg.training.max_train_steps = 2 + cfg.training.max_val_steps = 1 + cfg.dataloader.num_workers = cfg.val_dataloader.num_workers = 0 + if method == "ctmp": + cfg.training.p_epochs = 1 + cfg.experiment = { + "task": task, "method": method, "smoke": smoke, + "validation_in_sample": True, "checkpoint_selection": "final_epoch", + "num_candidates": 1, "evaluation_episodes": 2 if smoke else 50, + "episode_seed_base": 100000, "max_episode_steps": 300, + "legacy": False, "teacher_weights": "model", + "competing_gpu_workload": True, "paper_ready": False, + "memory_fraction": 0.65, "sleep_per_update": 0.02, + "adaptations": [ + "PushT state/keypoints/image task metadata; unchanged official models/losses", + "All 206 demonstrations and Table-2 normalization; no held-out selection", + "Local legacy=False simulator, matching the completed Table-2 teachers", + "Separate EDM teacher; NOT distilled from the completed DDPM checkpoint", + "Final-epoch rollouts only; not a Table-2 Last/Max checkpoint sweep", + "Isolated Torch 2.8/CUDA 12.8 runtime for RTX5090", + "In-sample diagnostic validation capped at two batches every 50 epochs", + ], + } + OmegaConf.resolve(cfg) + if OmegaConf.missing_keys(cfg): + raise ValueError(f"Unresolved configuration: {OmegaConf.missing_keys(cfg)}") + return cfg + + +def prepare(output, tasks=TASKS, seed=0, smoke=False): + output = Path(output).resolve() + if output.exists(): + raise FileExistsError(f"Choose a new output directory: {output}") + if not tasks or len(set(tasks)) != len(tasks) or any(t not in TASKS for t in tasks): + raise ValueError("Select distinct supported PushT tasks") + source = verify_source(DEFAULT_REPO) + configs = {} + for task in tasks: + dataset_manifest(task) + directory = output / task / f"seed{seed}" + configs[task] = {m: make_config(task, m, directory, seed, smoke) for m in ("edm", "ctmp")} + output.mkdir(parents=True) + for task, methods in configs.items(): + directory = output / task / f"seed{seed}" + (directory / "configs").mkdir(parents=True) + paths = {} + for method, cfg in methods.items(): + path = directory / "configs" / f"{method}.yaml" + OmegaConf.save(cfg, path, resolve=True) + paths[method] = {"path": str(path), "sha256": hashlib.sha256(path.read_bytes()).hexdigest()} + atomic_json(directory / "manifest.json", { + "upstream": source, "task": task, "seed": seed, "smoke": smoke, + "configs": paths, "existing_ddpm_reference": str(ROOT / "outputs/table2/chi_unet_20260909_v1/runs" / task / "seed0/model_1000000.pt"), + "dataset_manifest": str(dataset_manifest(task)), + "paper_ready": False, "training_started": False, + }) + atomic_json(output / "preparation.json", {"tasks": list(tasks), "seed": seed, + "smoke": smoke, "upstream": source, "paper_ready": False}) + print(json.dumps({"prepared": str(output), "tasks": list(tasks), "smoke": smoke})) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--output", required=True) + p.add_argument("--tasks", default=",".join(TASKS)) + p.add_argument("--seed", type=int, default=0) + p.add_argument("--smoke", action="store_true") + args = p.parse_args() + prepare(args.output, args.tasks.split(","), args.seed, args.smoke) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/run.sh b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/run.sh new file mode 100644 index 0000000000000000000000000000000000000000..68c9fdb1fc4758b0d687fb9bef2b7cebd4cdd123 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/run.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +set -euo pipefail +cp_task_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)" +export PYTHONPATH="$cp_task_root/outputs/consistency_pusht/runtime/packages:$cp_task_root/outputs/consistency_pusht/runtime/pytorch3d:$cp_task_root/third_party/consistency-policy:$cp_task_root${PYTHONPATH:+:$PYTHONPATH}" +# Only trusted local EDM/CP checkpoints use the legacy upstream torch.load API. +export TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 +export WANDB_MODE=disabled +export WANDB_SILENT=true +export OMP_NUM_THREADS=2 +export MKL_NUM_THREADS=2 +export HF_HUB_OFFLINE=1 +export DIFFUSERS_OFFLINE=1 +exec bash "$cp_task_root/experiments/table2/run.sh" "$@" diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/train.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/train.py new file mode 100644 index 0000000000000000000000000000000000000000..ef3ddf85b715c2a211b7eb8e34b581157b29e575 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/consistency_pusht/train.py @@ -0,0 +1,159 @@ +"""Run the pinned official workspace with finite-loss/provenance guards. + +The training loop, warm start, CTMP+DSM objective, dropout, and schedulers are +upstream. This wrapper only audits progress, caps memory, and saves the exact +final checkpoint after the upstream periodic checkpoint thread has finished. +""" +import argparse +import hashlib +import json +import os +from pathlib import Path +import subprocess +import time +import types + +import hydra +from omegaconf import OmegaConf +import torch + +from experiments.table2.common import ROOT, atomic_json +from tools.consistency_policy.evaluate import DEFAULT_REPO, verify_source + + +def sha256(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(8 * 1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def run(config_path, smoke=False): + cfg = OmegaConf.load(config_path) + if bool(cfg.experiment.smoke) != smoke: + raise ValueError("--smoke must agree with the prepared configuration") + source = verify_source(DEFAULT_REPO) + output = Path(cfg.training.output_dir) + if output.exists(): + raise FileExistsError(f"Refusing to overwrite training output: {output}") + if cfg.training.resume or cfg.training.gradient_accumulate_every != 1: + raise ValueError("This initial-run wrapper requires resume=False, accumulation=1") + if not torch.cuda.is_available(): + raise RuntimeError("Official PushT CP runs require host CUDA access") + torch.cuda.set_per_process_memory_fraction(float(cfg.experiment.memory_fraction), cfg.training.device) + os.nice(10) + output.mkdir(parents=True) + metadata = { + "upstream": source, "config": OmegaConf.to_container(cfg, resolve=True), + "config_sha256": sha256(config_path), "pid": os.getpid(), + "started_unix": time.time(), "torch": torch.__version__, + "gpu": torch.cuda.get_device_name(cfg.training.device), + "smoke": smoke, "paper_ready": False, + "base_commit": subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=ROOT, text=True).strip(), + "wrapper_source": {p.name: sha256(p) for p in Path(__file__).parent.glob("*.py")}, + } + if cfg.experiment.method == "ctmp": + teacher = Path(cfg.policy.teacher_path) + if not teacher.is_file() or str(teacher) != cfg.policy.edm: + raise ValueError("An exact local EDM teacher/warm-start checkpoint is required") + import dill + saved = torch.load(teacher, map_location="cpu", pickle_module=dill, weights_only=False) + teacher_cfg = saved["cfg"] + if teacher_cfg.experiment.task != cfg.experiment.task or teacher_cfg.experiment.method != "edm": + raise ValueError("Teacher task/method mismatch") + if bool(teacher_cfg.experiment.smoke) != smoke: + raise ValueError("Smoke teacher cannot seed a full run") + if teacher_cfg.task.dataset != cfg.task.dataset or teacher_cfg.policy.shape_meta != cfg.policy.shape_meta: + raise ValueError("Teacher and CP data/representation mismatch") + metadata["teacher_sha256"] = sha256(teacher) + del saved + atomic_json(output / "run.json", metadata) + (output / "environment.txt").write_text(subprocess.check_output( + [os.sys.executable, "-m", "pip", "freeze"], text=True)) + atomic_json(output / "status.json", {"state": "initializing", "pid": os.getpid()}) + try: + workspace = hydra.utils.get_class(cfg._target_)(cfg) + metadata["parameters"] = sum(p.numel() for p in workspace.model.parameters()) + metadata["denoiser_parameters"] = sum(p.numel() for p in workspace.model.model.parameters()) + atomic_json(output / "run.json", metadata) + original_loss = workspace.model.compute_loss + counter = {"updates": 0, "last_losses": {}, "hook": None, "pre_hook": None} + started = time.monotonic() + + def before_update(optimizer, args, kwargs): + parameters = [p for group in optimizer.param_groups for p in group["params"]] + # Diagnostic gate only: finite gradients are left unchanged. + torch.nn.utils.clip_grad_norm_(parameters, float("inf"), error_if_nonfinite=True) + + def after_update(optimizer, args, kwargs): + counter["updates"] += 1 + current = counter["updates"] + if current == 1 or current % 100 == 0: + elapsed = time.monotonic() - started + row = {"state": "training", "pid": os.getpid(), "updates": current, + "epoch": int(workspace.epoch), "losses": counter["last_losses"], + "elapsed_seconds": elapsed, "seconds_per_update": elapsed / current, + "gpu_peak_allocated_bytes": torch.cuda.max_memory_allocated()} + atomic_json(output / "status.json", row) + print("[CP_PROGRESS] " + json.dumps(row), flush=True) + time.sleep(float(cfg.experiment.sleep_per_update)) + + def guarded_loss(model, batch): + if counter["hook"] is None: + counter["hook"] = workspace.optimizer.register_step_post_hook(after_update) + counter["pre_hook"] = workspace.optimizer.register_step_pre_hook(before_update) + result = original_loss(batch) + losses = result if isinstance(result, dict) else {"edm": result} + for key, value in losses.items(): + if not torch.isfinite(value).all(): + raise FloatingPointError(f"Nonfinite {key} at update {counter['updates']}") + if torch.is_grad_enabled(): + counter["last_losses"] = {k: float(v.detach()) for k, v in losses.items()} + return result + + workspace.model.compute_loss = types.MethodType(guarded_loss, workspace.model) + workspace.run() + if workspace._saving_thread is not None: + workspace._saving_thread.join() + if int(workspace.epoch) != int(cfg.training.num_epochs): + raise RuntimeError("Workspace did not finish the declared epoch budget") + if any(not torch.isfinite(p).all() for p in workspace.model.parameters()): + raise FloatingPointError("Nonfinite final policy parameters") + # Official payload layout retained for strict official evaluation loading. + checkpoint = output / "checkpoints/final.ckpt" + # Upstream warm start inserts CPU zero-columns. The official threaded + # saver first copies tensors to CPU; a synchronous GPU payload would + # break that exact warm-start path. Join before atomically publishing. + temporary = output / "checkpoints/final.pending.ckpt" + workspace.save_checkpoint(path=temporary, use_thread=True) + workspace._saving_thread.join() + import dill + verified = torch.load(temporary, map_location="cpu", pickle_module=dill, weights_only=False) + if "model" not in verified["state_dicts"]: + raise RuntimeError("Incomplete final checkpoint") + del verified + temporary.replace(checkpoint) + completed = {"state": "smoke_complete" if smoke else "complete", + "pid": os.getpid(), "epochs": int(workspace.epoch), "updates": counter["updates"], + "checkpoint": str(checkpoint), "checkpoint_sha256": sha256(checkpoint), + "losses": counter["last_losses"], "elapsed_seconds": time.monotonic() - started, + "smoke": smoke, "paper_ready": False} + atomic_json(output / "status.json", completed) + print("[CP_COMPLETE] " + json.dumps(completed), flush=True) + except Exception as error: + atomic_json(output / "failure.json", {"error": repr(error), "pid": os.getpid()}) + atomic_json(output / "status.json", {"state": "failed", "error": repr(error), "pid": os.getpid()}) + raise + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", required=True, type=Path) + parser.add_argument("--smoke", action="store_true") + args = parser.parse_args() + run(args.config, args.smoke) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c24629b86106b85295f6c6d6d2858860a517763a --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/__init__.py @@ -0,0 +1 @@ +# Low-dim FDK distillation (keypoints / lift / square). diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/chi_unet_student.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/chi_unet_student.py new file mode 100644 index 0000000000000000000000000000000000000000..cc899facc165ed0cb6adebedb14be7e660903020 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/chi_unet_student.py @@ -0,0 +1,59 @@ +"""Chi_UNet student body (skips) with a VDV bottleneck. + +Imports ChiUNet1d from cleandiffuser; does not edit table-2 sources. +Teacher-matched width: model_dim=256, dim_mult=[1, 2, 2], obs as global cond. +""" +from __future__ import annotations + +import torch +from torch import nn + +from cleandiffuser.nn_diffusion.chiunet import ChiUNet1d + + +class ChiUNetSkipBackbone(ChiUNet1d): + """Same down/up as the Chi_UNet teacher, with mid features exposed.""" + + def __init__(self, in_dim: int, obs_dim: int, obs_steps: int): + super().__init__( + in_dim, + obs_dim, + obs_steps, + model_dim=256, + emb_dim=256, + dim_mult=[1, 2, 2], + obs_as_global_cond=True, + timestep_emb_type="positional", + ) + self.mid_channels = int(self.mids[0].conv1[0].out_channels) + + def _embed(self, t: torch.Tensor, obs: torch.Tensor): + emb = self.map_emb(self.map_noise(t.float())) + cond = self.global_cond_encoder(torch.flatten(obs, 1)) + return torch.cat([emb, cond], dim=-1) + + def down_mid(self, x: torch.Tensor, emb: torch.Tensor): + """x: [B, C, H] → mid [B, mid_channels, H/4], skips for the up path.""" + skips = [] + for resnet1, resnet2, downsample in self.downs: + x = resnet1(x, emb) + x = resnet2(x, emb) + skips.append(x) + x = downsample(x) + for mid in self.mids: + x = mid(x, emb) + return x, skips + + def up_final(self, x: torch.Tensor, skips: list, emb: torch.Tensor): + for resnet1, resnet2, upsample in self.ups: + x = torch.cat((x, skips.pop()), dim=1) + x = resnet1(x, emb) + x = resnet2(x, emb) + x = upsample(x) + return self.final_conv(x).permute(0, 2, 1) + + def encode_mid(self, traj: torch.Tensor, t: torch.Tensor, obs: torch.Tensor): + """Pool bottleneck features [B, H, C] → [B, mid_channels] (no skips).""" + emb = self._embed(t, obs) + mid, _ = self.down_mid(traj.permute(0, 2, 1), emb) + return mid.mean(dim=-1), emb diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/diagnose_square.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/diagnose_square.py new file mode 100644 index 0000000000000000000000000000000000000000..b54b9dfa896eac3b24605dcebb7d14073201efb0 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/diagnose_square.py @@ -0,0 +1,308 @@ +"""Square 0.36-model diagnostics: capacity vs Chi_UNet, per-dim MSE, fail/success videos.""" +from __future__ import annotations + +import argparse +from fractions import Fraction +import json +from pathlib import Path +from unittest.mock import patch + +import gym +import numpy as np +from omegaconf import OmegaConf +from PIL import Image, ImageDraw, ImageFont +import torch + +from cleandiffuser.env.utils import VideoRecorder +from cleandiffuser.nn_diffusion import ChiUNet1d +from experiments.lowdim_fdk_distill.run import ( + EXECUTE, + HORIZON, + LowdimFDKStudent, + SUPPORTED, + load_teacher, +) +from experiments.table2 import evaluate as evaluation +from experiments.table2.common import ROOT, atomic_json, build_dataset, seed_all +from experiments.teacher_videos.record_pusht_state import sha256, validate_video + + +BEST_ALIGN = ROOT / ( + "outputs/lowdim_fdk_distill/square_lowdim_best_align_filmgate_20260915_v1/" + "arm_vdv/best_ema.pt" +) +PAIRS = ROOT / "outputs/lowdim_fdk_distill/square_lowdim_fdk_20260914_v1/pairs.pt" +SUCCESS_SEEDS = (100001, 100007) +FAIL_SEEDS = (100000, 100002) + + +def count_params(module): + return int(sum(p.numel() for p in module.parameters())) + + +def capacity_report(): + fdk128 = LowdimFDKStudent(23, 10, latent_dim=128, teacher_steps=50, layout="tau_prior") + fdk256 = LowdimFDKStudent(23, 10, latent_dim=256, teacher_steps=50, layout="tau_prior") + fdk_w4 = LowdimFDKStudent( + 23, 10, latent_dim=512, teacher_steps=50, layout="tau_prior", backbone_width=4) + act128 = LowdimFDKStudent(23, 10, latent_dim=128, teacher_steps=50, layout="action_only") + chi = ChiUNet1d(10, 23, 2, model_dim=256, emb_dim=256, dim_mult=[1, 2, 2]) + return { + "fdk_tau_lat128": count_params(fdk128), + "fdk_tau_lat256": count_params(fdk256), + "fdk_tau_w4_lat512": count_params(fdk_w4), + "fdk_action_only_lat128": count_params(act128), + "chi_unet_teacher_body": count_params(chi), + "ratio_chi_over_fdk128": round(count_params(chi) / max(count_params(fdk128), 1), 2), + "ratio_chi_over_fdk_w4": round(count_params(chi) / max(count_params(fdk_w4), 1), 2), + "fdk_bottleneck": ( + "E compresses the whole H=16 trajectory into one 128-D vector; " + "K/D never see per-timestep skip connections. Chi_UNet keeps " + "multi-scale action features [B,H,10] with obs as FiLM/global cond." + ), + "fdk_input_tau": "16 x (23+10) = 528 numbers -> 128-D", + "chi_input": "16 x 10 actions, conditioned on 2 x 23 obs, channels 256/512", + } + + +@torch.no_grad() +def per_dim_mse(model, pairs, device): + start, end = EXECUTE + val = pairs["validation"] + model.eval() + if model.use_koopman: + model.koopman.cache_inference_basis() + pred, tgt = [], [] + for offset in range(0, len(val["obs"]), 256): + o = val["obs"][offset:offset + 256].to(device) + n = val["noise"][offset:offset + 256].to(device) + a = val["actions"][offset:offset + 256].to(device) + pred.append(model(n, o)) + tgt.append(a) + pred = torch.cat(pred) + tgt = torch.cat(tgt) + err = (pred - tgt).square() + chunk = err[:, start:end] + dim_mse = chunk.mean(dim=(0, 1)).cpu().tolist() + return { + "chunk_mse": float(chunk.mean()), + "xyz_mse": float(chunk[..., :3].mean()), + "rot6d_mse": float(chunk[..., 3:9].mean()), + "grip_mse": float(chunk[..., 9].mean()), + "per_dim": dim_mse, + "last_chunk_mse": float(err[:, end - 2:end].mean()), + "first_exec_mse": float(err[:, start].mean()), + } + + +def load_student(checkpoint, device): + saved = torch.load(Path(checkpoint).resolve(), map_location="cpu", weights_only=False) + meta = saved.get("metadata") or {} + model = LowdimFDKStudent( + meta.get("obs_dim", 23), + meta.get("act_dim", 10), + use_koopman=True, + latent_dim=int(meta.get("latent_dim") or 128), + teacher_steps=int(meta.get("teacher_steps") or 50), + context_mode=meta.get("context_mode") or "obs_film_gate", + layout=meta.get("layout") or "tau_prior", + backbone_width=int(meta.get("backbone_width") or 1), + ) + model.load_state_dict(saved["state_dict"]) + model.to(device).eval().requires_grad_(False) + return model, saved + + +def make_offscreen_env(task, cfg): + """Same as table2.evaluate.make_env but with offscreen RGB for low-dim square.""" + import collections + import collections.abc + for name in ("Mapping", "Iterable"): + if not hasattr(collections, name): + setattr(collections, name, getattr(collections.abc, name)) + import robomimic.utils.env_utils as EnvUtils + import robomimic.utils.file_utils as FileUtils + import robomimic.utils.obs_utils as ObsUtils + from cleandiffuser.env.robomimic.robomimic_lowdim_wrapper import RobomimicLowdimWrapper + from cleandiffuser.env.wrapper import MultiStepWrapper + meta = FileUtils.get_env_metadata_from_dataset(cfg.dataset_path) + ObsUtils.initialize_obs_modality_mapping_from_dict({"low_dim": list(cfg.obs_keys)}) + meta["env_kwargs"]["controller_configs"]["control_delta"] = False + env = EnvUtils.create_env_from_metadata( + meta, render=False, render_offscreen=True, use_image_obs=False) + env.env.hard_reset = False + env = RobomimicLowdimWrapper( + env, obs_keys=list(cfg.obs_keys), render_hw=(384, 384), + render_camera_name="agentview") + return MultiStepWrapper(env, cfg.obs_steps, cfg.action_steps, cfg.max_episode_steps) + + +class SquareRecorder(gym.Wrapper): + def __init__(self, env, directory, *, title, stem, banner=88): + super().__init__(env) + self.directory = Path(directory) + self.directory.mkdir(parents=True, exist_ok=True) + self.title, self.stem = title, stem + self.banner = int(banner) + self.fps = Fraction(20, 1) + font = Path("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf") + self.font = ImageFont.truetype(str(font), 16) if font.exists() else ImageFont.load_default() + self.writer = None + self.steps = self.frames = 0 + self.score = 0.0 + self.episode_seed = None + self.snapshots = [] + + def seed(self, seed=None): + self.episode_seed = int(seed) + return self.env.seed(seed) + + def reset(self): + observation = self.env.reset() + self.steps = self.frames = 0 + self.score = 0.0 + self.snapshots = [] + self.video_path = self.directory / f"{self.stem}_seed{self.episode_seed}.mp4" + if self.video_path.exists(): + raise FileExistsError(self.video_path) + self.writer = VideoRecorder.create_h264(fps=self.fps, crf=18, thread_count=1) + self.writer.start(str(self.video_path)) + self.write_frame() + return observation + + def step(self, action): + observation, reward, done, info = self.env.step(action) + self.steps += 1 + self.score = max(self.score, float(reward > 0)) + self.write_frame() + return observation, reward, done, info + + def write_frame(self): + frame = np.asarray(self.env.render(mode="rgb_array")) + if frame.dtype != np.uint8: + frame = np.clip(frame, 0, 255).astype(np.uint8) + canvas = Image.new("RGB", (frame.shape[1], frame.shape[0] + self.banner), (32, 36, 44)) + canvas.paste(Image.fromarray(frame), (0, self.banner)) + draw = ImageDraw.Draw(canvas) + draw.text((10, 8), self.title, font=self.font, fill=(255, 255, 255)) + draw.text( + (10, 32), + f"ep {self.episode_seed} | step {self.steps} | success {self.score:.0f}", + font=self.font, fill=(180, 220, 190), + ) + draw.text((10, 56), "agentview | 1-NFE FDK | binary insert", font=self.font, fill=(200, 205, 212)) + self.writer.write_frame(np.asarray(canvas)) + self.frames += 1 + self.last_frame = canvas + if self.steps in (0, 100, 200, 300): + path = self.directory / f"{self.stem}_seed{self.episode_seed}_t{self.steps:03d}.png" + canvas.save(path) + self.snapshots.append(str(path)) + + def close(self): + if self.writer is not None: + self.writer.stop() + self.writer = None + final = self.directory / f"{self.stem}_seed{self.episode_seed}_final.png" + self.last_frame.save(final) + self.snapshots.append(str(final)) + self.record = { + "episode_seed": self.episode_seed, + "video": str(self.video_path), + "snapshots": list(self.snapshots), + "steps": self.steps, + "frames": self.frames, + "score": self.score, + } + atomic_json(self.video_path.with_suffix(".json"), self.record) + self.env.close() + + +def record_seeds(checkpoint, out_dir, seeds, device, label): + from experiments.lowdim_fdk_distill.run import StudentAgent + model, saved = load_student(checkpoint, device) + agent = StudentAgent(model, model, "square_lowdim") + _, config, _, _ = load_teacher("square_lowdim", device) + config = OmegaConf.merge(config, {"sample_steps": 1, "device": device}) + dataset = build_dataset("square_lowdim", config) + rows = [] + for seed in seeds: + seed_dir = Path(out_dir) / f"{label}_seed{seed}" + seed_dir.mkdir(parents=True, exist_ok=True) + holders = [] + + def factory(task, cfg, _holders=holders, _dir=seed_dir): + env = make_offscreen_env(task, cfg) + rec = SquareRecorder( + env.env, _dir, + title=f"Square-lowdim | {label} | 1-NFE FDK", + stem=f"square_{label}", + ) + env.env = rec + _holders.append(rec) + return env + + seed_all(int(seed)) + with patch.object(evaluation, "make_env", factory): + result = evaluation.evaluate( + "square_lowdim", config, dataset, agent, + episodes=1, env_count=1, seed_base=int(seed), + ) + rec = holders[0] + rec.record["eval_score"] = float(result["score"]) + rec.record["verification"] = validate_video(rec.video_path, rec.frames) + rec.record["checkpoint"] = str(Path(checkpoint).resolve()) + rec.record["selected_step"] = saved.get("step") + atomic_json(rec.video_path.with_suffix(".json"), rec.record) + rows.append(rec.record) + print(f"[{label}] seed={seed} score={rec.record['score']:.0f} -> {rec.video_path}", flush=True) + return rows + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", default="outputs/lowdim_fdk_distill/square_diag_20260916_v1") + parser.add_argument("--checkpoint", default=str(BEST_ALIGN)) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--skip-videos", action="store_true") + args = parser.parse_args() + out = Path(args.output).resolve() + if not out.is_relative_to(ROOT / "outputs"): + raise ValueError("output under outputs/") + out.mkdir(parents=True, exist_ok=True) + + cap = capacity_report() + print("[CAPACITY]", json.dumps(cap, indent=2), flush=True) + + device = args.device + model, saved = load_student(args.checkpoint, device) + pairs = torch.load(PAIRS, map_location="cpu", weights_only=False) + dims = per_dim_mse(model, pairs, device) + print("[PER-DIM]", json.dumps(dims, indent=2), flush=True) + + videos = {"success": [], "fail": []} + if not args.skip_videos: + videos["success"] = record_seeds( + args.checkpoint, out / "videos", SUCCESS_SEEDS, device, "success") + videos["fail"] = record_seeds( + args.checkpoint, out / "videos", FAIL_SEEDS, device, "fail") + + report = { + "checkpoint": str(Path(args.checkpoint).resolve()), + "checkpoint_sha256": sha256(Path(args.checkpoint)), + "selected_step": saved.get("step"), + "capacity": cap, + "per_dim_val": dims, + "videos": videos, + "note": ( + "If fail videos reach the nut but miss the hole, bottleneck is " + "late-chunk precision (supports per-dim / last-chunk selection). " + "If they never approach, bottleneck is earlier policy." + ), + } + atomic_json(out / "diagnosis.json", report) + print(json.dumps({"output": str(out), "capacity": cap, "per_dim_val": dims}, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/generate_pairs.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/generate_pairs.py new file mode 100644 index 0000000000000000000000000000000000000000..156e4f2606c5363a858f1ad473b9412466433b51 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/generate_pairs.py @@ -0,0 +1,62 @@ +"""Generate lowdim FDK distillation pairs without training.""" +from __future__ import annotations + +import argparse +from pathlib import Path + +from experiments.lowdim_fdk_distill.run import ( + SUPPORTED, + generate_pairs, + load_teacher, +) +from experiments.table2.common import ROOT, atomic_json, build_dataset, seed_all +from experiments.table2.train import normalizer_state +from experiments.teacher_videos.record_pusht_state import restore_normalizer + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--task", required=True, choices=tuple(SUPPORTED)) + parser.add_argument("--output", required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--train-pairs", type=int, default=250000) + parser.add_argument("--validation-pairs", type=int, default=16384) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--batch-size", type=int, default=64) + args = parser.parse_args() + + out = Path(args.output).resolve() + if not out.is_relative_to(ROOT / "outputs"): + raise ValueError("output under outputs/") + out.mkdir(parents=True, exist_ok=False) + seed_all(args.seed) + + agent, config, meta, teacher_sha = load_teacher( + args.task, args.device, require_source_hash=False) + dataset = build_dataset(args.task, config) + if normalizer_state(dataset.normalizer) != meta["normalizer"]: + raise ValueError("Demonstration normalizer differs from frozen teacher") + dataset.normalizer = restore_normalizer(meta["normalizer"]) + print(f"[TEACHER] {args.task} steps={config.sample_steps} sha={teacher_sha[:12]}", flush=True) + + prepared = generate_pairs( + args.task, agent, config, dataset, out, + args.train_pairs, args.validation_pairs, args.seed, + batch_size=args.batch_size, + ) + prepared["metadata"]["teacher_sha256"] = teacher_sha + atomic_json(out / "plan.json", { + "task": args.task, + "train_pairs": args.train_pairs, + "validation_pairs": args.validation_pairs, + "seed": args.seed, + "teacher_sha256": teacher_sha, + "purpose": "shared distillation pairs for seq lat-weight sweep", + }) + n_train = len(prepared["train"]["obs"]) + n_val = len(prepared["validation"]["obs"]) + print(f"[PAIRS] wrote {out/'pairs.pt'} train={n_train} val={n_val}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/reeval_checkpoint.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/reeval_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..52a5c8143e8e406b7cda5ab8568ddad4bd98b55c --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/reeval_checkpoint.py @@ -0,0 +1,125 @@ +"""Re-evaluate frozen square FDK checkpoints at 1-NFE (no retraining).""" +from __future__ import annotations + +import argparse +import json +from copy import deepcopy +from pathlib import Path + +from omegaconf import OmegaConf +import torch + +from experiments.lowdim_fdk_distill.run import ( + LowdimFDKStudent, + StudentAgent, + SUPPORTED, + load_teacher, +) +from experiments.table2.common import ROOT, atomic_json, build_dataset, seed_all +from experiments.table2.evaluate import evaluate +from experiments.teacher_videos.record_pusht_state import restore_normalizer + + +DEFAULTS = { + "align": ROOT / ( + "outputs/lowdim_fdk_distill/square_lowdim_best_align_filmgate_20260915_v1/" + "arm_vdv/best_ema.pt" + ), +} + + +def load_student(checkpoint: Path, device: str): + saved = torch.load(checkpoint, map_location="cpu", weights_only=False) + meta = saved.get("metadata") or {} + task = meta.get("task") or "square_lowdim" + model = LowdimFDKStudent( + meta.get("obs_dim", SUPPORTED[task]["obs_dim"]), + meta.get("act_dim", SUPPORTED[task]["act_dim"]), + use_koopman=True, + latent_dim=int(meta.get("latent_dim") or 128), + teacher_steps=int(meta.get("teacher_steps") or SUPPORTED[task]["teacher_steps"]), + context_mode=meta.get("context_mode") or "obs_film_gate", + layout=meta.get("layout") or "tau_prior", + backbone_width=int(meta.get("backbone_width") or 1), + ) + model.load_state_dict(saved["state_dict"]) + model.to(device).eval().requires_grad_(False) + return model, saved, task + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", action="append", default=None, + help="Repeatable; default = align + wide best_ema") + parser.add_argument("--tag", action="append", default=None) + parser.add_argument("--output", required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--episodes", type=int, default=50) + parser.add_argument("--seed", type=int, default=0) + args = parser.parse_args() + + out = Path(args.output).resolve() + if not out.is_relative_to(ROOT / "outputs"): + raise ValueError("output under outputs/") + out.mkdir(parents=True, exist_ok=True) + seed_all(args.seed) + + if args.checkpoint: + ckpts = [Path(p).resolve() for p in args.checkpoint] + tags = args.tag or [p.stem for p in ckpts] + if len(tags) != len(ckpts): + raise ValueError("--tag count must match --checkpoint") + else: + ckpts = [DEFAULTS["align"]] + tags = ["align_0p36"] + + rows = [] + for tag, ckpt in zip(tags, ckpts): + print(f"[REEVAL] {tag} {ckpt}", flush=True) + model, saved, task = load_student(ckpt, args.device) + agent, config, meta, teacher_sha = load_teacher( + task, args.device, require_source_hash=False) + dataset = build_dataset(task, config) + dataset.normalizer = restore_normalizer(meta["normalizer"]) + del agent + if torch.cuda.is_available(): + torch.cuda.empty_cache() + adapter = StudentAgent(model, deepcopy(model), task) + student_config = OmegaConf.merge(config, {"sample_steps": 1, "device": args.device}) + result = evaluate( + task, student_config, dataset, adapter, + episodes=args.episodes, + env_count=min(10, args.episodes), + seed_base=100000, + ) + row = { + "tag": tag, + "checkpoint": str(ckpt), + "selected_step": saved.get("step"), + "prior_summary_score": None, + "reeval_score": result["score"], + "episodes": args.episodes, + "validation": saved.get("validation"), + "metadata": { + k: saved.get("metadata", {}).get(k) + for k in ( + "latent_dim", "backbone_width", "layout", "context_mode", + "teacher_sha256", + ) + }, + "teacher_sha256": teacher_sha, + } + # Attach prior reported score if sibling summary exists. + summary_path = ckpt.parents[1] / "summary.json" + if summary_path.is_file(): + row["prior_summary_score"] = json.loads(summary_path.read_text())["score"] + atomic_json(out / f"{tag}.json", {**row, "rollout": result}) + rows.append(row) + print(json.dumps(row, indent=2), flush=True) + + atomic_json(out / "summary.json", {"rows": rows}) + print(json.dumps({"rows": rows}, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/run.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/run.py new file mode 100644 index 0000000000000000000000000000000000000000..0710d729cb1795e9b4555a40825834eb67a02927 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/run.py @@ -0,0 +1,1093 @@ +"""Multi-task low-dim FDK distill: pusht_keypoints / lift_lowdim / square_lowdim. + +Aligned with PushT-state best recipe (fair 0.934): + τ-prior (pin o_t into τ[:,0,:obs_dim]), obs_film_gate, VDV^{-1}, + loss = chunk + 0.25*full + 0.5*first + 0.15*lat + 0.01*spec + 0.001*film + (pred/rec off). +""" +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import time +from copy import deepcopy +from pathlib import Path + +import numpy as np +import torch +from omegaconf import OmegaConf +from torch import nn +from torch.nn import functional as F + +from experiments.lowdim_fdk_distill.wide_backbone import ( + SequenceTemporalDecoder, + SequenceTemporalEncoder, + TemporalDecoder, + TemporalEncoder, +) +from experiments.pusht_state_distill.vdv_koopman import InvertibleDynamicKoopman +from experiments.table2.common import ( + ROOT, + atomic_json, + build_agent, + build_dataset, + config_for, + seed_all, + training_condition, +) +from experiments.table2.evaluate import evaluate +from experiments.table2.train import normalizer_state, source_manifest +from experiments.teacher_videos.record_pusht_state import restore_normalizer + + +SUPPORTED = { + "relay_kitchen": {"obs_dim": 60, "act_dim": 9, "teacher_steps": 50}, + "pusht_keypoints": {"obs_dim": 20, "act_dim": 2, "teacher_steps": 5}, + "lift_lowdim": {"obs_dim": 19, "act_dim": 10, "teacher_steps": 50}, + "square_lowdim": {"obs_dim": 23, "act_dim": 10, "teacher_steps": 50}, + "can_lowdim": {"obs_dim": 23, "act_dim": 10, "teacher_steps": 50}, + "tool_hang_lowdim": {"obs_dim": 53, "act_dim": 10, "teacher_steps": 50}, +} +# Final Chi_UNet teachers: Table-2 path by default; Can/Tool Hang live in table2_addon. +TEACHER_CKPT = { + "can_lowdim": ROOT / ( + "outputs/table2_addon/robomimic_20260916_v1/runs/can_lowdim/seed0/model_1000000.pt"), + "tool_hang_lowdim": ROOT / ( + "outputs/table2_addon/robomimic_20260916_v1/runs/tool_hang_lowdim/seed0/" + "model_1000000.pt"), +} +ADDON_TASKS = frozenset(TEACHER_CKPT) +HORIZON = 16 +OBS_STEPS = 2 +EXECUTE = (1, 9) +CONTEXT_DIM = 128 + + +def ensure_task_registered(task: str): + """Can/Tool Hang are process-local extensions (do not edit frozen table2.common).""" + if task in ADDON_TASKS: + from experiments.table2_addon.common import register_tasks + register_tasks() + + +def teacher_checkpoint(task: str) -> Path: + if task in TEACHER_CKPT: + return TEACHER_CKPT[task] + return ROOT / f"outputs/table2/chi_unet_20260909_v1/runs/{task}/seed0/model_1000000.pt" + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(8 * 1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def load_teacher(task: str, device: str, require_source_hash: bool = True): + if task not in SUPPORTED: + raise ValueError(f"task must be one of {tuple(SUPPORTED)}") + ensure_task_registered(task) + path = teacher_checkpoint(task) + if not path.is_file(): + raise FileNotFoundError(f"Missing frozen teacher: {path}") + saved = torch.load(path, map_location="cpu", weights_only=False) + meta = saved["metadata"] + if saved["step"] != 1000000 or meta["task"] != task or meta["seed"] != 0 or meta["smoke"]: + raise ValueError(f"Need final nonsmoke seed0 teacher for {task}") + # Addon teachers fingerprint table2_addon sources; Table-2 teachers use + # experiments.table2.train.source_manifest. Distill always allows mismatch. + if task not in ADDON_TASKS: + current = source_manifest()["sha256"] + expected = meta["source"]["sha256"] + if current != expected: + if require_source_hash: + raise ValueError("Table-2 source hash differs from frozen teacher") + print( + f"[WARN] source hash mismatch (teacher weights unchanged) " + f"got={current[:12]} expected={expected[:12]}", + flush=True, + ) + else: + print(f"[TEACHER] addon checkpoint {path}", flush=True) + config = config_for(task, 0, device) + spec = SUPPORTED[task] + if int(config.sample_steps) != spec["teacher_steps"]: + raise ValueError(f"Expected sample_steps={spec['teacher_steps']}, got {config.sample_steps}") + if int(config.obs_dim) != spec["obs_dim"] or int(config.action_dim) != spec["act_dim"]: + raise ValueError("obs/action dims mismatch config") + agent = build_agent(task, config) + agent.model.load_state_dict(saved["model"]) + agent.model_ema.load_state_dict(saved["model_ema"]) + agent.model.eval().requires_grad_(False) + agent.model_ema.eval().requires_grad_(False) + return agent, config, meta, sha256(path) + + +def episode_split(dataset, seed=20260914, fraction=0.1): + ends = np.asarray(dataset.replay_buffer.episode_ends) + if len(ends) < 2: + raise ValueError("Need >=2 episodes") + episode_ids = np.searchsorted(ends, dataset.sampler.indices[:, 0], side="right") + ids = np.random.default_rng(seed).permutation(len(ends)) + count = min(len(ids) - 1, max(1, round(len(ids) * fraction))) + held = set(ids[:count].tolist()) + mask = np.asarray([int(e) in held for e in episode_ids]) + return np.flatnonzero(~mask), np.flatnonzero(mask), { + "split_seed": seed, + "train_episodes": sorted(set(ids[count:].tolist())), + "validation_episodes": sorted(held), + } + + +def obs_batch(task, dataset, indices, config): + """Normalized obs [B, OBS_STEPS, obs_dim] using the same path as table2 training.""" + rows = [] + for index in indices: + sample = dataset[int(index)] + # training_condition expects a batch axis; dataset __getitem__ is unbatched. + batched = { + "obs": { + key: value.unsqueeze(0) if torch.is_tensor(value) else value + for key, value in sample["obs"].items() + } + } + flat = training_condition(task, batched, config) + rows.append(flat.reshape(OBS_STEPS, SUPPORTED[task]["obs_dim"])) + return torch.stack(rows, dim=0) + + +@torch.no_grad() +def generate_pairs(task, agent, config, dataset, out_dir, train_pairs, val_pairs, seed, batch_size): + out_dir = Path(out_dir) + device = config.device + act_dim = SUPPORTED[task]["act_dim"] + steps = int(config.sample_steps) + tr, va, split = episode_split(dataset, seed=seed + 7) + rng = np.random.default_rng(seed + 1001) + selected = { + "train": rng.choice(tr, train_pairs, replace=True), + "validation": rng.choice(va, val_pairs, replace=True), + } + collected = { + "train": {"obs": [], "noise": [], "actions": []}, + "validation": {"obs": [], "noise": [], "actions": []}, + } + for split_name, idxs in selected.items(): + for start in range(0, len(idxs), batch_size): + batch_idx = idxs[start:start + batch_size] + obs = obs_batch(task, dataset, batch_idx, config).to(device) + condition = obs.flatten(start_dim=1) + prior = torch.zeros(len(batch_idx), HORIZON, act_dim, device=device) + noise = torch.randn(len(batch_idx), HORIZON, act_dim, device=device) + actions, _ = agent.sample( + prior, + n_samples=len(batch_idx), + sample_steps=steps, + condition_cfg=condition, + w_cfg=1.0, + use_ema=True, + initial_noise=noise, + ) + actions = actions.clamp(-1, 1) + collected[split_name]["obs"].append(obs.cpu()) + collected[split_name]["noise"].append(noise.cpu()) + collected[split_name]["actions"].append(actions.cpu()) + if start == 0 or (start // batch_size) % 20 == 0: + print(f"[PAIRS] {split_name} {start + len(batch_idx)}/{len(idxs)}", flush=True) + prepared = { + "train": {k: torch.cat(v) for k, v in collected["train"].items()}, + "validation": {k: torch.cat(v) for k, v in collected["validation"].items()}, + "metadata": { + "task": task, + "format": "lowdim_fdk_pairs_v1", + "split": split, + "teacher_steps": steps, + "obs_dim": SUPPORTED[task]["obs_dim"], + "act_dim": act_dim, + "normalizer": normalizer_state(dataset.normalizer), + }, + } + torch.save(prepared, out_dir / "pairs.pt") + return prepared + + +def pack_tau_prior(obs: torch.Tensor, action_noise: torch.Tensor) -> torch.Tensor: + """τ_prior: pin current frame into step-0 state; actions = noise.""" + obs_dim, act_dim = obs.shape[-1], action_noise.shape[-1] + traj_dim = obs_dim + act_dim + prior = torch.zeros( + obs.shape[0], HORIZON, traj_dim, device=obs.device, dtype=obs.dtype) + prior[:, :, obs_dim:] = action_noise + prior[:, 0, :obs_dim] = obs[:, -1] + return prior + + +def pack_tau_clean(obs: torch.Tensor, actions: torch.Tensor) -> torch.Tensor: + clean = pack_tau_prior(obs, torch.zeros_like(actions)) + clean[:, :, obs.shape[-1]:] = actions + return clean + + +def unpack_actions(traj: torch.Tensor, obs_dim: int) -> torch.Tensor: + return traj[:, :, obs_dim:] + + +class LowdimFDKStudent(nn.Module): + """E(τ_prior) → gated FiLM → K(o_t) → D(τ̂_0). + + latent_mode: + - global: one pooled z (paper figure) + - sequence: z_{1:T} with T=H/4; shared K(c) applied in parallel (design A, still 1-NFE) + """ + + def __init__( + self, + obs_dim, + act_dim, + use_koopman=True, + latent_dim=128, + teacher_steps=5, + context_mode="obs_film_gate", + layout="tau_prior", + backbone_width=1, + latent_mode="global", + ): + super().__init__() + if context_mode not in ("obs_film", "obs_film_gate"): + raise ValueError("lowdim recipe supports obs_film / obs_film_gate only") + if layout not in ("tau_prior", "action_only"): + raise ValueError("layout must be tau_prior or action_only") + if latent_mode not in ("global", "sequence"): + raise ValueError("latent_mode must be global or sequence") + self.obs_dim = obs_dim + self.act_dim = act_dim + self.layout = layout + self.traj_dim = act_dim if layout == "action_only" else obs_dim + act_dim + self.use_koopman = use_koopman + self.latent_dim = latent_dim + self.context_mode = context_mode + self.backbone_width = int(backbone_width) + self.latent_mode = latent_mode + if latent_mode == "sequence": + self.encoder = SequenceTemporalEncoder( + self.traj_dim, latent_dim, HORIZON, width=self.backbone_width) + self.decoder = SequenceTemporalDecoder( + latent_dim, HORIZON, self.traj_dim, width=self.backbone_width) + self.latent_steps = HORIZON // 4 + else: + self.encoder = TemporalEncoder( + self.traj_dim, latent_dim, HORIZON, width=self.backbone_width) + self.decoder = TemporalDecoder( + latent_dim, HORIZON, self.traj_dim, width=self.backbone_width) + self.latent_steps = 1 + self.time_embed = nn.Embedding(teacher_steps, latent_dim) + nn.init.zeros_(self.time_embed.weight) + feat_dim = OBS_STEPS * obs_dim + self.obs_film = nn.Sequential( + nn.Linear(feat_dim, CONTEXT_DIM), + nn.SiLU(), + nn.Linear(CONTEXT_DIM, 2 * latent_dim), + ) + nn.init.zeros_(self.obs_film[-1].weight) + nn.init.zeros_(self.obs_film[-1].bias) + self.film_alpha = None + if context_mode == "obs_film_gate": + self.film_alpha = nn.Parameter(torch.tensor(-4.0)) + self.koopman = None + if use_koopman: + # γ on current frame only (same as PushT obs_film_gate). + self.koopman = InvertibleDynamicKoopman( + latent_dim, + gamma_mode="obs0", + obs_dim=obs_dim, + normalize_context=True, + inverse_mode="exact", + ) + + def encode(self, traj): + raw = self.encoder(traj.permute(0, 2, 1)) + # global: [B, L]; sequence: [B, T, L] + return F.normalize(raw, dim=-1, eps=1e-6 * math.sqrt(self.latent_dim)) * math.sqrt( + self.latent_dim) + + def _apply_film(self, z, obs): + scale, shift = self.obs_film(obs.reshape(obs.shape[0], -1)).chunk(2, dim=-1) + if z.dim() == 3: + scale, shift = scale[:, None, :], shift[:, None, :] + if self.film_alpha is not None: + alpha = torch.sigmoid(self.film_alpha) + z = z + alpha * (z * scale + shift) + else: + z = z * (1.0 + scale) + shift + return z, scale.squeeze(1) if scale.dim() == 3 else scale, ( + shift.squeeze(1) if shift.dim() == 3 else shift) + + def _apply_koopman(self, z, gamma_ctx): + """Shared K(c): one gamma(c); apply in parallel over time if z is [B,T,L].""" + if z.dim() == 2: + return self.koopman(z, gamma_ctx) + b, t, lat = z.shape + z_flat = z.reshape(b * t, lat) + ctx = gamma_ctx[:, None, :].expand(b, t, -1).reshape(b * t, -1) + z_out, gamma = self.koopman(z_flat, ctx) + # gamma is identical across t under obs0; keep [B, L] for logging/spec. + return z_out.view(b, t, lat), gamma.view(b, t, lat)[:, 0] + + def _pack(self, obs, actions): + if self.layout == "action_only": + return actions + return pack_tau_prior(obs, actions) + + def _unpack(self, traj, prior): + if self.layout == "action_only": + return traj + traj = traj.clone() + traj[:, 0, :self.obs_dim] = prior[:, 0, :self.obs_dim] + return unpack_actions(traj, self.obs_dim) + + def forward(self, noise, obs, clean=None, t=None): + prior = self._pack(obs, noise) + if t is None: + t = torch.full( + (noise.shape[0],), + self.time_embed.num_embeddings - 1, + device=noise.device, + dtype=torch.long, + ) + te = self.time_embed(t) + z = self.encode(prior) + z = z + (te[:, None, :] if z.dim() == 3 else te) + z, film_scale, film_shift = self._apply_film(z, obs) + gamma_ctx = obs[:, -1] + gamma = None + z_pred = z + if self.use_koopman: + z_pred, gamma = self._apply_koopman(z_pred, gamma_ctx) + traj = self.decoder(z_pred).view(-1, HORIZON, self.traj_dim) + actions = self._unpack(traj, prior) + if clean is None: + return actions + clean_traj = self._pack(obs, clean) + z_clean = self.encode(clean_traj) + z_clean = z_clean + (te[:, None, :] if z_clean.dim() == 3 else te) + recon_traj = self.decoder(z_clean).view(-1, HORIZON, self.traj_dim) + return { + "actions": actions, + "reconstruction": self._unpack(recon_traj, clean_traj), + "z_pred": z_pred, + "z_target": z_clean.detach(), + "gamma": gamma if gamma is not None else actions.new_zeros( + actions.shape[0], self.latent_dim), + "film_scale": film_scale, + "film_shift": film_shift, + } + + +class ChunkDiscriminator(nn.Module): + """KDM-style 4-layer discriminator adapted to action chunks (Conv1d). + + Official KDM uses Conv2d on images; see azencot-group/KDM Discriminator. + Here x0 / x̂0 are action trajectories [B, H, A], viewed as [B, A, H]. + Training matches KDM trainer: D-step then G-step with λ_adv=0.01. + """ + + def __init__(self, act_dim: int, advers_w: float = 1.0): + super().__init__() + self.advers_w = float(advers_w) + self.model = nn.Sequential( + nn.Conv1d(act_dim, 64, kernel_size=4, stride=2, padding=1), + nn.LeakyReLU(0.2, inplace=True), + nn.Conv1d(64, 128, kernel_size=4, stride=2, padding=1), + nn.BatchNorm1d(128), + nn.LeakyReLU(0.2, inplace=True), + nn.Conv1d(128, 256, kernel_size=4, stride=2, padding=1), + nn.BatchNorm1d(256), + nn.LeakyReLU(0.2, inplace=True), + nn.Conv1d(256, 1, kernel_size=2, stride=1, padding=0), + nn.Flatten(), + ) + + def forward(self, actions: torch.Tensor) -> torch.Tensor: + # [B, H, A] -> [B, A, H] + return self.model(actions.permute(0, 2, 1)) + + def disc_loss(self, real_actions: torch.Tensor, fake_actions: torch.Tensor): + d_real = self.forward(real_actions.detach()) + d_fake = self.forward(fake_actions.detach()) + real_loss = F.binary_cross_entropy_with_logits(d_real, torch.ones_like(d_real)) + fake_loss = F.binary_cross_entropy_with_logits(d_fake, torch.zeros_like(d_fake)) + adv = (real_loss + fake_loss) * self.advers_w + return adv, { + "adv_d": float(adv.detach()), + "adv_real": float(real_loss.detach()), + "adv_fake": float(fake_loss.detach()), + } + + def gen_loss(self, fake_actions: torch.Tensor) -> torch.Tensor: + # KDM: BCE(D(x̂0), 1) * 0.01 * advers_w + d_fake = self.forward(fake_actions) + return F.binary_cross_entropy_with_logits(d_fake, torch.ones_like(d_fake)) + + +def square_obs_parts(obs: torch.Tensor): + """square_lowdim flat obs: object(14)+eef_pos(3)+eef_quat(4)+grip(2)=23. + + Robosuite NutAssemblySquare object layout: nut_pos(3), nut_quat(4), + peg_pos(3), peg_quat(4). + """ + frame = obs[:, -1] if obs.dim() == 3 else obs + if frame.shape[-1] != 23: + raise ValueError(f"square geometry expects obs_dim=23, got {frame.shape[-1]}") + return { + "nut_pos": frame[:, 0:3], + "peg_pos": frame[:, 7:10], + "eef_pos": frame[:, 14:17], + } + + +def nut_peg_distance(obs: torch.Tensor) -> torch.Tensor: + parts = square_obs_parts(obs) + return (parts["nut_pos"] - parts["peg_pos"]).norm(dim=-1) + + +def geo_sample_probs(obs: torch.Tensor, temperature: float = 1.0) -> torch.Tensor: + """Oversample insertion-near windows via softmax(-dist / (std·T)).""" + dist = nut_peg_distance(obs) + scale = dist.std().clamp_min(1e-3) * max(float(temperature), 1e-3) + weight = torch.softmax(-dist / scale, dim=0) + return weight + + +def geo_proximity(obs: torch.Tensor, temperature: float = 1.0) -> torch.Tensor: + """Per-sample proximity in (0,1): closer nut–peg → larger.""" + dist = nut_peg_distance(obs) + scale = dist.std().clamp_min(1e-3) * max(float(temperature), 1e-3) + # sigmoid so batch stats stay stable under oversampling + return torch.sigmoid(-(dist - dist.median()) / scale) + + +def student_loss( + result, + target, + latent_weight=0.15, + spectral_weight=0.01, + full_weight=0.25, + first_weight=0.0, + film_reg_weight=0.0, + adv_weight=0.0, + late_weight=0.0, + late_steps=2, + sample_weights: torch.Tensor | None = None, + discriminator: ChunkDiscriminator | None = None, + act_weights: torch.Tensor | None = None, +): + start, end = EXECUTE + delta = (result["actions"] - target).square() + if act_weights is not None: + delta = delta * act_weights.view(1, 1, -1) + late_k = max(1, min(int(late_steps), end - start)) + per = delta.mean(dim=(1, 2)) # [B] + chunk_per = delta[:, start:end].mean(dim=(1, 2)) + late_per = delta[:, end - late_k:end].mean(dim=(1, 2)) + first_per = delta[:, start].mean(dim=-1) + if sample_weights is not None: + w = sample_weights.reshape(-1).to(dtype=delta.dtype, device=delta.device) + w = w / w.mean().clamp_min(1e-6) + chunk = (chunk_per * w).mean() + full = (per * w).mean() + first = (first_per * w).mean() + late = (late_per * w).mean() + else: + chunk = chunk_per.mean() + full = per.mean() + first = first_per.mean() + late = late_per.mean() + energy = result["z_target"].square().mean().clamp_min(1e-4) + lat_rel = F.mse_loss(result["z_pred"], result["z_target"]) / energy + lat_cos = 1.0 - F.cosine_similarity(result["z_pred"], result["z_target"], dim=-1).mean() + pred_rms = result["z_pred"].square().mean(-1).sqrt() + tgt_rms = result["z_target"].square().mean(-1).sqrt() + lat = lat_rel + 0.5 * lat_cos + 0.25 * (pred_rms - tgt_rms).square().mean() + spec = F.relu(result["gamma"].abs() - 1).mean() + film = result["actions"].new_zeros(()) + if float(film_reg_weight) > 0 and result.get("film_scale") is not None: + film = result["film_scale"].square().mean() + result["film_shift"].square().mean() + adv = result["actions"].new_zeros(()) + if float(adv_weight) > 0 and discriminator is not None: + # Paper / KDM code: λ_adv = 0.01; advers_w scales both D and G. + adv = discriminator.gen_loss(result["actions"]) * float(adv_weight) * discriminator.advers_w + loss = ( + chunk + full_weight * full + first_weight * first + + float(late_weight) * late + + latent_weight * lat + spectral_weight * spec + + float(film_reg_weight) * film + + adv + ) + return loss, { + "chunk": float(chunk.detach()), + "full": float(full.detach()), + "first": float(first.detach()), + "late": float(late.detach()), + "lat": float(lat.detach()), + "spec": float(spec.detach()), + "film": float(film.detach()), + "adv_g": float(adv.detach()), + } + + +@torch.no_grad() +def val_metrics(model, obs, noise, actions, device, batch_size=256): + was = model.training + model.eval() + if model.use_koopman: + model.koopman.cache_inference_basis() + start, end = EXECUTE + chunk_tot = last_tot = n_chunk = n_last = 0.0 + for offset in range(0, len(obs), batch_size): + o = obs[offset:offset + batch_size].to(device) + a_n = noise[offset:offset + batch_size].to(device) + a_c = actions[offset:offset + batch_size].to(device) + pred = model(a_n, o) + err = (pred - a_c).square() + chunk_tot += err[:, start:end].sum().item() + last_tot += err[:, end - 2:end].sum().item() + n_chunk += pred.shape[0] * (end - start) * pred.shape[-1] + n_last += pred.shape[0] * 2 * pred.shape[-1] + if model.use_koopman: + model.koopman.clear_inference_cache() + model.train(was) + return { + "chunk_mse": chunk_tot / max(n_chunk, 1), + "last_chunk_mse": last_tot / max(n_last, 1), + } + + +def val_chunk_mse(model, obs, noise, actions, device, batch_size=256): + return val_metrics(model, obs, noise, actions, device, batch_size)["chunk_mse"] + + +def action_dim_weights(act_dim: int, xyz_weight: float, rot_weight: float, grip_weight: float): + """Square/lift abs-action: [xyz(3), rot6d(6), grip(1)]. PushT 2D: uniform.""" + weights = torch.ones(act_dim) + if act_dim == 10: + weights[:3] = float(xyz_weight) + weights[3:9] = float(rot_weight) + weights[9] = float(grip_weight) + return weights + + +class StudentAgent: + def __init__(self, model, ema, task): + self.model = model + self.model_ema = ema + self.task = task + self.obs_dim = SUPPORTED[task]["obs_dim"] + self.act_dim = SUPPORTED[task]["act_dim"] + + @torch.no_grad() + def sample(self, prior, n_samples, sample_steps, condition_cfg, use_ema=True, **kwargs): + if sample_steps != 1 or n_samples != len(prior): + raise ValueError("one-step student only") + net = self.model_ema if use_ema else self.model + obs = condition_cfg.reshape(-1, OBS_STEPS, self.obs_dim) + noise = torch.randn(obs.shape[0], HORIZON, self.act_dim, device=obs.device, dtype=obs.dtype) + if net.use_koopman: + net.koopman.cache_inference_basis() + try: + return net(noise, obs), {} + finally: + if net.use_koopman: + net.koopman.clear_inference_cache() + + +def update_ema(target, source, beta=0.9999): + with torch.no_grad(): + for t_p, s_p in zip(target.parameters(), source.parameters()): + t_p.data.mul_(beta).add_(s_p.data, alpha=1 - beta) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--task", required=True, choices=tuple(SUPPORTED)) + parser.add_argument("--output", required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--reuse-pairs", default=None, + help="Existing pairs.pt (same task) to skip teacher sampling") + parser.add_argument("--train-pairs", type=int, default=100000) + parser.add_argument("--validation-pairs", type=int, default=8192) + parser.add_argument("--updates", type=int, default=150000) + parser.add_argument("--episodes", type=int, default=50) + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--validate-every", type=int, default=2000) + parser.add_argument("--latent-weight", type=float, default=0.15) + parser.add_argument("--first-weight", type=float, default=0.5) + parser.add_argument("--film-reg-weight", type=float, default=0.001) + parser.add_argument("--full-weight", type=float, default=0.25) + parser.add_argument("--spectral-weight", type=float, default=0.01) + parser.add_argument( + "--context-mode", + default="obs_film_gate", + choices=("obs_film", "obs_film_gate"), + ) + parser.add_argument("--latent-dim", type=int, default=128) + parser.add_argument( + "--backbone-width", + type=int, + default=1, + help="Temporal E/D channel multiplier (1 keeps 128/192/256; 4 ≈ 16x conv capacity)", + ) + parser.add_argument( + "--latent-mode", + default="global", + choices=("global", "sequence"), + help="global: one pooled z; sequence: z_1:T (T=H/4) + shared parallel K(c), still 1-NFE", + ) + parser.add_argument( + "--layout", + default="tau_prior", + choices=("tau_prior", "action_only"), + help="tau_prior packs obs into the trajectory; action_only encodes actions only", + ) + parser.add_argument("--xyz-weight", type=float, default=1.0) + parser.add_argument("--rot-weight", type=float, default=1.0) + parser.add_argument("--grip-weight", type=float, default=1.0) + parser.add_argument( + "--energy-sample", + action="store_true", + help="Oversample high-energy execute windows (contact / insertion proxy)", + ) + parser.add_argument( + "--geo-sample", + action="store_true", + help="square only: oversample small nut–peg distance windows", + ) + parser.add_argument( + "--geo-temp", + type=float, + default=1.0, + help="Temperature on dist.std() for geo softmax / proximity", + ) + parser.add_argument( + "--geo-loss-boost", + type=float, + default=0.0, + help="Per-sample weight = 1 + boost * proximity (closer nut–peg → larger)", + ) + parser.add_argument( + "--late-weight", + type=float, + default=0.0, + help="Extra MSE on the last late_steps of the execute window", + ) + parser.add_argument( + "--late-steps", + type=int, + default=2, + help="How many trailing execute steps enter the late term", + ) + parser.add_argument( + "--select", + default="chunk", + choices=("chunk", "last_chunk", "rollout"), + help="EMA checkpoint selection: val MSE or periodic env score", + ) + parser.add_argument( + "--rollout-every", + type=int, + default=10000, + help="When select=rollout, evaluate this often (steps)", + ) + parser.add_argument( + "--rollout-episodes", + type=int, + default=20, + help="Episodes per intermediate rollout selection (final uses --episodes)", + ) + parser.add_argument("--advers", action="store_true", + help="KDM-style adversarial refinement (λ_adv=0.01)") + parser.add_argument("--advers-w", type=float, default=1.0, + help="KDM advers_w scale on D and G adv terms") + parser.add_argument("--adv-weight", type=float, default=0.01, + help="Paper λ_adv on generator (KDM uses 0.01)") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--smoke", action="store_true") + args = parser.parse_args() + + if args.geo_sample or args.geo_loss_boost > 0: + if args.task != "square_lowdim": + raise ValueError("--geo-sample / --geo-loss-boost require --task square_lowdim") + if args.select == "rollout" and args.rollout_every < 1: + raise ValueError("--rollout-every must be positive") + if args.late_steps < 1: + raise ValueError("--late-steps must be >= 1") + + out = Path(args.output).resolve() + if not out.is_relative_to(ROOT / "outputs"): + raise ValueError("output under outputs/") + out.mkdir(parents=True, exist_ok=False) + seed_all(args.seed) + if args.smoke: + args.updates, args.validate_every, args.episodes = 200, 100, 2 + args.train_pairs, args.validation_pairs = 512, 128 + args.batch_size = 32 + args.rollout_every = 100 + args.rollout_episodes = 2 + + spec = SUPPORTED[args.task] + # The teacher is a frozen Table-2 artifact; this student module is allowed + # to evolve independently of the teacher's original source fingerprint. + agent, config, meta, teacher_sha = load_teacher( + args.task, args.device, require_source_hash=False) + dataset = build_dataset(args.task, config) + if normalizer_state(dataset.normalizer) != meta["normalizer"]: + raise ValueError("Demonstration normalizer differs from frozen teacher") + dataset.normalizer = restore_normalizer(meta["normalizer"]) + print(f"[TEACHER] {args.task} steps={config.sample_steps} sha={teacher_sha[:12]}", flush=True) + + if args.reuse_pairs: + reuse = Path(args.reuse_pairs).resolve() + prepared = torch.load(reuse, map_location="cpu", weights_only=False) + if prepared["metadata"]["task"] != args.task: + raise ValueError("reuse-pairs task mismatch") + print(f"[PAIRS] reuse {reuse}", flush=True) + # Keep a copy under this run for provenance. + torch.save(prepared, out / "pairs.pt") + else: + prepared = generate_pairs( + args.task, agent, config, dataset, out, + args.train_pairs, args.validation_pairs, args.seed, + batch_size=16 if args.smoke else 64, + ) + prepared["metadata"]["teacher_sha256"] = teacher_sha + loss_str = ( + f"chunk+{args.full_weight}*full+{args.first_weight}*first+" + f"{args.latent_weight}*lat+{args.spectral_weight}*spec" + + (f"+{args.film_reg_weight}*film" if args.film_reg_weight else "") + + (f"+{args.late_weight}*late{args.late_steps}" if args.late_weight else "") + + (f"+{args.adv_weight}*adv(w={args.advers_w})" if args.advers else "") + + "; pred/rec off; tau_prior+" + args.context_mode + + ("; KDM-style L_adv" if args.advers else "") + + f"; layout={args.layout}" + + (f"; w_xyz/rot/grip={args.xyz_weight}/{args.rot_weight}/{args.grip_weight}" + if args.xyz_weight != 1 or args.rot_weight != 1 or args.grip_weight != 1 else "") + + ("; energy_sample" if args.energy_sample else "") + + (f"; geo_sample(T={args.geo_temp})" if args.geo_sample else "") + + (f"; geo_boost={args.geo_loss_boost}" if args.geo_loss_boost else "") + + f"; select={args.select}" + + (f"; width={args.backbone_width}" if args.backbone_width != 1 else "") + + (f"; latent_mode={args.latent_mode}" if args.latent_mode != "global" else "") + ) + atomic_json(out / "plan.json", { + **{k: v for k, v in vars(args).items() if k != "reuse_pairs"}, + "reuse_pairs": args.reuse_pairs, + "teacher_sha256": teacher_sha, + "spec": spec, + "loss": loss_str, + "operator": "VDV^-1 + tau_prior + " + args.context_mode, + "aligned_with": "pusht_state filmgate_first05_freg_lat015_u150k fair=0.934", + "layout": args.layout, + "latent_mode": args.latent_mode, + "act_weights": { + "xyz": args.xyz_weight, "rot": args.rot_weight, "grip": args.grip_weight, + }, + "energy_sample": bool(args.energy_sample), + "geo_sample": bool(args.geo_sample), + "geo_temp": args.geo_temp, + "geo_loss_boost": args.geo_loss_boost, + "late_weight": args.late_weight, + "late_steps": args.late_steps, + "select": args.select, + "rollout_every": args.rollout_every, + "rollout_episodes": args.rollout_episodes, + "adversarial": { + "enabled": bool(args.advers), + "lambda_adv": args.adv_weight if args.advers else 0.0, + "advers_w": args.advers_w if args.advers else None, + "reference": "https://github.com/azencot-group/KDM Discriminator + trainer", + }, + }) + + # Free teacher weights before student training (pairs already generated / reused). + del agent + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + train, val = prepared["train"], prepared["validation"] + model = LowdimFDKStudent( + spec["obs_dim"], + spec["act_dim"], + use_koopman=True, + latent_dim=args.latent_dim, + teacher_steps=spec["teacher_steps"], + context_mode=args.context_mode, + layout=args.layout, + backbone_width=args.backbone_width, + latent_mode=args.latent_mode, + ).to(args.device) + ema = deepcopy(model).eval().requires_grad_(False) + opt = torch.optim.Adam(model.parameters(), lr=3e-4) + discriminator = None + opt_d = None + if args.advers: + discriminator = ChunkDiscriminator(spec["act_dim"], advers_w=args.advers_w).to(args.device) + opt_d = torch.optim.Adam(discriminator.parameters(), lr=3e-4, betas=(0.9, 0.999), eps=1e-8) + act_w = action_dim_weights( + spec["act_dim"], args.xyz_weight, args.rot_weight, args.grip_weight, + ).to(args.device) + sample_probs = None + if args.geo_sample: + sample_probs = geo_sample_probs(train["obs"], temperature=args.geo_temp) + print( + f"[GEO] sample probs mean={float(sample_probs.mean()):.3e} " + f"max/mean={float(sample_probs.max() / sample_probs.mean()):.2f}", + flush=True, + ) + elif args.energy_sample: + start, end = EXECUTE + energy_scores = train["actions"][:, start:end].square().mean(dim=(1, 2)).clamp_min(1e-6) + sample_probs = energy_scores / energy_scores.sum() + if args.select == "rollout": + best = {"score": float("-inf"), "step": 0, "chunk_mse": float("inf"), + "last_chunk_mse": float("inf"), "rollout_score": float("-inf")} + else: + best = {"score": float("inf"), "step": 0, "chunk_mse": float("inf"), + "last_chunk_mse": float("inf"), "rollout_score": None} + arm = out / "arm_vdv" + arm.mkdir() + tick = time.time() + model.train() + n_train = len(train["obs"]) + student_config = OmegaConf.merge(config, {"sample_steps": 1, "device": args.device}) + rollout_log = [] + for step in range(1, args.updates + 1): + if sample_probs is not None: + idx = torch.multinomial(sample_probs, args.batch_size, replacement=True) + else: + idx = torch.randint(0, n_train, (args.batch_size,)) + obs = train["obs"][idx].to(args.device) + noise = train["noise"][idx].to(args.device) + actions = train["actions"][idx].to(args.device) + result = model(noise, obs, clean=actions) + disc_parts = {} + if discriminator is not None: + # KDM order: update D on detached real/fake, then G with L_Koopman + λ L_adv. + opt_d.zero_grad(set_to_none=True) + d_loss, disc_parts = discriminator.disc_loss(actions, result["actions"]) + d_loss.backward() + torch.nn.utils.clip_grad_norm_(discriminator.parameters(), 1.0) + opt_d.step() + sample_w = None + if args.geo_loss_boost > 0: + sample_w = 1.0 + float(args.geo_loss_boost) * geo_proximity( + obs, temperature=args.geo_temp) + loss, parts = student_loss( + result, actions, + latent_weight=args.latent_weight, + spectral_weight=args.spectral_weight, + full_weight=args.full_weight, + first_weight=args.first_weight, + film_reg_weight=args.film_reg_weight, + adv_weight=args.adv_weight if args.advers else 0.0, + late_weight=args.late_weight, + late_steps=args.late_steps, + sample_weights=sample_w, + discriminator=discriminator, + act_weights=act_w, + ) + parts.update(disc_parts) + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + opt.step() + update_ema(ema, model) + if step % 100 == 0 or step == 1: + print( + f"[TRAIN] {args.task} {step}/{args.updates} " + f"loss={loss.item():.5f} {parts}", + flush=True, + ) + if step % args.validate_every == 0 or step == args.updates: + metrics = val_metrics(ema, val["obs"], val["noise"], val["actions"], args.device) + print( + f"[VAL] step={step} chunk_mse={metrics['chunk_mse']:.6f} " + f"last_chunk_mse={metrics['last_chunk_mse']:.6f}", + flush=True, + ) + take = False + select_val = None + rollout_score = None + if args.select == "rollout": + do_roll = ( + step % args.rollout_every == 0 + or step == args.updates + or step == args.validate_every + ) + if do_roll: + adapter = StudentAgent(ema, ema, args.task) + n_ep = args.rollout_episodes if step < args.updates else args.episodes + print(f"[ROLL] step={step} episodes={n_ep}", flush=True) + roll = evaluate( + args.task, student_config, dataset, adapter, + episodes=n_ep, + env_count=min(10, n_ep), + seed_base=100000, + max_steps=int(student_config.action_steps) if args.smoke else None, + ) + rollout_score = float(roll["score"]) + select_val = rollout_score + rollout_log.append({"step": step, "score": rollout_score, "episodes": n_ep}) + atomic_json(arm / "rollout_select_log.json", {"rows": rollout_log}) + print( + f"[ROLL] step={step} score={rollout_score:.4f} " + f"best={best['score']:.4f}", + flush=True, + ) + take = select_val > best["score"] + else: + select_val = ( + metrics["last_chunk_mse"] if args.select == "last_chunk" + else metrics["chunk_mse"] + ) + take = select_val < best["score"] + if take: + best = { + "score": select_val, + "step": step, + "chunk_mse": metrics["chunk_mse"], + "last_chunk_mse": metrics["last_chunk_mse"], + "rollout_score": rollout_score, + } + payload = { + "state_dict": ema.state_dict(), + "step": step, + "validation": { + "chunk_mse": best["chunk_mse"], + "last_chunk_mse": best["last_chunk_mse"], + "select": args.select, + "select_value": best["score"], + "rollout_score": rollout_score, + }, + "metadata": { + "task": args.task, + "use_koopman": True, + "operator": "vdv", + "context_mode": args.context_mode, + "layout": args.layout, + "backbone_width": args.backbone_width, + "latent_mode": args.latent_mode, + "obs_dim": spec["obs_dim"], + "act_dim": spec["act_dim"], + "latent_dim": args.latent_dim, + "teacher_steps": spec["teacher_steps"], + "teacher_sha256": teacher_sha, + "tau_prior": True, + "first_weight": args.first_weight, + "film_reg_weight": args.film_reg_weight, + "latent_weight": args.latent_weight, + "late_weight": args.late_weight, + "late_steps": args.late_steps, + "geo_sample": bool(args.geo_sample), + "geo_loss_boost": args.geo_loss_boost, + "advers": bool(args.advers), + "adv_weight": args.adv_weight if args.advers else 0.0, + "advers_w": args.advers_w if args.advers else 0.0, + }, + } + if discriminator is not None: + payload["discriminator"] = discriminator.state_dict() + torch.save(payload, arm / "best_ema.pt") + + if not (arm / "best_ema.pt").exists(): + # select=rollout never triggered a save (should not happen); fall back to last EMA. + torch.save({ + "state_dict": ema.state_dict(), + "step": args.updates, + "validation": {"chunk_mse": float("nan"), "select": args.select}, + "metadata": {"task": args.task, "latent_dim": args.latent_dim}, + }, arm / "best_ema.pt") + + saved = torch.load(arm / "best_ema.pt", map_location="cpu", weights_only=False) + model.load_state_dict(saved["state_dict"]) + model.eval().requires_grad_(False) + ema_roll = deepcopy(model) + adapter = StudentAgent(model, ema_roll, args.task) + print(f"[EVAL] {args.task} {args.episodes} eps", flush=True) + result = evaluate( + args.task, student_config, dataset, adapter, + episodes=args.episodes, + env_count=min(10, args.episodes), + seed_base=100000, + max_steps=int(student_config.action_steps) if args.smoke else None, + ) + result.update( + selected_step=saved["step"], + validation=saved["validation"], + task=args.task, + nfe=1, + teacher_sha256=teacher_sha, + wall_seconds=time.time() - tick, + context_mode=args.context_mode, + advers=bool(args.advers), + select=args.select, + ) + atomic_json(arm / "rollout.json", result) + summary = { + "task": args.task, + "score": result["score"], + "episodes": args.episodes, + "selected_step": saved["step"], + "chunk_mse": saved["validation"]["chunk_mse"], + "last_chunk_mse": saved["validation"].get("last_chunk_mse"), + "select": args.select, + "select_value": saved["validation"].get("select_value"), + "layout": args.layout, + "backbone_width": args.backbone_width, + "latent_mode": args.latent_mode, + "xyz_weight": args.xyz_weight, + "rot_weight": args.rot_weight, + "grip_weight": args.grip_weight, + "energy_sample": bool(args.energy_sample), + "geo_sample": bool(args.geo_sample), + "geo_temp": args.geo_temp, + "geo_loss_boost": args.geo_loss_boost, + "late_weight": args.late_weight, + "late_steps": args.late_steps, + "teacher_sha256": teacher_sha, + "latent_weight": args.latent_weight, + "first_weight": args.first_weight, + "film_reg_weight": args.film_reg_weight, + "latent_dim": args.latent_dim, + "context_mode": args.context_mode, + "tau_prior": args.layout == "tau_prior", + "advers": bool(args.advers), + "adv_weight": args.adv_weight if args.advers else 0.0, + "advers_w": args.advers_w if args.advers else 0.0, + "updates": args.updates, + "smoke": args.smoke, + "loss": loss_str, + "baseline_best_align_square": 0.36, + "baseline_wide_square": 0.50, + "baseline_chiunet_square": 0.72, + "baseline_onedp_square": 0.86, + "prev_action_only_score": { + "lift_lowdim": 0.98, + "square_lowdim": 0.28, + "pusht_keypoints": 0.811, + }.get(args.task), + } + teacher_eval = ROOT / ( + f"outputs/table2/chi_unet_20260909_v1/runs/{args.task}/seed0/eval_1000000.json" + ) + if teacher_eval.exists(): + summary["teacher_baseline_eval"] = float( + json.loads(teacher_eval.read_text())["score"] + ) + atomic_json(out / "summary.json", summary) + print(json.dumps(summary, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/wide_backbone.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/wide_backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..2de959ab6c86a8e01ae1451dc4a2fcd27917889b --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/lowdim_fdk_distill/wide_backbone.py @@ -0,0 +1,185 @@ +"""Kitchen temporal_v2 E/D kept out of cleandiffuser (Table-2 source hash). + +width=1 matches pinned FDK checkpoints (stem/blocks/head). width>1 scales +128/192/256 channels. Still pools H → one latent (no UNet skips). +""" +from __future__ import annotations + +import math + +import torch +from torch import nn +from torch.nn import functional as F + + +def _channels(width: int): + width = int(width) + if width < 1: + raise ValueError(f"backbone width must be >= 1, got {width}") + return 128 * width, 192 * width, 256 * width + + +class TemporalResidualBlock(nn.Module): + """Residual Conv1d block (state-dict key: *.net.*).""" + + def __init__(self, channels: int, dilation: int = 1): + super().__init__() + groups = math.gcd(channels, 8) + self.net = nn.Sequential( + nn.GroupNorm(groups, channels), + nn.SiLU(), + nn.Conv1d( + channels, channels, + kernel_size=3, padding=dilation, dilation=dilation, + ), + nn.GroupNorm(groups, channels), + nn.SiLU(), + nn.Conv1d(channels, channels, kernel_size=3, padding=1), + ) + + def forward(self, x): + return self.net(x) + x + + +class TemporalEncoder(nn.Module): + """[B, C, H] → [B, latent]. Strided Conv1d pool, no skips.""" + + def __init__(self, input_dim: int, latent_dim: int, horizon: int, width: int = 1): + super().__init__() + if horizon < 4: + raise ValueError("temporal_v2 backbone requires horizon >= 4") + c1, c2, c3 = _channels(width) + self.horizon = horizon + self.width = int(width) + self.reduced_horizon = horizon // 4 + self.stem = nn.Conv1d(input_dim, c1, kernel_size=5, padding=2) + self.pos_embedding = nn.Parameter(torch.zeros(1, c1, horizon)) + nn.init.normal_(self.pos_embedding, std=0.02) + self.blocks = nn.Sequential( + TemporalResidualBlock(c1), + nn.Conv1d(c1, c2, kernel_size=4, stride=2, padding=1), + TemporalResidualBlock(c2, dilation=2), + nn.Conv1d(c2, c3, kernel_size=4, stride=2, padding=1), + TemporalResidualBlock(c3, dilation=2), + ) + self.projection = nn.Sequential( + nn.Flatten(), + nn.Linear(c3 * self.reduced_horizon, latent_dim), + ) + + def forward(self, x): + if x.shape[-1] != self.horizon: + raise ValueError(f"expected horizon {self.horizon}, got {x.shape[-1]}") + x = self.stem(x) + self.pos_embedding + return self.projection(self.blocks(x)) + + +class TemporalDecoder(nn.Module): + """[B, latent] → flat [B, H*C] (caller .view's to [B,H,C]).""" + + def __init__(self, latent_dim: int, horizon: int, input_dim: int, width: int = 1): + super().__init__() + if horizon < 4: + raise ValueError("temporal_v2 backbone requires horizon >= 4") + c1, c2, c3 = _channels(width) + self.horizon = horizon + self.input_dim = input_dim + self.width = int(width) + self.reduced_horizon = horizon // 4 + self.c3 = c3 + self.projection = nn.Linear(latent_dim, c3 * self.reduced_horizon) + self.up1 = nn.Sequential( + TemporalResidualBlock(c3, dilation=2), + nn.ConvTranspose1d(c3, c2, kernel_size=4, stride=2, padding=1), + TemporalResidualBlock(c2), + ) + self.up2 = nn.ConvTranspose1d(c2, c1, kernel_size=4, stride=2, padding=1) + self.pos_embedding = nn.Parameter(torch.zeros(1, c1, horizon)) + nn.init.normal_(self.pos_embedding, std=0.02) + gn_groups = math.gcd(c1, 8) + self.head = nn.Sequential( + TemporalResidualBlock(c1), + nn.GroupNorm(gn_groups, c1), + nn.SiLU(), + nn.Conv1d(c1, input_dim, kernel_size=3, padding=1), + ) + + def forward(self, z): + x = self.projection(z).view(z.shape[0], self.c3, self.reduced_horizon) + x = self.up2(self.up1(x)) + if x.shape[-1] != self.horizon: + x = F.interpolate(x, size=self.horizon, mode="linear", align_corners=False) + x = self.head(x + self.pos_embedding) + return x.permute(0, 2, 1).reshape(z.shape[0], self.horizon * self.input_dim) + + +class SequenceTemporalEncoder(nn.Module): + """[B, C, H] → [B, T, latent] with T=H/4. Same stem/blocks as global E; no pool.""" + + def __init__(self, input_dim: int, latent_dim: int, horizon: int, width: int = 1): + super().__init__() + if horizon < 4 or horizon % 4 != 0: + raise ValueError("sequence backbone requires horizon divisible by 4") + c1, c2, c3 = _channels(width) + self.horizon = horizon + self.width = int(width) + self.reduced_horizon = horizon // 4 + self.stem = nn.Conv1d(input_dim, c1, kernel_size=5, padding=2) + self.pos_embedding = nn.Parameter(torch.zeros(1, c1, horizon)) + nn.init.normal_(self.pos_embedding, std=0.02) + self.blocks = nn.Sequential( + TemporalResidualBlock(c1), + nn.Conv1d(c1, c2, kernel_size=4, stride=2, padding=1), + TemporalResidualBlock(c2, dilation=2), + nn.Conv1d(c2, c3, kernel_size=4, stride=2, padding=1), + TemporalResidualBlock(c3, dilation=2), + ) + self.projection = nn.Conv1d(c3, latent_dim, kernel_size=1) + + def forward(self, x): + if x.shape[-1] != self.horizon: + raise ValueError(f"expected horizon {self.horizon}, got {x.shape[-1]}") + h = self.blocks(self.stem(x) + self.pos_embedding) + return self.projection(h).permute(0, 2, 1) # [B, T, L] + + +class SequenceTemporalDecoder(nn.Module): + """[B, T, latent] → flat [B, H*C]. Shared upsampling head; T must be H/4.""" + + def __init__(self, latent_dim: int, horizon: int, input_dim: int, width: int = 1): + super().__init__() + if horizon < 4 or horizon % 4 != 0: + raise ValueError("sequence backbone requires horizon divisible by 4") + c1, c2, c3 = _channels(width) + self.horizon = horizon + self.input_dim = input_dim + self.width = int(width) + self.reduced_horizon = horizon // 4 + self.projection = nn.Conv1d(latent_dim, c3, kernel_size=1) + self.up1 = nn.Sequential( + TemporalResidualBlock(c3, dilation=2), + nn.ConvTranspose1d(c3, c2, kernel_size=4, stride=2, padding=1), + TemporalResidualBlock(c2), + ) + self.up2 = nn.ConvTranspose1d(c2, c1, kernel_size=4, stride=2, padding=1) + self.pos_embedding = nn.Parameter(torch.zeros(1, c1, horizon)) + nn.init.normal_(self.pos_embedding, std=0.02) + gn_groups = math.gcd(c1, 8) + self.head = nn.Sequential( + TemporalResidualBlock(c1), + nn.GroupNorm(gn_groups, c1), + nn.SiLU(), + nn.Conv1d(c1, input_dim, kernel_size=3, padding=1), + ) + + def forward(self, z): + # z: [B, T, L] + if z.shape[1] != self.reduced_horizon: + raise ValueError( + f"expected T={self.reduced_horizon}, got {z.shape[1]}") + x = self.projection(z.permute(0, 2, 1)) + x = self.up2(self.up1(x)) + if x.shape[-1] != self.horizon: + x = F.interpolate(x, size=self.horizon, mode="linear", align_corners=False) + x = self.head(x + self.pos_embedding) + return x.permute(0, 2, 1).reshape(z.shape[0], self.horizon * self.input_dim) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a233d6f02312c6a3360f11224f6c8b6bd74070d2 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/__init__.py @@ -0,0 +1 @@ +"""OneDP-S published-formulation adaptation for native CleanDiffuser DDPM teachers.""" diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/core.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/core.py new file mode 100644 index 0000000000000000000000000000000000000000..891bc055b83dc0fd72757388f3acd617fbbb75fe --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/core.py @@ -0,0 +1,242 @@ +"""Standalone OneDP-S kernel using an explicit epsilon-to-x0 generator. + +The generator, auxiliary score network and frozen teacher own independent +observation encoders. Student encoders are trainable unless explicitly frozen; +eval mode fixes preprocessing/crops without disabling their weight gradients. + +The native teacher cumulative-alpha schedule and a valid fixed t_init are supplied +by the caller. Action shape is configurable, including (16, 2) and (16, 10). +There are no dataset, environment or other experimental-package dependencies. + +Generator choice: G(z,O) = (z - sigma_init * F_theta(z,t_init,O)) / alpha_init. +Teacher and auxiliary networks still predict epsilon. The final paper does not +uniquely establish this internal parameterization; neither matching tensor shapes +nor high scores demonstrate author-code equivalence. Never load a direct-action +head checkpoint as this wrapped generator without explicit semantic validation. +""" +from collections.abc import Mapping +import copy + +import torch +from torch import nn + + +PARAMETERIZATION = "epsilon_to_x0" +CORE_SCHEMA_VERSION = 1 + + +def _coefficients(reference, k, bar_alpha): + if reference.ndim < 2 or reference.shape[0] == 0: + raise ValueError("Expected a nonempty batched action tensor") + if (not isinstance(k, torch.Tensor) or k.dtype != torch.long + or tuple(k.shape) != (reference.shape[0],)): + raise ValueError("k must be a long tensor with one timestep per sample") + if bar_alpha.ndim != 1 or not bar_alpha.numel(): + raise ValueError("bar_alpha must be a nonempty vector") + if k.device != bar_alpha.device or reference.device != bar_alpha.device: + raise ValueError("Actions, timesteps and schedule must share a device") + if torch.any(k < 0) or torch.any(k >= bar_alpha.numel()): + raise ValueError("Timestep outside the actual teacher schedule") + shape = (reference.shape[0],) + (1,) * (reference.ndim - 1) + cumulative = bar_alpha[k].to(dtype=reference.dtype).reshape(shape) + return cumulative.sqrt(), (1 - cumulative).sqrt() + + +def forward_diffuse(actions, k, epsilon, bar_alpha): + """A_k = alpha_k A + sigma_k epsilon, preserving the gradient through A.""" + if (epsilon.shape != actions.shape or epsilon.device != actions.device + or epsilon.dtype != actions.dtype): + raise ValueError("Forward noise must match action shape, dtype and device") + alpha, sigma = _coefficients(actions, k, bar_alpha) + return alpha * actions + sigma * epsilon + + +def ddpm_score(epsilon_prediction, k, bar_alpha): + """Raw epsilon-to-score conversion s(A_k,k) = -epsilon_hat / sigma_k.""" + _, sigma = _coefficients(epsilon_prediction, k, bar_alpha) + return -epsilon_prediction / sigma + + +def auxiliary_score_loss(epsilon_prediction, epsilon): + """DDPM auxiliary denoising loss with lambda(k)=1, NOT sigma_k squared.""" + if epsilon_prediction.shape != epsilon.shape: + raise ValueError("Score prediction and noise target shapes must match") + return (epsilon_prediction - epsilon.detach()).square().mean() + + +def generator_surrogate(actions, k, epsilon, epsilon_teacher, epsilon_student, bar_alpha): + """Weighted reverse-KL gradient surrogate; the scalar is NOT a KL estimate. + + With w(k)=sigma_k^2, w(s_psi-s_phi)=sigma_k(eps_phi-eps_psi). + Detach that score difference, but NOT A_k: its alpha_k chain factor is + required. Mean reduction divides by all batch/action elements, matching the + implementation's auxiliary MSE convention. + """ + if epsilon_teacher.shape != actions.shape or epsilon_student.shape != actions.shape: + raise ValueError("Both raw epsilon predictions must match generated actions") + noisy = forward_diffuse(actions, k, epsilon, bar_alpha) + _, sigma = _coefficients(actions, k, bar_alpha) + direction = (sigma * (epsilon_teacher - epsilon_student)).detach() + return (direction * noisy).mean() + + +def _fresh_observations(observations): + # CleanDiffuser's image encoder replaces dictionary values after flattening + # history. A fresh dictionary prevents this from affecting another encoder. + if isinstance(observations, Mapping): + return {key: _fresh_observations(value) for key, value in observations.items()} + if not isinstance(observations, torch.Tensor): + raise TypeError("Normalized observations must be tensors or a tensor dictionary") + return observations + + +def _observation_batch(observations): + if isinstance(observations, Mapping): + if not observations: + raise ValueError("Empty observation dictionary") + sizes = {_observation_batch(value) for value in observations.values()} + if len(sizes) != 1: + raise ValueError("Observation fields have inconsistent batch sizes") + return sizes.pop() + if not isinstance(observations, torch.Tensor) or observations.ndim < 1: + raise ValueError("Observations must have a batch dimension") + return observations.shape[0] + + +class OneDPS(nn.Module): + """Teacher-initialized OneDP-S with isolated trainable generator and psi. + + ``teacher`` is the EMA ``nn.ModuleDict(diffusion=..., condition=...)``; + the caller's module is never changed. ``bar_alpha`` is its exact native + schedule, not a new linear/cosine reconstruction. All modules can be moved + together with ``.to(device)`` after construction. + """ + parameterization = PARAMETERIZATION + PARAMETERIZATION = PARAMETERIZATION + core_schema_version = CORE_SCHEMA_VERSION + + def __init__(self, teacher, bar_alpha, t_init, action_shape=(16, 2), frozen_vision=False): + super().__init__() + if not isinstance(teacher, nn.ModuleDict) or set(teacher) != {"diffusion", "condition"}: + raise TypeError("teacher must be a ModuleDict containing diffusion and condition") + schedule = torch.as_tensor(bar_alpha).detach().clone() + if (schedule.ndim != 1 or schedule.numel() < 1 or not schedule.is_floating_point() + or not torch.isfinite(schedule).all() or torch.any(schedule <= 0) + or torch.any(schedule >= 1) or torch.any(schedule[1:] >= schedule[:-1])): + raise ValueError("Require finite, strictly decreasing teacher bar_alpha in (0,1)") + if isinstance(t_init, bool) or not isinstance(t_init, int) or not 0 <= t_init < schedule.numel(): + raise ValueError("t_init must index the actual teacher schedule") + if len(action_shape) != 2 or any(int(value) != value or value <= 0 for value in action_shape): + raise ValueError("action_shape must contain positive horizon and action dimension") + self.register_buffer("bar_alpha", schedule) + self._t_init = t_init + self.action_shape = tuple(int(value) for value in action_shape) + self.frozen_vision = bool(frozen_vision) + + # Deepcopy preserves requires_grad=False from an EMA teacher. Explicitly + # unfreeze BOTH students; train() alone cannot fix frozen parameters. + self.teacher = copy.deepcopy(teacher).requires_grad_(False) + self.generator = copy.deepcopy(teacher).requires_grad_(True) + self.score_network = copy.deepcopy(teacher).requires_grad_(True) + if self.frozen_vision: + self.generator["condition"].requires_grad_(False) + self.score_network["condition"].requires_grad_(False) + self.train(True) + + @property + def t_init(self): + return self._t_init + + def train(self, mode=True): + super().train(mode) + self.teacher.requires_grad_(False).eval() + # Deterministic center crops/disabled condition dropout while preserving + # student encoder weight gradients. Explicit augmentation ablations can + # be implemented separately, not silently introduced into this pilot. + self.generator["condition"].eval() + self.score_network["condition"].eval() + return self + + def generator_parameters(self): + return tuple(parameter for parameter in self.generator.parameters() if parameter.requires_grad) + + def score_parameters(self): + return tuple(parameter for parameter in self.score_network.parameters() if parameter.requires_grad) + + def _predict(self, network, observations, noisy_actions, k): + # Reassert deterministic vision preprocessing if a caller toggled a + # child module directly. eval() does not turn off autograd. + network["condition"].eval() + condition = network["condition"](_fresh_observations(observations)) + prediction = network["diffusion"](noisy_actions, k, condition) + if prediction.shape != noisy_actions.shape: + raise ValueError("Teacher/student epsilon output has the wrong action shape") + return prediction + + def generate(self, observations, z): + """One-pass normalized action generation, with caller-supplied z~N(0,I).""" + if (z.ndim != 3 or tuple(z.shape[1:]) != self.action_shape + or z.shape[0] != _observation_batch(observations)): + raise ValueError("z must have shape [observation_batch, horizon, action_dim]") + k = torch.full((z.shape[0],), self.t_init, device=z.device, dtype=torch.long) + alpha, sigma = _coefficients(z, k, self.bar_alpha) + prediction = self._predict(self.generator, observations, z, k) + return (z - sigma * prediction) / alpha + + def score_loss(self, observations, z, k, epsilon): + """Train psi on detached G samples; never backpropagate this loss into G.""" + with torch.no_grad(): + actions = self.generate(observations, z) + noisy = forward_diffuse(actions, k, epsilon, self.bar_alpha) + prediction = self._predict(self.score_network, observations, noisy, k) + return auxiliary_score_loss(prediction, epsilon) + + def generator_loss(self, observations, z, k, epsilon): + """Fresh score evaluations for G, using psi AFTER its optimizer step.""" + actions = self.generate(observations, z) + with torch.no_grad(): + noisy = forward_diffuse(actions.detach(), k, epsilon, self.bar_alpha) + self.teacher.requires_grad_(False).eval() + teacher_epsilon = self._predict(self.teacher, observations, noisy, k) + student_epsilon = self._predict(self.score_network, observations, noisy, k) + return generator_surrogate(actions, k, epsilon, teacher_epsilon, student_epsilon, + self.bar_alpha) + + +def _gradient_norm(parameters, max_grad_norm): + if max_grad_norm is not None and max_grad_norm <= 0: + raise ValueError("max_grad_norm must be positive or None") + norm = nn.utils.clip_grad_norm_(parameters, + float("inf") if max_grad_norm is None else max_grad_norm, error_if_nonfinite=True) + return float(norm.detach()) + + +def alternating_step(model, observations, z, k, epsilon, generator_optimizer, + score_optimizer, max_grad_norm=10.0): + """One psi update THEN one G update with freshly recomputed psi predictions. + + Scores in the G surrogate cannot receive gradients. Clear psi gradients + after its own update so that the post-step state also exposes that invariant. + Returns scalar diagnostics, not a claim that the surrogate value is a KL. + """ + model.train() + generator_optimizer.zero_grad(set_to_none=True) + score_optimizer.zero_grad(set_to_none=True) + loss_score = model.score_loss(observations, z, k, epsilon) + if not torch.isfinite(loss_score): + raise FloatingPointError("Nonfinite auxiliary score loss") + loss_score.backward() + grad_score = _gradient_norm(model.score_parameters(), max_grad_norm) + score_optimizer.step() + score_optimizer.zero_grad(set_to_none=True) + + loss_generator = model.generator_loss(observations, z, k, epsilon) + if not torch.isfinite(loss_generator): + raise FloatingPointError("Nonfinite generator surrogate") + loss_generator.backward() + grad_generator = _gradient_norm(model.generator_parameters(), max_grad_norm) + generator_optimizer.step() + return {"loss_score": float(loss_score.detach()), + "loss_generator": float(loss_generator.detach()), + "grad_score": grad_score, "grad_generator": grad_generator} + diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/data.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/data.py new file mode 100644 index 0000000000000000000000000000000000000000..4c67ec93040d62e33124855e5bda9de8032ef3c6 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/data.py @@ -0,0 +1,167 @@ +"""Observation-only whole-episode student split; teacher statistics stay fixed.""" +import hashlib +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import Dataset + +from .teacher import task_spec, sha256 + + +def build_dataset(task, config): + task_spec(task) + common = dict(horizon=config.horizon, pad_before=config.obs_steps - 1, + pad_after=config.action_steps - 1) + if task == "relay_kitchen": + from unittest.mock import patch + from cleandiffuser.dataset import kitchen_dataset + parser, audit = kitchen_dataset.parse_mjl_logs, [] + def audited_parser(path, **kwargs): + row = {"path": str(Path(path).relative_to(Path(config.dataset_path)))} + try: + result = parser(path, **kwargs) + row["status"] = "parsed" + return result + except Exception as error: + row.update(status="skipped", error=repr(error)) + raise + finally: + audit.append(row) + # Same original MJL order, skipamount=40, and observation-noise seed42 + # as the teacher dataset; exceptions remain visible in the manifest. + with patch.object(kitchen_dataset, "parse_mjl_logs", audited_parser): + dataset = kitchen_dataset.KitchenMjlDataset(config.dataset_path, + abs_action=config.abs_action, **common) + dataset.parse_audit = audit + if sum(row["status"] == "parsed" for row in audit) != dataset.replay_buffer.n_episodes: + raise RuntimeError("Kitchen preprocessing dropped a successfully parsed demonstration") + if not len(dataset): + raise ValueError("Empty demonstration dataset") + return dataset + common.update(abs_action=config.abs_action, obs_keys=list(config.obs_keys)) + if task.startswith("pusht"): + from cleandiffuser.dataset.pusht_dataset import PushTStateDataset, PushTKeypointDataset, PushTImageDataset + cls = {"pusht": PushTStateDataset, "pusht_keypoints": PushTKeypointDataset, + "pusht_image": PushTImageDataset}[task] + else: + from cleandiffuser.dataset.robomimic_dataset import RobomimicDataset + cls = RobomimicDataset + dataset = cls(config.dataset_path, **common) + if not len(dataset): + raise ValueError("Empty demonstration dataset") + return dataset + + +def split_episodes(dataset, seed=20260914, validation_fraction=.1): + if not 0 < validation_fraction < 1: + raise ValueError("Validation fraction must be strictly between zero and one") + ends = np.asarray(dataset.replay_buffer.episode_ends) + if len(ends) < 2 or np.any(np.diff(ends) <= 0): + raise ValueError("Need at least two nonempty episodes") + indices = np.asarray(dataset.sampler.indices) + episode_ids = np.searchsorted(ends, indices[:, 0], side="right") + if np.any(episode_ids >= len(ends)) or np.any(indices[:, 1] > ends[episode_ids]): + raise ValueError("Invalid or cross-episode sequence window") + ids = np.random.default_rng(seed).permutation(len(ends)) + count = min(len(ids) - 1, max(1, round(len(ids) * validation_fraction))) + mask = np.isin(episode_ids, ids[:count]) + train, validation = np.flatnonzero(~mask), np.flatnonzero(mask) + if not len(train) or not len(validation): + raise ValueError("Empty student training or validation window split") + return train, validation, dict(split_seed=seed, validation_fraction=validation_fraction, + train_episodes=sorted(ids[count:].tolist()), validation_episodes=sorted(ids[:count].tolist()), + note="Student-context episode holdout; teacher was trained on all demonstrations.") + + +class ObservationDataset(Dataset): + """Do not return demonstration actions: OneDP learns online from teacher scores.""" + def __init__(self, dataset, indices, task): + task_spec(task) + self.dataset, self.indices, self.task = dataset, np.asarray(indices), task + + def __len__(self): + return len(self.indices) + + def __getitem__(self, index): + obs = self.dataset[int(self.indices[index])]["obs"] + if self.task == "pusht_image": + expected = {"image": (2, 3, 96, 96), "agent_pos": (2, 2)} + elif self.task == "pusht_keypoints": + expected = {"keypoint": (2, 18), "agent_pos": (2, 2)} + else: + expected = {"state": (2, task_spec(self.task)["obs_dim"])} + if set(obs) != set(expected): + raise ValueError("Unexpected observation fields; no privileged state is allowed") + result = {key: torch.as_tensor(value[:2], dtype=torch.float32) for key, value in obs.items()} + if any(tuple(result[key].shape) != shape or not torch.isfinite(result[key]).all() + for key, shape in expected.items()): + raise ValueError("Wrong shape or nonfinite normalized observation") + if self.task == "pusht_image": + return result + value = (torch.cat((result["keypoint"], result["agent_pos"]), dim=-1) + if self.task == "pusht_keypoints" else result["state"]) + return value.flatten() + + +def dataset_fingerprint(path): + """Content digest of a standalone HDF5 file or a complete Zarr tree.""" + path = Path(path).resolve(strict=True) + if path.is_file(): + return {"kind": "file", "sha256": sha256(path), "bytes": path.stat().st_size} + digest, count, total = hashlib.sha256(), 0, 0 + for member in sorted(path.rglob("*")): + if not member.is_file(): + continue + relative = member.relative_to(path).as_posix() + size = member.stat().st_size + digest.update(f"{relative}\0{size}\0{sha256(member)}\n".encode()) + count, total = count + 1, total + size + if not count: + raise ValueError("Empty dataset directory") + return {"kind": "directory", "tree_sha256": digest.hexdigest(), "files": count, "bytes": total} + + +def validate_kitchen_preprocessing(dataset, bundle): + """Keep teacher's exact parsed episode order and normalization statistics.""" + expected = bundle.metadata.get("dataset_parse_audit", []) + if not expected or len(expected) != len(dataset.parse_audit): + raise ValueError("Kitchen parse count differs from saved teacher audit") + # Teacher paths are repository relative; our paths are dataset relative. + for teacher_row, actual in zip(expected, dataset.parse_audit): + if (not teacher_row["path"].endswith("/" + actual["path"]) + or teacher_row["status"] != actual["status"]): + raise ValueError("Kitchen parsed episode order/status differs from teacher") + for computed, saved in ((dataset.normalizer["obs"]["state"], bundle.normalizer["obs"]["state"]), + (dataset.normalizer["action"], bundle.normalizer["action"])): + for key in ("min", "max", "range", "mask"): + if not np.array_equal(getattr(computed, key), getattr(saved, key)): + raise ValueError("Kitchen preprocessing statistics differ from teacher: " + key) + + +def prepare_data(bundle, dataset_path=None): + path = Path(bundle.config.dataset_path if dataset_path is None else dataset_path).expanduser().resolve(strict=True) + # Evaluation needs this path to read Robomimic environment metadata as well. + bundle.config.dataset_path = str(path) + fingerprint = dataset_fingerprint(path) + dataset = build_dataset(bundle.task, bundle.config) + if bundle.task == "relay_kitchen": + validate_kitchen_preprocessing(dataset, bundle) + dataset.normalizer = bundle.normalizer + train, validation, split = split_episodes(dataset) + metadata = dict(task=bundle.task, dataset_path=str(path), dataset_fingerprint=fingerprint, + train_windows=len(train), validation_windows=len(validation), total_windows=len(dataset), + normalizer_source="teacher_checkpoint", teacher_checkpoint_sha256=bundle.checkpoint_sha256, + demonstration_action_targets=False, **split) + if bundle.task == "relay_kitchen": + metadata.update(parse_audit=dataset.parse_audit, observation_shape=[2, 60], + action_shape=[16, 9], padded_observation_coordinates=list(range(30, 60)), + padded_future_states_supervised=False, teacher_preprocessing_bitwise_match=True) + return (ObservationDataset(dataset, train, bundle.task), + ObservationDataset(dataset, validation, bundle.task), metadata) + + +def to_device(observations, device): + if isinstance(observations, dict): + return {key: value.to(device=device, dtype=torch.float32) for key, value in observations.items()} + return observations.to(device=device, dtype=torch.float32) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/evaluation.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/evaluation.py new file mode 100644 index 0000000000000000000000000000000000000000..34c9cfceab4f0e853702d694dae880ce69b61c2e --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/evaluation.py @@ -0,0 +1,192 @@ +"""Native one-candidate closed-loop protocol, without old experiment imports.""" +import copy +import math +import random +import time + +import numpy as np +import torch + +from .teacher import task_spec, normalize_observations + + +def make_env(task, config): + """Same observations, absolute controller and action chunk as teacher evaluation.""" + import collections + import collections.abc + for name in ("Mapping", "Iterable"): + if not hasattr(collections, name): + setattr(collections, name, getattr(collections.abc, name)) + import gym + from cleandiffuser.env.wrapper import MultiStepWrapper + task_spec(task) + if task == "relay_kitchen": + from cleandiffuser.env import kitchen # noqa: F401 -- environment registration + from cleandiffuser.env.kitchen.kitchen_lowdim_wrapper import KitchenLowdimWrapper + env = KitchenLowdimWrapper(gym.make(config.env_name, use_abs_action=True), + init_qpos=None, init_qvel=None, render_hw=(240, 360)) + elif task.startswith("pusht"): + from cleandiffuser.env import pusht # noqa: F401 -- environment registration + env = gym.make(config.env_name, legacy=False) + else: + import robomimic.utils.file_utils as FileUtils + import robomimic.utils.env_utils as EnvUtils + import robomimic.utils.obs_utils as ObsUtils + from cleandiffuser.env.robomimic.robomimic_lowdim_wrapper import RobomimicLowdimWrapper + meta = FileUtils.get_env_metadata_from_dataset(config.dataset_path) + if meta["env_kwargs"].get("reward_shaping", False): + raise ValueError("Robomimic success requires sparse, unshaped rewards") + ObsUtils.initialize_obs_modality_mapping_from_dict({"low_dim": list(config.obs_keys)}) + meta["env_kwargs"]["controller_configs"]["control_delta"] = False + env = EnvUtils.create_env_from_metadata(meta, render=False, render_offscreen=False, + use_image_obs=False) + env.env.hard_reset = False + env = RobomimicLowdimWrapper(env, obs_keys=list(config.obs_keys)) + return MultiStepWrapper(env, config.obs_steps, config.action_steps, config.max_episode_steps) + + +def undo_absolute_action(action): + """Single-arm position(3)+rotation6D(6)+gripper(1) -> axis-angle action(7).""" + from cleandiffuser.dataset.dataset_utils import RotationTransformer + if action.shape[-1] != 10: + raise ValueError("Require 10D absolute single-arm action") + rotation = RotationTransformer(from_rep="axis_angle", to_rep="rotation_6d") + return np.concatenate((action[..., :3], rotation.inverse(action[..., 3:9]), action[..., [-1]]), axis=-1) + + +def action_chunk(bundle, prediction): + config = bundle.config + if prediction.ndim != 3 or tuple(prediction.shape[1:]) != (config.horizon, config.action_dim): + raise ValueError("Incorrect policy action tensor shape") + if not torch.isfinite(prediction).all(): + raise FloatingPointError("Nonfinite normalized policy action") + normalized = prediction.detach().cpu().numpy().clip(-1, 1) + actions = bundle.normalizer["action"].unnormalize(normalized) + if bundle.task.endswith("_lowdim"): + actions = undo_absolute_action(actions) + actions = actions[:, config.obs_steps - 1:config.obs_steps - 1 + config.action_steps] + if not np.isfinite(actions).all(): + raise FloatingPointError("Nonfinite physical policy action") + return actions + + +def episode_score(task, rewards, completed_count=0): + task_spec(task) + if task == "relay_kitchen": + return min(completed_count, 4) / 4. + return float(max(rewards, default=0.)) if task.startswith("pusht") else float(any(r > 0 for r in rewards)) + + +def evaluate_policy(bundle, model, args): + """Evaluate EMA teacher (model=None) or the actual ONLINE student generator. + + RNG and model modes are restored even on exceptions, so an optional zero- + update evaluation does not perturb training. The formal protocol is 50 + episodes in batches of ten environments. Smoke is explicitly truncated. + """ + task, config = bundle.task, copy.deepcopy(bundle.config) + task_spec(task) + if getattr(args, "task", task) != task: + raise ValueError("Evaluation task and teacher task differ") + smoke = bool(getattr(args, "smoke", False)) + episodes = int(getattr(args, "episodes", 50)) + if episodes <= 0 or (not smoke and episodes != 50): + raise ValueError("Formal evaluation requires exactly 50 full episodes") + config.seed = (int(bundle.metadata["seed"]) if model is None else + int(getattr(args, "seed", bundle.metadata["seed"]))) + device = torch.device(config.device) + if torch.device(getattr(args, "device", device)) != device: + raise ValueError("Evaluation must use the loaded teacher/student device") + devices = [device.index or 0] if device.type == "cuda" else [] + env_count, seed_base = min(2 if smoke else 10, episodes), int(getattr(args, "seed_base", 100000)) + offset = getattr(args, "policy_seed_offset", None) + # Fresh Relay comparisons reuse identical rollout RNG for every student + # training seed, matching the fixed-checkpoint FDK evaluation protocol. + if offset is None: + offset = 0 if task == "relay_kitchen" else config.seed + policy_rng_seed = seed_base + int(offset) + nfe = int(config.sample_steps) if model is None else 1 + numpy_state, python_state = np.random.get_state(), random.getstate() + modules = [bundle.agent.model, bundle.agent.model_ema] if model is None else [model] + modes = [module.training for module in modules] + envs, rows, latencies = [], [], [] + try: + with torch.random.fork_rng(devices=devices), torch.no_grad(): + torch.random.default_generator.manual_seed(policy_rng_seed) + if devices: + torch.cuda.default_generators[devices[0]].manual_seed(policy_rng_seed) + if task == "relay_kitchen": + random.seed(seed_base) + np.random.seed(seed_base) + for module in modules: + module.eval() + # Append separately so partial construction is also cleaned up. + for _ in range(env_count): + envs.append(make_env(task, config)) + for start in range(0, episodes, env_count): + count = min(env_count, episodes - start) + observations = [] + for index, env in enumerate(envs[:count]): + env.seed(seed_base + start + index) + observations.append(env.reset()) + finished, rewards = [False] * count, [[] for _ in range(count)] + completed = [0] * count + limit = min(8, config.max_episode_steps) if smoke else config.max_episode_steps + for _ in range(math.ceil(limit / config.action_steps)): + if all(finished): + break + if devices: + torch.cuda.synchronize(device) + tick = time.perf_counter() + condition = normalize_observations(bundle, observations) + prior = torch.zeros(count, config.horizon, config.action_dim, device=device) + if model is None: + prediction, _ = bundle.agent.sample(prior, n_samples=count, + sample_steps=nfe, condition_cfg=condition, w_cfg=1., use_ema=True) + else: + prediction = model.generate(condition, torch.randn_like(prior)) + actions = action_chunk(bundle, prediction) + if devices: + torch.cuda.synchronize(device) + latencies.append((time.perf_counter() - tick) * 1000) + for index, env in enumerate(envs[:count]): + if finished[index]: + continue + observation, reward, done, info = env.step(actions[index]) + observations[index] = observation + rewards[index].append(float(reward)) + finished[index] = bool(done) + if task == "relay_kitchen" and "completed_tasks" in info: + completed[index] = max(completed[index], len(info["completed_tasks"][-1])) + for index in range(count): + rows.append(dict(episode_seed=seed_base + start + index, + score=episode_score(task, rewards[index], completed[index]), + completed_tasks=completed[index] if task == "relay_kitchen" else None)) + print(f"[EVAL] {task} {len(rows)}/{episodes} mean={np.mean([row['score'] for row in rows]):.6f}", flush=True) + finally: + try: + for env in envs: + env.close() + finally: + np.random.set_state(numpy_state) + random.setstate(python_state) + for module, mode in zip(modules, modes): + module.train(mode) + scores = np.asarray([row["score"] for row in rows]) + return dict(task=task, training_seed=config.seed, episodes=rows, score=float(scores.mean()), + num_episodes=episodes, sample_steps=nfe, nfe=nfe, num_candidates=1, + weights="EMA" if model is None else "online", seed_base=seed_base, + policy_rng_seed=policy_rng_seed, + horizon=int(config.horizon), obs_steps=int(config.obs_steps), action_steps=int(config.action_steps), + action_shape=[int(config.horizon), int(config.action_dim)], + maximum_episode_steps=int(config.max_episode_steps), + pusht_legacy=False if task.startswith("pusht") else None, + absolute_action=bool(config.abs_action), smoke=smoke, truncated_smoke=smoke, + teacher_checkpoint_sha256=bundle.checkpoint_sha256, + episode_std=float(scores.std()), worst_episode=float(scores.min()), + lower_10_percent_mean=float(np.sort(scores)[:max(1, math.ceil(.1 * len(scores)))].mean()), + training_seed_std=None, decision_ms=float(np.mean(latencies)), + decision_ms_p50=float(np.percentile(latencies, 50)), decision_ms_p95=float(np.percentile(latencies, 95)), + latency_batch_envs=env_count, latency_concurrent_diagnostic=True, + latency_scope="normalization, transfer, encoder, sampler, action unnormalization and D2H; includes first call", + paper_ready=False) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/queue.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/queue.py new file mode 100644 index 0000000000000000000000000000000000000000..4d24ed8eacfd58019ec72909df7c299b49b563ec --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/queue.py @@ -0,0 +1,187 @@ +"""Sequential OneDP-S jobs, optionally detached; never manage other processes.""" +import argparse +import hashlib +import json +import os +from pathlib import Path +import subprocess +import sys +import tempfile +import time + + +TASKS = ("pusht", "pusht_keypoints", "pusht_image", "lift_lowdim", "square_lowdim") + + +def write_json(path, value): + path = Path(path) + descriptor, temporary = tempfile.mkstemp(prefix=path.name + ".", dir=path.parent) + try: + with os.fdopen(descriptor, "w") as stream: + json.dump(value, stream, indent=2, allow_nan=False) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + finally: + if os.path.exists(temporary): + os.unlink(temporary) + + +def read_json(path): + return json.loads(Path(path).read_text()) + + +def sources(): + folder = Path(__file__).resolve().parent + return {str(p.relative_to(folder)): hashlib.sha256(p.read_bytes()).hexdigest() + for p in sorted(folder.rglob("*.py"))} + + +def dataset_overrides(entries, tasks): + result = {} + for entry in entries: + task, separator, value = entry.partition("=") + if not separator or not value or task not in tasks or task in result: + raise ValueError("Each --dataset must uniquely specify a queued TASK=PATH") + path = Path(value).expanduser().resolve(strict=True) + result[task] = str(path) + return result + + +def build_jobs(args, output): + if (not args.tasks or len(set(args.tasks)) != len(args.tasks) + or any(task not in TASKS for task in args.tasks) or not args.teacher_root): + raise ValueError("Require unique supported tasks and an explicit teacher root") + if min(args.epochs, args.batch_size, args.episodes) < 1 or min(args.seed, args.teacher_seed) < 0: + raise ValueError("Positive budgets and nonnegative seeds required") + teachers = Path(args.teacher_root).expanduser().resolve(strict=True) + datasets = dataset_overrides(args.dataset, args.tasks) + jobs = [] + for task in args.tasks: + teacher = (teachers / task / f"seed{args.teacher_seed}" / "model_1000000.pt").resolve(strict=True) + if not teacher.is_file() or teacher.stat().st_size == 0: + raise ValueError("Teacher checkpoint must be a nonempty regular file") + destination = output / task / f"seed{args.seed}" + command = [sys.executable, "-m", "experiments.onedp_s.run", "--task", task, + "--teacher-checkpoint", str(teacher), "--output", str(destination), + "--device", args.device, "--seed", str(args.seed), "--epochs", str(args.epochs), + "--batch-size", str(args.batch_size), "--episodes", str(args.episodes)] + if task in datasets: + command += ["--dataset", datasets[task]] + jobs.append({"task": task, "seed": args.seed, "teacher_seed": args.teacher_seed, + "teacher_checkpoint": str(teacher), "teacher_bytes": teacher.stat().st_size, + "teacher_mtime_ns": teacher.stat().st_mtime_ns, + "output": str(destination), "command": command}) + return jobs + + +def prepare(args): + output = Path(args.output).expanduser().resolve() + if output.exists(): + raise FileExistsError(output) + jobs = build_jobs(args, output) # Validate every input before creating anything. + output.mkdir(parents=True, exist_ok=False) + plan = {"schema": 1, "package": "onedp_s", "created_unix": time.time(), "paper_ready": False, + "args": vars(args).copy(), "jobs": jobs, "sources": sources(), + "working_directory": str(Path(__file__).resolve().parents[2])} + write_json(output / "plan.json", plan) + write_json(output / "status.json", {"state": "prepared", "completed": [], "paper_ready": False}) + return output, plan + + +def validate_result(job, plan): + folder = Path(job["output"]) + status, summary = read_json(folder / "status.json"), read_json(folder / "summary.json") + if (status.get("state") != "complete" or summary.get("task") != job["task"] + or summary.get("training_seed") != job["seed"] or summary.get("method") != "OneDP-S" + or summary.get("smoke") is not False or summary.get("paper_ready") is not False + or summary.get("checkpoint_selection") != "final" or summary.get("weights") != "online" + or summary.get("epochs") != plan["args"]["epochs"] + or summary.get("batch_size") != plan["args"]["batch_size"] + or summary.get("evaluation", {}).get("num_episodes") != plan["args"]["episodes"] + or not isinstance(summary.get("updates"), int) or summary["updates"] < 1): + raise RuntimeError("Child exited without a matching completed experiment") + checkpoint = Path(summary["checkpoint"]).resolve(strict=True) + if checkpoint != (folder / "checkpoints/final.pt").resolve() or checkpoint.stat().st_size == 0: + raise RuntimeError("Missing expected final checkpoint") + return {"task": job["task"], "seed": job["seed"], "score": summary["evaluation"]["score"], + "updates": summary["updates"], "checkpoint_sha256": summary["checkpoint_sha256"]} + + +def worker(output): + output = Path(output).resolve(strict=True) + plan = read_json(output / "plan.json") + # An exclusive receipt refuses a second worker or accidental restart. + with (output / "worker.json").open("x") as stream: + json.dump({"pid": os.getpid(), "started_unix": time.time()}, stream) + stream.flush() + os.fsync(stream.fileno()) + completed = [] + try: + for job in plan["jobs"]: + if sources() != plan["sources"]: + raise RuntimeError("Queued package sources changed") + teacher = Path(job["teacher_checkpoint"]) + if teacher.stat().st_size != job["teacher_bytes"] or teacher.stat().st_mtime_ns != job["teacher_mtime_ns"]: + raise RuntimeError("Queued teacher checkpoint changed") + if Path(job["output"]).exists(): + raise FileExistsError(job["output"]) + log_path = output / (job["task"] + ".log") + with log_path.open("xb") as log: + child = subprocess.Popen(job["command"], cwd=plan["working_directory"], + stdin=subprocess.DEVNULL, stdout=log, stderr=subprocess.STDOUT) + write_json(output / "status.json", {"state": "running", "pid": os.getpid(), + "child_pid": child.pid, "task": job["task"], "completed": completed, + "log": str(log_path), "paper_ready": False}) + code = child.wait() + if code: + raise RuntimeError(f"{job['task']} exited with code {code}; see {log_path}") + completed.append(validate_result(job, plan)) + write_json(output / "status.json", {"state": "complete", "pid": os.getpid(), + "child_pid": None, "completed": completed, "paper_ready": False}) + return 0 + except BaseException as error: + write_json(output / "status.json", {"state": "failed", "pid": os.getpid(), + "completed": completed, "error": repr(error), "paper_ready": False}) + raise + + +def launch(args): + output, plan = prepare(args) + if not args.detach: + return worker(output) + command = [sys.executable, "-m", "experiments.onedp_s.queue", "--worker", "--output", str(output)] + try: + with (output / "queue.log").open("xb") as log: + child = subprocess.Popen(command, cwd=plan["working_directory"], stdin=subprocess.DEVNULL, + stdout=log, stderr=subprocess.STDOUT, start_new_session=True) + write_json(output / "launcher.json", {"pid": child.pid, "command": command, + "started_unix": time.time(), "detached": True, "paper_ready": False}) + except BaseException as error: + write_json(output / "status.json", {"state": "failed", "error": repr(error), "paper_ready": False}) + raise + print(json.dumps({"queue_pid": child.pid, "output": str(output), "status": str(output / "status.json")}), flush=True) + return 0 + + +def parse_args(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tasks", nargs="+", choices=TASKS) + parser.add_argument("--teacher-root", help="Root containing TASK/seedN/model_1000000.pt") + parser.add_argument("--teacher-seed", type=int, default=0) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--dataset", action="append", default=[], help="Optional TASK=PATH; repeat for each override") + parser.add_argument("--output", required=True) + parser.add_argument("--epochs", type=int, default=20) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--episodes", type=int, default=50) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--detach", action="store_true", help="Start only this queue in a new background session") + parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS) + return parser.parse_args(argv) + + +if __name__ == "__main__": + arguments = parse_args() + raise SystemExit(worker(arguments.output) if arguments.worker else launch(arguments)) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/repeat_eval.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/repeat_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..37f4ce7b7f353ff88c36d939a93f4c9a8412f5a8 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/repeat_eval.py @@ -0,0 +1,120 @@ +"""Fresh fixed-checkpoint Relay-Kitchen test banks, shared with FDK and CP. + +Default banks are disjoint from the old FDK checkpoint-selection episodes. +This measures evaluation variation, conditional on the declared training seed. +""" +import argparse +import statistics +import time +from pathlib import Path +from types import SimpleNamespace + +import torch + +from .core import OneDPS +from .evaluation import evaluate_policy +from .run import atomic_save, module_sha256, restore_generator, source_fingerprints +from .teacher import load_teacher, sha256 + +SEED_BASES = (400000, 500000, 600000) + + +def aggregate_groups(groups): + if len(groups) != 3 or len({g["training_seed"] for g in groups}) != 1: + raise ValueError("Require three groups of one fixed training checkpoint") + seeds = [row["episode_seed"] for group in groups for row in group["episodes"]] + if len(seeds) != 150 or len(set(seeds)) != 150: + raise ValueError("Require 150 distinct episode seeds") + values = [] + for group in groups: + scores = [row["score"] for row in group["episodes"]] + if (len(scores) != 50 or any(not 0 <= score <= 1 for score in scores) + or abs(statistics.mean(scores) - group["score"]) > 1e-12 + or group["policy_rng_seed"] != group["seed_base"]): + raise ValueError("Invalid raw scores or unmatched policy RNG") + values.extend(scores) + means = [group["score"] for group in groups] + return dict(group_scores=means, mean=statistics.mean(means), + evaluation_group_std=statistics.stdev(means), evaluation_group_std_ddof=1, + best_group_mean=max(means), worst_group_mean=min(means), worst_episode=min(values), + lower_10_percent_mean=statistics.mean(sorted(values)[:15]), + total_episodes=150, num_groups=3, episodes_per_group=50, + training_seed=groups[0]["training_seed"], num_training_seeds=1, training_seed_std=None, + uncertainty_scope="evaluation variation conditional on one fixed checkpoint") + + +def run(args): + bases = list(args.seed_bases) + if (len(bases) != 3 or min(bases) < 0 + or len({seed for base in bases for seed in range(base, base + 50)}) != 150): + raise ValueError("Specify three disjoint 50-episode seed banks") + out = Path(args.output).resolve() + out.mkdir(parents=True, exist_ok=False) + started = time.time() + try: + atomic_save(out / "status.json", dict(state="loading")) + checkpoint = Path(args.checkpoint).resolve(strict=True) + digest = sha256(checkpoint) + if args.teacher_only: + bundle = load_teacher("relay_kitchen", checkpoint, args.device) + model, seed, step, weights = None, bundle.metadata["seed"], 1_000_000, "EMA" + method, before = "DDPM teacher", module_sha256(bundle.agent.model_ema) + else: + payload = torch.load(checkpoint, map_location="cpu", weights_only=False, mmap=True) + if (payload.get("task") != "relay_kitchen" or payload.get("smoke") is not False + or payload.get("step", 0) <= 0): + raise ValueError("Require a completed, nonsmoke Relay-Kitchen student") + bundle = load_teacher("relay_kitchen", payload["teacher_checkpoint"], args.device) + model = OneDPS(bundle.agent.model_ema, bundle.agent.bar_alpha, payload["t_init"], + action_shape=(16, 9)).to(args.device) + restore_generator(model, payload, "relay_kitchen", bundle.checkpoint_sha256) + model.requires_grad_(False).eval() + seed, step, weights = payload["args"]["seed"], payload["step"], "online" + method, before = "OneDP-S", module_sha256(model.generator) + del payload + plan = dict(method=method, task="relay_kitchen", checkpoint=str(checkpoint), + checkpoint_sha256=digest, training_seed=seed, checkpoint_step=step, + teacher_checkpoint_sha256=bundle.checkpoint_sha256, weights=weights, + episode_seed_bases=bases, policy_rng_seeds=bases, episodes_per_group=50, + source_fingerprints=source_fingerprints(), training_updates=0, + paper_original_recipe_reproduced=False, fresh_evaluation=True) + atomic_save(out / "plan.json", plan) + groups = [] + for index, base in enumerate(bases): + atomic_save(out / "status.json", dict(state="evaluating", group=index, seed_base=base)) + options = SimpleNamespace(task="relay_kitchen", episodes=50, smoke=False, device=args.device, + seed=seed, seed_base=base, policy_seed_offset=0) + result = evaluate_policy(bundle, model, options) + result.update(group_index=index, checkpoint_sha256=digest, method=method, + checkpoint_step=step) + atomic_save(out / f"group{index}.json", result) + groups.append(result) + actual = module_sha256(bundle.agent.model_ema if model is None else model.generator) + if actual != before: + raise RuntimeError("Evaluation modified policy weights") + summary = {**aggregate_groups(groups), **plan, "source_fingerprints": plan["source_fingerprints"], + "weights_unchanged": True, "elapsed_seconds": time.time() - started} + atomic_save(out / "summary.json", summary) + atomic_save(out / "status.json", dict(state="complete", total_episodes=150, + elapsed_seconds=time.time() - started)) + print(f"[COMPLETE] {method} seed={seed} groups={summary['group_scores']} " + f"mean={summary['mean']:.6f} std={summary['evaluation_group_std']:.6f}", flush=True) + return summary + except BaseException as error: + atomic_save(out / "status.json", dict(state="failed", error=repr(error))) + raise + + +def parse_args(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--output", required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--seed-bases", nargs=3, type=int, default=SEED_BASES) + parser.add_argument("--teacher-only", action="store_true", + help="Checkpoint is the frozen DDPM teacher rather than OneDP-S") + return parser.parse_args(argv) + + +if __name__ == "__main__": + run(parse_args()) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/run.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/run.py new file mode 100644 index 0000000000000000000000000000000000000000..84b112f68adca4b0a44afa24a3f49f1947303723 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/run.py @@ -0,0 +1,375 @@ +"""Portable, explicitly epsilon-to-x0-parameterized OneDP-S DDPM adaptation. + +Teacher weights and the dataset are user-supplied local inputs. No downloads, +process management, or original-paper reproduction claims are implicit here. +""" +import argparse +from contextlib import contextmanager +import hashlib +import json +import math +import os +from pathlib import Path +import platform +import random +import tempfile +import time + +import numpy as np +import torch +from torch.utils.data import DataLoader + +from .core import CORE_SCHEMA_VERSION, PARAMETERIZATION, OneDPS, alternating_step + + +TASKS = ("pusht", "pusht_keypoints", "pusht_image", "lift_lowdim", "square_lowdim", "relay_kitchen") +CHECKPOINT_SCHEMA = 1 + + +def file_sha256(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for chunk in iter(lambda: stream.read(8 * 1024**2), b""): + digest.update(chunk) + return digest.hexdigest() + + +def module_sha256(module): + digest = hashlib.sha256() + for name, value in sorted(module.state_dict().items()): + value = value.detach().cpu().contiguous() + digest.update(name.encode()) + digest.update(str((value.dtype, tuple(value.shape))).encode()) + digest.update(value.numpy().tobytes()) + return digest.hexdigest() + + +def plain(value): + """JSON-compatible provenance without silently dropping configuration fields.""" + from omegaconf import OmegaConf + if OmegaConf.is_config(value): + return plain(OmegaConf.to_container(value, resolve=True)) + if isinstance(value, dict): + return {str(k): plain(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [plain(v) for v in value] + if isinstance(value, (torch.Tensor, np.ndarray)): + return value.detach().cpu().tolist() if isinstance(value, torch.Tensor) else value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, Path): + return str(value) + return value + + +def atomic_save(path, value, *, checkpoint=False): + """Publish only a fully flushed file; temporary files stay inside the run.""" + path = Path(path) + descriptor, temporary = tempfile.mkstemp(prefix=path.name + ".", suffix=".tmp", dir=path.parent) + try: + with os.fdopen(descriptor, "wb") as stream: + if checkpoint: + torch.save(value, stream) + else: + stream.write((json.dumps(plain(value), indent=2, allow_nan=False) + "\n").encode()) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + finally: + if os.path.exists(temporary): + os.unlink(temporary) + + +def source_fingerprints(): + root = Path(__file__).resolve().parents[2] + paths = list(Path(__file__).parent.rglob("*.py")) + paths += list((root / "cleandiffuser").rglob("*.py")) + return {str(p.relative_to(root)): file_sha256(p) for p in sorted(paths)} + + +def tags(task): + return {"package": "onedp_s", "method": "OneDP-S", "task": task, + "parameterization": PARAMETERIZATION, "core_schema": CORE_SCHEMA_VERSION, + "adaptation": "native_teacher_DDPM_epsilon_to_x0", "paper_ready": False} + + +@contextmanager +def preserve_rng(loader_rng, device): + numpy_state, python_state = np.random.get_state(), random.getstate() + loader_state = loader_rng.get_state().clone() + devices = list(range(torch.cuda.device_count())) if torch.device(device).type == "cuda" else [] + try: + with torch.random.fork_rng(devices=devices): + yield + finally: + np.random.set_state(numpy_state) + random.setstate(python_state) + loader_rng.set_state(loader_state) + + +def configure_noise(args, bundle): + """Use the actual trained grid. score_min/max are inclusive native indices.""" + count = len(bundle.agent.bar_alpha) + expected = 5 if args.task.startswith("pusht") else 50 + if (count != expected or bundle.agent.diffusion_steps != count + or int(bundle.config.sample_steps) != count or not bundle.agent.predict_noise): + raise ValueError("Teacher schedule and metadata must match the supported native DDPM grid") + defaults = (2, 0, 3) if count == 5 else (32, 1, 47) + for name, default in zip(("t_init", "score_min", "score_max"), defaults): + if getattr(args, name) is None: + setattr(args, name, default) + if not (0 <= args.t_init < count and 0 <= args.score_min <= args.score_max < count - 1): + raise ValueError("Noise indices must be valid; exclude the terminal score timestep") + return {"diffusion_steps": count, "bar_alpha": plain(bundle.agent.bar_alpha), + "t_init": args.t_init, "score_timesteps": list(range(args.score_min, args.score_max + 1)), + "mapping_note": "DDPM5 retained; DDPM50 defaults scale paper indices, not an official reproduction"} + + +def batch_count(observations): + return len(next(iter(observations.values()))) if isinstance(observations, dict) else len(observations) + + +def conditioning_module(network): + if any(True for _ in network["condition"].parameters()): + return network["condition"] + return getattr(network["diffusion"], "global_cond_encoder", network["condition"]) + + +@torch.no_grad() +def check_parameters_finite(model): + for name, value in model.named_parameters(): + if not torch.isfinite(value).all(): + raise FloatingPointError("Nonfinite model parameter: " + name) + + +@torch.no_grad() +def diagnostics(model, observations, z): + actions, alternative = model.generate(observations, z), model.generate(observations, -z) + if not torch.isfinite(actions).all() or not torch.isfinite(alternative).all(): + raise FloatingPointError("Nonfinite generated actions") + return {"generated_action_rms": float(actions.square().mean().sqrt()), + "generated_action_absmax": float(actions.abs().max()), + "fraction_outside_action_bounds": float((actions.abs() > 1).float().mean()), + "fixed_observation_noise_sensitivity_mse": float((actions - alternative).square().mean())} + + +def checkpoint_payload(model, optimizers, args, bundle, epoch, step, data, loader_rng, source): + return {**tags(args.task), "schema": CHECKPOINT_SCHEMA, + "generator": model.generator.state_dict(), "score_network": model.score_network.state_dict(), + "generator_optimizer": optimizers[0].state_dict(), "score_optimizer": optimizers[1].state_dict(), + "bar_alpha": model.bar_alpha.detach().cpu(), "t_init": model.t_init, + "action_shape": list(model.action_shape), "epoch": epoch, "step": step, + "args": vars(args).copy(), "teacher_checkpoint_sha256": bundle.checkpoint_sha256, + "teacher_checkpoint": str(bundle.checkpoint_path), "teacher_metadata": plain(bundle.metadata), + "teacher_config": plain(bundle.config), "data": plain(data), "source_fingerprints": source, + "torch_rng": torch.get_rng_state(), "numpy_rng": np.random.get_state(), + "python_rng": random.getstate(), "loader_rng": loader_rng.get_state(), + "cuda_rng": torch.cuda.get_rng_state_all() if torch.device(args.device).type == "cuda" else [], + "weights": "online", "checkpoint_selection": "final", "smoke": args.smoke} + + +def restore_generator(model, payload, task, teacher_sha256): + expected = {**tags(task), "schema": CHECKPOINT_SCHEMA, "weights": "online", "checkpoint_selection": "final"} + if type(model) is not OneDPS: + raise TypeError("Require the declared epsilon-to-x0 OneDPS class") + for key, value in expected.items(): + if payload.get(key) != value: + raise ValueError("Checkpoint identity mismatch: " + key) + if (payload.get("teacher_checkpoint_sha256") != teacher_sha256 + or payload.get("t_init") != model.t_init + or payload.get("action_shape") != list(model.action_shape) + or not isinstance(payload.get("bar_alpha"), torch.Tensor) + or not torch.equal(payload["bar_alpha"].cpu(), model.bar_alpha.cpu())): + raise ValueError("Checkpoint teacher, action shape or native noise schedule differs") + model.generator.load_state_dict(payload["generator"], strict=True) + + +def verify_checkpoint(path, model, observations, z, task, teacher_sha256): + was_training = model.training + model.eval() + try: + with torch.no_grad(): + expected = model.generate(observations, z).clone() + payload = torch.load(path, map_location="cpu", weights_only=False) + restore_generator(model, payload, task, teacher_sha256) + with torch.no_grad(): + error = float((model.generate(observations, z) - expected).abs().max()) + finally: + model.train(was_training) + if not math.isfinite(error) or error > 1e-6: + raise RuntimeError("Checkpoint output round-trip mismatch") + return {"checkpoint_round_trip_max_abs_error": error, "strict_typed_load": True} + + +def run(args): + from .teacher import load_teacher, teacher_parity + from .data import prepare_data, to_device + from .evaluation import evaluate_policy + if args.task not in TASKS or min(args.epochs, args.batch_size, args.episodes) < 1 or args.workers < 0: + raise ValueError("Supported task and positive budgets required") + if args.smoke: + smoke_batch = int(getattr(args, "smoke_batch_size", 1)) + if smoke_batch < 1: + raise ValueError("Smoke batch size must be positive") + args.epochs, args.batch_size, args.episodes, args.workers = 1, smoke_batch, 2, 0 + output = Path(args.output).expanduser().resolve() + output.mkdir(parents=True, exist_ok=False) + (output / "checkpoints").mkdir() + started = time.time() + source = source_fingerprints() + def status(state, **fields): + atomic_save(output / "status.json", {**tags(args.task), "state": state, "smoke": args.smoke, **fields}) + try: + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + status("loading_teacher") + bundle = load_teacher(args.task, args.teacher_checkpoint, args.device) + noise_config = configure_noise(args, bundle) + train, validation, data = prepare_data(bundle, dataset_path=args.dataset) + if len(train) < 1 or len(validation) < 1: + raise ValueError("Nonempty whole-episode train and validation splits required") + loader_rng = torch.Generator().manual_seed(args.seed + 6001) + loader = DataLoader(train, batch_size=args.batch_size, shuffle=True, drop_last=False, + num_workers=args.workers, generator=loader_rng, + pin_memory=torch.device(args.device).type == "cuda", persistent_workers=args.workers > 0) + expected_updates = 2 if args.smoke else args.epochs * len(loader) + if args.smoke and len(loader) < 2: + raise ValueError("Smoke requires at least two training contexts") + probe = to_device(next(iter(DataLoader(validation, batch_size=1 if args.smoke else 4))), args.device) + with preserve_rng(loader_rng, args.device): + parity = teacher_parity(bundle, probe) + atomic_save(output / "teacher_parity.json", parity) + atomic_save(output / "data_manifest.json", data) + plan = {**tags(args.task), **vars(args), **noise_config, "expected_updates": expected_updates, + "train_windows": len(train), "validation_windows": len(validation), + "generator_lr": 1e-6, "score_lr": 2e-5, "adam_betas": [0., .999], + "score_lambda": 1., "gradient_clip_norm": 10., "teacher_metadata": bundle.metadata, + "teacher_checkpoint_sha256": bundle.checkpoint_sha256, "source_fingerprints": source, + "teacher_config": bundle.config, "runtime": {"python": platform.python_version(), + "torch": str(torch.__version__), "numpy": np.__version__, "cuda": torch.version.cuda}, + "checkpoint_selection": "final", "student_weights": "online", "num_candidates": 1} + atomic_save(output / "plan.json", plan) + shape = (int(bundle.config.horizon), int(bundle.config.action_dim)) + model = OneDPS(bundle.agent.model_ema, bundle.agent.bar_alpha, args.t_init, action_shape=shape).to(args.device) + optimizers = (torch.optim.Adam(model.generator_parameters(), lr=1e-6, betas=(0., .999)), + torch.optim.Adam(model.score_parameters(), lr=2e-5, betas=(0., .999))) + teacher_hash = module_sha256(model.teacher) + student_hashes = {name: module_sha256(getattr(model, name)) for name in ("generator", "score_network")} + condition_hashes = {name: module_sha256(conditioning_module(getattr(model, name))) for name in student_hashes} + probe_rng = torch.Generator(device=args.device).manual_seed(args.seed + 6002) + probe_z = torch.randn(batch_count(probe), *shape, device=args.device, generator=probe_rng) + with torch.no_grad(): + observations = dict(probe) if isinstance(probe, dict) else probe + condition = bundle.agent.model_ema["condition"](observations) + k = torch.full((len(probe_z),), args.t_init, device=args.device, dtype=torch.long) + epsilon = bundle.agent.model_ema["diffusion"](probe_z, k, condition) + alpha, sigma = model.bar_alpha[args.t_init].sqrt(), (1 - model.bar_alpha[args.t_init]).sqrt() + warm_error = float((model.generate(probe, probe_z) - (probe_z - sigma * epsilon) / alpha).abs().max()) + if not math.isfinite(warm_error) or warm_error > 1e-5: + raise RuntimeError("Teacher initialization output mismatch") + verification = {"warm_start_max_abs_error": warm_error, "initial_diagnostics": diagnostics(model, probe, probe_z)} + evaluations = {} + for stage, policy in (("teacher", None), ("zero_update", model)): + status("evaluating_" + stage, step=0) + with preserve_rng(loader_rng, args.device): + result = evaluate_policy(bundle, policy, args) + result.update(evaluation_stage=stage, parameterization="native_ddpm" if policy is None else PARAMETERIZATION, + paper_ready=False) + evaluations[stage] = result + atomic_save(output / ("teacher_evaluation.json" if stage == "teacher" else "zero_update_evaluation.json"), result) + if module_sha256(model.teacher) != teacher_hash or any(module_sha256(getattr(model, n)) != h for n, h in student_hashes.items()): + raise RuntimeError("Initial evaluation changed teacher/student parameters") + verification["initial_evaluations_weights_unchanged"] = True + verification["initial_evaluations_rng_restored"] = True + atomic_save(output / "verification.json", verification) + step, train_started = 0, time.time() + with (output / "train.jsonl").open("x") as log: + for epoch in range(1, args.epochs + 1): + for index, batch in enumerate(loader): + if args.smoke and index >= 2: + break + observations = to_device(batch, args.device) + z = torch.randn(batch_count(observations), *shape, device=args.device) + k = torch.randint(args.score_min, args.score_max + 1, (len(z),), device=args.device) + metrics = alternating_step(model, observations, z, k, torch.randn_like(z), *optimizers, max_grad_norm=10.) + if not all(math.isfinite(float(v)) for v in metrics.values()): + raise FloatingPointError("Nonfinite training diagnostics") + step += 1 + log.write(json.dumps({"epoch": epoch, "step": step, **metrics, + "elapsed_seconds": time.time() - train_started}, allow_nan=False) + "\n") + if step == 1 or step % 50 == 0 or args.smoke: + log.flush() + status("training", step=step, expected_updates=expected_updates, epoch=epoch) + print(f"[TRAIN] {args.task} step={step}/{expected_updates} {metrics}", flush=True) + if step == 1: + checks = {"teacher_bitwise_unchanged": module_sha256(model.teacher) == teacher_hash, + "teacher_no_gradients": all(p.grad is None for p in model.teacher.parameters())} + for name, digest in student_hashes.items(): + checks[name + "_changed"] = module_sha256(getattr(model, name)) != digest + child = conditioning_module(getattr(model, name)) + if any(True for _ in child.parameters()): + checks[name + "_conditioning_changed"] = module_sha256(child) != condition_hashes[name] + if not all(checks.values()): + raise RuntimeError("Teacher/student update isolation check failed") + verification["first_step"] = checks + atomic_save(output / "verification.json", verification) + check_parameters_finite(model) + if module_sha256(model.teacher) != teacher_hash: + raise RuntimeError("Teacher changed during training") + atomic_save(output / f"diagnostics_epoch{epoch:03d}.json", diagnostics(model, probe, probe_z)) + if step != expected_updates: + raise RuntimeError("Training did not complete its declared update budget") + final = output / "checkpoints/final.pt" + atomic_save(final, checkpoint_payload(model, optimizers, args, bundle, args.epochs, step, data, loader_rng, source), checkpoint=True) + verification.update(verify_checkpoint(final, model, probe, probe_z, args.task, bundle.checkpoint_sha256)) + verification["final_teacher_bitwise_unchanged"] = module_sha256(model.teacher) == teacher_hash + atomic_save(output / "verification.json", verification) + status("evaluating_final", step=step) + evaluation = evaluate_policy(bundle, model, args) + evaluation.update(evaluation_stage="final", parameterization=PARAMETERIZATION, paper_ready=False) + atomic_save(output / "evaluation.json", evaluation) + summary = {**tags(args.task), "smoke": args.smoke, "training_seed": args.seed, + "epochs": args.epochs, "updates": step, "generator_optimizer_steps": step, + "score_optimizer_steps": step, "batch_size": args.batch_size, "weights": "online", + "checkpoint_selection": "final", "checkpoint": str(final), "checkpoint_sha256": file_sha256(final), + "teacher_checkpoint_sha256": bundle.checkpoint_sha256, "teacher_evaluation": evaluations["teacher"], + "zero_update_evaluation": evaluations["zero_update"], "evaluation": evaluation, + "verification": verification, "source_fingerprints": source, "data": data, + "train_windows": len(train), "validation_windows": len(validation), + "total_seconds": time.time() - started} + atomic_save(output / "summary.json", summary) + status("complete", step=step, formal_reproduction_verified=False) + return summary + except BaseException as error: + status("failed", error=repr(error)) + raise + + +def parse_args(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--task", choices=TASKS, required=True) + parser.add_argument("--teacher-checkpoint", required=True, help="Trusted local native DDPM checkpoint") + parser.add_argument("--dataset", default=None, help="Explicit dataset path; otherwise use saved teacher configuration") + parser.add_argument("--output", required=True, help="Fresh output directory; existing paths are never overwritten") + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--epochs", type=int, default=20) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--episodes", type=int, default=50) + parser.add_argument("--seed-base", type=int, default=100000, + help="Diagnostic episode seed bank; repeat_eval provides fresh test banks") + parser.add_argument("--policy-seed-offset", type=int, default=None, + help="Default: 0 for Relay-Kitchen, student training seed for other tasks") + parser.add_argument("--workers", type=int, default=0) + parser.add_argument("--t-init", type=int, default=None) + parser.add_argument("--score-min", type=int, default=None, help="Inclusive native score timestep") + parser.add_argument("--score-max", type=int, default=None, help="Inclusive native score timestep; exclude terminal") + parser.add_argument("--smoke", action="store_true", help="Two updates and truncated evaluation; not performance") + parser.add_argument("--smoke-batch-size", type=int, default=1, + help="Use the formal batch size to audit GPU memory before a long run") + return parser.parse_args(argv) + + +if __name__ == "__main__": + run(parse_args()) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/teacher.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/teacher.py new file mode 100644 index 0000000000000000000000000000000000000000..1053f0f4debc1f3ba334a14d0701a8a828f9c5c6 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/onedp_s/teacher.py @@ -0,0 +1,279 @@ +"""Explicit, trusted local CleanDiffuser teachers for the portable OneDP-S package. + +Checkpoint pickle files are code-bearing: only pass your own trusted training +artifacts. No checkpoint is downloaded and no class is selected by its metadata. +""" +from dataclasses import dataclass +import hashlib +from pathlib import Path +from typing import Any + +import numpy as np +import torch + + +ROBO_KEYS = ("object", "robot0_eef_pos", "robot0_eef_quat", "robot0_gripper_qpos") +TASKS = { + "relay_kitchen": dict(obs_dim=60, action_dim=9, steps=50, env="kitchen-all-v0", + max_steps=280, obs_keys=("state",)), + "pusht": dict(obs_dim=5, action_dim=2, steps=5, env="pusht-v0", max_steps=300, + obs_keys=("state", "action")), + "pusht_keypoints": dict(obs_dim=20, action_dim=2, steps=5, env="pusht-keypoints-v0", + max_steps=300, obs_keys=("keypoint", "state", "action")), + "pusht_image": dict(obs_dim=None, action_dim=2, steps=5, env="pusht-image-v0", + max_steps=300, obs_keys=("img", "state", "action")), + "lift_lowdim": dict(obs_dim=19, action_dim=10, steps=50, env="lift", max_steps=400, + obs_keys=ROBO_KEYS), + "square_lowdim": dict(obs_dim=23, action_dim=10, steps=50, env="square", max_steps=500, + obs_keys=ROBO_KEYS), +} + + +@dataclass +class TeacherBundle: + agent: Any + config: Any + metadata: dict + normalizer: dict + checkpoint_path: Path + checkpoint_sha256: str + task: str + + +def task_spec(task): + if task not in TASKS: + raise ValueError(f"Unsupported task {task!r}; choose from {tuple(TASKS)}") + return TASKS[task] + + +def sha256(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def restore_normalizers(task, payload): + """Restore only known saved statistics; never recompute on a student split.""" + from cleandiffuser.dataset.dataset_utils import MinMaxNormalizer, ImageNormalizer + spec = task_spec(task) + expected = ({"keypoint": 2, "agent_pos": 2} if task == "pusht_keypoints" else + {"image": None, "agent_pos": 2} if task == "pusht_image" else + {"state": spec["obs_dim"]}) + if set(payload) != {"obs", "action"} or set(payload["obs"]) != set(expected): + raise ValueError("Saved normalizer fields do not match the requested task") + + def restore(value, dimension): + prefix = "cleandiffuser.dataset.dataset_utils." + if task == "relay_kitchen": + # Kitchen's utility normalizer masks constant coordinates. Using + # dataset_utils.MinMaxNormalizer would turn its zero padding into -1. + from cleandiffuser.utils.normalizers import MinMaxNormalizer as KitchenMinMax + if (value.get("class") != "cleandiffuser.utils.normalizers.MinMaxNormalizer" + or set(value.get("state", {})) != {"min", "max", "range", "mask"}): + raise ValueError("Require the saved masked Kitchen MinMaxNormalizer") + arrays = {key: np.asarray(array, dtype=np.float32) for key, array in value["state"].items()} + if any(a.shape != (dimension,) or not np.isfinite(a).all() for a in arrays.values()): + raise ValueError("Invalid Kitchen normalization shapes or values") + span = arrays["max"] - arrays["min"] + if (np.any(span < 0) or not np.array_equal(arrays["mask"], (span != 0).astype(np.float32)) + or not np.allclose(arrays["range"], np.where(span != 0, span, 1.))): + raise ValueError("Inconsistent Kitchen mask/range statistics") + normalizer = KitchenMinMax.__new__(KitchenMinMax) + for key, array in arrays.items(): + setattr(normalizer, key, array) + return normalizer + if dimension is None: + if value.get("class") != prefix + "ImageNormalizer" or value.get("state") != {}: + raise ValueError("Expected stateless ImageNormalizer for RGB input") + return ImageNormalizer() + if value.get("class") != prefix + "MinMaxNormalizer": + raise ValueError("Only saved MinMaxNormalizer statistics are supported") + if set(value["state"]) != {"min", "max", "range"}: + raise ValueError("Require min, max and range statistics") + arrays = {key: np.asarray(array, dtype=np.float32) for key, array in value["state"].items()} + if any(array.shape != (dimension,) or not np.isfinite(array).all() for array in arrays.values()): + raise ValueError(f"Expected finite {dimension}D normalization statistics") + if np.any(arrays["range"] <= 0) or np.any(arrays["max"] < arrays["min"]): + raise ValueError("Invalid min/max or nonpositive normalization range") + normalizer = MinMaxNormalizer.__new__(MinMaxNormalizer) + for key, array in arrays.items(): + setattr(normalizer, key, array) + return normalizer + + return {"obs": {key: restore(payload["obs"][key], dimension) for key, dimension in expected.items()}, + "action": restore(payload["action"], spec["action_dim"])} + + +def validate_checkpoint(task, payload): + from omegaconf import OmegaConf + spec = task_spec(task) + metadata = payload["metadata"] + seed = metadata.get("seed") + if (metadata.get("task") != task or type(seed) is not int or seed < 0 or metadata.get("smoke") is not False + or payload.get("step") != 1_000_000 or "model_ema" not in payload or "model" not in payload): + raise ValueError("Require the requested final 1M-update nonsmoke EMA teacher with an explicit seed") + config = OmegaConf.create(metadata["config"]) + if type(config.seed) is not int or config.seed != seed: + raise ValueError("Checkpoint config and metadata training seeds differ") + if (config.diffusion != "ddpm" or config.nn != "chi_unet" or config.sample_steps != spec["steps"] + or config.obs_steps != 2 or config.horizon != 16 or config.action_steps != 8 + or config.action_dim != spec["action_dim"] or config.env_name != spec["env"] + or config.max_episode_steps != spec["max_steps"] or float(config.get("temperature", 1.)) != 1.): + raise ValueError("Teacher config does not match the native DDPM task protocol") + if task == "relay_kitchen": + if config.abs_action is not True: + raise ValueError("Relay-Kitchen requires the native absolute-joint controller") + elif task.endswith("_lowdim"): + if tuple(config.obs_keys) != ROBO_KEYS or config.abs_action is not True or config.env_type != "ph": + raise ValueError("Require PH low-dimensional data, canonical observation order and absolute actions") + elif set(config.obs_keys) != set(spec["obs_keys"]) or config.abs_action is not False: + raise ValueError("Incorrect PushT observation keys or action convention") + if spec["obs_dim"] is not None and config.obs_dim != spec["obs_dim"]: + raise ValueError("Incorrect observation dimension") + if task == "pusht_image": + obs = config.shape_meta.obs + if (set(obs) != {"image", "agent_pos"} or list(obs.image.shape) != [3, 96, 96] + or obs.image.type != "rgb" or list(obs.agent_pos.shape) != [2] + or config.rgb_model != "resnet18" or list(config.crop_shape) != [84, 84] + or config.resize_shape is not None or not config.random_crop + or not config.use_group_norm or not config.use_seq): + raise ValueError("Incorrect native PushT image architecture or observation channels") + return config + + +def build_agent(task, config): + """Explicit native architecture; no Hydra targets or experimental imports.""" + from cleandiffuser.nn_diffusion import ChiUNet1d + from cleandiffuser.nn_condition import IdentityCondition, MultiImageObsCondition + from cleandiffuser.diffusion.ddpm import DDPM + task_spec(task) + image = task == "pusht_image" + condition = (MultiImageObsCondition(shape_meta=config.shape_meta, emb_dim=256, + rgb_model_name=config.rgb_model, resize_shape=config.resize_shape, + crop_shape=config.crop_shape, random_crop=config.random_crop, + use_group_norm=config.use_group_norm, use_seq=config.use_seq) + if image else IdentityCondition(dropout=0.)) + diffusion = ChiUNet1d(config.action_dim, 256 if image else config.obs_dim, + config.obs_steps, model_dim=256, emb_dim=256, dim_mult=[1, 2, 2], + obs_as_global_cond=True, timestep_emb_type="positional") + maximum = torch.ones((1, config.horizon, config.action_dim), device=config.device) + return DDPM(diffusion, condition, device=config.device, diffusion_steps=config.sample_steps, + x_max=maximum, x_min=-maximum, optim_params={"lr": config.lr}) + + +def load_teacher(task, checkpoint, device="cpu"): + """Load a caller-selected trusted local checkpoint, preserving initialization RNG.""" + path = Path(checkpoint).expanduser().resolve(strict=True) + digest = sha256(path) + payload = torch.load(path, map_location="cpu", weights_only=False, mmap=True) + config = validate_checkpoint(task, payload) + normalizer = restore_normalizers(task, payload["metadata"]["normalizer"]) + config.device = "cpu" + with torch.random.fork_rng(devices=[]): + agent = build_agent(task, config) + agent.model.load_state_dict(payload["model"], strict=True) + agent.model_ema.load_state_dict(payload["model_ema"], strict=True) + agent.device = config.device = str(device) + agent.model.to(device).requires_grad_(False).eval() + agent.model_ema.to(device).requires_grad_(False).eval() + for name in ("alpha", "beta", "bar_alpha", "x_min", "x_max", "fix_mask", "loss_weight"): + value = getattr(agent, name) + if isinstance(value, torch.Tensor): + setattr(agent, name, value.to(device)) + if task != "pusht_image" and any(True for _ in agent.model_ema["condition"].parameters()): + raise ValueError("Low-dimensional teachers must use parameter-free IdentityCondition") + return TeacherBundle(agent, config, payload["metadata"], normalizer, path, digest, task) + + +def normalize_observations(bundle, observations): + """Raw environment observations only; dataset observations are already normalized.""" + spec = task_spec(bundle.task) + normalizer = bundle.normalizer["obs"] + if bundle.task == "pusht_image": + raw = ({key: np.stack([obs[key] for obs in observations]) for key in observations[0]} + if not isinstance(observations, dict) else observations) + if set(raw) != {"image", "agent_pos"}: + raise ValueError("Image policy receives only RGB and agent position") + expected = {"image": (2, 3, 96, 96), "agent_pos": (2, 2)} + result = {} + for key, value in raw.items(): + array = np.asarray(value) + if array.ndim != len(expected[key]) + 1 or tuple(array.shape[1:]) != expected[key]: + raise ValueError(f"Incorrect raw {key} shape") + result[key] = torch.as_tensor(normalizer[key].normalize(array), device=bundle.agent.device, + dtype=torch.float32) + if len({value.shape[0] for value in result.values()}) != 1: + raise ValueError("Inconsistent image observation batch sizes") + if any(not torch.isfinite(value).all() for value in result.values()): + raise ValueError("Nonfinite image observations") + return result + raw = np.asarray(observations, dtype=np.float32) + if raw.ndim != 3 or tuple(raw.shape[1:]) != (2, spec["obs_dim"]) or not len(raw): + raise ValueError(f"Expected raw observations [B,2,{spec['obs_dim']}]") + if not np.isfinite(raw).all(): + raise ValueError("Nonfinite observations") + if bundle.task == "pusht_keypoints": + keypoints = normalizer["keypoint"].normalize(raw[:, :, :18].reshape(-1, 2)).reshape(len(raw), 2, 18) + normalized = np.concatenate((keypoints, normalizer["agent_pos"].normalize(raw[:, :, 18:])), axis=-1) + else: + normalized = normalizer["state"].normalize(raw) + return torch.as_tensor(normalized, device=bundle.agent.device, dtype=torch.float32).flatten(start_dim=1) + + +def _fresh(condition): + return dict(condition) if isinstance(condition, dict) else condition + + +@torch.no_grad() +def teacher_parity(bundle, observations, seed=73913): + """Bitwise native/manual DDPM5 or DDPM50 check, including sampling RNG state.""" + agent = bundle.agent + steps = task_spec(bundle.task)["steps"] + if (agent.diffusion_steps != steps or not agent.predict_noise or agent.classifier is not None + or torch.as_tensor(agent.fix_mask).count_nonzero().item()): + raise ValueError("Require native epsilon DDPM without classifier or fixed coordinates") + device = torch.device(agent.device) + devices = [device.index or 0] if device.type == "cuda" else [] + count = len(next(iter(observations.values()))) if isinstance(observations, dict) else len(observations) + prior = torch.zeros(count, bundle.config.horizon, bundle.config.action_dim, device=device) + modes = agent.model.training, agent.model_ema.training + + def seed_selected(): + torch.random.default_generator.manual_seed(seed) + if devices: + torch.cuda.default_generators[devices[0]].manual_seed(seed) + + def states(): + return [torch.get_rng_state().clone()] + ([torch.cuda.get_rng_state(device).clone()] if devices else []) + + try: + agent.model.eval() + agent.model_ema.eval() + with torch.random.fork_rng(devices=devices): + seed_selected() + expected, _ = agent.sample(prior, n_samples=count, sample_steps=steps, + condition_cfg=_fresh(observations), w_cfg=1., use_ema=True) + expected_rng = states() + seed_selected() + xt = torch.randn_like(prior) + context = agent.model_ema["condition"](_fresh(observations)) + for t in range(steps - 1, -1, -1): + t_batch = torch.full((count,), t, device=device, dtype=torch.long) + cumulative = agent.bar_alpha[t] + previous = agent.bar_alpha[t - 1] if t else torch.tensor(1., device=device) + prediction, _ = agent.predict_function(xt, t_batch, cumulative, use_ema=True, + requires_grad=False, condition_vec_cfg=context, w_cfg=1., w_cg=0.) + xt = 1 / agent.alpha[t].sqrt() * (xt - agent.beta[t] / (1 - cumulative).sqrt() * prediction) + if t: + xt = xt + (agent.beta[t] * (1 - previous) / (1 - cumulative)).sqrt() * torch.randn_like(xt) + error = float((expected - xt).abs().max()) + rng_equal = all(torch.equal(left, right) for left, right in zip(expected_rng, states())) + finally: + agent.model.train(modes[0]) + agent.model_ema.train(modes[1]) + if not np.isfinite(error) or error > 1e-5 or not rng_equal: + raise RuntimeError(f"Native teacher parity failed: error={error}, RNG={rng_equal}") + return dict(task=bundle.task, max_abs_error=error, rng_equal=rng_equal, teacher_steps=steps, + sampler="ddpm", weights="EMA", action_shape=list(prior.shape)) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/relay_baselines/cp.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/relay_baselines/cp.py new file mode 100644 index 0000000000000000000000000000000000000000..7a006c6c174e44e920c760f686a5bf96f7762622 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/relay_baselines/cp.py @@ -0,0 +1,226 @@ +"""Relay-Kitchen official EDM/CP preparation and matched repeat evaluation. + +Run through experiments/consistency_pusht/run.sh. Training delegates to the +unchanged official wrapper; all three CP seeds share one seed-0 EDM teacher. +Fresh test banks are disjoint from FDK's earlier checkpoint-selection bank. +""" +import argparse +from collections import Counter +import json +from pathlib import Path +import statistics +import subprocess +import sys + +from omegaconf import OmegaConf + +from experiments.consistency_manipulation.prepare import prepare as official_prepare +from experiments.table2.common import atomic_json +from tools.consistency_policy.evaluate import ( + DEFAULT_REPO, checkpoint_hash, load_policy, set_sampling, verify_source, +) + +TASK = "relay_kitchen" +BASES = (400000, 500000, 600000) + + +def run_directory(output, seed): + return Path(output).resolve() / f"seed{seed}" / TASK / f"seed{seed}" + + +def duplicate_manifest_audit(path): + members = json.loads(Path(path).read_text())["members"] + logs = [entry for entry in members if entry["path"].endswith(".mjl")] + counts = Counter(entry["sha256"] for entry in logs) + return dict(archive_mjl_entries=len(logs), unique_mjl_hashes=len(counts), + duplicate_extra_entries=sum(count - 1 for count in counts.values()), + note="Archive-level audit; actual parse successes are recorded by the dataset adapter.") + + +def prepare(output, seeds=(0, 1, 2), smoke=False, device="cuda:0"): + output = Path(output).resolve() + seeds = tuple(seeds) + if not seeds or 0 not in seeds or len(set(seeds)) != len(seeds) or min(seeds) < 0: + raise ValueError("Distinct nonnegative seeds including teacher seed 0 are required") + if output.exists(): + raise FileExistsError(f"Use a fresh CP preparation directory: {output}") + teacher = run_directory(output, 0) / "edm/checkpoints/final.ckpt" + paths = {} + for seed in seeds: + official_prepare(output / f"seed{seed}", (TASK,), seed=seed, smoke=smoke) + directory = run_directory(output, seed) + manifest_path = directory / "manifest.json" + manifest = json.loads(manifest_path.read_text()) + for method in ("edm", "ctmp"): + path = directory / "configs" / f"{method}.yaml" + cfg = OmegaConf.load(path) + cfg.training.device = device + if method == "ctmp": + cfg.policy.teacher_path = cfg.policy.edm = str(teacher) + OmegaConf.save(cfg, path, resolve=True) + manifest["configs"][method]["sha256"] = checkpoint_hash(path) + manifest.update(shared_edm_teacher=str(teacher), shared_teacher_training_seed=0, + train_edm_for_this_seed=seed == 0, + seed_semantics="CP optimization/dataloader/dropout seed conditional on fixed EDM teacher") + atomic_json(manifest_path, manifest) + paths[str(seed)] = {key: str(directory / "configs" / f"{key}.yaml") + for key in ("edm", "ctmp")} + manifest = json.loads((run_directory(output, 0) / "manifest.json").read_text()) + report = dict(task=TASK, seeds=list(seeds), smoke=smoke, shared_teacher=str(teacher), + shared_teacher_training_seed=0, configs=paths, + teacher_budget_epochs=1 if smoke else 400, + student_budget_epochs=1 if smoke else 500, + data_protocol={ + "normalization": "Table-2 Kitchen exact affine, 30 masked goal channels stay zero", + "observation_history": 2, "observation_dim": 60, + "action_horizon": 16, "action_dim": 9, "execution_slice": [1, 9], + "train_split": "All parsed demonstrations, including archive duplicates", + "validation": "In-sample diagnostic only; final-epoch checkpoint", + "duplicate_audit": duplicate_manifest_audit(manifest["dataset_manifest"]), + "fairness_caveat": "CP retains official CTM+DSM supervision on demonstration actions. " + "FDK and OneDP use teacher-supervised student-context splits. " + "Matched environment/evaluation does not imply identical student data objectives/splits.", + }, + evaluation={"seed_bases": list(BASES), "episodes_per_group": 50, + "policy_rng": "group seed base", "n_envs": 10, + "max_episode_steps": 280, "num_candidates": 1, + "cp_nfe": [1, 3], "edm_bins": 80, "edm_nfe": 158}, + state="prepared", paper_ready=False) + atomic_json(output / "preparation.json", report) + return report + + +def validate_result(row, *, checkpoint_sha256, method, nfe, base, seed, smoke): + expected = dict(task=TASK, checkpoint_sha256=checkpoint_sha256, method=method, + weights="model", seed_base=base, eval_seed=base, training_seed=seed, + num_episodes=2 if smoke else 50, num_candidates=1, + horizon=16, obs_steps=2, action_steps=8, + max_episode_steps=8 if smoke else 280, truncated_smoke=smoke) + if any(row.get(key) != value for key, value in expected.items()): + raise ValueError("Existing/new CP evaluation has mismatched provenance or protocol") + if row.get("nfe", {}).get("min") != nfe or row.get("nfe", {}).get("max") != nfe: + raise ValueError("Measured denoiser NFE differs from official sampling") + episodes = row.get("episodes", []) + if [episode["episode_seed"] for episode in episodes] != list(range(base, base + expected["num_episodes"])): + raise ValueError("Evaluation episode seeds are incomplete or different") + scores = [episode["score"] for episode in episodes] + if any(not 0 <= score <= 1 for score in scores): + raise ValueError("Invalid Kitchen score") + if abs(row["score"] - statistics.mean(scores)) > 1e-10: + raise ValueError("Reported score does not match episode scores") + + +def evaluate_checkpoint(checkpoint, output, method, device="cuda:0", seed_bases=BASES, smoke=False): + from experiments.consistency_manipulation.evaluate import evaluate, validate_policy + from experiments.consistency_pusht.evaluate import runtime_provenance + + if method not in ("edm", "cp"): + raise ValueError("Expected edm or cp") + bases = tuple(seed_bases) + count = 2 if smoke else 50 + if len(bases) != 3 or len(set(bases)) != 3 or min(bases) < 0: + raise ValueError("Exactly three distinct nonnegative evaluation seed bases are required") + if any(abs(left - right) < count for i, left in enumerate(bases) for right in bases[i + 1:]): + raise ValueError("Evaluation seed banks must not overlap") + source = verify_source(DEFAULT_REPO) + if source["repo"] not in sys.path: + sys.path.insert(0, source["repo"]) + if not OmegaConf.has_resolver("eval"): + OmegaConf.register_new_resolver("eval", eval) + checkpoint = Path(checkpoint).resolve(strict=True) + digest = checkpoint_hash(checkpoint) + policy, cfg, weights = load_policy(checkpoint, method, "model", device) + validate_policy(policy, TASK, cfg.policy.shape_meta) + if cfg.experiment.task != TASK or bool(cfg.experiment.smoke) != smoke or weights != "model": + raise ValueError("Checkpoint task, smoke flag or raw weights mismatch") + seed = int(cfg.training.seed) + output = Path(output).resolve() + output.mkdir(parents=True, exist_ok=True) + runtime = runtime_provenance(device) + summaries = {} + for steps in ((80,) if method == "edm" else (1, 3)): + times = set_sampling(policy, method, steps, "official") + nfe = 2 * steps - 2 if method == "edm" else steps + rows, files = [], [] + for base in bases: + path = output / f"{method}_{steps}_group{base}.json" + if path.exists(): + row = json.loads(path.read_text()) + else: + row = evaluate(policy, TASK, episodes=count, envs=2 if smoke else 10, + seed_base=base, eval_seed=base, max_steps=8 if smoke else None, warmup=3) + row.update(schema_version=1, method=method, weights=weights, + checkpoint=str(checkpoint), checkpoint_sha256=digest, + training_seed=seed, source=source, runtime=runtime, + shared_teacher=str(cfg.policy.teacher_path) if method == "cp" else None, + sampling={"steps": steps, "nfe": nfe, "chaining_times": times, + "chain_mode": "official", "solver": policy.noise_scheduler.solver}, + training=OmegaConf.to_container(cfg.training, resolve=False), + seed_semantics="Fixed checkpoint; episode and policy RNG both use group base") + validate_result(row, checkpoint_sha256=digest, method=method, nfe=nfe, + base=base, seed=seed, smoke=smoke) + atomic_json(path, row) + validate_result(row, checkpoint_sha256=digest, method=method, nfe=nfe, + base=base, seed=seed, smoke=smoke) + rows.append(row) + files.append(str(path)) + scores = [row["score"] for row in rows] + summaries[str(nfe)] = dict(nfe=nfe, group_scores=scores, mean=statistics.mean(scores), + evaluation_group_std=statistics.stdev(scores), + std_ddof=1, worst_group_mean=min(scores), + worst_episode=min(ep["score"] for row in rows for ep in row["episodes"]), + total_episodes=3 * count, files=files) + report = dict(task=TASK, method=method, checkpoint=str(checkpoint), checkpoint_sha256=digest, + training_seed=seed, weights=weights, seed_bases=list(bases), smoke=smoke, + state="smoke_complete" if smoke else "complete", by_nfe=summaries, + teacher_variance_included=False, training_seed_std=None, + comparison_note="Three evaluation groups of one fixed checkpoint; not three trainings.", + latency_paper_ready=False, paper_ready=False) + atomic_json(output / "summary.json", report) + print(json.dumps(report, indent=2), flush=True) + return report + + +def smoke(output, device="cuda:0", seed_bases=BASES): + prepare(output, seeds=(0,), smoke=True, device=device) + directory = run_directory(output, 0) + for method in ("edm", "ctmp"): + subprocess.run([sys.executable, "-m", "experiments.consistency_pusht.train", "--config", + str(directory / "configs" / f"{method}.yaml"), "--smoke"], check=True) + eval_method = "edm" if method == "edm" else "cp" + subprocess.run([sys.executable, "-m", "experiments.relay_baselines.cp", "evaluate", + "--checkpoint", str(directory / method / "checkpoints/final.ckpt"), + "--output", str(directory / f"eval_{eval_method}"), "--method", eval_method, + "--device", device, "--seed-bases", ",".join(map(str, seed_bases)), "--smoke"], check=True) + atomic_json(Path(output) / "smoke_summary.json", dict(task=TASK, state="smoke_complete", paper_ready=False)) + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="command", required=True) + prep = sub.add_parser("prepare") + prep.add_argument("--output", type=Path, required=True) + prep.add_argument("--seeds", default="0,1,2") + prep.add_argument("--smoke", action="store_true") + prep.add_argument("--device", default="cuda:0") + for command in ("evaluate", "smoke"): + entry = sub.add_parser(command) + entry.add_argument("--output", type=Path, required=True) + entry.add_argument("--device", default="cuda:0") + entry.add_argument("--seed-bases", default=",".join(map(str, BASES))) + if command == "evaluate": + entry.add_argument("--checkpoint", type=Path, required=True) + entry.add_argument("--method", choices=("cp", "edm"), required=True) + entry.add_argument("--smoke", action="store_true") + args = parser.parse_args(argv) + if args.command == "prepare": + prepare(args.output, tuple(map(int, args.seeds.split(","))), args.smoke, args.device) + elif args.command == "evaluate": + evaluate_checkpoint(args.checkpoint, args.output, args.method, args.device, + tuple(map(int, args.seed_bases.split(","))), args.smoke) + else: + smoke(args.output, args.device, tuple(map(int, args.seed_bases.split(",")))) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/relay_baselines/fdk.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/relay_baselines/fdk.py new file mode 100644 index 0000000000000000000000000000000000000000..e40af20dcabb83dfb1e4fdb8ea14b0710a2c5b03 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/relay_baselines/fdk.py @@ -0,0 +1,268 @@ +"""Frozen-recipe Relay FDK training and independent test-bank evaluation. + +The legacy training selector is deliberately retained, including ten validation +episodes at intermediate checkpoints and fifty at the last checkpoint. All +formal test banks here are disjoint from that selector's 100000--100049 bank. +""" +import argparse +import hashlib +import json +import math +from pathlib import Path +import statistics +import subprocess +import sys +import time + + +ROOT = Path(__file__).resolve().parents[2] +TEACHER = ROOT / "checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/ddpm_teacher/model_1000000.pt" +TEACHER_SHA = "a489c3ae549771cebe14cab0a83a1e480ffd60ffab9d5ba9bd9b5b8174da9686" +PAIRS = ROOT / "outputs/lowdim_fdk_distill/relay_kitchen_filmgate_lat015_u150k_20260916_v1/pairs.pt" +PAIRS_SHA = "8a9f054b7897babe289a00bc46298c0c53d95d91cd56d2015409d6409bd05893" +STUDENT0_SHA = "04873fd942bc52dfa2829ac5db8db65628aaf07e18fbe33f4d6b41b0a8240add" +PLAN0_SHA = "76dd3f80e2a1b488e9aaaa8f3c7c31fc4dc4139afc346c86074cee9f4ea6fe97" +TEST_BASES = (400000, 500000, 600000) +RECIPE = dict(task="relay_kitchen", updates=150000, episodes=50, batch_size=128, + train_pairs=100000, validation_pairs=8192, + validate_every=10000, latent_weight=.15, first_weight=0., + film_reg_weight=.001, full_weight=.25, spectral_weight=.01, + context_mode="obs_film_gate", latent_dim=128, backbone_width=1, + latent_mode="global", layout="tau_prior", select="rollout", + rollout_every=40000, rollout_episodes=10) +LIMITATION = ("Legacy validation selection: 10 episodes at steps 10000, 40000, " + "80000 and 120000; 50 episodes at step 150000; all use bank 100000. " + "Final-test banks are disjoint. Student seeds share one fixed " + "teacher and one fixed distillation dataset.") + + +def digest(path): + h = hashlib.sha256() + with Path(path).open("rb") as stream: + for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +def save_json(path, value): + path = Path(path) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n") + temporary.replace(path) + + +def checked_file(path, expected): + path = Path(path).resolve(strict=True) + if digest(path) != expected: + raise ValueError(f"Pinned artifact hash mismatch: {path}") + return path + + +def validate_banks(bases): + if len(bases) != 3 or any(type(b) is not int or b < 0 for b in bases): + raise ValueError("Require three nonnegative integer test-bank bases") + banks = [set(range(base, base + 50)) for base in bases] + if len(set.union(*banks)) != 150 or any(b & set(range(100000, 100050)) for b in banks): + raise ValueError("Test banks must be disjoint from each other and selection bank") + return tuple(bases) + + +def training_command(seed, output, pairs, device): + if type(seed) is not int or seed not in (1, 2): + raise ValueError("This wrapper trains additional student seeds 1 or 2") + command = [sys.executable, "-m", "experiments.lowdim_fdk_distill.run"] + for key, value in RECIPE.items(): + command.extend(("--" + key.replace("_", "-"), str(value))) + command.extend(("--seed", str(seed), "--device", device, + "--output", str(output), "--reuse-pairs", str(pairs))) + return command + + +def validate_student(payload, plan, checkpoint_sha, plan_sha): + seed = plan.get("seed") + if type(seed) is not int or seed not in (0, 1, 2): + raise ValueError("Source plan must identify student training seed 0, 1 or 2") + for key, expected in RECIPE.items(): + actual = plan.get(key, "global" if key == "latent_mode" else None) + if actual != expected: + raise ValueError(f"Source training plan differs from frozen recipe: {key}") + if plan.get("teacher_sha256") != TEACHER_SHA or plan.get("smoke") is not False: + raise ValueError("Require the pinned nonsmoke DDPM teacher") + for flag in ("energy_sample", "geo_sample", "advers"): + if plan.get(flag, False): + raise ValueError(f"Additional training objective/sampling is not frozen recipe: {flag}") + for weight in ("xyz_weight", "rot_weight", "grip_weight"): + if plan.get(weight, 1.) != 1.: + raise ValueError(f"Changed action-coordinate weighting: {weight}") + if plan.get("geo_loss_boost", 0.) != 0. or plan.get("late_weight", 0.) != 0.: + raise ValueError("Additional action loss differs from frozen recipe") + if seed == 0 and (checkpoint_sha != STUDENT0_SHA or plan_sha != PLAN0_SHA): + raise ValueError("Training seed0 must be the pinned historical student and plan") + meta = payload.get("metadata", {}) + expected = dict(task="relay_kitchen", use_koopman=True, operator="vdv", obs_dim=60, + act_dim=9, teacher_steps=50, teacher_sha256=TEACHER_SHA, + context_mode=plan["context_mode"], layout=plan["layout"], + latent_dim=plan["latent_dim"], first_weight=0., latent_weight=.15, + film_reg_weight=.001) + for key, value in expected.items(): + if meta.get(key) != value: + raise ValueError(f"Checkpoint metadata differs from source plan: {key}") + if (meta.get("backbone_width", 1) != plan["backbone_width"] + or meta.get("latent_mode", "global") != plan.get("latent_mode", "global") + or ("seed" in meta and meta["seed"] != seed)): + raise ValueError("Checkpoint architecture or training seed disagrees with source plan") + if (type(payload.get("step")) is not int + or payload["step"] not in (10000, 40000, 80000, 120000, 150000) + or payload.get("validation", {}).get("select") != "rollout" + or "state_dict" not in payload): + raise ValueError("Expected validation-selected EMA student") + return seed + + +def summarize(groups): + if len(groups) != 3 or len({g["training_seed"] for g in groups}) != 1: + raise ValueError("Require three test groups for one student checkpoint") + if len({g["checkpoint_sha256"] for g in groups}) != 1: + raise ValueError("Test groups contain different checkpoints") + validate_banks([g["episode_seed_base"] for g in groups]) + scores = [] + for group in groups: + rows, base = group["episodes"], group["episode_seed_base"] + if (len(rows) != 50 or {r["episode_seed"] for r in rows} != set(range(base, base + 50)) + or group["policy_rng_seed"] != base): + raise ValueError("Incorrect episode bank or policy RNG seed") + values = [r["score"] for r in rows] + if any(not math.isfinite(x) or not 0 <= x <= 1 for x in values): + raise ValueError("Nonfinite or out-of-range episode score") + if not math.isclose(statistics.mean(values), group["score"], abs_tol=1e-12): + raise ValueError("Group score disagrees with raw episodes") + scores.extend(values) + means = [g["score"] for g in groups] + return dict(task="relay_kitchen", method="FDK", training_seed=groups[0]["training_seed"], + checkpoint_sha256=groups[0]["checkpoint_sha256"], + group_scores=means, mean=statistics.mean(means), + evaluation_group_std=statistics.stdev(means), evaluation_group_std_ddof=1, + best_group_mean=max(means), worst_group_mean=min(means), + pooled_episode_mean=statistics.mean(scores), pooled_episode_std=statistics.stdev(scores), + pooled_episode_std_ddof=1, worst_episode=min(scores), + lower_10_percent_mean=statistics.mean(sorted(scores)[:15]), + num_groups=3, episodes_per_group=50, total_episodes=150, + num_training_seeds=1, training_seed_std=None, + uncertainty_scope="evaluation variation conditional on one fixed student checkpoint") + + +def train(args): + output = Path(args.output).resolve() + if output.exists(): + raise FileExistsError("Training requires a fresh output directory") + checked_file(TEACHER, TEACHER_SHA) + # The legacy training program resolves this original source path. + checked_file(ROOT / "outputs/table2/chi_unet_20260909_v1/runs/relay_kitchen/seed0/model_1000000.pt", TEACHER_SHA) + pairs = checked_file(args.reuse_pairs, PAIRS_SHA) + command = training_command(args.seed, output, pairs, args.device) + source_hash = digest(ROOT / "experiments/lowdim_fdk_distill/run.py") + started = time.time() + completed = subprocess.run(command, cwd=ROOT, check=False) + if output.is_dir(): + save_json(output / "relay_comparison_provenance.json", dict(command=command, + training_seed=args.seed, teacher_sha256=TEACHER_SHA, pairs_sha256=PAIRS_SHA, + legacy_source_sha256=source_hash, recipe=RECIPE, limitation=LIMITATION, + elapsed_seconds=time.time() - started, exit_code=completed.returncode)) + if completed.returncode: + raise subprocess.CalledProcessError(completed.returncode, command) + + +def evaluate(args): + # Delayed imports let the protocol/provenance tests run without simulators. + import random + from types import SimpleNamespace + import numpy as np + import torch + from experiments.onedp_s.teacher import restore_normalizers + from experiments.table2.common import config_for + from experiments.table2.evaluate import evaluate as rollout + from experiments.lowdim_fdk_distill.run import LowdimFDKStudent, StudentAgent + + bases = validate_banks(args.seed_bases) + checkpoint = Path(args.checkpoint).resolve(strict=True) + source_plan = Path(args.source_plan).resolve(strict=True) + plan = json.loads(source_plan.read_text()) + output = Path(args.output).resolve() + if output.exists(): + raise FileExistsError("Evaluation requires a fresh output directory") + checkpoint_sha, plan_sha = digest(checkpoint), digest(source_plan) + payload = torch.load(checkpoint, map_location="cpu", weights_only=False) + training_seed = validate_student(payload, plan, checkpoint_sha, plan_sha) + if training_seed != 0 and checkpoint != (source_plan.parent / "arm_vdv/best_ema.pt").resolve(): + raise ValueError("New student checkpoint must belong to its source plan directory") + checked_file(plan["reuse_pairs"], PAIRS_SHA) + teacher_path = checked_file(TEACHER, TEACHER_SHA) + teacher = torch.load(teacher_path, map_location="cpu", weights_only=False) + normalizer = restore_normalizers("relay_kitchen", teacher["metadata"]["normalizer"]) + config = config_for("relay_kitchen", 0, args.device) + for key in ("horizon", "obs_steps", "action_steps", "action_dim", "obs_dim", "max_episode_steps", "env_name", "abs_action"): + if config[key] != teacher["metadata"]["config"][key]: + raise ValueError(f"Runtime config differs from pinned teacher: {key}") + config.sample_steps = 1 + meta = payload["metadata"] + model = LowdimFDKStudent(meta["obs_dim"], meta["act_dim"], + use_koopman=meta["use_koopman"], latent_dim=meta["latent_dim"], + teacher_steps=meta["teacher_steps"], context_mode=meta["context_mode"], + layout=meta["layout"], backbone_width=meta.get("backbone_width", 1), + latent_mode=meta.get("latent_mode", "global")).to(args.device) + model.load_state_dict(payload["state_dict"], strict=True) + model.eval().requires_grad_(False) + agent = StudentAgent(model, model, "relay_kitchen") + selected_step = payload["step"] + del teacher, payload + output.mkdir(parents=True, exist_ok=False) + provenance = dict(method="FDK", task="relay_kitchen", training_seed=training_seed, + checkpoint=str(checkpoint), checkpoint_sha256=checkpoint_sha, + checkpoint_step=selected_step, source_plan=str(source_plan), source_plan_sha256=plan_sha, + teacher_sha256=TEACHER_SHA, pairs_sha256=PAIRS_SHA, seed_bases=list(bases), + policy_rng_seeds=list(bases), episodes_per_group=50, nfe=1, weights="EMA", + checkpoint_selection="legacy_rollout_validation", limitation=LIMITATION, + teacher_fail_filter=False, latency_comparison_ready=False) + save_json(output / "plan.json", provenance) + started, groups = time.time(), [] + for i, base in enumerate(bases): + random.seed(base) + np.random.seed(base) + result = rollout("relay_kitchen", config, SimpleNamespace(normalizer=normalizer), agent, + episodes=50, env_count=10, seed_base=base, max_steps=None) + result.update(training_seed=training_seed, group_index=i, episode_seed_base=base, + policy_rng_seed=base, nfe=1, method="FDK", checkpoint_sha256=checkpoint_sha, + latency_comparison_ready=False) + result["latency_scope"] = "normalization, transfers, FDK, action unnormalization; includes first call; batched" + save_json(output / f"group{i}.json", result) + groups.append(result) + summary = summarize(groups) + summary.update(checkpoint_step=provenance["checkpoint_step"], checkpoint=str(checkpoint), + elapsed_seconds=time.time() - started, limitation=LIMITATION) + if digest(checkpoint) != checkpoint_sha: + raise RuntimeError("Checkpoint bytes changed during evaluation") + save_json(output / "summary.json", summary) + save_json(output / "status.json", dict(state="complete", total_episodes=150)) + print(json.dumps(summary, indent=2)) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="mode", required=True) + training = sub.add_parser("train") + training.add_argument("--seed", type=int, choices=(1, 2), required=True) + training.add_argument("--output", required=True) + training.add_argument("--reuse-pairs", default=str(PAIRS)) + training.add_argument("--device", default="cuda:0") + evaluation = sub.add_parser("evaluate") + evaluation.add_argument("--checkpoint", required=True) + evaluation.add_argument("--source-plan", required=True) + evaluation.add_argument("--output", required=True) + evaluation.add_argument("--seed-bases", type=int, nargs=3, default=TEST_BASES) + evaluation.add_argument("--device", default="cuda:0") + args = parser.parse_args() + (train if args.mode == "train" else evaluate)(args) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/relay_baselines/launch.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/relay_baselines/launch.py new file mode 100644 index 0000000000000000000000000000000000000000..b3eef15b640a65e34250123892e88598b1954f3d --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/relay_baselines/launch.py @@ -0,0 +1,319 @@ +"""Durable Relay-Kitchen priority queue with identity-checked DDPM recovery. + +Only the active Robomimic DDPM continuation and its descendants are paused. +SIGSTOP preserves their optimizer state in memory. A separate watchdog resumes +exactly these processes after the Relay queue exits, including on failure. +""" +import argparse +import ctypes +import fcntl +import hashlib +import json +import os +from pathlib import Path +import signal +import subprocess +import sys +import time + +ROOT = Path(__file__).resolve().parents[2] +DDPM = ROOT / "outputs/robomimic_ddpm_continue/robomimic_20260916_v1" +BASE = ROOT / "outputs/relay_baselines" +TEACHER = ROOT / "checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/ddpm_teacher/model_1000000.pt" +FDK = ROOT / "checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/fdk_student/model_step80000_best_ema.pt" +FDK_PLAN = ROOT / "outputs/lowdim_fdk_distill/relay_fw0_rollout_20260917_v1/plan.json" +BASES = [400000, 500000, 600000] + + +def read(path): + return json.loads(Path(path).read_text()) + + +def write(path, payload): + path = Path(path) + tmp = path.with_name(f".{path.name}.{os.getpid()}.tmp") + with tmp.open("w") as stream: + json.dump(payload, stream, indent=2, allow_nan=False) + stream.flush() + os.fsync(stream.fileno()) + tmp.replace(path) + + +def sha(path): + h = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(8 * 1024**2), b""): + h.update(block) + return h.hexdigest() + + +def identity(pid): + p = Path(f"/proc/{int(pid)}") + stat = (p / "stat").read_text().rsplit(") ", 1)[1].split() + return dict(pid=int(pid), state=stat[0], ppid=int(stat[1]), pgid=int(stat[2]), + start_ticks=stat[19], args=[s.decode() for s in (p / "cmdline").read_bytes().split(b"\0") if s]) + + +def same(expected): + actual = identity(expected["pid"]) + if any(actual[k] != expected[k] for k in ("pid", "start_ticks", "args")): + raise RuntimeError("Process identity changed; refusing to signal reused PID") + return actual + + +def signal_process(expected, signum): + # The pinned Python runtime lacks os.pidfd_open; host glibc provides it. + libc = ctypes.CDLL(None, use_errno=True) + open_fd = libc.pidfd_open + open_fd.argtypes, open_fd.restype = (ctypes.c_int, ctypes.c_uint), ctypes.c_int + send = libc.pidfd_send_signal + send.argtypes = (ctypes.c_int, ctypes.c_int, ctypes.c_void_p, ctypes.c_uint) + send.restype = ctypes.c_int + fd = open_fd(expected["pid"], 0) + if fd < 0: + error = ctypes.get_errno() + raise OSError(error, os.strerror(error)) + try: + same(expected) + if send(fd, signum, None, 0) < 0: + error = ctypes.get_errno() + raise OSError(error, os.strerror(error)) + finally: + os.close(fd) + + +def children(pid): + result = [] + for path in Path("/proc").iterdir(): + if path.name.isdigit(): + try: + row = identity(int(path.name)) + if row["ppid"] == pid and row["state"] not in ("Z", "X"): + result.append(row) + except (FileNotFoundError, ProcessLookupError, PermissionError): + pass + return sorted(result, key=lambda row: row["pid"]) + + +def pause_ddpm(directory): + status = read(DDPM / "status.json") + target = identity(status["pid"]) + if (status["state"] != "training" or target["state"] in ("T", "t", "Z", "X") + or "experiments.robomimic_ddpm_continue.launch" not in target["args"] + or str(DDPM) not in target["args"]): + raise RuntimeError("DDPM is not the expected active continuation; re-audit before pause") + receipt = dict(state="pausing", previous_status=status, stopped=[], created_unix=time.time()) + write(directory / "handoff.json", receipt) + + def stop_tree(row): + if row["state"] in ("T", "t", "Z", "X"): + raise RuntimeError("Refuse to take ownership of a previously stopped process") + receipt["stopped"].append(dict(process=row, resumed=False)) + write(directory / "handoff.json", receipt) + signal_process(row, signal.SIGSTOP) + for _ in range(100): + if same(row)["state"] in ("T", "t"): + break + time.sleep(.05) + else: + raise RuntimeError("DDPM did not stop") + for child in children(row["pid"]): + stop_tree(child) + + stop_tree(target) + receipt.update(state="paused_for_relay", paused_unix=time.time(), + note="Optimizer and RNG preserved in RAM/GPU; no checkpoint overwritten") + write(directory / "handoff.json", receipt) + + +def restore_ddpm(directory): + path = directory / "handoff.json" + if not path.exists(): + return + receipt = read(path) + for item in reversed(receipt["stopped"]): + if item["resumed"]: + continue + try: + row = same(item["process"]) + if row["state"] in ("T", "t"): + signal_process(item["process"], signal.SIGCONT) + item["resumed"] = True + except (FileNotFoundError, ProcessLookupError): + item.update(resumed=True, already_exited=True) + write(path, receipt) + receipt.update(state="restored", restored_unix=time.time()) + write(path, receipt) + + +def command(module, *args, cp=False): + wrapper = "consistency_pusht" if cp else "table2" + return ["bash", str(ROOT / f"experiments/{wrapper}/run.sh"), "-m", module, *map(str, args)] + + +def make_jobs(directory): + jobs = [] + def add(name, cmd, artifact): + jobs.append(dict(name=name, command=cmd, artifact=str(artifact))) + # Both actual architectures and simulators must pass before long training. + out = directory / "smoke/onedp" + add("onedp_smoke", command("experiments.onedp_s.run", "--task", "relay_kitchen", + "--teacher-checkpoint", TEACHER, "--output", out, "--smoke", "--smoke-batch-size", 64), out / "summary.json") + out = directory / "smoke/cp" + add("cp_smoke", command("experiments.relay_baselines.cp", "smoke", "--output", out, cp=True), out / "smoke_summary.json") + out = directory / "evaluations/fdk/seed0" + add("fdk_seed0_fresh_test", command("experiments.relay_baselines.fdk", "evaluate", "--checkpoint", FDK, + "--source-plan", FDK_PLAN, "--output", out), out / "summary.json") + # Three students share the same frozen DDPM teacher and context split. + for seed in range(3): + out = directory / f"onedp/seed{seed}" + add(f"onedp_seed{seed}_train", command("experiments.onedp_s.run", "--task", "relay_kitchen", + "--teacher-checkpoint", TEACHER, "--output", out, "--seed", seed, + "--epochs", 20, "--batch-size", 64, "--episodes", 50), out / "summary.json") + ev = directory / f"evaluations/onedp/seed{seed}" + add(f"onedp_seed{seed}_test", command("experiments.onedp_s.repeat_eval", + "--checkpoint", out / "checkpoints/final.pt", "--output", ev), ev / "summary.json") + ddpm_eval = directory / "evaluations/ddpm/seed0" + add("ddpm_fresh_test", command("experiments.onedp_s.repeat_eval", "--teacher-only", + "--checkpoint", TEACHER, "--output", ddpm_eval), ddpm_eval / "summary.json") + # Official EDM-to-CP route: a fixed shared seed0 EDM teacher, three CP seeds. + cp_root = directory / "cp" + edm = cp_root / "seed0/relay_kitchen/seed0/edm" + add("edm_train", command("experiments.consistency_pusht.train", "--config", + edm.parent / "configs/edm.yaml", cp=True), edm / "status.json") + ev = directory / "evaluations/edm/seed0" + add("edm_fresh_test", command("experiments.relay_baselines.cp", "evaluate", "--method", "edm", + "--checkpoint", edm / "checkpoints/final.ckpt", "--output", ev, cp=True), ev / "summary.json") + for seed in range(3): + out = cp_root / f"seed{seed}/relay_kitchen/seed{seed}/ctmp" + add(f"cp_seed{seed}_train", command("experiments.consistency_pusht.train", "--config", + out.parent / "configs/ctmp.yaml", cp=True), out / "status.json") + ev = directory / f"evaluations/cp/seed{seed}" + add(f"cp_seed{seed}_test", command("experiments.relay_baselines.cp", "evaluate", "--method", "cp", + "--checkpoint", out / "checkpoints/final.ckpt", "--output", ev, cp=True), ev / "summary.json") + for seed in (1, 2): + out = directory / f"fdk/seed{seed}" + add(f"fdk_seed{seed}_train", command("experiments.relay_baselines.fdk", "train", + "--seed", seed, "--output", out), out / "summary.json") + ev = directory / f"evaluations/fdk/seed{seed}" + add(f"fdk_seed{seed}_test", command("experiments.relay_baselines.fdk", "evaluate", + "--checkpoint", out / "arm_vdv/best_ema.pt", "--source-plan", out / "plan.json", + "--output", ev), ev / "summary.json") + return jobs + + +def sources(): + folders = ["experiments/relay_baselines", "experiments/onedp_s", "experiments/consistency_manipulation", + "experiments/consistency_pusht", "experiments/table2", "experiments/lowdim_fdk_distill", + "cleandiffuser", "tools/consistency_policy"] + return {str(p.relative_to(ROOT)): sha(p) for folder in folders for p in sorted((ROOT / folder).rglob("*")) + if p.is_file() and p.suffix in (".py", ".sh", ".yaml") and not p.name.startswith("test_")} + + +def verify(plan): + for name, expected in plan["sources"].items(): + if sha(ROOT / name) != expected: + raise RuntimeError("Frozen source changed: " + name) + + +def worker(directory): + plan = read(directory / "plan.json") + state = dict(state="starting", pid=os.getpid(), completed=[], started_unix=time.time()) + child = None + try: + verify(plan) + pause_ddpm(directory) + for index, job in enumerate(plan["jobs"]): + verify(plan) + log = directory / "logs" / f"{index:02d}_{job['name']}.log" + with log.open("x") as stream: + child = subprocess.Popen(job["command"], cwd=ROOT, stdin=subprocess.DEVNULL, + stdout=stream, stderr=subprocess.STDOUT) + state.update(state="running", job=job["name"], position=index, total_jobs=len(plan["jobs"]), + child_pid=child.pid, log=str(log), job_started_unix=time.time()) + write(directory / "status.json", state) + code = child.wait() + if code: + raise RuntimeError(f"{job['name']} exited {code}; see {log}") + result = read(job["artifact"]) + if result.get("state") == "failed": + raise RuntimeError("Child reports failed despite exit0") + state["completed"].append(dict(name=job["name"], artifact=job["artifact"], + sha256=sha(job["artifact"]), finished_unix=time.time())) + state.update(child_pid=None) + write(directory / "status.json", state) + from .report import summarize + summarize(directory) + state.update(state="complete", finished_unix=time.time()) + write(directory / "status.json", state) + except BaseException as error: + state.update(state="failed", error=repr(error), failed_unix=time.time()) + write(directory / "status.json", state) + raise + finally: + # On interruption a GPU child may outlive this worker. The watchdog + # waits for that process group before restoring DDPM. + if child is None or child.poll() is not None: + restore_ddpm(directory) + + +def watchdog(directory): + # This separate CPU process survives unexpected worker exits. + with (BASE / "priority.lock").open("a") as lock: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + child = subprocess.Popen([sys.executable, "-m", __package__ + ".launch", "--output", str(directory), + "--worker"], cwd=ROOT, start_new_session=True) + write(directory / "launcher.json", dict(watchdog_pid=os.getpid(), worker=identity(child.pid))) + child.wait() + # If the worker was killed, allow its already-running GPU child to finish. + while True: + living = [] + for path in Path("/proc").iterdir(): + if path.name.isdigit(): + try: + row = identity(int(path.name)) + if row["pgid"] == child.pid and row["state"] not in ("Z", "X"): + living.append(row["pid"]) + except (FileNotFoundError, ProcessLookupError, PermissionError): + pass + if not living: + break + time.sleep(10) + restore_ddpm(directory) + write(directory / "watchdog.json", dict(state="exited", worker_exit_code=child.returncode, + finished_unix=time.time())) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--worker", action="store_true") + p.add_argument("--watchdog", action="store_true") + args = p.parse_args() + directory = args.output.resolve() + if directory.parent != BASE: + raise ValueError("Use a fresh direct child of outputs/relay_baselines") + if args.worker: + return worker(directory) + if args.watchdog: + return watchdog(directory) + directory.mkdir(parents=True, exist_ok=False) + (directory / "logs").mkdir() + subprocess.run(command("experiments.relay_baselines.cp", "prepare", "--output", directory / "cp", cp=True), + cwd=ROOT, check=True) + plan = dict(task="relay_kitchen", student_training_seeds=[0, 1, 2], test_seed_bases=BASES, + episodes_per_group=50, sources=sources(), teacher_checkpoint=str(TEACHER), + teacher_sha256=sha(TEACHER), historical_fdk_checkpoint=str(FDK), historical_fdk_sha256=sha(FDK), + jobs=make_jobs(directory), created_unix=time.time(), + selection_note="100000 was used for FDK selection; fresh 400000/500000/600000 test banks", + variability_note="Independent student training conditional on fixed teachers; not full pipeline variance") + write(directory / "plan.json", plan) + with (directory / "supervisor.log").open("x") as log: + process = subprocess.Popen([sys.executable, "-m", __package__ + ".launch", "--output", str(directory), + "--watchdog"], cwd=ROOT, stdin=subprocess.DEVNULL, stdout=log, + stderr=subprocess.STDOUT, start_new_session=True) + print(json.dumps(dict(watchdog_pid=process.pid, output=str(directory)), indent=2)) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/relay_baselines/report.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/relay_baselines/report.py new file mode 100644 index 0000000000000000000000000000000000000000..3b48ba8b34ebfe9beee468df42243134da39cf55 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/relay_baselines/report.py @@ -0,0 +1,94 @@ +"""Separate fixed-checkpoint evaluation SD from independent-training SD.""" +import argparse +import json +from pathlib import Path +import statistics + +BASES = (400000, 500000, 600000) + + +def scores_from_groups(groups, seed): + if len(groups) != 3: + raise ValueError("Require all three evaluation groups") + result = [] + for row, base in zip(groups, BASES): + if row["training_seed"] != seed: + raise ValueError("Training seed mismatch") + episodes = row["episodes"] + if len(episodes) != 50 or [e["episode_seed"] for e in episodes] != list(range(base, base + 50)): + raise ValueError("Wrong test bank or episode count") + values = [ep["score"] for ep in episodes] + if any(not 0 <= value <= 1 for value in values): + raise ValueError("Invalid score") + mean = statistics.mean(values) + if abs(mean - row["score"]) > 1e-10: + raise ValueError("Raw episode scores disagree with reported mean") + result.append(mean) + return result + + +def aggregate_training(rows): + """Never treat nine groups as nine independently trained students.""" + ordered = sorted(rows, key=lambda row: row["training_seed"]) + seeds = [row["training_seed"] for row in ordered] + if len(set(seeds)) != len(seeds): + raise ValueError("Duplicate training seeds") + if len({row["checkpoint_sha256"] for row in ordered}) != len(ordered): + raise ValueError("Same checkpoint cannot represent different training seeds") + complete = seeds == [0, 1, 2] + result = dict(complete=complete, completed_training_seeds=seeds, + num_training_seeds=len(seeds), training_seed_std=None) + if complete: + primary = [row["group_scores"][0] for row in ordered] + all_groups = [statistics.mean(row["group_scores"]) for row in ordered] + def stats(values, episodes): + return dict(seed_scores=values, mean=statistics.mean(values), + training_seed_std=statistics.stdev(values), std_ddof=1, + worst_training_seed_score=min(values), episodes_per_training_seed=episodes) + result.update(primary_same_50_episodes=stats(primary, 50), + secondary_150_episodes_per_seed=stats(all_groups, 150), + training_seed_std=statistics.stdev(primary)) + return result + + +def summarize(directory): + directory = Path(directory) + rows = {} + for method in ("fdk", "onedp", "cp", "edm", "ddpm"): + for seed in range(3): + folder = directory / f"evaluations/{method}/seed{seed}" + path = folder / "summary.json" + if not path.exists(): + continue + summary = json.loads(path.read_text()) + if summary.get("smoke", False): + raise ValueError("Smoke is not a performance measurement") + variants = summary.get("by_nfe", {"50" if method == "ddpm" else "1": None}) + for nfe, variant in variants.items(): + files = variant["files"] if variant else [folder / f"group{i}.json" for i in range(3)] + groups = [json.loads(Path(file).read_text()) for file in files] + scores = scores_from_groups(groups, seed) + row = dict(method=method, nfe=int(nfe), training_seed=seed, + checkpoint_sha256=summary["checkpoint_sha256"], group_scores=scores, + mean=statistics.mean(scores), evaluation_group_std=statistics.stdev(scores), + std_ddof=1, worst_group_mean=min(scores), source=str(path)) + rows.setdefault(f"{method}_nfe{nfe}", []).append(row) + report = dict(task="relay_kitchen", test_seed_bases=list(BASES), episodes_per_group=50, + fixed_checkpoint_evaluation_groups=rows, + independent_student_training={key: aggregate_training(value) for key, value in rows.items() + if not key.startswith(("ddpm_", "edm_"))}, + teacher_training_variance_included=False, + historical_seed100000_overlap="excluded; used in FDK checkpoint selection", + latency_paper_ready=False, + note="Baselines use final online weights; FDK uses validation-selected EMA. " + "CP uses its own EDM teacher and official CTM+DSM demonstration objective; " + "FDK/OneDP use a fixed DDPM teacher. Student budgets/data splits differ.") + from .launch import write + write(directory / "comparison.json", report) + return report + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + print(json.dumps(summarize(parser.parse_args().output), indent=2)) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/common.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/common.py new file mode 100644 index 0000000000000000000000000000000000000000..a19c323e132502bdd3cb318ec77a7afe5f9b3a1d --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/common.py @@ -0,0 +1,145 @@ +"""Table-2 task definitions; reuse this checkout's Chi_UNet and datasets.""" +import json +import random +from pathlib import Path + +import numpy as np +import torch +from omegaconf import OmegaConf + +ROOT = Path(__file__).resolve().parents[2] +TASKS = { + "relay_kitchen": ("kitchen", "kitchen_abs", False), + "pusht": ("pusht", "pusht", False), + "pusht_keypoints": ("pusht", "pusht_keypoint", False), + "pusht_image": ("pusht", "pusht_image", True), + "lift_lowdim": ("robomimic", "lift_abs", False), + "square_lowdim": ("robomimic", "square_abs", False), + "lift_image": ("robomimic", "lift_abs", True), + "square_image": ("robomimic", "square_abs", True), +} + + +def atomic_json(path, payload): + path = Path(path) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2, allow_nan=False) + "\n") + temporary.replace(path) + + +def config_for(task, seed=0, device="cuda:0"): + family, name, image = TASKS[task] + config = OmegaConf.load(ROOT / f"configs/dp/{family}/chi_unet/{name}.yaml") + if family == "robomimic" and image: + # Image pipeline's shape metadata is supplied by the multimodal config. + config = OmegaConf.load(ROOT / f"configs/dp/robomimic_multi_modal/chi_unet/{name}.yaml") + config.seed, config.device = seed, device + config.save_video = False + config.dataset_path = str((ROOT / config.dataset_path).resolve()) + # Logging paths are explicit; Hydra's date resolver isn't used here. + for key in ("hydra", "work_dir"): + if key in config: + del config[key] + return config + + +def seed_all(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + +def build_dataset(task, config): + family, _, image = TASKS[task] + common = dict(horizon=config.horizon, pad_before=config.obs_steps - 1, + pad_after=config.action_steps - 1) + if family == "kitchen": + from unittest.mock import patch + from cleandiffuser.dataset import kitchen_dataset + parser = kitchen_dataset.parse_mjl_logs + audit = [] + def audited_parser(path, **kwargs): + record = {"path": str(Path(path).relative_to(ROOT))} + try: + value = parser(path, **kwargs) + record["status"] = "parsed" + return value + except Exception as error: + record.update(status="skipped", error=repr(error)) + raise + finally: + audit.append(record) + # Preserve the author's skip-on-parse-error behavior, but make it visible. + with patch.object(kitchen_dataset, "parse_mjl_logs", audited_parser): + dataset = kitchen_dataset.KitchenMjlDataset(config.dataset_path, **common) + dataset.table2_parse_audit = audit + if sum(r["status"] == "parsed" for r in audit) != dataset.replay_buffer.n_episodes: + raise RuntimeError("Kitchen preprocessing dropped a successfully parsed demonstration") + elif family == "pusht": + from cleandiffuser.dataset.pusht_dataset import PushTStateDataset, PushTKeypointDataset, PushTImageDataset + cls = PushTImageDataset if image else (PushTKeypointDataset if task == "pusht_keypoints" else PushTStateDataset) + dataset = cls(config.dataset_path, obs_keys=list(config.obs_keys), abs_action=config.abs_action, **common) + else: + from cleandiffuser.dataset.robomimic_dataset import RobomimicDataset, RobomimicImageDataset + if image: + dataset = RobomimicImageDataset(config.dataset_path, shape_meta=config.shape_meta, + n_obs_steps=config.obs_steps, abs_action=config.abs_action, **common) + else: + dataset = RobomimicDataset(config.dataset_path, obs_keys=list(config.obs_keys), + abs_action=config.abs_action, **common) + if len(dataset) == 0: + raise ValueError("Empty demonstration dataset") + return dataset + + +def build_agent(task, config): + from cleandiffuser.nn_diffusion import ChiUNet1d + from cleandiffuser.nn_condition import IdentityCondition, MultiImageObsCondition + from cleandiffuser.diffusion.ddpm import DDPM + image = TASKS[task][2] + if image: + condition = MultiImageObsCondition(shape_meta=config.shape_meta, emb_dim=256, + rgb_model_name=config.rgb_model, resize_shape=config.resize_shape, + crop_shape=config.crop_shape, random_crop=config.random_crop, + use_group_norm=config.use_group_norm, use_seq=config.use_seq) + else: + condition = IdentityCondition(dropout=0.0) + diffusion = ChiUNet1d(config.action_dim, 256 if image else config.obs_dim, + config.obs_steps, model_dim=256, emb_dim=256, dim_mult=[1, 2, 2], + obs_as_global_cond=True, timestep_emb_type="positional") + # Same defaults as the local dp_* pipelines: cosine noise schedule, predict + # noise, EMA .995, AdamW(lr=1e-4), per-task 5/50 training diffusion steps. + maximum = torch.ones((1, config.horizon, config.action_dim), device=config.device) + return DDPM(diffusion, condition, device=config.device, diffusion_steps=config.sample_steps, + x_max=maximum, x_min=-maximum, optim_params={"lr": config.lr}) + + +def training_condition(task, batch, config): + if TASKS[task][2]: + return {k: v[:, :config.obs_steps].to(config.device).float() for k, v in batch["obs"].items()} + if task == "pusht_keypoints": + state = torch.cat([batch["obs"]["keypoint"], batch["obs"]["agent_pos"]], dim=-1) + return state[:, :config.obs_steps].to(config.device).float().flatten(start_dim=1) + return batch["obs"]["state"][:, :config.obs_steps].to(config.device).float().flatten(start_dim=1) + + +def eval_checkpoints(total_steps, image): + if total_steps != 1000000: + return [total_steps] + return [800000, 900000, 1000000] if image else list(range(100000, 1000001, 100000)) + + +def metric_from_episode(task, rewards, completed_count=0): + if task == "relay_kitchen": + return min(completed_count, 4) / 4.0 + return float(max(rewards, default=0.0)) if task.startswith("pusht") else float(any(r > 0 for r in rewards)) + + +def dataset_manifest(task): + family, _, image = TASKS[task] + name = ("kitchen" if family == "kitchen" else "pusht" if family == "pusht" + else "robomimic_image" if image else "robomimic_lowdim") + path = ROOT / f"outputs/table2/downloads/{name}_manifest.json" + if not path.is_file(): + raise FileNotFoundError(f"Dataset has not passed archive verification: {path}") + return path diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/evaluate.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/evaluate.py new file mode 100644 index 0000000000000000000000000000000000000000..d672144392a3ef2b78123bf4903f7f5680e25a5d --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/evaluate.py @@ -0,0 +1,146 @@ +"""Seeded one-candidate DP rollouts without vector-autoreset contamination.""" +import math +import random +import time + +import numpy as np +import torch + +from experiments.table2.common import TASKS, metric_from_episode + + +def make_env(task, config): + # Legacy adept_envs / robosuite imports on Python 3.10. + import collections + import collections.abc + for name in ("Mapping", "Iterable"): + if not hasattr(collections, name): + setattr(collections, name, getattr(collections.abc, name)) + import gym + from cleandiffuser.env.wrapper import MultiStepWrapper + family, _, image = TASKS[task] + if family == "kitchen": + from cleandiffuser.env import kitchen + from cleandiffuser.env.kitchen.kitchen_lowdim_wrapper import KitchenLowdimWrapper + env = KitchenLowdimWrapper(gym.make(config.env_name, use_abs_action=True), + init_qpos=None, init_qvel=None, render_hw=(240, 360)) + elif family == "pusht": + from cleandiffuser.env import pusht + # Keep this checkout's default legacy=False. Record it explicitly. + env = gym.make(config.env_name, legacy=False) + else: + import robomimic.utils.file_utils as FileUtils + import robomimic.utils.env_utils as EnvUtils + import robomimic.utils.obs_utils as ObsUtils + meta = FileUtils.get_env_metadata_from_dataset(config.dataset_path) + if meta["env_kwargs"].get("reward_shaping", False): + raise ValueError("Robomimic success metric requires sparse, unshaped rewards") + if image: + mapping = {"low_dim": [], "rgb": []} + for key, spec in config.shape_meta.obs.items(): + mapping[spec.get("type", "low_dim")].append(key) + meta["env_kwargs"]["use_object_obs"] = False + else: + mapping = {"low_dim": list(config.obs_keys)} + ObsUtils.initialize_obs_modality_mapping_from_dict(mapping) + meta["env_kwargs"]["controller_configs"]["control_delta"] = False + env = EnvUtils.create_env_from_metadata(meta, render=False, render_offscreen=image, + use_image_obs=image) + env.env.hard_reset = False + if image: + from cleandiffuser.env.robomimic.robomimic_image_wrapper import RobomimicImageWrapper + env = RobomimicImageWrapper(env, config.shape_meta, render_obs_key=config.render_obs_key) + else: + from cleandiffuser.env.robomimic.robomimic_lowdim_wrapper import RobomimicLowdimWrapper + env = RobomimicLowdimWrapper(env, obs_keys=list(config.obs_keys)) + return MultiStepWrapper(env, config.obs_steps, config.action_steps, config.max_episode_steps) + + +def condition_from_obs(task, observations, dataset, config): + normalizer = dataset.normalizer["obs"] + if TASKS[task][2]: + return {k: torch.as_tensor(normalizer[k].normalize(np.stack([o[k] for o in observations])), + device=config.device, dtype=torch.float32) for k in observations[0]} + obs = np.stack(observations).astype(np.float32) + if task == "pusht_keypoints": + batch, history = obs.shape[:2] + keypoints = normalizer["keypoint"].normalize(obs[:, :, :18].reshape(-1, 2)).reshape(batch, history, 18) + agent = normalizer["agent_pos"].normalize(obs[:, :, 18:20]) + normalized = np.concatenate([keypoints, agent], axis=-1) + else: + normalized = normalizer["state"].normalize(obs) + return torch.as_tensor(normalized, device=config.device, dtype=torch.float32).flatten(start_dim=1) + + +def evaluate(task, config, dataset, agent, episodes=50, env_count=10, seed_base=100000, max_steps=None): + """RNG/mode preserving evaluation; seeds reused across checkpoints/methods.""" + device = torch.device(config.device) + devices = [device.index or 0] if device.type == "cuda" else [] + numpy_state, python_state = np.random.get_state(), random.getstate() + model_training, ema_training = agent.model.training, agent.model_ema.training + envs, rows, latencies = [], [], [] + try: + with torch.random.fork_rng(devices=devices), torch.no_grad(): + torch.manual_seed(seed_base + int(config.seed)) + agent.model.eval() + agent.model_ema.eval() + envs = [make_env(task, config) for _ in range(min(env_count, episodes))] + for start in range(0, episodes, len(envs)): + count = min(len(envs), episodes - start) + observations = [] + for idx, env in enumerate(envs[:count]): + env.seed(seed_base + start + idx) + observations.append(env.reset()) + finished = [False] * count + rewards, completed = [[] for _ in range(count)], [0] * count + limit = config.max_episode_steps if max_steps is None else min(max_steps, config.max_episode_steps) + for _ in range(math.ceil(limit / config.action_steps)): + if all(finished): + break + if device.type == "cuda": + torch.cuda.synchronize() + tick = time.perf_counter() + condition = condition_from_obs(task, observations, dataset, config) + prior = torch.zeros(count, config.horizon, config.action_dim, device=device) + prediction, _ = agent.sample(prior, n_samples=count, sample_steps=config.sample_steps, + condition_cfg=condition, w_cfg=1.0, use_ema=True) + normalized = prediction.cpu().numpy().clip(-1, 1) + actions = dataset.normalizer["action"].unnormalize(normalized) + if TASKS[task][0] == "robomimic": + actions = dataset.undo_transform_action(actions) + actions = actions[:, config.obs_steps - 1:config.obs_steps - 1 + config.action_steps] + if not np.isfinite(actions).all(): + raise FloatingPointError("Nonfinite policy action") + if device.type == "cuda": + torch.cuda.synchronize() + latencies.append((time.perf_counter() - tick) * 1000) + for idx, env in enumerate(envs[:count]): + if finished[idx]: + continue + obs, reward, done, info = env.step(actions[idx]) + observations[idx] = obs + rewards[idx].append(float(reward)) + finished[idx] = bool(done) + if task == "relay_kitchen" and "completed_tasks" in info: + completed[idx] = max(completed[idx], len(info["completed_tasks"][-1])) + for idx in range(count): + rows.append({"episode_seed": seed_base + start + idx, + "score": metric_from_episode(task, rewards[idx], completed[idx]), + "completed_tasks": completed[idx] if task == "relay_kitchen" else None}) + print(f"[EVAL] {task} {len(rows)}/{episodes} mean={np.mean([r['score'] for r in rows]):.6f}", flush=True) + finally: + for env in envs: + env.close() + np.random.set_state(numpy_state) + random.setstate(python_state) + agent.model.train(model_training) + agent.model_ema.train(ema_training) + return {"task": task, "training_seed": int(config.seed), "episodes": rows, + "score": float(np.mean([r["score"] for r in rows])), "num_episodes": episodes, + "sample_steps": int(config.sample_steps), "num_candidates": 1, + "action_steps": int(config.action_steps), "obs_steps": int(config.obs_steps), + "horizon": int(config.horizon), "pusht_legacy": False if task.startswith("pusht") else None, + "weights": "EMA", "seed_base": seed_base, "truncated_smoke": max_steps is not None, + "decision_ms": float(np.mean(latencies)), "latency_batch_envs": min(env_count, episodes), + "latency_scope": "normalization, transfer, observation encoder, DDPM, action unnormalization and D2H; includes first call", + "paper_ready": False} diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/launch_queue.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/launch_queue.py new file mode 100644 index 0000000000000000000000000000000000000000..c853de5f2dfa4fd7e250887be69fa729aca0f3f2 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/launch_queue.py @@ -0,0 +1,83 @@ +"""Persistent serial GPU queue: validate each task, then train three seeds.""" +import argparse +import fcntl +import json +import os +from pathlib import Path +import subprocess +import sys +import time + +from experiments.table2.common import ROOT, TASKS, atomic_json, dataset_manifest +from experiments.table2.train import source_manifest + + +def worker(directory): + plan = json.loads((directory / "plan.json").read_text()) + lock_path = ROOT / "outputs/table2/gpu.lock" + with lock_path.open("a") as lock: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + atomic_json(directory / "status.json", {"state": "preflight", "pid": os.getpid()}) + stages = [(task, 0, True) for task in plan["tasks"]] + stages += [(task, seed, False) for seed in plan["seeds"] for task in plan["tasks"]] + for task, seed, smoke in stages: + if source_manifest()["sha256"] != plan["source"]["sha256"]: + raise RuntimeError("Source files changed; queue stopped before next job") + output = directory / ("preflight" if smoke else "runs") / task / f"seed{seed}" + output.mkdir(parents=True, exist_ok=True) + command = [sys.executable, "experiments/table2/train.py", "--task", task, + "--seed", str(seed), "--output", str(output), "--workers", "0" if smoke else "4"] + if smoke: + command += ["--smoke", "--steps", "2", "--episodes", "2", "--envs", "2"] + else: + command += ["--episodes", "50", "--envs", "10"] + with (output / "train.log").open("x") as log: + process = subprocess.Popen(command, cwd=ROOT, stdout=log, stderr=subprocess.STDOUT) + status = {"state": "preflight" if smoke else "training", "task": task, + "seed": seed, "pid": os.getpid(), "child_pid": process.pid, + "log": str(output / "train.log"), "started_unix": time.time()} + atomic_json(directory / "status.json", status) + print(json.dumps(status), flush=True) + code = process.wait() + if code: + raise RuntimeError(f"{task} seed{seed} smoke={smoke} exited {code}; see {output / 'train.log'}") + atomic_json(directory / "status.json", {"state": "all_seeds_complete", "pid": os.getpid()}) + subprocess.run([sys.executable, "experiments/table2/summarize.py", str(directory)], check=True, cwd=ROOT) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", required=True) + parser.add_argument("--tasks", nargs="+", choices=TASKS, default=list(TASKS)) + parser.add_argument("--seeds", nargs="+", type=int, default=[0, 1, 2]) + parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS) + args = parser.parse_args() + directory = Path(args.root).resolve() + if not directory.is_relative_to(ROOT / "outputs/table2"): + raise ValueError("Queue output must be under outputs/table2") + if args.worker: + try: + worker(directory) + except Exception as error: + atomic_json(directory / "status.json", {"state": "failed", "error": repr(error), "pid": os.getpid()}) + raise + return + if len(set(args.tasks)) != len(args.tasks) or sorted(args.seeds) != [0, 1, 2]: + raise ValueError("Use unique tasks and training seeds 0,1,2") + for task in args.tasks: + dataset_manifest(task) + directory.mkdir(parents=True, exist_ok=False) + plan = {"tasks": args.tasks, "seeds": args.seeds, "source": source_manifest(), + "updates_per_seed": 1000000, "episodes_per_checkpoint": 50, + "serial_gpu_jobs": True, "created_unix": time.time(), + "note": "Chi_UNet representative Robomimic ph tasks only; not the full Table-2 average"} + atomic_json(directory / "plan.json", plan) + with (directory / "queue.log").open("x") as log: + process = subprocess.Popen([sys.executable, "experiments/table2/launch_queue.py", "--root", str(directory), + "--worker"], cwd=ROOT, stdin=subprocess.DEVNULL, stdout=log, stderr=subprocess.STDOUT, + start_new_session=True) + print(json.dumps({"queue_pid": process.pid, "root": str(directory), "training_jobs": len(args.tasks) * 3}), flush=True) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/prepare_data.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/prepare_data.py new file mode 100644 index 0000000000000000000000000000000000000000..c32bbfa8fa8b55efbf3274014c43e05d76ae70ef --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/prepare_data.py @@ -0,0 +1,100 @@ +"""Fetch author-hosted IL datasets, verify ZIP CRCs, record provenance. + +Image archive supports selected-member extraction over HTTP Range to avoid +downloading unrelated tasks. Existing files are never overwritten silently. +""" +import argparse +import hashlib +import json +from pathlib import Path +import shutil +import stat +import zipfile + +import requests + +ROOT = Path(__file__).resolve().parents[2] +BASE = "https://diffusion-policy.cs.columbia.edu/data/training/" + + +def sha256(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(8 * 1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def extract(archive, members, destination): + destination = destination.resolve() + records = [] + for info in members: + path = (destination / info.filename).resolve() + if not path.is_relative_to(destination) or stat.S_ISLNK(info.external_attr >> 16): + raise ValueError(f"Unsafe archive member {info.filename}") + if info.is_dir(): + continue + if path.exists(): + raise FileExistsError(path) + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(path.name + ".partial") + with archive.open(info) as source, temporary.open("xb") as target: + shutil.copyfileobj(source, target, length=8 * 1024 * 1024) + temporary.rename(path) + records.append({"path": str(path.relative_to(ROOT)), "size": info.file_size, + "crc32": info.CRC, "sha256": sha256(path)}) + return records + + +def prepare(name, image_tasks): + directory = ROOT / "outputs/table2/downloads" + directory.mkdir(parents=True, exist_ok=True) + manifest = directory / f"{name}_manifest.json" + if manifest.exists(): + print(f"[SKIP] Existing completed dataset manifest: {manifest}", flush=True) + return + url = BASE + name + ".zip" + if name == "robomimic_image": + from remotezip import RemoteZip + with RemoteZip(url) as archive: + selected = [i for i in archive.infolist() if any( + f"/{task}/ph/image_abs.hdf5" in "/" + i.filename for task in image_tasks)] + if len(selected) != len(image_tasks): + raise ValueError(f"Missing/duplicate image members: {[i.filename for i in selected]}") + print("[SELECTED]", [(i.filename, i.file_size) for i in selected], flush=True) + records = extract(archive, selected, ROOT / "dev") + metadata = {"source": url, "selection": image_tasks, "members": records, + "archive_sha256": None, "verification": "ZIP member CRC and locally computed SHA256"} + else: + path = directory / (name + ".zip") + if not path.exists(): + part = path.with_suffix(".zip.partial") + if part.exists(): + raise FileExistsError(f"Incomplete download requires explicit resume: {part}") + with requests.get(url, stream=True, timeout=(30, 120)) as response: + response.raise_for_status() + print(f"[DOWNLOAD] {url}: {response.headers.get('Content-Length', '?')} bytes", flush=True) + with part.open("xb") as stream: + for block in response.iter_content(8 * 1024 * 1024): + stream.write(block) + part.rename(path) + with zipfile.ZipFile(path) as archive: + print(f"[EXTRACT] {name}, {len(archive.infolist())} members", flush=True) + records = extract(archive, archive.infolist(), ROOT / "dev") + metadata = {"source": url, "archive_sha256": sha256(path), "members": records, + "verification": "ZIP member CRC and locally computed SHA256; no published reference hash"} + manifest.write_text(json.dumps(metadata, indent=2) + "\n") + print(f"[READY] {manifest}", flush=True) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("datasets", nargs="+", choices=["pusht", "kitchen", "robomimic_lowdim", "robomimic_image"]) + parser.add_argument("--image-tasks", nargs="+", default=["lift", "square"]) + args = parser.parse_args() + for name in args.datasets: + prepare(name, args.image_tasks) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/run.sh b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/run.sh new file mode 100644 index 0000000000000000000000000000000000000000..1f97e1c33a827f94f554504b7cd9709fe0c2b657 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/run.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +set -euo pipefail +table2_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)" +export PYTHONPATH="$table2_root:$table2_root/pipelines${PYTHONPATH:+:$PYTHONPATH}" +export PATH="$table2_root/outputs/table2/runtime/venv/bin:$PATH" +export LD_LIBRARY_PATH="$table2_root/outputs/runtime/cleandiffuser-kitchen/lib:$table2_root/outputs/runtime/mujoco210/bin${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" +export MUJOCO_PY_MUJOCO_PATH="$table2_root/outputs/runtime/mujoco210" +export CC="${CC:-/usr/bin/gcc-13}" +export CPATH="$table2_root/outputs/table2/runtime/sysroot/usr/include${CPATH:+:$CPATH}" +export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu" +export MUJOCO_GL="${MUJOCO_GL:-egl}" +export SDL_VIDEODRIVER=dummy +export PYGAME_HIDE_SUPPORT_PROMPT=1 +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}" +export MKL_NUM_THREADS="${MKL_NUM_THREADS:-4}" +export PYTHONUNBUFFERED=1 +export MPLCONFIGDIR="$table2_root/outputs/table2/runtime/matplotlib" +export XDG_CACHE_HOME="$table2_root/outputs/table2/runtime/cache" +cd -- "$table2_root" +exec "$table2_root/outputs/table2/runtime/venv/bin/python" "$@" diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/summarize.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/summarize.py new file mode 100644 index 0000000000000000000000000000000000000000..07b019c502c36f57351b6366ce6d1aee2eea1553 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/summarize.py @@ -0,0 +1,56 @@ +"""Report Last/Max only after the prespecified checkpoint windows complete.""" +import argparse +import json +from pathlib import Path + +import numpy as np + +from experiments.table2.common import TASKS, atomic_json, eval_checkpoints + + +def summarize(directory): + directory = Path(directory) + plan = json.loads((directory / "plan.json").read_text()) + rows = [] + for task in plan["tasks"]: + scores, missing = [], [] + window = eval_checkpoints(1000000, TASKS[task][2]) + for seed in [0, 1, 2]: + sequence = [] + for step in window: + path = directory / "runs" / task / f"seed{seed}" / f"eval_{step:07d}.json" + if not path.exists(): + missing.append(str(path)) + continue + result = json.loads(path.read_text()) + if (result["smoke"] or result["truncated_smoke"] or result["num_episodes"] != 50 + or result["step"] != step or result["training_seed"] != seed + or result["source_sha256"] != plan["source"]["sha256"] + or result["num_candidates"] != 1 + or [r["episode_seed"] for r in result["episodes"]] != list(range(100000, 100050))): + raise ValueError(f"Protocol mismatch: {path}") + measured = np.mean([r["score"] for r in result["episodes"]]) + if not np.isfinite(measured) or abs(measured - result["score"]) > 1e-10: + raise ValueError(f"Inconsistent episode scores: {path}") + sequence.append(measured) + scores.append(sequence) + if missing: + rows.append({"task": task, "state": "incomplete", "missing_count": len(missing)}) + continue + values = np.asarray(scores) + rows.append({"task": task, "state": "protocol_complete", "training_seeds": [0, 1, 2], + "last_mean": float(values[:, -1].mean()), "last_seed_std": float(values[:, -1].std(ddof=1)), + "max_mean": float(values.max(axis=1).mean()), "max_seed_std": float(values.max(axis=1).std(ddof=1)), + "max_of_seed_means_alternative": float(values.mean(axis=0).max()), + "max_definition": "mean across seeds of each seed's best checkpoint in the fixed window", + "checkpoint_steps": window, "seed_checkpoint_scores": scores, "paper_ready": False, + "note": "Complete local protocol; runtime deviations and agreement with paper still require review."}) + atomic_json(directory / "table2_results.json", {"rows": rows, "full_table2_average": None}) + print(json.dumps(rows, indent=2)) + return rows + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("directory") + summarize(parser.parse_args().directory) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/train.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/train.py new file mode 100644 index 0000000000000000000000000000000000000000..b7d2b28d4c1ecf62f6adec75d5d7cc0bae5d29f2 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/experiments/table2/train.py @@ -0,0 +1,228 @@ +"""Auditable Chi_UNet Table-2 training using the existing IL implementations.""" +import argparse +import hashlib +import json +import math +import os +from pathlib import Path +import random +import subprocess +import time + +import numpy as np +from omegaconf import OmegaConf +import torch +from torch.utils.data import DataLoader, Sampler + +from experiments.table2.common import (ROOT, TASKS, atomic_json, build_agent, + build_dataset, config_for, dataset_manifest, eval_checkpoints, seed_all, + training_condition) +from experiments.table2.evaluate import evaluate + + +class ResumableBatches(Sampler): + """Epoch-wise shuffle, including the last partial batch; no RNG coupling. + + Resume is based on completed optimizer updates, not prefetched batches. + """ + def __init__(self, size, batch_size, seed, start, stop): + if size < 1 or batch_size < 1 or start < 0 or stop < start: + raise ValueError("Invalid batch sampler bounds") + self.size, self.batch_size, self.seed = size, batch_size, seed + self.start, self.stop = start, stop + + def __len__(self): + return self.stop - self.start + + def __iter__(self): + batches = math.ceil(self.size / self.batch_size) + old_epoch, indices = None, None + for step in range(self.start, self.stop): + epoch, offset = divmod(step, batches) + if epoch != old_epoch: + generator = torch.Generator().manual_seed(self.seed + epoch * 1009) + indices = torch.randperm(self.size, generator=generator).tolist() + old_epoch = epoch + yield indices[offset * self.batch_size:(offset + 1) * self.batch_size] + + +def source_manifest(): + paths = [] + for directory in ("cleandiffuser", "configs/dp", "experiments/table2"): + paths.extend(p for p in (ROOT / directory).rglob("*") + if p.suffix in (".py", ".yaml", ".sh", ".patch") or p.name == "requirements.txt") + # The isolated compatibility patch is not represented by pip freeze alone. + legacy = ROOT / "outputs/table2/runtime/venv/lib/python3.10/site-packages/mujoco_py" + paths.extend(p for p in (legacy / "__init__.py", legacy / "builder.py") if p.is_file()) + files = {str(p.relative_to(ROOT)): hashlib.sha256(p.read_bytes()).hexdigest() + for p in sorted(paths) if p.is_file()} + return {"commit": subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=ROOT, text=True).strip(), + "files": files, "sha256": hashlib.sha256(json.dumps(files, sort_keys=True).encode()).hexdigest()} + + +def normalizer_state(value): + if isinstance(value, dict): + return {k: normalizer_state(v) for k, v in value.items()} + return {"class": type(value).__module__ + "." + type(value).__name__, + "state": {k: v.tolist() if isinstance(v, np.ndarray) else v + for k, v in vars(value).items()}} + + +def save_checkpoint(path, agent, scheduler, step, metadata, resumable): + checkpoint = {"model": agent.model.state_dict(), "model_ema": agent.model_ema.state_dict(), + "step": step, "metadata": metadata} + if resumable: + checkpoint.update(optimizer=agent.optimizer.state_dict(), scheduler=scheduler.state_dict(), + rng={"torch": torch.get_rng_state(), "cuda": torch.cuda.get_rng_state_all() + if torch.cuda.is_available() else [], "numpy": np.random.get_state(), + "python": random.getstate()}) + temporary = path.with_suffix(".tmp") + torch.save(checkpoint, temporary) + temporary.replace(path) + + +def train(args): + if args.steps != 1000000 and not args.smoke: + raise ValueError("Noncanonical update counts require --smoke") + if args.steps < 1 or args.episodes < 1 or args.envs < 1: + raise ValueError("Steps, episodes, and envs must be positive") + if not args.smoke and args.episodes != 50: + raise ValueError("Table-2 evaluation requires 50 episodes") + config = config_for(args.task, args.seed, args.device) + directory = Path(args.output).resolve() + directory.mkdir(parents=True, exist_ok=True) + if (directory / "run.json").exists() and not args.resume: + raise FileExistsError(f"Existing run: {directory}; use --resume explicitly") + seed_all(args.seed) + manifest_path = dataset_manifest(args.task) + manifest = json.loads(manifest_path.read_text()) + source = source_manifest() + dataset = build_dataset(args.task, config) + agent = build_agent(args.task, config) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(agent.optimizer, T_max=args.steps) + metadata = {"task": args.task, "seed": args.seed, "total_steps": args.steps, + "smoke": args.smoke, "config": OmegaConf.to_container(config, resolve=True), + "source": source, "dataset_manifest": str(manifest_path), + "dataset_manifest_sha256": hashlib.sha256(manifest_path.read_bytes()).hexdigest(), + "dataset_source": manifest["source"], "dataset_sequences": len(dataset), + "dataset_episodes": dataset.replay_buffer.n_episodes, + "dataset_parse_audit": getattr(dataset, "table2_parse_audit", None), + "normalizer": normalizer_state(dataset.normalizer), + "parameters": sum(p.numel() for p in agent.model.parameters()), + "evaluation_steps": eval_checkpoints(args.steps, TASKS[args.task][2]), + "evaluation_episodes": args.episodes, "evaluation_envs": args.envs, + "episode_seed_base": 100000, "pid": os.getpid(), + "torch": torch.__version__, "device": str(config.device), + "gpu": torch.cuda.get_device_name(config.device) if "cuda" in config.device else None, + "precision": "float32; no AMP", "paper_ready": False} + step = 0 + if args.resume: + # Only our locally generated trusted checkpoint; it contains RNG arrays. + saved = torch.load(directory / "resume.pt", map_location="cpu", weights_only=False) + for key in ("task", "seed", "total_steps", "smoke", "config", "normalizer", + "dataset_manifest_sha256", "dataset_parse_audit", "evaluation_episodes", "evaluation_envs"): + if saved["metadata"][key] != metadata[key]: + raise ValueError(f"Resume metadata mismatch: {key}") + if saved["metadata"]["source"]["sha256"] != source["sha256"]: + raise ValueError("Source changed since checkpoint; use a fresh run") + agent.model.load_state_dict(saved["model"]) + agent.model_ema.load_state_dict(saved["model_ema"]) + agent.optimizer.load_state_dict(saved["optimizer"]) + scheduler.load_state_dict(saved["scheduler"]) + torch.set_rng_state(saved["rng"]["torch"]) + if saved["rng"]["cuda"]: + torch.cuda.set_rng_state_all(saved["rng"]["cuda"]) + np.random.set_state(saved["rng"]["numpy"]) + random.setstate(saved["rng"]["python"]) + step = saved["step"] + atomic_json(directory / "run.json", metadata) + frozen = subprocess.check_output([os.sys.executable, "-m", "pip", "freeze"], text=True) + (directory / "environment.txt").write_text(frozen) + print(f"[START] {args.task} seed={args.seed} parameters={metadata['parameters']:,} " + f"demonstrations={metadata['dataset_episodes']} step={step}/{args.steps} smoke={args.smoke}", flush=True) + sampler = ResumableBatches(len(dataset), config.batch_size, args.seed, step, args.steps) + loader = DataLoader(dataset, batch_sampler=sampler, num_workers=args.workers, + pin_memory="cuda" in config.device, persistent_workers=args.workers > 0, + generator=torch.Generator().manual_seed(args.seed + 98765)) + start_time, start_step, loss_sum, loss_count = time.monotonic(), step, 0.0, 0 + evaluation_steps = metadata["evaluation_steps"] + + def checkpoint_evaluation(current): + target = directory / f"eval_{current:07d}.json" + if target.exists(): + return + save_checkpoint(directory / f"model_{current:07d}.pt", agent, scheduler, current, metadata, False) + result = evaluate(args.task, config, dataset, agent, args.episodes, args.envs, + max_steps=config.action_steps if args.smoke else None) + result.update(step=current, checkpoint=str(directory / f"model_{current:07d}.pt"), + source_sha256=source["sha256"], smoke=args.smoke) + atomic_json(target, result) + + # A crash after saving resume.pt but during rollout must not drop evaluation. + if step in evaluation_steps: + checkpoint_evaluation(step) + for batch in loader: + action = batch["action"].to(config.device).float() + condition = training_condition(args.task, batch, config) + agent.optimizer.zero_grad(set_to_none=True) + loss = agent.loss(action, condition) + if not torch.isfinite(loss): + raise FloatingPointError(f"Nonfinite loss at update {step + 1}") + loss.backward() + # Infinity leaves finite gradients unchanged; this is a diagnostic gate. + grad_norm = torch.nn.utils.clip_grad_norm_(agent.model.parameters(), float("inf"), error_if_nonfinite=True) + agent.optimizer.step() + agent.ema_update() + scheduler.step() + step += 1 + loss_sum += loss.item() + loss_count += 1 + if step % args.log_every == 0 or step == 1 or step == args.steps: + record = {"step": step, "loss": loss_sum / loss_count, "grad_norm": float(grad_norm), + "lr": scheduler.get_last_lr()[0], "elapsed_seconds": time.monotonic() - start_time, + "seconds_per_update": (time.monotonic() - start_time) / (step - start_step)} + print("[TRAIN] " + json.dumps(record), flush=True) + with (directory / "metrics.jsonl").open("a") as stream: + stream.write(json.dumps(record, allow_nan=False) + "\n") + atomic_json(directory / "status.json", dict(record, state="training", pid=os.getpid())) + loss_sum, loss_count = 0.0, 0 + if step % args.save_every == 0 or step in evaluation_steps: + save_checkpoint(directory / "resume.pt", agent, scheduler, step, metadata, True) + if step in evaluation_steps: + checkpoint_evaluation(step) + scores = [json.loads((directory / f"eval_{s:07d}.json").read_text())["score"] for s in evaluation_steps] + atomic_json(directory / "summary.json", {"task": args.task, "seed": args.seed, + "step": step, "last": scores[-1], "max": max(scores), "checkpoint_scores": scores, + "evaluation_steps": evaluation_steps, "episodes_per_checkpoint": args.episodes, + "smoke": args.smoke, "paper_ready": False, + "note": "One training seed only; Table 2 requires aggregation across seeds 0,1,2."}) + atomic_json(directory / "status.json", {"state": "smoke_complete" if args.smoke else "seed_complete", "step": step}) + print(f"[COMPLETE] {directory}", flush=True) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--task", required=True, choices=TASKS) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--output", required=True) + parser.add_argument("--steps", type=int, default=1000000) + parser.add_argument("--episodes", type=int, default=50) + parser.add_argument("--envs", type=int, default=10) + parser.add_argument("--workers", type=int, default=4) + parser.add_argument("--log-every", type=int, default=100) + parser.add_argument("--save-every", type=int, default=10000) + parser.add_argument("--smoke", action="store_true") + parser.add_argument("--resume", action="store_true") + args = parser.parse_args() + try: + train(args) + except Exception as error: + directory = Path(args.output) + if directory.exists(): + atomic_json(directory / "failure.json", {"error": repr(error), "pid": os.getpid()}) + raise + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/__init__.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e83d23441fba2f43f17cb79425d9c7f10b73c74e --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/__init__.py @@ -0,0 +1 @@ +"""Reproducible comparison with the official RSS 2024 Consistency Policy.""" diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/bootstrap.sh b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/bootstrap.sh new file mode 100644 index 0000000000000000000000000000000000000000..9169f2e5ac319763caba49c08f8a3f4fa39805c6 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/bootstrap.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +set -euo pipefail + +task_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)" +cp_repo="${1:-${task_root}/third_party/consistency-policy}" +cp_revision=eed0c465d3c1aa6dcf22f10e784df07dec06beac +if [[ ! -e "$cp_repo" ]]; then + git clone --no-checkout https://github.com/Aaditya-Prasad/consistency-policy.git "$cp_repo" + git -C "$cp_repo" checkout --detach "$cp_revision" +fi +if [[ "$(git -C "$cp_repo" rev-parse HEAD)" != "$cp_revision" ]]; then + echo "Expected revision $cp_revision. Choose a new destination for a clean checkout." >&2 + exit 1 +fi +if [[ -n "$(git -C "$cp_repo" status --porcelain --untracked-files=no)" ]]; then + echo "The official checkout has tracked modifications; choose a clean destination." >&2 + exit 1 +fi +echo "Official Consistency Policy: $cp_repo @ $cp_revision" +echo "Use the upstream dependency environment; this script does not install packages." diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/evaluate.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/evaluate.py new file mode 100644 index 0000000000000000000000000000000000000000..7dee3774abaf8b150a3682a99b983a41ad05a941 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/evaluate.py @@ -0,0 +1,315 @@ +"""Evaluate pinned official CP / EDM / DDIM checkpoints on image tasks. + +Run in the official dependency environment. Timing covers predict_action on +device-resident observations, including normalization and the vision encoder. +""" +import argparse +import copy +import hashlib +import inspect +import json +import math +import platform +import random +import subprocess +import sys +import time +from pathlib import Path + +PINNED_COMMIT = "eed0c465d3c1aa6dcf22f10e784df07dec06beac" +DEFAULT_REPO = Path(__file__).resolve().parents[2] / "third_party/consistency-policy" +TARGETS = { + "cp": "consistency_policy.student.ctm_policy.CTMPPUnetHybridImagePolicy", + "edm": "consistency_policy.teacher.edm_policy.KarrasUnetHybridImagePolicy", + "ddim": "diffusion_policy.policy.diffusion_unet_hybrid_image_policy.DiffusionUnetHybridImagePolicy", +} +LABELS = {"cp": "Consistency Policy", "edm": "EDM teacher", "ddim": "DDIM"} + + +def verify_source(repo): + repo = Path(repo).resolve() + revision = subprocess.check_output( + ["git", "-C", str(repo), "rev-parse", "HEAD"], text=True).strip() + if revision != PINNED_COMMIT: + raise ValueError("Expected official revision " + PINNED_COMMIT + "; got " + revision) + changes = subprocess.check_output( + ["git", "-C", str(repo), "status", "--porcelain", "--untracked-files=no"], text=True) + if changes: + raise ValueError("Official source has tracked modifications; use a clean pinned checkout.") + return {"repo": str(repo), "commit": revision, + "url": "https://github.com/Aaditya-Prasad/consistency-policy"} + + +def parse_steps(value): + steps = [int(v) for v in value.split(",")] + if not steps or min(steps) < 1 or len(set(steps)) != len(steps): + raise ValueError("steps must be distinct positive integers") + return steps + + +def cp_chaining_times(steps, mode, bins): + if steps == 1: + return [] + if mode == "official": + if steps != 3: + raise ValueError("The official recipe supports 1 or 3 steps; use --chain-mode bins for other counts.") + # Upstream interprets D as raw noise magnitudes, despite README wording. + return ["D", 27, 54] + if steps >= bins: + raise ValueError("CP steps must be fewer than scheduler bins") + # C is the upstream code's actual bin-index-to-sigma conversion switch. + return ["C"] + [round(i * bins / steps) for i in range(1, steps)] + + +def set_sampling(policy, method, steps, chain_mode): + if method == "cp": + times = cp_chaining_times(steps, chain_mode, policy.noise_scheduler.bins) + policy.disable_chaining() + if times: + policy.chaining_times = times + policy.enable_chaining() + return times + if method == "ddim": + if steps > policy.noise_scheduler.config.num_train_timesteps: + raise ValueError("DDIM steps exceed training timesteps") + policy.num_inference_steps = steps + else: + if steps < 2: + raise ValueError("EDM requires at least 2 bins; bins are not denoiser NFE") + policy.noise_scheduler.bins = steps + policy.noise_scheduler.bins_max = steps + return [] + + +def load_policy(checkpoint, method, weights, device): + import dill + import hydra + import torch + from omegaconf import OmegaConf + + # Checkpoints are locally trained, trusted pickle files. This also supports + # CPU loading of CUDA-trained checkpoints and Torch's newer load default. + kwargs = {"map_location": "cpu", "pickle_module": dill} + if "weights_only" in inspect.signature(torch.load).parameters: + kwargs["weights_only"] = False + payload = torch.load(str(checkpoint), **kwargs) + cfg = copy.deepcopy(payload["cfg"]) + if cfg.policy._target_ != TARGETS[method]: + raise ValueError("Checkpoint policy does not match --method " + method) + if method == "cp" and cfg.policy.get("use_kde", False): + raise ValueError("This one-candidate protocol does not support CP use_kde=True checkpoints") + if method == "ddim" and "DDIMScheduler" not in cfg.policy.noise_scheduler._target_: + raise ValueError("DDIM baseline requires a DDIMScheduler checkpoint") + if weights == "auto": + weights = "ema_model" if method != "cp" and cfg.training.use_ema else "model" + if method == "cp" and weights != "model": + raise ValueError("Official CP predicts with state_dicts.model; choose --weights model") + if weights not in payload["state_dicts"]: + raise ValueError("Checkpoint has no " + weights + " weights") + policy_cfg = copy.deepcopy(cfg.policy) + if method in ("cp", "edm"): + policy_cfg.inference_mode = True + if method == "cp": + # Missing training-only OmegaConf values must not block inference. + policy_cfg.teacher_path = None + policy_cfg.edm = None + policy = hydra.utils.instantiate(policy_cfg) + policy.load_state_dict(payload["state_dicts"][weights], strict=True) + if hasattr(policy, "drop_teacher"): + policy.drop_teacher() + if method == "cp": + # The target EMA network is only used in the training loss. + policy.model_ema = None + policy.requires_grad_(False) + policy.eval().to(device) + del payload + return policy, cfg, weights + + +class TimedPolicy: + """Duck-typed upstream policy wrapper; no changes to rollout behavior.""" + def __init__(self, policy, warmup): + self.policy = policy + self.warmup = warmup + self.durations = [] + self.nfes = [] + self.calls = 0 + self.warmed = False + self.hook = policy.model.register_forward_hook(self._count_forward) + + def __getattr__(self, name): + return getattr(self.policy, name) + + def _count_forward(self, *args): + self.calls += 1 + + def _sync(self): + import torch + if self.policy.device.type == "cuda": + torch.cuda.synchronize(self.policy.device) + + def predict_action(self, obs): + import torch + if not self.warmed: + # Warmup must not consume the rollout's random action samples. + devices = [self.policy.device.index or 0] if self.policy.device.type == "cuda" else [] + with torch.random.fork_rng(devices=devices), torch.no_grad(): + for _ in range(self.warmup): + self.policy.predict_action(obs) + self._sync() + self.warmed = True + self.calls = 0 + self._sync() + start = time.perf_counter() + result = self.policy.predict_action(obs) + self._sync() + self.durations.append((time.perf_counter() - start) * 1000) + self.nfes.append(self.calls) + return result + + def close(self): + self.hook.remove() + + +def checkpoint_hash(path): + digest = hashlib.sha256() + with open(path, "rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repo", type=Path, default=DEFAULT_REPO) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--method", choices=TARGETS, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--dataset", type=Path, help="Relocate the Robomimic rollout dataset") + parser.add_argument("--steps", help="CP:1,3; DDIM:1,2,4,8,16,32,50; EDM bins:20,40,80") + parser.add_argument("--chain-mode", choices=("official", "bins"), default="official") + parser.add_argument("--weights", choices=("auto", "model", "ema_model"), default="auto") + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--eval-seed", type=int, default=0) + parser.add_argument("--episodes", type=int, default=50) + parser.add_argument("--test-start-seed", type=int, default=100000) + parser.add_argument("--max-steps", type=int, help="Override episode limit for all compared methods") + parser.add_argument("--warmup", type=int, default=10) + args = parser.parse_args(argv) + if args.episodes < 1 or args.warmup < 0 or (args.max_steps is not None and args.max_steps < 1): + parser.error("episodes/max-steps must be positive; warmup must be nonnegative") + args.checkpoint = args.checkpoint.resolve(strict=True) + args.output = args.output.resolve() + source = verify_source(args.repo) + sys.path.insert(0, source["repo"]) + import hydra + import numpy as np + import torch + from omegaconf import OmegaConf + + steps_list = parse_steps(args.steps or {"cp": "1,3", "ddim": "1,2,4,8,16,32,50", "edm": "20,40,80"}[args.method]) + policy, cfg, weights = load_policy(args.checkpoint, args.method, args.weights, args.device) + runner_cfg = copy.deepcopy(cfg.task.env_runner) + allowed_runners = { + "diffusion_policy.env_runner.pusht_image_runner.PushTImageRunner", + "diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner", + } + if runner_cfg._target_ not in allowed_runners: + raise ValueError("Only official PushT/Robomimic image runners are supported") + is_pusht = "PushT" in runner_cfg._target_ + if args.dataset is not None and not is_pusht: + runner_cfg.dataset_path = str(args.dataset.resolve(strict=True)) + if not is_pusht: + runner_cfg.dataset_path = str(Path(runner_cfg.dataset_path).resolve(strict=True)) + runner_cfg.n_envs = 1 + runner_cfg.n_train = runner_cfg.n_train_vis = runner_cfg.n_test_vis = 0 + runner_cfg.n_test = args.episodes + runner_cfg.test_start_seed = args.test_start_seed + if args.max_steps is not None: + runner_cfg.max_steps = args.max_steps + if (runner_cfg.n_obs_steps != policy.n_obs_steps or + runner_cfg.n_action_steps != policy.n_action_steps): + raise ValueError("Runner observation/action chunk lengths do not match checkpoint") + args.output.mkdir(parents=True, exist_ok=True) + # Validate every requested sweep point before starting expensive rollouts. + for steps in steps_list: + set_sampling(policy, args.method, steps, args.chain_mode) + filename = args.output / (args.method + "_" + str(steps) + "_" + args.chain_mode + ".json") + if filename.exists(): + raise FileExistsError("Use a new output directory; result already exists: " + str(filename)) + dataset = str(cfg.task.dataset.get("zarr_path", cfg.task.dataset.get("dataset_path", ""))) + protocol = { + "horizon": policy.horizon, "n_obs_steps": policy.n_obs_steps, + "n_action_steps": policy.n_action_steps, "n_test": args.episodes, + "test_start_seed": args.test_start_seed, "max_steps": runner_cfg.max_steps, + "n_envs": 1, "n_candidates": 1, + "abs_action": bool(runner_cfg.get("abs_action", False)), + "dataset": dataset if is_pusht else runner_cfg.dataset_path, + "shape_meta": OmegaConf.to_container(cfg.policy.shape_meta, resolve=True), + "crop_shape": OmegaConf.to_container(cfg.policy.crop_shape) if cfg.policy.crop_shape is not None else None, + "legacy_test": bool(runner_cfg.get("legacy_test", False)), + "dataset_seed": cfg.task.dataset.get("seed"), + "val_ratio": cfg.task.dataset.get("val_ratio"), + "max_train_episodes": cfg.task.dataset.get("max_train_episodes"), + } + device = policy.device + hardware = { + "device": str(device), "device_name": torch.cuda.get_device_name(device) if device.type == "cuda" else (platform.processor() or platform.machine()), + "torch": torch.__version__, "cuda": torch.version.cuda, "dtype": str(policy.dtype), + "threads": torch.get_num_threads(), "cudnn_benchmark": torch.backends.cudnn.benchmark, + "tf32_matmul": torch.backends.cuda.matmul.allow_tf32, + "tf32_cudnn": torch.backends.cudnn.allow_tf32, + "cudnn_deterministic": torch.backends.cudnn.deterministic, + } + digest = checkpoint_hash(args.checkpoint) + for steps in steps_list: + random.seed(args.eval_seed) + np.random.seed(args.eval_seed) + torch.manual_seed(args.eval_seed) + if device.type == "cuda": + torch.cuda.manual_seed_all(args.eval_seed) + times = set_sampling(policy, args.method, steps, args.chain_mode) + run_dir = args.output / (args.method + "_" + str(steps) + "_" + args.chain_mode) + run_dir.mkdir() + runner = hydra.utils.instantiate(runner_cfg, output_dir=str(run_dir)) + timed = TimedPolicy(policy, args.warmup) + try: + with torch.no_grad(): + metrics = runner.run(timed) + finally: + timed.close() + runner.env.close() + episodes = {key: float(value) for key, value in metrics.items() + if key.startswith("test/sim_max_reward_")} + if len(episodes) != args.episodes or not timed.durations: + raise ValueError("Incomplete rollout: expected all episodes and latency samples") + score = float(metrics["test/mean_score"]) + if not math.isfinite(score) or any(not math.isfinite(v) for v in episodes.values()): + raise ValueError("Non-finite rollout score") + result = { + "schema_version": 1, "method": args.method, "method_label": LABELS[args.method], + "sampling": {"steps": steps, "chain_mode": args.chain_mode if args.method == "cp" else None, + "chaining_times": times, "steps_unit": "bins" if args.method == "edm" else "denoiser_calls"}, + "checkpoint": str(args.checkpoint), "checkpoint_sha256": digest, + "weights": weights, "train_seed": cfg.training.seed, "eval_seed": args.eval_seed, + "task": str(cfg.task.name), + "score_metric": "mean_max_normalized_coverage" if is_pusht else "success_rate", + "score": score, "episode_scores": episodes, "protocol": protocol, "hardware": hardware, + "latency": {"scope": "predict_action_device_resident", "mean_ms": float(np.mean(timed.durations)), + "p50_ms": float(np.percentile(timed.durations, 50)), + "p95_ms": float(np.percentile(timed.durations, 95)), + "samples": len(timed.durations), "warmup_calls": args.warmup, + "raw_ms": timed.durations}, + "nfe": {"mean": float(np.mean(timed.nfes)), "min": min(timed.nfes), "max": max(timed.nfes)}, + "source": source, + "architecture": OmegaConf.to_container(cfg.policy, resolve=False), + "training": OmegaConf.to_container(cfg.training, resolve=False), + } + path = args.output / (run_dir.name + ".json") + path.write_text(json.dumps(result, indent=2, allow_nan=False) + "\n") + print(f"{LABELS[args.method]} {steps}: score={score:.4f}, " + f"predict_action={result['latency']['mean_ms']:.3f} ms, NFE={result['nfe']['mean']:.1f} -> {path}") + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/locomotion_eval.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/locomotion_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..3c874c6e493bc9e67be703c21a13ee2ee503ae0e --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/locomotion_eval.py @@ -0,0 +1,398 @@ +"""Matched locomotion evaluation for a frozen Consistency Policy student. + +The actor interface is predict(repeated_normalized_observations, steps=1), returning +flat [num_envs * num_candidates, action_dim] actions in environment units. This is +an adaptation to the paper's frozen IQL@64 locomotion protocol, not the original +Consistency Policy manipulation benchmark. Training and detaching jobs are owned +by the caller; each completed evaluation seed is persisted independently. +""" + +import hashlib +import inspect +import json +import os +from pathlib import Path +import random +import time + +import numpy as np + + +BENCHMARK = "consistency_policy_locomotion_iql" +PROVENANCE_KEYS = ("student_checkpoint", "student_sha256", "teacher_checkpoint", "teacher_sha256", "training_seed") + + +def sha256_file(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _array(value): + if hasattr(value, "detach"): + value = value.detach().cpu().numpy() + return np.asarray(value, dtype=np.float32) + + +def normalizer_digest(normalizer): + digest = hashlib.sha256() + for key in ("mean", "std"): + attribute = key if hasattr(normalizer, key) else key + "s" + value = _array(getattr(normalizer, attribute)) + digest.update(key.encode()) + digest.update(str(value.shape).encode()) + digest.update(value.tobytes()) + return digest.hexdigest() + + +def _atomic_json(path, value): + temporary = path.with_suffix(path.suffix + ".tmp") + with temporary.open("w") as stream: + json.dump(value, stream, indent=2, sort_keys=True, allow_nan=False) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + temporary.replace(path) + + +def _atomic_npz(path, arrays): + temporary = path.with_suffix(path.suffix + ".tmp") + with temporary.open("wb") as stream: + np.savez_compressed(stream, **arrays) + stream.flush() + os.fsync(stream.fileno()) + temporary.replace(path) + + +def return_statistics(norm_returns): + """Separate rollout-group statistics used by the table from episode statistics.""" + values = np.asarray(norm_returns, dtype=np.float32) + if values.ndim != 2 or not values.size or not np.isfinite(values).all(): + raise ValueError("norm_returns must be a finite [groups, environments] array") + group_means = values.mean(axis=1) + aggregate = { + "normalized_return_mean_over_eps": float(group_means.mean()), + "normalized_return_std_over_eps": float(group_means.std(ddof=0)), + "normalized_return_mean_over_all_env_eps": float(values.mean()), + "normalized_return_std_over_all_env_eps": float(values.std(ddof=0)), + "worst_episode_mean_return": float(group_means.min()), + "worst_group_mean_return": float(group_means.min()), + "individual_episode_std": float(values.std(ddof=0)), + "individual_episode_min": float(values.min()), + "group_mean_returns": group_means.tolist(), + } + return aggregate + + +def aggregate_seed_results(results): + """Match existing table conventions while retaining alternative definitions.""" + if not results: + raise ValueError("no completed evaluation seeds") + results = sorted(results, key=lambda item: item["seed"]) + if len({item["seed"] for item in results}) != len(results): + raise ValueError("duplicate evaluation seeds") + reference = results[0] + for item in results: + for key in ("env", "provenance", "protocol", "hardware", "sampling_steps"): + if item[key] != reference[key]: + raise ValueError("cannot aggregate evaluations with differing " + key) + if not item.get("complete"): + raise ValueError("cannot aggregate incomplete evaluation") + means = np.asarray([item["normalized_return_mean_over_eps"] for item in results], dtype=np.float32) + latency = np.asarray([item["decision_time_mean_ms"] for item in results], dtype=np.float32) + aggregate = { + "schema_version": 1, "benchmark": BENCHMARK, "method": "consistency_policy", + "env": reference["env"], "evaluation_seeds": [item["seed"] for item in results], + "num_evaluation_seeds": len(results), "num_training_seeds": 1, + "sampling_steps": reference["sampling_steps"], "provenance": reference["provenance"], + "protocol": reference["protocol"], "hardware": reference["hardware"], + "legacy": { + "return_mean": float(means.mean()), "return_std_across_seeds": float(means.std(ddof=0)) if len(results) > 1 else None, + "sigma_group": float(np.mean([item["normalized_return_std_over_eps"] for item in results])), + "worst_group_mean": min(item["worst_group_mean_return"] for item in results), + "latency_mean_ms": float(latency.mean()), + "latency_within_run_std_ms_mean": float(np.mean([item["decision_time_std_ms"] for item in results])), + }, + "alternative_statistics": { + "return_sample_std_across_eval_seeds": float(means.std(ddof=1)) if len(results) > 1 else None, + "worst_seed_mean": float(means.min()), + "individual_episode_std_mean": float(np.mean([item["individual_episode_std"] for item in results])), + "worst_individual_episode": min(item["individual_episode_min"] for item in results), + "latency_std_across_eval_seed_means_ms": float(latency.std(ddof=0)) if len(results) > 1 else None, + }, + "metric_definitions": { + "return": "Mean of the five evaluation-seed means; population SD (ddof=0). One frozen trained student.", + "sigma_group": "Mean across seeds of population SD of ten group means, each averaging fifty environment episodes.", + "worst_group_mean": "Minimum over evaluation seed and rollout group of the fifty-environment mean return.", + "latency_uncertainty": "Mean across seeds of within-run population SD of decision latencies; not SD of seed means.", + }, + "run_files": [item["result_json"] for item in results], + } + success_keys = sorted(key for key in reference if key.startswith("success_at_least_")) + if success_keys: + aggregate["success_rates"] = { + key: float(np.mean([item[key] for item in results])) for key in success_keys} + return aggregate + + +def _run_rollouts(actor, vector_env, score_env, normalizer, q_target, value_net, device, + num_envs, num_episodes, num_candidates, max_steps, steps, seed, + reseed_each_group=False, direct_argmax=False): + """Run first episodes only, counting their terminal rewards before masking.""" + import torch + + obs_dim = int(vector_env.single_observation_space.shape[0]) + act_dim = int(vector_env.single_action_space.shape[0]) + use_cuda = torch.device(device).type == "cuda" + raw_returns, episode_lengths, episode_done = [], [], [] + episode_times, decision_times, group_step_counts = [], [], [] + + def sync(): + if use_cuda: + torch.cuda.synchronize(device) + + for group in range(num_episodes): + if reseed_each_group: + from cleandiffuser.utils.eval_utils import reset_seeded_env + group_seeds = [seed * 100000 + group * num_envs + index for index in range(num_envs)] + obs = reset_seeded_env(vector_env, group_seeds) + else: + obs = vector_env.reset() + if isinstance(obs, tuple): + obs = obs[0] + rewards = np.zeros(num_envs, dtype=np.float32) + lengths = np.zeros(num_envs, dtype=np.int32) + cumulative_done = np.zeros(num_envs, dtype=bool) + group_start = time.perf_counter() + control_steps = 0 + while not cumulative_done.all() and control_steps < max_steps: + alive = ~cumulative_done + sync() + started = time.perf_counter() + with torch.no_grad(): + obs_t = torch.tensor(normalizer.normalize(obs), device=device, dtype=torch.float32) + repeated = obs_t.unsqueeze(1).repeat(1, num_candidates, 1).view(-1, obs_dim) + actions = actor.predict(repeated, steps=steps) + if actions.shape != (num_envs * num_candidates, act_dim): + raise ValueError("actor.predict returned incorrect candidate shape: " + str(actions.shape)) + q = q_target(repeated, actions) + v = value_net(repeated) + advantage = (q - v).view(num_envs, num_candidates) + if direct_argmax: + indices = advantage.argmax(dim=1) + else: + probabilities = torch.softmax(advantage * 100.0, dim=1) + indices = probabilities.argmax(dim=1) + candidates = actions.view(num_envs, num_candidates, act_dim) + selected = candidates[torch.arange(num_envs, device=device), indices].cpu().numpy() + sync() + decision_times.append((time.perf_counter() - started) * 1000.0) + if not np.isfinite(selected).all(): + raise ValueError("policy selected a nonfinite action") + obs, reward, done, _ = vector_env.step(selected) + control_steps += 1 + # Alive is captured before step: count terminal reward once and ignore + # all rewards from Gym's automatic reset for finished environments. + rewards += np.asarray(reward, dtype=np.float32) * alive + lengths += alive.astype(np.int32) + cumulative_done |= np.asarray(done, dtype=bool) + raw_returns.append(rewards) + episode_lengths.append(lengths) + episode_done.append(cumulative_done) + group_step_counts.append(control_steps) + episode_times.append(time.perf_counter() - group_start) + print("[CP eval seed {} group {}/{}] wall={:.2f}s mean_raw_return={:.3f}".format( + seed, group + 1, num_episodes, episode_times[-1], rewards.mean() + ), flush=True) + + raw_returns = np.asarray(raw_returns, dtype=np.float32) + norm_returns = np.asarray([ + [score_env.get_normalized_score(float(reward)) for reward in group] for group in raw_returns + ], dtype=np.float32) + decision_times = np.asarray(decision_times, dtype=np.float32) + if not decision_times.size or not np.isfinite(decision_times).all(): + raise ValueError("missing or nonfinite timing samples") + return { + "raw_returns": raw_returns, "norm_returns": norm_returns, + "per_ep_mean": norm_returns.mean(axis=1), "per_ep_std": norm_returns.std(axis=1), + "episode_lengths": np.asarray(episode_lengths, dtype=np.int32), + "environment_done": np.asarray(episode_done, dtype=bool), + "group_step_counts": np.asarray(group_step_counts, dtype=np.int32), + "episode_times": np.asarray(episode_times, dtype=np.float32), "decision_times_ms": decision_times, + } + + +def evaluate_actor(actor, env_name, device, seeds, num_envs=50, num_episodes=10, + num_candidates=64, max_steps=1000, output_dir=None, iql_checkpoint=None, + steps=1, normalizer=None, metadata=None): + """Evaluate and resume completed seeds; return the aggregate summary. + + metadata must identify student/teacher checkpoint paths and SHA256 hashes plus + training_seed. A provided normalizer must be the training GaussianNormalizer + with means/stds (or mean/std) attributes; otherwise the existing D4RL dataset loader is used. + """ + import d4rl + import gym + import torch + from cleandiffuser.utils import IDQLQNet, IDQLVNet + + supported_envs = ( + "halfcheetah-medium-expert-v2", "walker2d-medium-expert-v2", + "kitchen-partial-v0", "kitchen-mixed-v0", + ) + if env_name not in supported_envs: + raise ValueError(f"unsupported D4RL state task {env_name!r}") + seeds = list(seeds) + if not seeds or len(seeds) != len(set(seeds)) or any(isinstance(seed, bool) or not isinstance(seed, int) or seed < 0 for seed in seeds): + raise ValueError("seeds must be unique nonnegative integers") + if any(isinstance(value, bool) or not isinstance(value, int) or value < 1 + for value in (num_envs, num_episodes, num_candidates, max_steps, steps)): + raise ValueError("evaluation dimensions and steps must be positive integers") + if output_dir is None or iql_checkpoint is None: + raise ValueError("output_dir and the frozen iql_checkpoint are required") + if metadata is None or any(key not in metadata for key in PROVENANCE_KEYS): + raise ValueError("metadata requires " + ", ".join(PROVENANCE_KEYS)) + provenance = dict(metadata) + for prefix in ("student", "teacher"): + checkpoint = Path(provenance[prefix + "_checkpoint"]).resolve() + if sha256_file(checkpoint) != provenance[prefix + "_sha256"]: + raise ValueError(prefix + " checkpoint SHA256 does not match metadata") + provenance[prefix + "_checkpoint"] = str(checkpoint) + output = Path(output_dir).resolve() + output.mkdir(parents=True, exist_ok=True) + iql_checkpoint = Path(iql_checkpoint).resolve() + provenance.update(iql_checkpoint=str(iql_checkpoint), iql_sha256=sha256_file(iql_checkpoint)) + + score_env = gym.make(env_name) + try: + is_kitchen = env_name.startswith("kitchen-") + if normalizer is None: + if env_name.startswith("kitchen-"): + from cleandiffuser.dataset.d4rl_kitchen_dataset import D4RLKitchenTDDataset + dataset = D4RLKitchenTDDataset(d4rl.qlearning_dataset(score_env)) + normalizer_name = "D4RLKitchenTDDataset.GaussianNormalizer" + else: + from cleandiffuser.dataset.d4rl_mujoco_dataset import D4RLMuJoCoTDDataset + dataset = D4RLMuJoCoTDDataset(d4rl.qlearning_dataset(score_env), True) + normalizer_name = "D4RLMuJoCoTDDataset.GaussianNormalizer" + normalizer = dataset.get_normalizer() + del dataset + else: + normalizer_name = type(normalizer).__name__ + provenance["normalizer_sha256"] = normalizer_digest(normalizer) + obs_dim = score_env.observation_space.shape[0] + act_dim = score_env.action_space.shape[0] + q_target = IDQLQNet(obs_dim, act_dim).to(device).requires_grad_(False).eval() + value_net = IDQLVNet(obs_dim).to(device).requires_grad_(False).eval() + load_kwargs = {"map_location": device} + if "weights_only" in inspect.signature(torch.load).parameters: + load_kwargs["weights_only"] = False + critic = torch.load(iql_checkpoint, **load_kwargs) + if is_kitchen and "normalizer_mean" in critic: + from cleandiffuser.utils.selector_normalizer import checkpoint_normalizer + selector_normalizer = checkpoint_normalizer(critic, obs_dim) + if normalizer_digest(selector_normalizer) != normalizer_digest(normalizer): + raise ValueError("Kitchen CP and frozen IQL TD-data normalizers differ") + provenance["selector_normalizer_sha256"] = normalizer_digest(selector_normalizer) + q_target.load_state_dict(critic["iql_q_target"], strict=True) + value_net.load_state_dict(critic["iql_v"], strict=True) + del critic + if hasattr(actor, "eval"): + actor.eval() + protocol = { + "num_envs": num_envs, "num_episodes": num_episodes, "num_candidates": num_candidates, + "max_steps": max_steps, "select": "argmax", "weight_temperature": 100.0, + "prediction_horizon": 1, "observation_horizon": 1, "executed_actions": 1, + "observation_normalizer": normalizer_name, + "action_normalization": "none; policy actions in environment units", + "candidate_initial_noise_std": 1.0, + "candidate_initial_noise_note": "Released CP uses unit Gaussian initial samples while the trajectory map starts at sigma_max=80; no additional sigma_max multiplier.", + "environment_seed_formula": ( + "100000 * evaluation_seed + group * num_envs + vector_environment_index" + if is_kitchen else + "1000 * evaluation_seed + vector_environment_index; seed once before first reset"), + "done_handling": "alive mask before step includes terminal reward and excludes autoreset rewards", + "latency_scope": ( + "observation normalization, H2D, candidate generation, frozen IQL Q/V scoring, " + + ("argmax" if is_kitchen else "softmax/argmax") + ", selected-action D2H"), + "latency_batch": num_envs * num_candidates, "latency_warmup_calls": 0, + "latency_includes_initial_calls": True, "cuda_synchronization": torch.device(device).type == "cuda", + "historical_seed_caveat": "Original sp_consistency_policy.py did not explicitly seed vector environments; exact original initial states cannot be reconstructed.", + } + hardware = { + "device": str(device), "device_name": torch.cuda.get_device_name(device) if torch.device(device).type == "cuda" else "CPU", + "torch": torch.__version__, "cuda": torch.version.cuda, "dtype": "torch.float32", + "threads": torch.get_num_threads(), "cudnn_benchmark": torch.backends.cudnn.benchmark, + "tf32_matmul": torch.backends.cuda.matmul.allow_tf32, + "tf32_cudnn": torch.backends.cudnn.allow_tf32, + "cudnn_deterministic": torch.backends.cudnn.deterministic, + } + results = [] + for seed in seeds: + path = output / "cp_seed{}.json".format(seed) + npz_path = path.with_suffix(".npz") + if path.exists(): + previous = json.loads(path.read_text()) + expected = {"provenance": provenance, "protocol": protocol, "hardware": hardware, + "env": env_name, "seed": seed, "sampling_steps": steps, "complete": True} + if not npz_path.is_file() or any(previous.get(key) != value for key, value in expected.items()): + raise ValueError("existing result does not match this evaluation; choose another output directory: " + str(path)) + results.append(previous) + print("[CP eval] resume completed seed {}".format(seed), flush=True) + continue + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.device(device).type == "cuda": + torch.cuda.manual_seed_all(seed) + vector_env = gym.vector.make(env_name, num_envs) + if is_kitchen: + environment_seeds = [ + [seed * 100000 + group * num_envs + index for index in range(num_envs)] + for group in range(num_episodes)] + else: + environment_seeds = [1000 * seed + index for index in range(num_envs)] + try: + if not is_kitchen: + vector_env.seed(environment_seeds) + arrays = _run_rollouts(actor, vector_env, score_env, normalizer, q_target, value_net, + device, num_envs, num_episodes, num_candidates, max_steps, steps, seed, + reseed_each_group=is_kitchen, direct_argmax=is_kitchen) + finally: + vector_env.close() + metrics = return_statistics(arrays["norm_returns"]) + timings = arrays["decision_times_ms"] + result = { + "schema_version": 1, "benchmark": BENCHMARK, "method": "consistency_policy", "model": "cp", + "env": env_name, "seed": seed, "evaluation_seed": seed, + "sampling_steps": steps, "num_envs": num_envs, "num_episodes": num_episodes, + "num_candidates": num_candidates, "select": "argmax", "weight_temperature": 100.0, + "steps_cap": max_steps, "environment_seeds": environment_seeds, + "complete": True, "provenance": provenance, "protocol": protocol, "hardware": hardware, + "result_json": str(path), "raw_npz": str(npz_path), + "raw_shape": list(arrays["norm_returns"].shape), + "episode_wall_time_mean_s": float(arrays["episode_times"].mean()), + "episode_wall_time_std_s": float(arrays["episode_times"].std(ddof=0)), + "decision_time_mean_ms": float(timings.mean()), "decision_time_std_ms": float(timings.std(ddof=0)), + "decision_time_p95_ms": float(np.percentile(timings, 95)), "decision_time_samples": int(timings.size), + **metrics, + } + if env_name.startswith("kitchen-"): + for completed in (1, 2, 3, 4): + result[f"success_at_least_{completed}_tasks"] = float( + (arrays["raw_returns"] >= completed).mean()) + _atomic_npz(npz_path, arrays) + _atomic_json(path, result) + results.append(result) + partial = aggregate_seed_results(results) + partial["complete"] = len(results) == len(seeds) + partial["requested_evaluation_seeds"] = seeds + _atomic_json(output / "summary.json", partial) + summary = aggregate_seed_results(results) + summary.update(complete=True, requested_evaluation_seeds=seeds) + _atomic_json(output / "summary.json", summary) + return summary + finally: + score_env.close() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/locomotion_policy.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/locomotion_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..38933abd0d2203a658d04a44f6fecabcbcc2c918 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/locomotion_policy.py @@ -0,0 +1,217 @@ +"""Released Consistency Policy objectives adapted to the D4RL IDQL MLP. + +The trajectory map, time sampling, loss and unit-variance initial noise follow +the MIT-licensed official implementation at commit +eed0c465d3c1aa6dcf22f10e784df07dec06beac. The existing CleanDiffuser EDM teacher +retains its native preconditioning and time embedding. Both student and target +are warm-started, avoiding the upstream first-batch random target artifact. +""" + +import copy +import inspect +import math +from pathlib import Path + +import torch +from torch import nn + +from cleandiffuser.nn_diffusion.idqlmlp import IDQLMlp + + +class TwoTimeIDQLMlp(IDQLMlp): + """IDQL MLP with shared start/stop embedding and a zero-column warm start.""" + + def __init__(self, obs_dim, act_dim, emb_dim=64, hidden_dim=256, + n_blocks=3, dropout=0.2): + super().__init__(obs_dim, act_dim, emb_dim=emb_dim, + hidden_dim=hidden_dim, n_blocks=n_blocks, + dropout=dropout, timestep_emb_type="untrainable_fourier") + self.act_dim = act_dim + self.emb_dim = emb_dim + self.affine_in = nn.Linear(obs_dim + act_dim + 2 * emb_dim, hidden_dim) + + def warm_start(self, teacher): + state = copy.deepcopy(teacher.state_dict()) + old = state["affine_in.weight"] + split = self.act_dim + self.emb_dim + state["affine_in.weight"] = torch.cat( + (old[:, :split], old.new_zeros(old.shape[0], self.emb_dim), + old[:, split:]), dim=1) + self.load_state_dict(state, strict=True) + + def forward(self, x, noise, stop_noise, condition=None): + if condition is None: + condition = x.new_zeros(x.shape[0], self.obs_dim) + start = self.time_mlp(self.map_noise(noise)) + stop = self.time_mlp(self.map_noise(stop_noise)) + features = torch.cat((x, start, stop, condition), dim=-1) + return self.affine_out(self.ln_resnet(self.affine_in(features))) + + +class LocomotionConsistencyPolicy(nn.Module): + """CP with native D4RL observations/actions and the pretrained EDM teacher.""" + + time_min = 0.02 + time_max = 80.0 + data_std = 0.5 + rho = 7.0 + bins = 80 + + def __init__(self, obs_dim, act_dim, teacher_checkpoint: Path, + device="cuda:0", dropout=0.2): + super().__init__() + self.obs_dim, self.act_dim = obs_dim, act_dim + self.dropout = dropout + self.teacher_checkpoint = str(Path(teacher_checkpoint).resolve()) + kwargs = {"map_location": "cpu"} + if "weights_only" in inspect.signature(torch.load).parameters: + # This is the user's locally trained, trusted checkpoint. + kwargs["weights_only"] = False + checkpoint = torch.load(self.teacher_checkpoint, **kwargs) + state = {key[len("diffusion."):]: value + for key, value in checkpoint["model_ema"].items() + if key.startswith("diffusion.")} + emb_dim = state["time_mlp.0.weight"].shape[1] + hidden_dim = state["affine_in.weight"].shape[0] + n_blocks = len({key.split(".")[1] for key in state + if key.startswith("ln_resnet.")}) + if state["affine_in.weight"].shape[1] != obs_dim + act_dim + emb_dim: + raise ValueError("Teacher observation/action dimensions do not match") + architecture = dict(emb_dim=emb_dim, hidden_dim=hidden_dim, n_blocks=n_blocks) + self.teacher = IDQLMlp(obs_dim, act_dim, dropout=0.1, + timestep_emb_type="untrainable_fourier", **architecture) + self.teacher.load_state_dict(state, strict=True) + self.teacher.requires_grad_(False).eval() + self.model = TwoTimeIDQLMlp(obs_dim, act_dim, dropout=dropout, **architecture) + self.model.warm_start(self.teacher) + self.target = copy.deepcopy(self.model).requires_grad_(False) + self._architecture = architecture + self.to(device) + self.train() + + @property + def config(self): + return { + "method": "Consistency Policy (D4RL IDQL MLP adaptation)", + "upstream_commit": "eed0c465d3c1aa6dcf22f10e784df07dec06beac", + "upstream_url": "https://github.com/Aaditya-Prasad/consistency-policy", + "obs_dim": self.obs_dim, "act_dim": self.act_dim, + **self._architecture, "dropout": self.dropout, + "time_min": self.time_min, "time_max": self.time_max, + "rho": self.rho, "bins": self.bins, "data_std": self.data_std, + "teacher_checkpoint": self.teacher_checkpoint, + "teacher_weights": "model_ema", "teacher_preconditioning": "native_edm", + "time_embedding": "log_sigma_div_4_shared_start_stop", + "losses": {"ctm": 1.0, "dsm": 1.0}, + "loss_formula": "sqrt(residual**4 + delta**2) - delta", + "target_decay": 0.0, "target_dropout_during_training": True, + "initial_noise_std": 1.0, "inference_weights": "model", + "chaining_times_3": ["D", 27, 54], + "adaptations": ["IDQL MLP backbone and normalized D4RL states", + "native pretrained EDM teacher and Fourier scale", + "student and target both initialized from teacher", + "vectorized equivalent official time sampler"], + } + + def train(self, mode=True): + super().train(mode) + self.teacher.eval() + return self + + def sigma(self, indices): + return (self.time_max ** (1 / self.rho) + indices / (self.bins - 1) + * (self.time_min ** (1 / self.rho) + - self.time_max ** (1 / self.rho))).pow(self.rho).clamp( + self.time_min, self.time_max) + + def sample_ctm_indices(self, batch, device): + # Integer distributions exactly match upstream, including equal times + # and the extra minimum-time bin. Avoid GPU-to-CPU scalar syncs. + start = torch.randint(self.bins, (batch,), device=device) + stop = start + (torch.rand(batch, device=device) + * (self.bins + 1 - start)).floor().long() + intermediate = start + (torch.rand(batch, device=device) + * (stop - start + 1)).floor().long() + intermediate = torch.minimum(intermediate, start + 1) + return start, stop, intermediate + + def trajectory(self, network, x, start, stop, obs, clamp=False): + start = torch.as_tensor(start, device=x.device, dtype=x.dtype).expand(x.shape[0]) + stop = torch.as_tensor(stop, device=x.device, dtype=x.dtype).expand(x.shape[0]) + denom = (start.square() + self.data_std ** 2).sqrt() + c_in = denom.reciprocal()[:, None] + c_skip = (self.data_std ** 2 / + ((start - self.time_min).square() + self.data_std ** 2))[:, None] + c_out = ((start - self.time_min) * self.data_std / denom)[:, None] + raw = network(x * c_in, start.log() / 4, stop.log() / 4, obs) + denoised = c_skip * x + c_out * raw + ratio = (stop / start)[:, None] + result = ratio * x + (1 - ratio) * denoised + return result.clamp(-1, 1) if clamp else result + + @torch.no_grad() + def teacher_denoise(self, x, sigma, obs): + denom = (sigma.square() + self.data_std ** 2).sqrt() + raw = self.teacher(x / denom[:, None], sigma.log() / 4, obs) + return ((self.data_std ** 2 / denom.square())[:, None] * x + + (sigma * self.data_std / denom)[:, None] * raw) + + @torch.no_grad() + def teacher_step(self, x, start, stop, obs): + step = (stop - start)[:, None] + slope = (x - self.teacher_denoise(x, start, obs)) / start[:, None] + euler = x + step * slope + next_slope = (euler - self.teacher_denoise(euler, stop, obs)) / stop[:, None] + return x + step * (slope + next_slope) / 2 + + def released_loss(self, prediction, target, weights=None): + delta = math.sqrt(self.act_dim) * 0.00054 + # Preserve the released code's fourth-power residual, including DSM. + loss = ((prediction - target).square().square() + delta ** 2).sqrt() - delta + if weights is not None: + loss = loss * weights[:, None] + return loss.mean() + + def compute_loss(self, actions, obs): + batch = actions.shape[0] + i, j, k = self.sample_ctm_indices(batch, actions.device) + start, stop, intermediate = self.sigma(i), self.sigma(j), self.sigma(k) + noisy = actions + start[:, None] * torch.randn_like(actions) + transported = self.teacher_step(noisy, start, intermediate, obs) + pred = self.trajectory(self.model, noisy, start, stop, obs) + # Target weights are frozen, but this projection needs input gradients. + pred = self.trajectory(self.target, pred, stop, self.time_min, obs) + with torch.no_grad(): + target = self.trajectory(self.target, transported, intermediate, stop, obs) + target = self.trajectory(self.target, target, stop, self.time_min, obs) + ctm = self.released_loss(pred, target) + + high_count = (batch + 1) // 2 + fractions = torch.rand(high_count, device=actions.device) * 0.7 + high = (self.time_max ** (1 / self.rho) + fractions + * (self.time_min ** (1 / self.rho) + - self.time_max ** (1 / self.rho))).pow(self.rho).clamp( + self.time_min, self.time_max) + low = (torch.randn(batch // 2, device=actions.device) * 1.2 - 1.2).exp() + sigma = torch.cat((high, low)) + noisy = actions + sigma[:, None] * torch.randn_like(actions) + prediction = self.trajectory(self.model, noisy, sigma, self.time_min, obs) + weights = ((sigma.square() + self.data_std ** 2) + / (sigma * self.data_std).square()) + return {"ctm": ctm, "dsm": self.released_loss(prediction, actions, weights)} + + @torch.no_grad() + def update_target(self): + self.target.load_state_dict(self.model.state_dict(), strict=True) + + @torch.no_grad() + def predict(self, obs, steps=1): + if steps not in (1, 3): + raise ValueError("Released CP supports 1 or 3 steps in this benchmark") + noise = torch.randn(obs.shape[0], self.act_dim, device=obs.device, dtype=obs.dtype) + result = self.trajectory(self.model, noise, self.time_max, self.time_min, obs, clamp=True) + if steps == 3: + for sigma in (27.0, 54.0): + noise = result + sigma * torch.randn_like(result) + result = self.trajectory(self.model, noise, sigma, self.time_min, obs, clamp=True) + return result diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/prepare.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/prepare.py new file mode 100644 index 0000000000000000000000000000000000000000..ce5082724cb9dcb35c2cc38c69fa0767e53f3696 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/prepare.py @@ -0,0 +1,300 @@ +"""Prepare matched vision configurations for the official Consistency Policy code. + +This module only reads YAML and writes configuration/provenance files. Run its +printed commands in the upstream environment to train EDM, then CTMP, and DDIM. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +import shlex +import subprocess +import sys +from typing import Any + +from omegaconf import DictConfig, OmegaConf + + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +UPSTREAM_URL = "https://github.com/Aaditya-Prasad/Consistency-Policy" +UPSTREAM_COMMIT = "eed0c465d3c1aa6dcf22f10e784df07dec06beac" +# Match the existing CleanDiffuser vision task episode limits and crop sizes. +TASKS = { + "pusht_image": ("pusht_image.yaml", 300, 84), + "lift_image": ("lift_image_abs.yaml", 400, 76), + "can_image": ("can_image_abs.yaml", 400, 76), + "square_image": ("square_image_abs.yaml", 500, 76), + "tool_hang_image": ("tool_hang_image_abs.yaml", 700, 216), +} +METHODS = ("edm", "ctmp", "ddim") + + +def positive_int(value: str) -> int: + number = int(value) + if number <= 0: + raise argparse.ArgumentTypeError("must be positive") + return number + + +def nonnegative_int(value: str) -> int: + number = int(value) + if number < 0: + raise argparse.ArgumentTypeError("must be nonnegative") + return number + + +def parser() -> argparse.ArgumentParser: + result = argparse.ArgumentParser(description=__doc__) + result.add_argument("--repo", type=Path, + default=PROJECT_ROOT / "third_party/consistency-policy") + result.add_argument("--task", choices=TASKS, required=True) + result.add_argument("--dataset", type=Path, required=True, + help="Existing PushT Zarr directory/zip or absolute-action Robomimic image HDF5") + result.add_argument("--output", type=Path, required=True, + help="Write configs/, manifest.json, and method output directories here") + result.add_argument("--teacher-checkpoint", type=Path, + help="EDM teacher and warm start; default: OUTPUT/edm/checkpoints/latest.ckpt") + result.add_argument("--seed", type=nonnegative_int, default=0) + result.add_argument("--dataset-seed", type=nonnegative_int, default=42) + result.add_argument("--device", default="cuda:0") + result.add_argument("--epochs", type=positive_int, + help="Set all three methods to this epoch count; otherwise retain upstream budgets") + result.add_argument("--batch-size", type=positive_int, default=64) + result.add_argument("--num-workers", type=nonnegative_int, default=8) + return result + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _flatten(value: Any, prefix: str = "") -> dict[str, Any]: + if isinstance(value, dict): + result = {} + for key, child in value.items(): + result.update(_flatten(child, f"{prefix}.{key}" if prefix else key)) + return result + return {prefix: value} + + +def _changes(before: dict, after: dict) -> list[dict[str, Any]]: + original, prepared = _flatten(before), _flatten(after) + changes = [] + for key in sorted(original.keys() | prepared.keys()): + if key not in original or key not in prepared or original[key] != prepared[key]: + change = {"path": key} + if key in original: + change["upstream"] = original[key] + if key in prepared: + change["prepared"] = prepared[key] + changes.append(change) + return changes + + +def _make_config(args: argparse.Namespace, method: str, template: DictConfig, + task_template: DictConfig) -> DictConfig: + cfg = OmegaConf.create(OmegaConf.to_container(template, resolve=False)) + # Replace the entire task, including nested shape metadata, to avoid retaining + # Square camera/action dimensions when constructing PushT or Tool Hang. + cfg.task = OmegaConf.create(OmegaConf.to_container(task_template, resolve=False)) + cfg.task.name = args.task + cfg.task_name = args.task + cfg.shape_meta = cfg.task.shape_meta + cfg.policy.shape_meta = cfg.task.shape_meta + cfg.name = f"consistency_policy_{method}_{args.task}" + cfg.exp_name = "cleandiffuser_comparison" + cfg.horizon = cfg.policy.horizon = cfg.task.dataset.horizon = 16 + cfg.n_obs_steps = cfg.dataset_obs_steps = cfg.policy.n_obs_steps = 2 + cfg.n_action_steps = cfg.policy.n_action_steps = 8 + cfg.n_latency_steps = 0 + cfg.past_action_visible = False + cfg.task.dataset.pad_before = 1 + cfg.task.dataset.pad_after = 7 + cfg.task.dataset.seed = args.dataset_seed + cfg.task.dataset.val_ratio = 0.02 + cfg.policy.crop_shape = [TASKS[args.task][2]] * 2 + cfg.policy.down_dims = [512, 1024, 2048] + cfg.policy.diffusion_step_embed_dim = 128 + cfg.policy.eval_fixed_crop = True + + dataset_path = str(args.dataset) + if args.task == "pusht_image": + cfg.task.dataset.zarr_path = dataset_path + # Use the same full training split for the three new methods. The + # upstream PushT task otherwise caps training at only 90 episodes. + cfg.task.dataset.max_train_episodes = None + else: + cfg.task.dataset_path = dataset_path + cfg.task.dataset.dataset_path = dataset_path + cfg.task.env_runner.dataset_path = dataset_path + cfg.task.dataset.n_obs_steps = 2 + + runner = cfg.task.env_runner + runner.max_steps = TASKS[args.task][1] + runner.n_obs_steps, runner.n_action_steps = 2, 8 + runner.n_train = runner.n_train_vis = runner.n_test_vis = 0 + runner.n_test, runner.n_envs = 50, 1 + runner.test_start_seed = 100000 + runner.past_action = False + + for key in ("dataloader", "val_dataloader"): + cfg[key].batch_size = args.batch_size + cfg[key].num_workers = args.num_workers + cfg[key].persistent_workers = False + + method_output = args.output / method + cfg.training.device = args.device + cfg.training.seed = args.seed + cfg.training.online_rollouts = False + cfg.training.inference_mode = False + cfg.training.resume = False + # Upstream compares this sentinel with the literal string "None". + cfg.training.resume_path = "None" + cfg.training.output_dir = str(method_output) + if args.epochs is not None: + cfg.training.num_epochs = args.epochs + if method == "ctmp": + cfg.training.p_epochs = min(int(cfg.training.p_epochs), args.epochs) + if method == "ctmp": + cfg.policy.teacher_path = str(args.teacher_checkpoint) + cfg.policy.edm = str(args.teacher_checkpoint) + if method in ("ctmp", "edm"): + cfg.policy.inference_mode = False + + cfg.logging.mode = "disabled" + cfg.logging.name = cfg.name + cfg.logging.resume = False + cfg.logging.id = None + if "tags" in cfg.logging: + cfg.logging.tags = [method, args.task, cfg.exp_name] + cfg.checkpoint.topk.k = 0 + cfg.checkpoint.save_last_ckpt = True + cfg.checkpoint.save_last_snapshot = False + cfg.multi_run.run_dir = str(method_output) + cfg.multi_run.wandb_name_base = cfg.name + cfg.hydra = {"run": {"dir": str(method_output)}, "job": {"chdir": False}} + # All upstream eval expressions are replaced above; resolution does not + # import Hydra, Torch, simulator packages, or install an eval resolver. + OmegaConf.resolve(cfg) + missing = OmegaConf.missing_keys(cfg) + if missing: + raise ValueError(f"Unresolved required configuration values: {sorted(missing)}") + return cfg + + +def prepare(args: argparse.Namespace) -> dict[str, Any]: + args.repo = args.repo.expanduser().resolve() + args.dataset = args.dataset.expanduser().resolve() + args.output = args.output.expanduser().resolve() + args.teacher_checkpoint = ( + args.teacher_checkpoint.expanduser().resolve() if args.teacher_checkpoint + else args.output / "edm/checkpoints/latest.ckpt" + ) + # Check every destination before creating directories or writing any file. + destinations = [args.output / "manifest.json"] + [ + args.output / "configs" / f"{method}.yaml" for method in METHODS + ] + occupied = [str(path) for path in destinations if path.exists() or path.is_symlink()] + if occupied: + raise ValueError("Refusing to overwrite prepared files: " + ", ".join(occupied)) + directories = [args.output, args.output / "configs"] + [ + args.output / method for method in METHODS + ] + for directory in directories: + if directory.exists() and not directory.is_dir(): + raise ValueError(f"Output directory is an existing file: {directory}") + if not args.dataset.exists(): + raise ValueError(f"Dataset does not exist: {args.dataset}") + if args.seed >= 2**32 or args.dataset_seed >= 2**32: + raise ValueError("Seeds must be smaller than 2**32") + try: + commit = subprocess.check_output( + ["git", "-C", str(args.repo), "rev-parse", "HEAD"], text=True, + stderr=subprocess.PIPE, + ).strip() + tracked_changes = subprocess.check_output( + ["git", "-C", str(args.repo), "status", "--porcelain", "--untracked-files=no"], + text=True, stderr=subprocess.PIPE, + ) + except (OSError, subprocess.CalledProcessError) as error: + raise ValueError(f"Cannot inspect upstream checkout {args.repo}; run bootstrap.sh first") from error + if commit != UPSTREAM_COMMIT: + raise ValueError(f"Expected upstream revision {UPSTREAM_COMMIT}, found {commit}") + if tracked_changes: + raise ValueError("Official source has tracked modifications; use a clean pinned checkout.") + task_source = args.repo / "diffusion_policy/config/task" / TASKS[args.task][0] + task_template = OmegaConf.load(task_source) + template_suffix = "th" if args.task == "tool_hang_image" else "square" + configs = {} + manifest: dict[str, Any] = { + "schema_version": 1, + "kind": "configuration_provenance", + "upstream": {"url": UPSTREAM_URL, "commit": commit, "repo": str(args.repo)}, + "task": args.task, + "task_template": {"path": str(task_source), "sha256": _sha256(task_source)}, + "dataset": {"path": str(args.dataset), "seed": args.dataset_seed, + "val_ratio": 0.02, "contents_validated": False}, + "seed": args.seed, + "teacher_checkpoint": str(args.teacher_checkpoint), + "adaptation_notes": [ + "These are upstream reproduction adaptations, not published performance results.", + "Horizon, observation/action steps, crops, U-Net widths, embedding dimension, batch size, dataset split, and rollout protocol match across the three methods.", + "This does not establish matched compute or architecture with CleanDiffuser's existing pipelines; report parameters, optimizer updates, teacher cost, and hardware separately.", + "Without --epochs, upstream epoch budgets remain EDM=400, CTMP=500, DDIM=8000. CTMP additionally requires the EDM teacher training cost.", + "With --epochs, CTMP p_epochs is capped at that budget; other method-specific optimizer/scheduler settings remain upstream values.", + "PushT uses the full training split, removing the upstream 90-episode cap. Robomimic requires absolute-action image data with the declared camera shapes.", + "Dataset existence is checked; contents, demonstrations, and image shapes are not read or validated.", + "Training rollouts and W&B logging are disabled. Latest checkpoints retain the upstream periodic save cadence; latest need not be the final epoch.", + "Offline training still needs upstream policy/encoder dependencies, including Robomimic; disabling rollouts only avoids constructing simulators.", + "CTMP teacher_path and edm warm start refer to the same checkpoint, which must exist before CTMP training starts.", + ], + "methods": {}, + } + for method in METHODS: + source = args.repo / "configs" / f"{method}_{template_suffix}.yaml" + template = OmegaConf.load(source) + cfg = _make_config(args, method, template, task_template) + config_path = args.output / "configs" / f"{method}.yaml" + command = [sys.executable, str(args.repo / "train.py"), + "--config-dir", str(config_path.parent), "--config-name", method] + manifest["methods"][method] = { + "template": str(source), "template_sha256": _sha256(source), + "config": str(config_path), "output_dir": str(args.output / method), + "epochs": cfg.training.num_epochs, "train_command": command, + "changes_from_method_template": _changes( + OmegaConf.to_container(template, resolve=False), + OmegaConf.to_container(cfg, resolve=True), + ), + } + configs[method] = cfg + + (args.output / "configs").mkdir(parents=True, exist_ok=True) + for method, cfg in configs.items(): + (args.output / method).mkdir(exist_ok=True) + path = Path(manifest["methods"][method]["config"]) + with path.open("x") as stream: + OmegaConf.save(cfg, stream, resolve=True) + manifest["methods"][method]["config_sha256"] = _sha256(path) + with (args.output / "manifest.json").open("x") as stream: + stream.write(json.dumps(manifest, indent=2) + "\n") + return manifest + + +def main() -> None: + cli = parser() + try: + manifest = prepare(cli.parse_args()) + except (ValueError, OSError) as error: + cli.error(str(error)) + print("Prepared configurations and provenance; no training has been launched.") + print("Use the upstream environment. Train EDM before CTMP; DDIM is independent.") + for method, info in manifest["methods"].items(): + print(f"{method} ({info['epochs']} epochs): {shlex.join(info['train_command'])}") + print(f"Manifest: {Path(manifest['methods']['edm']['config']).parent.parent / 'manifest.json'}") + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/record_locomotion.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/record_locomotion.py new file mode 100644 index 0000000000000000000000000000000000000000..b23a2b6f0c1de16721199b23dd823418771ffe31 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/record_locomotion.py @@ -0,0 +1,248 @@ +"""Record canonical CP locomotion students in the existing website video format. + +One predetermined seed per task, single-step CP, frozen IQL argmax over 64 +candidates. A demo episode is illustrative, not a replacement benchmark score. +Use --background to survive terminal disconnection; status and logs are durable. +""" + +import argparse +import fcntl +import json +import os +from pathlib import Path +import random +import shutil +import signal +import subprocess +import sys +import time +import traceback + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT)) +from tools.consistency_policy.run_locomotion import atomic_json, sha256 + +TASKS = ("halfcheetah", "walker2d") + + +def status(args, **fields): + path = args.run_dir / "status.json" + previous = json.loads(path.read_text()) if path.exists() else {} + previous.update(fields, pid=os.getpid(), updated_at=time.strftime("%Y-%m-%dT%H:%M:%S%z")) + atomic_json(path, previous) + + +def publish_local_video(source, destination): + destination.parent.mkdir(parents=True, exist_ok=True) + if destination.exists(): + if sha256(source) != sha256(destination): + raise FileExistsError(f"Refusing to replace a different website video: {destination}") + return + temporary = destination.with_suffix(".partial.mp4") + shutil.copy2(source, temporary) + os.replace(temporary, destination) + + +def record_task(args, task): + import d4rl + import gym + import imageio.v2 as imageio + import numpy as np + import torch + from cleandiffuser.utils import IDQLQNet, IDQLVNet + from tools.consistency_policy.locomotion_policy import LocomotionConsistencyPolicy + + env_name = task + "-medium-expert-v2" + directory = ROOT / "checkpoints/final_paper_main" / env_name + student_path = directory / "cp_latest.pt" + teacher_path = directory / "edm_latest.pt" + iql_path = directory / "iql_latest.pt" + video = ROOT / (task + "_consistency_policy.mp4") + web_video = ROOT / "website/assets/videos" / video.name + manifest = args.run_dir / (task + ".json") + source_hashes = {key: sha256(path) for key, path in + (("student", student_path), ("teacher", teacher_path), ("iql", iql_path))} + if manifest.exists() and video.exists(): + old = json.loads(manifest.read_text()) + if (old.get("complete") and old["video_sha256"] == sha256(video) + and old["checkpoint_sha256"] == source_hashes + and old["seed"] == args.seed): + publish_local_video(video, web_video) + print(f"Already recorded and verified: {video}", flush=True) + return old + raise ValueError("Existing recording differs; preserve it and use a new output name") + if video.exists(): + raise FileExistsError(f"Video exists without matching recording provenance: {video}") + + reference = ROOT / (task + "_diffusion.mp4") + reader = imageio.get_reader(str(reference)) + try: + reference_meta = reader.get_meta_data() + max_steps = reader.count_frames() + finally: + reader.close() + width, height = reference_meta["size"] + fps = reference_meta["fps"] + assert max_steps == 1000 and (width, height) == (512, 512) + checkpoint = torch.load(student_path, map_location=args.device, weights_only=False) + config = checkpoint["config"] + if checkpoint["method"] != "consistency_policy_d4rl_idql_adaptation" or checkpoint["step"] != 1_000_000: + raise ValueError("Expected the fixed final one-million-update CP checkpoint") + for key in ("teacher", "iql"): + if config[key + "_sha256"] != source_hashes[key]: + raise ValueError(f"{key} checkpoint differs from the student's training provenance") + policy_source = ROOT / "tools/consistency_policy/locomotion_policy.py" + expected_source = config["provenance"]["adapter_sha256"]["locomotion_policy.py"] + if sha256(policy_source) != expected_source: + raise ValueError("CP policy implementation differs from the frozen evaluated source") + obs_dim, act_dim = config["obs_dim"], config["act_dim"] + actor = LocomotionConsistencyPolicy(obs_dim, act_dim, teacher_path, + device=args.device, dropout=checkpoint["policy_config"]["dropout"]) + actor.load_state_dict(checkpoint["model_state"], strict=True) + actor.requires_grad_(False).eval() + mean = np.asarray(checkpoint["normalizer_mean"], dtype=np.float32) + std = np.asarray(checkpoint["normalizer_std"], dtype=np.float32) + assert mean.shape == std.shape == (obs_dim,) and np.all(std > 0) + q = IDQLQNet(obs_dim, act_dim).to(args.device).requires_grad_(False).eval() + v = IDQLVNet(obs_dim).to(args.device).requires_grad_(False).eval() + critic = torch.load(iql_path, map_location=args.device, weights_only=False) + q.load_state_dict(critic["iql_q_target"], strict=True) + v.load_state_dict(critic["iql_v"], strict=True) + del critic, checkpoint + + # Match the evaluation seed ordering: constructors first, then seed RNGs. + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + env = gym.make(env_name) + assert env.observation_space.shape == (obs_dim,) and env.action_space.shape == (act_dim,) + env.seed(1000 * args.seed) + obs = env.reset() + temporary_video = video.with_suffix(".partial.mp4") + writer = imageio.get_writer(str(temporary_video), fps=fps, codec="libx264", + pixelformat="yuv420p", quality=8, macro_block_size=16, + ffmpeg_params=["-movflags", "+faststart"]) + observations, actions_saved, rewards = [], [], [] + done = False + started = time.monotonic() + try: + for step in range(max_steps): + observations.append(np.asarray(obs).copy()) + with torch.no_grad(): + normalized = (obs[None, :] - mean) / std + repeated = torch.tensor(normalized, device=args.device, dtype=torch.float32).repeat(64, 1) + actions = actor.predict(repeated, steps=1) + assert actions.shape == (64, act_dim) + advantage = (q(repeated, actions) - v(repeated)).view(1, 64, 1) + weights = torch.softmax(100.0 * advantage, dim=1) + probabilities = weights / weights.sum(dim=1, keepdim=True) + index = probabilities.squeeze(-1).argmax(dim=1) + action = actions[index[0]].cpu().numpy() + if not np.isfinite(action).all(): + raise ValueError("Nonfinite CP action") + obs, reward, done, info = env.step(action) + actions_saved.append(action.copy()) + rewards.append(float(reward)) + frame = env.render(mode="rgb_array", width=width, height=height) + assert frame.shape == (height, width, 3) and frame.dtype == np.uint8 + writer.append_data(frame) + if step in (0, 149, 499, 999): + imageio.imwrite(str(args.run_dir / f"{task}_frame{step + 1:04d}.png"), frame) + if (step + 1) % 100 == 0 or done: + status(args, phase="recording", env=env_name, frames=step + 1, + max_frames=max_steps, elapsed_seconds=time.monotonic() - started) + print(f"{task}: {step + 1}/{max_steps} frames, cumulative reward {sum(rewards):.2f}", flush=True) + if done: + break + raw_return = float(sum(rewards)) + normalized_return = float(env.get_normalized_score(raw_return)) + simulation_dt = float(env.unwrapped.dt) + finally: + writer.close() + env.close() + + # Decode every frame and verify dimensions, frame count and stream format. + reader = imageio.get_reader(str(temporary_video)) + try: + encoded = reader.get_meta_data() + decoded_frames = 0 + for frame in reader: + if frame.shape != (height, width, 3): + raise ValueError("Encoded video dimensions differ from the reference") + decoded_frames += 1 + assert decoded_frames == len(rewards) and abs(encoded["fps"] - fps) < 0.01 + assert encoded["codec"] == "h264" and encoded["pix_fmt"].startswith("yuv420p") + finally: + reader.close() + os.replace(temporary_video, video) + trace = args.run_dir / (task + "_rollout.npz") + np.savez_compressed(trace, observations=np.asarray(observations), + actions=np.asarray(actions_saved), rewards=np.asarray(rewards)) + result = { + "complete": True, "method": "Consistency Policy (D4RL adaptation)", + "env": env_name, "seed": args.seed, "environment_seed": 1000 * args.seed, + "training_seed": config["training_seed"], "training_updates": 1_000_000, + "sampling_steps": 1, "num_candidates": 64, "selector": "frozen IQL softmax(100*advantage), argmax", + "normalizer": "saved CP training mean/std; no refit", "inference_weights": "model", + "checkpoint_paths": {"student": str(student_path), "teacher": str(teacher_path), "iql": str(iql_path)}, + "checkpoint_sha256": source_hashes, "policy_source_sha256": sha256(policy_source), + "recorder_sha256": sha256(Path(__file__)), "reference": str(reference), + "reference_sha256": sha256(reference), "video": str(video), "website_video": str(web_video), + "video_sha256": sha256(video), "width": width, "height": height, "fps": fps, + "frames": len(rewards), "duration_seconds": len(rewards) / fps, + "codec": "h264", "pixel_format": "yuv420p", "faststart": True, + "environment_done": bool(done), "simulation_dt_seconds": simulation_dt, + "raw_return": raw_return, "normalized_return": normalized_return, + "rollout_trace": str(trace), "rollout_trace_sha256": sha256(trace), + "note": "One preselected seed-0 single-environment illustrative rollout, not an aggregate benchmark score or exact replay of the 50-environment evaluation RNG stream. One frame per step, reference playback fps retained.", + "hardware": {"torch": torch.__version__, "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(args.device), "cpu_threads": torch.get_num_threads()}, + } + atomic_json(manifest, result) + publish_local_video(video, web_video) + print(f"Saved {video}: {len(rewards)} frames, normalized demo return {normalized_return:.4f}", flush=True) + return result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-dir", type=Path, default=ROOT / "outputs/consistency_policy/videos_20260908") + parser.add_argument("--background", action="store_true") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--device", default="cuda:0") + args = parser.parse_args() + if not 0 <= args.seed < 2**32 // 1000: + parser.error("Seed is out of range") + args.run_dir = args.run_dir.expanduser().resolve() + args.run_dir.mkdir(parents=True, exist_ok=True) + if args.background: + environment = dict(os.environ, PYTHONUNBUFFERED="1", D4RL_SUPPRESS_IMPORT_ERROR="1", + OMP_NUM_THREADS="1", MKL_NUM_THREADS="1") + command = [sys.executable, str(Path(__file__).resolve()), "--run-dir", str(args.run_dir), + "--seed", str(args.seed), "--device", args.device] + with (args.run_dir / "background.log").open("ab", buffering=0) as log: + child = subprocess.Popen(command, cwd=str(ROOT), stdin=subprocess.DEVNULL, + stdout=log, stderr=subprocess.STDOUT, env=environment, + start_new_session=True, close_fds=True) + (args.run_dir / "launcher.pid").write_text(str(child.pid) + "\n") + print(f"Detached recorder PID {child.pid}; log: {args.run_dir / 'background.log'}") + return + import torch + torch.set_num_threads(1) + if not torch.cuda.is_available(): + raise RuntimeError("A CUDA GPU is required for this recording configuration") + signal.signal(signal.SIGHUP, signal.SIG_IGN) + with (args.run_dir / "run.lock").open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + try: + status(args, phase="starting", complete=False) + results = [record_task(args, task) for task in TASKS] + status(args, phase="complete", complete=True, videos=[r["video"] for r in results]) + except BaseException as error: + status(args, phase="failed", complete=False, error=str(error), traceback=traceback.format_exc()) + raise + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/run_kitchen.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/run_kitchen.py new file mode 100644 index 0000000000000000000000000000000000000000..1ff87938947328391b313151fdbe2d6ad2c85805 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/run_kitchen.py @@ -0,0 +1,198 @@ +"""Train and evaluate the official-objective CP adaptation on D4RL Kitchen.""" + +import argparse +import inspect +import json +import math +from pathlib import Path +import random +import time + +import numpy as np + +from tools.consistency_policy.locomotion_eval import evaluate_actor, sha256_file +from tools.consistency_policy.locomotion_policy import LocomotionConsistencyPolicy +from tools.consistency_policy.run_locomotion import atomic_json, code_provenance, save_checkpoint + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--env", choices=("kitchen-partial-v0", "kitchen-mixed-v0"), + default="kitchen-partial-v0") + parser.add_argument("--run-dir", type=Path, required=True) + parser.add_argument("--teacher-checkpoint", type=Path) + parser.add_argument("--iql-checkpoint", type=Path) + parser.add_argument("--gradient-steps", type=int, default=1_000_000) + parser.add_argument("--checkpoint-every", type=int, default=10_000) + parser.add_argument("--batch-size", type=int, default=256) + parser.add_argument("--training-seed", type=int, default=0) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--lr", type=float, default=1e-4) + parser.add_argument("--threads", type=int, default=1) + parser.add_argument("--eval-only", action="store_true") + return parser.parse_args() + + +def main(): + import d4rl + import gym + import torch + from cleandiffuser.dataset.d4rl_kitchen_dataset import D4RLKitchenTDDataset + + args = parse_args() + if min(args.gradient_steps, args.checkpoint_every, args.batch_size, args.threads) < 1: + raise ValueError("training counts must be positive") + if args.lr <= 0 or not 0 <= args.training_seed < 2**32: + raise ValueError("invalid learning rate or training seed") + + root = Path(__file__).resolve().parents[2] + args.run_dir = args.run_dir.expanduser().resolve() + teacher_path = (args.teacher_checkpoint or ( + root / f"tutorials/results/sp_consistency_policy/{args.env}/edm_ckpt_latest.pt" + )).expanduser().resolve(strict=True) + iql_path = (args.iql_checkpoint or ( + root / f"results/idql_d4rl_kitchen/{args.env}/iql_ckpt_1000000.pt" + )).expanduser().resolve(strict=True) + args.run_dir.mkdir(parents=True, exist_ok=True) + torch.set_num_threads(args.threads) + + provenance = code_provenance() + provenance["kitchen_adapter_sha256"] = sha256_file(Path(__file__)) + config = { + "env": args.env, + "gradient_steps": args.gradient_steps, + "batch_size": args.batch_size, + "training_seed": args.training_seed, + "lr": args.lr, + "optimizer": "AdamW", + "betas": [0.95, 0.999], + "weight_decay": 1e-6, + "teacher_checkpoint": str(teacher_path), + "teacher_sha256": sha256_file(teacher_path), + "iql_checkpoint": str(iql_path), + "iql_sha256": sha256_file(iql_path), + "evaluation": { + "seeds": [0, 1, 2, 3, 4], "num_envs": 50, "groups": 10, + "num_candidates": 64, "max_steps": 280, "sampling_steps": 1, + }, + "provenance": provenance, + } + config_path = args.run_dir / "run_config.json" + if config_path.exists() and json.loads(config_path.read_text()) != config: + raise ValueError("Existing run configuration differs; use a new run directory") + atomic_json(config_path, config) + + random.seed(args.training_seed) + np.random.seed(args.training_seed) + torch.manual_seed(args.training_seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(args.training_seed) + + env = gym.make(args.env) + try: + dataset = D4RLKitchenTDDataset(d4rl.qlearning_dataset(env)) + finally: + env.close() + normalizer = dataset.get_normalizer() + obs_dim, act_dim = dataset.o_dim, dataset.a_dim + observations = dataset.obs.to(args.device) + actions = dataset.act.to(args.device) + del dataset + + actor = LocomotionConsistencyPolicy( + obs_dim, act_dim, teacher_path, device=args.device, dropout=0.2) + optimizer = torch.optim.AdamW( + actor.model.parameters(), lr=args.lr, betas=(0.95, 0.999), weight_decay=1e-6) + + def lr_factor(step): + if step < 500: + return float(step) / 500 + progress = (step - 500) / max(1, args.gradient_steps - 500) + return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress))) + + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_factor) + latest_path = args.run_dir / "latest.pt" + final_path = args.run_dir / "final.pt" + resume_path = final_path if final_path.exists() else latest_path + start_step, train_seconds = 0, 0.0 + if resume_path.exists(): + load_kwargs = {"map_location": args.device} + if "weights_only" in inspect.signature(torch.load).parameters: + load_kwargs["weights_only"] = False + checkpoint = torch.load(resume_path, **load_kwargs) + if checkpoint["config"] != config: + raise ValueError("Checkpoint configuration/provenance differs") + np.testing.assert_array_equal(checkpoint["normalizer_mean"], normalizer.mean) + np.testing.assert_array_equal(checkpoint["normalizer_std"], normalizer.std) + actor.load_state_dict(checkpoint["model_state"], strict=True) + optimizer.load_state_dict(checkpoint["optimizer"]) + scheduler.load_state_dict(checkpoint["scheduler"]) + start_step, train_seconds = checkpoint["step"], checkpoint["train_seconds"] + random.setstate(checkpoint["rng"]["python"]) + np.random.set_state(checkpoint["rng"]["numpy"]) + torch.set_rng_state(checkpoint["rng"]["torch"].cpu()) + if checkpoint["rng"]["cuda"]: + torch.cuda.set_rng_state_all([state.cpu() for state in checkpoint["rng"]["cuda"]]) + print(f"Resuming at update {start_step}", flush=True) + + if not args.eval_only and start_step < args.gradient_steps: + actor.train() + clock_start = time.monotonic() + totals = {"ctm": 0.0, "dsm": 0.0} + report_start = start_step + with (args.run_dir / "training.jsonl").open("a", buffering=1) as stream: + for step in range(start_step + 1, args.gradient_steps + 1): + indices = torch.randint(observations.shape[0], (args.batch_size,), device=args.device) + losses = actor.compute_loss(actions[indices], observations[indices]) + loss = losses["ctm"] + losses["dsm"] + if not torch.isfinite(loss).item(): + raise FloatingPointError(f"Non-finite CP loss at update {step}") + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + scheduler.step() + actor.update_target() + for key in totals: + totals[key] += losses[key].item() + if step % 1000 == 0 or step == args.gradient_steps: + count = step - report_start + record = { + "step": step, "target_steps": args.gradient_steps, + "ctm": totals["ctm"] / count, "dsm": totals["dsm"] / count, + "lr": scheduler.get_last_lr()[0], + } + stream.write(json.dumps(record, allow_nan=False) + "\n") + print(json.dumps(record), flush=True) + totals = {"ctm": 0.0, "dsm": 0.0} + report_start = step + if step % args.checkpoint_every == 0 or step == args.gradient_steps: + elapsed = train_seconds + time.monotonic() - clock_start + save_checkpoint( + latest_path, actor, optimizer, scheduler, step, + normalizer, config, elapsed) + if not final_path.exists(): + latest_path.replace(final_path) + + if not final_path.exists(): + raise FileNotFoundError("No final CP checkpoint; train first or remove --eval-only") + actor.eval() + metadata = { + "student_checkpoint": str(final_path), + "student_sha256": sha256_file(final_path), + "teacher_checkpoint": str(teacher_path), + "teacher_sha256": sha256_file(teacher_path), + "training_seed": args.training_seed, + "gradient_steps": args.gradient_steps, + "provenance": provenance, + "method_note": "Official CP CTM+DSM objective adapted to Kitchen state/action IDQL MLP", + } + summary = evaluate_actor( + actor, args.env, args.device, seeds=list(range(5)), num_envs=50, + num_episodes=10, num_candidates=64, max_steps=280, + output_dir=args.run_dir / "evaluation", iql_checkpoint=iql_path, + steps=1, normalizer=normalizer, metadata=metadata) + print(json.dumps(summary["legacy"], indent=2, sort_keys=True), flush=True) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/run_locomotion.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/run_locomotion.py new file mode 100644 index 0000000000000000000000000000000000000000..11f8fb9a12ef5db652533e15b101d70b123b890d --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/run_locomotion.py @@ -0,0 +1,338 @@ +"""Durable background CP distillation, D4RL evaluation, and paper update. + +The default experiment trains one frozen student per task for one million +gradient updates, then evaluates seeds 0..4. It does not select on test return. +""" +import argparse +import fcntl +import hashlib +import json +import math +import os +from pathlib import Path +import random +import shutil +import signal +import subprocess +import sys +import time +import traceback + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT)) +ENVS = ("halfcheetah-medium-expert-v2", "walker2d-medium-expert-v2") +STOP_REQUESTED = False + + +def sha256(path): + h = hashlib.sha256() + with Path(path).open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +def atomic_json(path, value): + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n") + os.replace(temporary, path) + + +def update_status(run_dir, **fields): + path = run_dir / "status.json" + state = json.loads(path.read_text()) if path.exists() else {} + state.update(fields) + state.update(pid=os.getpid(), updated_at=time.strftime("%Y-%m-%dT%H:%M:%S%z")) + atomic_json(path, state) + + +def signal_stop(signum, frame): + global STOP_REQUESTED + STOP_REQUESTED = True + print("Stop requested; saving at the next safe training boundary.", flush=True) + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-dir", type=Path, required=True) + parser.add_argument("--background", action="store_true", + help="Detach from terminal and write stdout/stderr to background.log") + parser.add_argument("--gradient-steps", type=int, default=1_000_000) + parser.add_argument("--checkpoint-every", type=int, default=10_000) + parser.add_argument("--batch-size", type=int, default=256) + parser.add_argument("--training-seed", type=int, default=0) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--lr", type=float, default=1e-4) + parser.add_argument("--threads", type=int, default=1) + parser.add_argument("--paper-output", type=Path, default=ROOT / "paper/consistency_policy_update") + return parser.parse_args() + + +def launch_background(args): + args.run_dir.mkdir(parents=True, exist_ok=True) + pid_file = args.run_dir / "launcher.pid" + if pid_file.exists(): + previous_pid = int(pid_file.read_text().strip()) + proc_cmd = Path("/proc") / str(previous_pid) / "cmdline" + if proc_cmd.exists(): + command = proc_cmd.read_bytes().replace(b"\x00", b" ").decode(errors="replace") + if "run_locomotion.py" in command and str(args.run_dir) in command: + print(f"Already running: PID {previous_pid}; {args.run_dir / 'status.json'}") + return + command = [sys.executable, str(Path(__file__).resolve()), + "--run-dir", str(args.run_dir), "--paper-output", str(args.paper_output), + "--gradient-steps", str(args.gradient_steps), + "--checkpoint-every", str(args.checkpoint_every), + "--batch-size", str(args.batch_size), "--training-seed", str(args.training_seed), + "--device", args.device, "--lr", str(args.lr), "--threads", str(args.threads)] + environment = os.environ.copy() + environment.update(PYTHONUNBUFFERED="1", D4RL_SUPPRESS_IMPORT_ERROR="1", + OMP_NUM_THREADS=str(args.threads), MKL_NUM_THREADS=str(args.threads)) + with (args.run_dir / "background.log").open("ab", buffering=0) as log: + process = subprocess.Popen(command, cwd=str(ROOT), stdin=subprocess.DEVNULL, + stdout=log, stderr=subprocess.STDOUT, + env=environment, start_new_session=True, close_fds=True) + pid_file.write_text(str(process.pid) + "\n") + print(f"Launched detached PID {process.pid}") + print(f"Log: {args.run_dir / 'background.log'}") + print(f"Status: {args.run_dir / 'status.json'}") + + +def code_provenance(): + from tools.consistency_policy.evaluate import verify_source, DEFAULT_REPO + source = verify_source(DEFAULT_REPO) + names = ("run_locomotion.py", "locomotion_policy.py", "locomotion_eval.py", "update_locomotion_paper.py") + return {"official_source": source, + "adapter_sha256": {name: sha256(Path(__file__).parent / name) for name in names}} + + +def training_config(args, provenance): + return {"gradient_steps": args.gradient_steps, "batch_size": args.batch_size, + "training_seed": args.training_seed, "lr": args.lr, + "optimizer": "AdamW", "betas": [0.95, 0.999], "weight_decay": 1e-6, + "lr_schedule": "500-step linear warmup then cosine", + "student_dropout": 0.2, "teacher_weights": "model_ema", + "student_weights_at_evaluation": "model", "target_decay": 0.0, + "checkpoint_selection": "fixed final gradient step; no test-return selection", + "evaluation_seeds": [0, 1, 2, 3, 4], "num_envs": 50, + "num_rollout_groups": 10, "num_candidates": 64, "max_steps": 1000, + "sampling_steps": 1, "provenance": provenance} + + +def save_checkpoint(path, actor, optimizer, scheduler, step, normalizer, config, train_seconds): + import numpy as np + import torch + payload = { + "schema_version": 1, "method": "consistency_policy_d4rl_idql_adaptation", + "step": step, "model_state": actor.state_dict(), + "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(), + "normalizer_mean": normalizer.mean, "normalizer_std": normalizer.std, + "config": config, "policy_config": actor.config, + "train_seconds": train_seconds, + "rng": {"python": random.getstate(), "numpy": np.random.get_state(), + "torch": torch.get_rng_state(), + "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else []}, + } + temporary = path.with_suffix(".tmp") + torch.save(payload, temporary) + os.replace(temporary, path) + + +def train_task(args, env_name, config): + import d4rl + import gym + import numpy as np + import torch + from cleandiffuser.dataset.d4rl_mujoco_dataset import D4RLMuJoCoTDDataset + from tools.consistency_policy.locomotion_policy import LocomotionConsistencyPolicy + + task_dir = args.run_dir / "training" / env_name + task_dir.mkdir(parents=True, exist_ok=True) + teacher_path = ROOT / "checkpoints/final_paper_main" / env_name / "edm_latest.pt" + critic_path = ROOT / "checkpoints/final_paper_main" / env_name / "iql_latest.pt" + task_config = dict(config, env=env_name, teacher_checkpoint=str(teacher_path), + teacher_sha256=sha256(teacher_path), iql_checkpoint=str(critic_path), + iql_sha256=sha256(critic_path)) + update_status(args.run_dir, phase="loading_dataset", env=env_name, completed_steps=0) + random.seed(args.training_seed) + np.random.seed(args.training_seed) + torch.manual_seed(args.training_seed) + env = gym.make(env_name) + try: + dataset = D4RLMuJoCoTDDataset(d4rl.qlearning_dataset(env), True) + finally: + env.close() + normalizer = dataset.get_normalizer() + obs_dim, act_dim = dataset.o_dim, dataset.a_dim + observations = dataset.obs.to(args.device) + actions = dataset.act.to(args.device) + task_config["dataset_transitions"] = len(dataset) + task_config["obs_dim"], task_config["act_dim"] = obs_dim, act_dim + del dataset + actor = LocomotionConsistencyPolicy(obs_dim, act_dim, teacher_path, device=args.device, dropout=0.2) + optimizer = torch.optim.AdamW(actor.model.parameters(), lr=args.lr, betas=(0.95, 0.999), + weight_decay=1e-6) + + def lr_factor(step): + if step < 500: + return float(step) / 500 + progress = (step - 500) / max(1, args.gradient_steps - 500) + return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress))) + + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_factor) + latest_path = task_dir / "latest.pt" + final_path = task_dir / "final.pt" + resume_path = final_path if final_path.exists() else latest_path + start_step, train_seconds = 0, 0.0 + if resume_path.exists(): + checkpoint = torch.load(resume_path, map_location=args.device, weights_only=False) + if checkpoint["config"] != task_config: + raise ValueError("Checkpoint experiment/code provenance differs; use a new run directory") + if not (np.array_equal(checkpoint["normalizer_mean"], normalizer.mean) and + np.array_equal(checkpoint["normalizer_std"], normalizer.std)): + raise ValueError("Dataset normalizer changed since the checkpoint") + actor.load_state_dict(checkpoint["model_state"], strict=True) + optimizer.load_state_dict(checkpoint["optimizer"]) + scheduler.load_state_dict(checkpoint["scheduler"]) + start_step, train_seconds = checkpoint["step"], checkpoint["train_seconds"] + random.setstate(checkpoint["rng"]["python"]) + np.random.set_state(checkpoint["rng"]["numpy"]) + torch.set_rng_state(checkpoint["rng"]["torch"].cpu()) + if checkpoint["rng"]["cuda"]: + torch.cuda.set_rng_state_all([state.cpu() for state in checkpoint["rng"]["cuda"]]) + del checkpoint + print(f"Resuming {env_name} at update {start_step}", flush=True) + atomic_json(task_dir / "config.json", dict(task_config, policy_config=actor.config)) + actor.train() + clock_start = time.monotonic() + last_report, last_save = clock_start, clock_start + report_step = start_step + totals = {"ctm": 0.0, "dsm": 0.0} + with (task_dir / "training.jsonl").open("a", buffering=1) as stream: + for step in range(start_step + 1, args.gradient_steps + 1): + indices = torch.randint(observations.shape[0], (args.batch_size,), device=args.device) + losses = actor.compute_loss(actions[indices], observations[indices]) + total = losses["ctm"] + losses["dsm"] + if not torch.isfinite(total).item(): + raise FloatingPointError(f"Non-finite CP loss at update {step}") + optimizer.zero_grad(set_to_none=True) + total.backward() + optimizer.step() + scheduler.step() + actor.update_target() + for key in totals: + totals[key] += losses[key].item() + now = time.monotonic() + if step % 1000 == 0 or now - last_report >= 60 or step == args.gradient_steps: + n = step - report_step + rate = n / max(now - last_report, 1e-9) + record = {"env": env_name, "step": step, "target_steps": args.gradient_steps, + "ctm": totals["ctm"] / n, "dsm": totals["dsm"] / n, + "lr": scheduler.get_last_lr()[0], "updates_per_second": rate, + "remaining_train_hours_estimate": (args.gradient_steps - step) / rate / 3600, + "train_seconds": train_seconds + now - clock_start} + stream.write(json.dumps(record, allow_nan=False) + "\n") + print(json.dumps(record), flush=True) + update_status(args.run_dir, phase="training", env=env_name, completed_steps=step, + target_steps=args.gradient_steps, last_training_record=record) + totals = {"ctm": 0.0, "dsm": 0.0} + last_report, report_step = now, step + if step % args.checkpoint_every == 0 or now - last_save >= 600 or STOP_REQUESTED: + save_checkpoint(latest_path, actor, optimizer, scheduler, step, normalizer, + task_config, train_seconds + now - clock_start) + last_save = time.monotonic() + if STOP_REQUESTED: + raise InterruptedError("Stopped after saving resumable checkpoint") + if not final_path.exists(): + save_checkpoint(final_path, actor, optimizer, scheduler, args.gradient_steps, normalizer, + task_config, train_seconds + time.monotonic() - clock_start) + del observations, actions, optimizer, scheduler + if torch.cuda.is_available(): + torch.cuda.empty_cache() + actor.eval() + return actor, normalizer, final_path, task_config + + +def run(args): + import torch + from tools.consistency_policy.locomotion_eval import evaluate_actor + + args.run_dir.mkdir(parents=True, exist_ok=True) + with (args.run_dir / "run.lock").open("w") as lock: + try: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + raise RuntimeError("Another worker already holds this run directory") + try: + signal.signal(signal.SIGTERM, signal_stop) + signal.signal(signal.SIGINT, signal_stop) + signal.signal(signal.SIGHUP, signal.SIG_IGN) + torch.set_num_threads(args.threads) + provenance = code_provenance() + config = training_config(args, provenance) + config_path = args.run_dir / "run_config.json" + if config_path.exists() and json.loads(config_path.read_text()) != config: + raise ValueError("Existing run configuration differs; use a new directory") + atomic_json(config_path, config) + snapshot_dir = args.run_dir / "source" + snapshot_dir.mkdir(exist_ok=True) + for name in provenance["adapter_sha256"]: + target = snapshot_dir / name + if not target.exists(): + shutil.copy2(Path(__file__).parent / name, target) + update_status(args.run_dir, phase="starting", complete=False, training_seed=args.training_seed) + for env_name in ENVS: + if STOP_REQUESTED: + raise InterruptedError("Stop requested before next task") + actor, normalizer, checkpoint, task_config = train_task(args, env_name, config) + if STOP_REQUESTED: + raise InterruptedError("Stop requested before evaluation") + update_status(args.run_dir, phase="evaluating", env=env_name, + completed_steps=args.gradient_steps) + metadata = {"student_checkpoint": str(checkpoint), "student_sha256": sha256(checkpoint), + "teacher_checkpoint": task_config["teacher_checkpoint"], + "teacher_sha256": task_config["teacher_sha256"], + "training_seed": args.training_seed, "gradient_steps": args.gradient_steps, + "provenance": provenance, + "method_note": "Official CP CTM+DSM adapted to the pretrained D4RL IDQL MLP backbone"} + evaluate_actor(actor, env_name, args.device, seeds=list(range(5)), num_envs=50, + num_episodes=10, num_candidates=64, max_steps=1000, + output_dir=args.run_dir / "evaluation" / env_name, + iql_checkpoint=task_config["iql_checkpoint"], steps=1, + normalizer=normalizer, metadata=metadata) + if STOP_REQUESTED: + raise InterruptedError("Stop requested after evaluation") + del actor + if torch.cuda.is_available(): + torch.cuda.empty_cache() + update_status(args.run_dir, phase="updating_paper") + subprocess.run([sys.executable, str(Path(__file__).parent / "update_locomotion_paper.py"), + "--results-dir", str(args.run_dir / "evaluation"), + "--output-dir", str(args.paper_output)], check=True, cwd=str(ROOT)) + update_status(args.run_dir, phase="complete", complete=True, + paper_output=str(args.paper_output), completed_at=time.strftime("%Y-%m-%dT%H:%M:%S%z")) + print("Both five-seed evaluations and paper update completed.", flush=True) + except BaseException as error: + update_status(args.run_dir, phase="interrupted" if isinstance(error, InterruptedError) else "failed", + complete=False, error=str(error), traceback=traceback.format_exc()) + raise + + +def main(): + args = parse_args() + args.run_dir = args.run_dir.expanduser().resolve() + args.paper_output = args.paper_output.expanduser().resolve() + if min(args.gradient_steps, args.checkpoint_every, args.batch_size, args.threads) < 1 or args.lr <= 0: + raise ValueError("Training counts and learning rate must be positive") + if not 0 <= args.training_seed < 2**32: + raise ValueError("Training seed out of range") + if args.background: + launch_background(args) + else: + run(args) + + +if __name__ == "__main__": + main() diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/summarize.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/summarize.py new file mode 100644 index 0000000000000000000000000000000000000000..fc560c64f5c386198e83d254b7f22028cb18fdfc --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/summarize.py @@ -0,0 +1,378 @@ +#!/usr/bin/env python3 +r"""Aggregate measured official Consistency Policy comparisons across train seeds. + +Example: + python tools/consistency_policy/summarize.py \ + --input 'results/consistency_policy/**/*.json' --output results/cp_summary --plot + +Only evaluator result JSON files should be supplied. Evaluation-seed repeats are +not independent training runs and must be combined before using this tool. +""" + +import argparse +import csv +import glob +import hashlib +import json +import math +import re +import statistics +from pathlib import Path + + +METHOD_LABELS = {"cp": "Consistency Policy", "ddim": "DDIM", "edm": "EDM"} +LATENCY_KEYS = ("mean_ms", "p50_ms", "p95_ms") +FIELDNAMES = [ + "task", "method", "method_label", "weights", "sampling_steps", "sampling_steps_unit", "chain_mode", + "chaining_times", "n_train_seeds", "train_seeds", "eval_seeds", "score_metric", + "score_mean", "score_std", "latency_mean_ms_mean", "latency_mean_ms_std", + "latency_p50_ms_mean", "latency_p50_ms_std", "latency_p95_ms_mean", + "latency_p95_ms_std", "latency_samples", "latency_warmup_calls", "latency_scope", + "nfe_mean", "nfe_std", "nfe_min", "nfe_max", "comparison_id", "checkpoints", + "checkpoint_sha256s", "result_files", "source_json", "protocol_json", "hardware_json", + "architecture_json", "training_json", +] + + +def canonical(value): + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def finite_number(value, name, minimum=None): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + raise ValueError("{} must be a finite number".format(name)) + if minimum is not None and value < minimum: + raise ValueError("{} must be at least {}".format(name, minimum)) + return value + + +def integer(value, name, minimum=0): + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise ValueError("{} must be an integer >= {}".format(name, minimum)) + return value + + +def nonempty_string(value, name): + if not isinstance(value, str) or not value: + raise ValueError("{} must be a nonempty string".format(name)) + return value + + +def validate_result(result): + """Reject incomplete results rather than create plausible comparison rows.""" + if not isinstance(result, dict) or result.get("schema_version") != 1: + raise ValueError("expected evaluator JSON with schema_version=1") + if result.get("method") not in METHOD_LABELS: + raise ValueError("method must be cp, ddim, or edm") + for key in ("task", "checkpoint", "method_label", "score_metric"): + nonempty_string(result[key], key) + if result["score_metric"] not in ("mean_max_coverage", "mean_max_normalized_coverage", "success_rate"): + raise ValueError("unknown score_metric: {}".format(result["score_metric"])) + if result.get("weights") not in ("model", "ema_model"): + raise ValueError("weights must be model or ema_model") + integer(result["train_seed"], "train_seed") + integer(result["eval_seed"], "eval_seed") + finite_number(result["score"], "score", 0) + if result["score"] > 1: + raise ValueError("score must be in [0, 1]") + + sampling = result["sampling"] + integer(sampling["steps"], "sampling.steps", 1) + if sampling["chain_mode"] not in (None, "official", "bins"): + raise ValueError("sampling.chain_mode must be null, official, or bins") + if not isinstance(sampling["chaining_times"], list): + raise ValueError("sampling.chaining_times must be a list") + times = sampling["chaining_times"] + if times: + if times[0] not in ("D", "C"): + raise ValueError("sampling.chaining_times must begin with D or C") + for value in times[1:]: + finite_number(value, "sampling.chaining_times entry", 0) + if "steps_unit" in sampling and sampling["steps_unit"] not in ("bins", "denoiser_calls"): + raise ValueError("sampling.steps_unit must be bins or denoiser_calls") + if "checkpoint_sha256" in result and not re.fullmatch(r"[0-9a-fA-F]{64}", str(result["checkpoint_sha256"])): + raise ValueError("checkpoint_sha256 must contain 64 hexadecimal characters") + for key in ("architecture", "training"): + if key in result and not isinstance(result[key], dict): + raise ValueError(key + " must be an object") + + protocol = result["protocol"] + for key in ("horizon", "n_obs_steps", "n_action_steps", "n_test", "max_steps", "n_envs"): + integer(protocol[key], "protocol." + key, 1) + integer(protocol["test_start_seed"], "protocol.test_start_seed") + nonempty_string(protocol["dataset"], "protocol.dataset") + if not isinstance(protocol["abs_action"], bool): + raise ValueError("protocol.abs_action must be boolean") + if not isinstance(protocol["shape_meta"], dict) or not protocol["shape_meta"]: + raise ValueError("protocol.shape_meta must be a nonempty object") + if "crop_shape" not in protocol: + raise ValueError("missing protocol.crop_shape") + if result["task"].startswith("pusht") and not isinstance(protocol.get("legacy_test"), bool): + raise ValueError("PushT results require boolean protocol.legacy_test") + + hardware = result["hardware"] + for key in ("device", "device_name", "torch", "dtype"): + nonempty_string(hardware[key], "hardware." + key) + if "cuda" not in hardware: + raise ValueError("missing hardware.cuda") + latency = result["latency"] + nonempty_string(latency["scope"], "latency.scope") + integer(latency["samples"], "latency.samples", 1) + integer(latency["warmup_calls"], "latency.warmup_calls") + for key in LATENCY_KEYS: + finite_number(latency[key], "latency." + key, 0) + if latency["p95_ms"] < latency["p50_ms"]: + raise ValueError("latency.p95_ms must be >= latency.p50_ms") + + nfe = result["nfe"] + finite_number(nfe["mean"], "nfe.mean", 1) + integer(nfe["min"], "nfe.min", 1) + integer(nfe["max"], "nfe.max", 1) + if not nfe["min"] <= nfe["mean"] <= nfe["max"]: + raise ValueError("nfe must satisfy min <= mean <= max") + for key in ("repo", "commit"): + nonempty_string(result["source"][key], "source." + key) + canonical(result) # Reject NaN values hidden inside provenance metadata. + + +def resolve_inputs(patterns): + files = set() + for pattern in patterns: + matches = glob.glob(str(Path(pattern).expanduser()), recursive=True) + if not matches: + raise ValueError("input matched no files: {}".format(pattern)) + for match in matches: + path = Path(match) + candidates = path.rglob("*.json") if path.is_dir() else [path] + files.update(p.resolve() for p in candidates if p.is_file()) + if not files: + raise ValueError("no result JSON files found") + return sorted(files) + + +def load_results(paths): + records = [] + for path in paths: + try: + with Path(path).open() as stream: + result = json.load(stream) + validate_result(result) + except (KeyError, TypeError, ValueError, OSError) as error: + raise ValueError("{}: {}".format(path, error)) from error + records.append((str(path), result)) + return records + + +def comparison_metadata(result): + return { + "protocol": result["protocol"], + "hardware": result["hardware"], + "source": result["source"], + "score_metric": result["score_metric"], + "latency_scope": result["latency"]["scope"], + "latency_warmup_calls": result["latency"]["warmup_calls"], + } + + +def group_metadata(result): + # Teacher locations and output paths vary by seed without changing the model. + return { + "architecture": {key: value for key, value in result.get("architecture", {}).items() + if key not in ("teacher_path", "edm")}, + "training": {key: value for key, value in result.get("training", {}).items() + if key not in ("seed", "output_dir", "resume_path")}, + } + + +def mean_std(values): + return statistics.mean(values), statistics.stdev(values) if len(values) > 1 else None + + +def aggregate_results(records): + """Return rows with equal weight per independent training seed. + + All sampling settings may be compared within one task, but protocol, source, + hardware, metric, and latency measurement scope must match exactly. + """ + if not records: + raise ValueError("no result records supplied") + task_metadata = {} + independent_seeds = {} + independent_checkpoints = {} + labels = {} + groups = {} + group_provenance = {} + for path, result in records: + validate_result(result) + task = result["task"] + metadata = comparison_metadata(result) + if task in task_metadata: + previous, previous_path = task_metadata[task] + differences = [key for key in metadata if canonical(previous[key]) != canonical(metadata[key])] + if differences: + raise ValueError( + "incomparable {} results: {} differ between {} and {}".format( + task, ", ".join(differences), previous_path, path + ) + ) + else: + task_metadata[task] = (metadata, path) + label_key = (task, result["method"]) + if label_key in labels and labels[label_key] != result["method_label"]: + raise ValueError("inconsistent method_label for {}".format(label_key)) + labels[label_key] = result["method_label"] + group = (task, result["method"], canonical(result["sampling"]), result["weights"]) + provenance = group_metadata(result) + if group in group_provenance: + previous, previous_path = group_provenance[group] + differences = [key for key in provenance if canonical(previous[key]) != canonical(provenance[key])] + if differences: + raise ValueError("incomparable training runs: {} differ between {} and {}".format( + ", ".join(differences), previous_path, path + )) + else: + group_provenance[group] = (provenance, path) + seed_key = group + (result["train_seed"],) + if seed_key in independent_seeds: + raise ValueError( + "duplicate training seed {} for {}: {} and {}; evaluation repeats are not " + "independent training runs".format(result["train_seed"], group, independent_seeds[seed_key], path) + ) + checkpoint_identity = ("sha256", result["checkpoint_sha256"].lower()) if result.get("checkpoint_sha256") else ( + "path", str(Path(result["checkpoint"]).expanduser().resolve()) + ) + checkpoint_key = group + checkpoint_identity + if checkpoint_key in independent_checkpoints: + raise ValueError("same checkpoint cannot represent independent training seeds: {} and {}".format( + independent_checkpoints[checkpoint_key], path + )) + independent_seeds[seed_key] = path + independent_checkpoints[checkpoint_key] = path + groups.setdefault(group, []).append((path, result)) + + rows = [] + for group in sorted(groups): + runs = sorted(groups[group], key=lambda entry: entry[1]["train_seed"]) + first = runs[0][1] + results = [result for _, result in runs] + sampling = first["sampling"] + metadata = comparison_metadata(first) + row = { + "task": first["task"], "method": first["method"], "method_label": first["method_label"], + "weights": first["weights"], "sampling_steps": sampling["steps"], + "sampling_steps_unit": sampling.get("steps_unit", "bins" if first["method"] == "edm" else "denoiser_calls"), + "chain_mode": sampling["chain_mode"], "chaining_times": canonical(sampling["chaining_times"]), + "n_train_seeds": len(runs), "train_seeds": canonical([r["train_seed"] for r in results]), + "eval_seeds": canonical([r["eval_seed"] for r in results]), "score_metric": first["score_metric"], + "latency_samples": canonical([r["latency"]["samples"] for r in results]), + "latency_warmup_calls": first["latency"]["warmup_calls"], + "latency_scope": first["latency"]["scope"], + "nfe_min": min(r["nfe"]["min"] for r in results), + "nfe_max": max(r["nfe"]["max"] for r in results), + "comparison_id": hashlib.sha256(canonical(metadata).encode()).hexdigest()[:16], + "checkpoints": canonical([r["checkpoint"] for r in results]), + "checkpoint_sha256s": canonical([r.get("checkpoint_sha256") for r in results]), + "architecture_json": canonical([r.get("architecture") for r in results]), + "training_json": canonical([r.get("training") for r in results]), + "result_files": canonical([p for p, _ in runs]), + "source_json": canonical(first["source"]), "protocol_json": canonical(first["protocol"]), + "hardware_json": canonical(first["hardware"]), + } + row["score_mean"], row["score_std"] = mean_std([r["score"] for r in results]) + row["nfe_mean"], row["nfe_std"] = mean_std([r["nfe"]["mean"] for r in results]) + for key in LATENCY_KEYS: + row["latency_" + key + "_mean"], row["latency_" + key + "_std"] = mean_std([ + r["latency"][key] for r in results + ]) + rows.append(row) + return rows + + +def pareto_front(rows): + """Return nondominated measured means: lower latency, higher score.""" + return sorted([ + row for row in rows + if not any( + other["latency_mean_ms_mean"] <= row["latency_mean_ms_mean"] + and other["score_mean"] >= row["score_mean"] + and (other["latency_mean_ms_mean"] < row["latency_mean_ms_mean"] + or other["score_mean"] > row["score_mean"]) + for other in rows + ) + ], key=lambda row: row["latency_mean_ms_mean"]) + + +def plot_results(rows, output): + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + tasks = sorted({row["task"] for row in rows}) + paths = [] + for task in tasks: + task_rows = [row for row in rows if row["task"] == task] + fig, axis = plt.subplots(figsize=(8, 5)) + colors = {"cp": "#0072B2", "ddim": "#D55E00", "edm": "#009E73"} + markers = {"cp": "o", "ddim": "s", "edm": "^"} + for row in task_rows: + mode = " ({})".format(row["chain_mode"]) if row["chain_mode"] else "" + times = " times={}".format(row["chaining_times"]) if row["chaining_times"] != "[]" else "" + unit = "bin" if row["sampling_steps_unit"] == "bins" else "step" + label = "{}, {} {}{}{}{}, NFE {:g}, {}, {} train seed{}".format( + row["method_label"], row["sampling_steps"], unit, "s" if row["sampling_steps"] != 1 else "", + mode, times, row["nfe_mean"], row["weights"], row["n_train_seeds"], + "s" if row["n_train_seeds"] != 1 else "", + ) + axis.errorbar( + row["latency_mean_ms_mean"], row["score_mean"], + xerr=row["latency_mean_ms_std"], yerr=row["score_std"], + fmt=markers[row["method"]], color=colors[row["method"]], capsize=3, label=label, + ) + front = pareto_front(task_rows) + if len(front) > 1: + axis.plot([r["latency_mean_ms_mean"] for r in front], [r["score_mean"] for r in front], + "--", color="0.55", linewidth=1, label="Pareto frontier of measured means") + axis.set_xlabel("Mean policy latency (ms; {})".format(task_rows[0]["latency_scope"])) + axis.set_ylabel(task_rows[0]["score_metric"].replace("_", " ").capitalize()) + axis.set_title(task + " — error bars: sample SD across training seeds") + axis.set_ylim(0, 1.02) + axis.set_xlim(left=0) + axis.grid(alpha=0.2) + axis.legend(fontsize=8) + fig.tight_layout() + safe_task = re.sub(r"[^A-Za-z0-9_.-]", "_", task) + filename = "pareto.png" if len(tasks) == 1 else "pareto_{}_{}.png".format( + safe_task, hashlib.sha256(task.encode()).hexdigest()[:8] + ) + path = output / filename + fig.savefig(path, dpi=180) + plt.close(fig) + paths.append(path) + return paths + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--input", nargs="+", required=True, help="Result JSON paths, directories, or recursive globs") + parser.add_argument("--output", type=Path, required=True, help="Directory for comparison.csv and optional plots") + parser.add_argument("--plot", action="store_true", help="Write a Pareto plot per task (requires matplotlib)") + args = parser.parse_args(argv) + try: + rows = aggregate_results(load_results(resolve_inputs(args.input))) + args.output.mkdir(parents=True, exist_ok=True) + csv_path = args.output / "comparison.csv" + with csv_path.open("w", newline="") as stream: + writer = csv.DictWriter(stream, fieldnames=FIELDNAMES) + writer.writeheader() + writer.writerows(rows) + print(csv_path) + if args.plot: + for path in plot_results(rows, args.output): + print(path) + except (ValueError, OSError, ImportError) as error: + parser.error(str(error)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/update_locomotion_paper.py b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/update_locomotion_paper.py new file mode 100644 index 0000000000000000000000000000000000000000..28942b273db75bb51a0d8b45640bb9d58c5454a3 --- /dev/null +++ b/checkpoints/relay_kitchen_fdk_fw0_step80k_20260917/comparison_20260918_cp_eval3/source_original_verified/tools/consistency_policy/update_locomotion_paper.py @@ -0,0 +1,454 @@ +#!/usr/bin/env python3 +"""Update the editable locomotion paper only from ten complete CP evaluations. + +The original attachment is never modified. Existing generated files may be +updated only if their contents still match the last recorded generation; manual +edits are preserved by refusing to overwrite them. +""" + +from __future__ import annotations + +import argparse +from datetime import datetime, timezone +import hashlib +import json +import math +import os +from pathlib import Path +import re +import statistics +import tempfile + + +ROOT = Path(__file__).resolve().parents[2] +DEFAULT_SOURCE = Path("/home/zack/.codex/attachments/8d9f4c3d-1bf6-47ad-b5f7-1de0943da480/pasted-text.txt") +DEFAULT_OUTPUT = ROOT / "paper/consistency_policy_update" +GENERATOR = "consistency_policy_locomotion_paper_v1" +TASKS = {"walker2d-medium-expert-v2": "Walker2d", "halfcheetah-medium-expert-v2": "HalfCheetah"} +SEEDS = set(range(5)) +GENERATED = ("manuscript.tex", "table.tex", "notes.md") +CP_LABEL = r"Consistency Policy (D4RL adaptation)~\cite{prasad2024consistency}" + + +def sha256(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def number(value, name, minimum=None): + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + raise ValueError(f"{name} must be finite") + if minimum is not None and value < minimum: + raise ValueError(f"{name} must be >= {minimum}") + return float(value) + + +def validate_run(path): + """Recompute table statistics from raw arrays, rejecting smoke protocols.""" + import numpy as np + + record = json.loads(path.read_text()) + required = { + "schema_version": 1, "benchmark": "consistency_policy_locomotion_iql", + "method": "consistency_policy", "complete": True, "sampling_steps": 1, + "num_envs": 50, "num_episodes": 10, "num_candidates": 64, + "select": "argmax", "steps_cap": 1000, + } + for key, expected in required.items(): + if record.get(key) != expected: + raise ValueError(f"{key}: expected {expected!r}, got {record.get(key)!r}") + env = record["env"] + if env not in TASKS: + raise ValueError(f"Unexpected task: {env}") + seed = record["evaluation_seed"] + if isinstance(seed, bool) or not isinstance(seed, int) or seed not in SEEDS or record["seed"] != seed: + raise ValueError("Evaluation seed must be one of 0,1,2,3,4 and match seed") + if path.name != f"cp_seed{seed}.json": + raise ValueError("Result filename does not match evaluation seed") + if number(record["weight_temperature"], "weight_temperature") != 100.0: + raise ValueError("Expected weight_temperature=100") + provenance = record["provenance"] + for key in ("student", "teacher", "iql"): + checkpoint = provenance.get(key + "_checkpoint") + if not isinstance(checkpoint, str) or not checkpoint: + raise ValueError(f"Missing {key} checkpoint provenance") + for key in ("student_sha256", "teacher_sha256", "iql_sha256", "normalizer_sha256"): + if not re.fullmatch(r"[0-9a-f]{64}", str(provenance.get(key, ""))): + raise ValueError(f"Missing valid {key} provenance") + training_seed = provenance.get("training_seed") + if isinstance(training_seed, bool) or not isinstance(training_seed, int) or training_seed < 0: + raise ValueError("Missing training_seed provenance") + raw_path = path.with_suffix(".npz") + with np.load(raw_path, allow_pickle=False) as archive: + returns = np.asarray(archive["norm_returns"], dtype=np.float64) + times = np.asarray(archive["decision_times_ms"], dtype=np.float64) + if returns.shape != (10, 50) or not np.isfinite(returns).all(): + raise ValueError("Expected 10 groups x 50 finite normalized episode returns") + if times.ndim != 1 or times.size == 0 or not np.isfinite(times).all() or np.any(times <= 0): + raise ValueError("Expected finite positive raw decision times") + group_means = returns.mean(axis=1) + if not np.allclose(np.asarray(archive["per_ep_mean"]), group_means, rtol=1e-5, atol=1e-6): + raise ValueError("Saved group means disagree with raw episode returns") + if "episode_lengths" in archive: + lengths = np.asarray(archive["episode_lengths"]) + if lengths.shape != (10, 50) or np.any(lengths <= 0) or np.any(lengths > 1000): + raise ValueError("Invalid raw episode lengths") + computed = { + "normalized_return_mean_over_eps": float(returns.mean()), + "normalized_return_std_over_eps": float(group_means.std(ddof=0)), + "worst_episode_mean_return": float(group_means.min()), + "decision_time_mean_ms": float(times.mean()), + "decision_time_std_ms": float(times.std(ddof=0)), + "decision_time_p95_ms": float(np.percentile(times, 95)), + } + for key, expected in computed.items(): + observed = number(record[key], key) + if not math.isclose(observed, expected, rel_tol=1e-5, abs_tol=1e-6): + raise ValueError(f"{key} disagrees with raw arrays: {observed} vs {expected}") + computed["individual_episode_std"] = float(returns.std(ddof=0)) + computed["individual_episode_min"] = float(returns.min()) + if not isinstance(record.get("protocol"), dict) or not record["protocol"]: + raise ValueError("Missing evaluation protocol") + if not isinstance(record.get("hardware"), dict) or not record["hardware"]: + raise ValueError("Missing evaluation hardware") + return { + "env": env, "evaluation_seed": seed, "provenance": provenance, + "protocol": record["protocol"], "hardware": record["hardware"], + "metrics": computed, "json": str(path), "json_sha256": sha256(path), + "npz": str(raw_path), "npz_sha256": sha256(raw_path), + } + + +def collect_results(results_dir): + runs = {env: {} for env in TASKS} + errors = [] + for path in sorted(results_dir.rglob("cp_seed*.json")): + try: + run = validate_run(path) + if run["evaluation_seed"] in runs[run["env"]]: + raise ValueError(f"Duplicate evaluation seed {run['evaluation_seed']} for {run['env']}") + runs[run["env"]][run["evaluation_seed"]] = run + except (OSError, ValueError, KeyError, TypeError) as error: + errors.append(f"{path}: {error}") + summaries = {} + for env, task_runs in runs.items(): + if set(task_runs) != SEEDS: + continue + ordered = [task_runs[seed] for seed in sorted(SEEDS)] + first = ordered[0]["provenance"] + identity_keys = ("student_sha256", "teacher_sha256", "iql_sha256", "normalizer_sha256", "training_seed") + if any(any(run["provenance"][key] != first[key] for key in identity_keys) for run in ordered): + errors.append(f"{env}: evaluation seeds do not use one frozen student, teacher, IQL selector and normalizer") + continue + if any(run["protocol"] != ordered[0]["protocol"] or run["hardware"] != ordered[0]["hardware"] for run in ordered): + errors.append(f"{env}: evaluation protocol or hardware differs across seeds") + continue + values = [run["metrics"] for run in ordered] + means = [v["normalized_return_mean_over_eps"] for v in values] + latency = [v["decision_time_mean_ms"] for v in values] + summaries[env] = { + "return_mean": statistics.mean(means), "return_std_across_evaluation_seeds": statistics.pstdev(means), + "latency_mean_ms": statistics.mean(latency), + "latency_within_run_std_ms_mean": statistics.mean(v["decision_time_std_ms"] for v in values), + "sigma_group": statistics.mean(v["normalized_return_std_over_eps"] for v in values), + "worst_group_mean": min(v["worst_episode_mean_return"] for v in values), + "latency_std_across_evaluation_seeds_ms": statistics.pstdev(latency), + "individual_episode_std_mean": statistics.mean(v["individual_episode_std"] for v in values), + "minimum_evaluation_seed_mean": min(means), "provenance": first, + } + if set(summaries) == set(TASKS): + devices = [runs[env][0]["hardware"] for env in TASKS] + if devices[0] != devices[1]: + errors.append("Hardware differs between task evaluations") + complete = not errors and set(summaries) == set(TASKS) + return complete, summaries, runs, errors + + +def extract_table(source): + matches = [m for m in re.finditer(r"\\begin\{table\}.*?\\end\{table\}", source, re.DOTALL) + if r"\label{tab:sim_main_results}" in m.group()] + if len(matches) != 1: + raise ValueError("Expected exactly one sim_main_results table") + return matches[0] + + +def strip_bold(cell): + # Balance braces instead of assuming that a formatted numeric cell has no + # nested TeX commands. Other formatting and every digit remain untouched. + command = r"\textbf{" + while command in cell: + start = cell.index(command) + cursor = start + len(command) + depth, end = 1, cursor + while end < len(cell) and depth: + depth += (cell[end] == "{") - (cell[end] == "}") + end += 1 + if depth: + raise ValueError("Unbalanced textbf in table cell") + cell = cell[:start] + cell[cursor:end-1] + cell[end:] + return cell + + +def numeric_cell(cell): + match = re.search(r"[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?", strip_bold(cell)) + if not match: + raise ValueError(f"Expected a numeric table cell: {cell}") + return float(match.group()) + + +def update_table(table, summaries, complete): + table = table.replace(r"\multirow{7}", r"\multirow{8}") + table = table.replace(r"\sigma_{\mathrm{ep}}", r"\sigma_{\mathrm{group}}") + old_caption = r"Return and latency are reported as mean \(\pm\) std over five seeds. Bold entries denote the best result." + new_caption = ( + r"Return is mean \(\pm\) population standard deviation across five evaluation seeds; " + r"latency is the mean of run means \(\pm\) the mean within-run decision-time standard deviation. " + r"\(\sigma_{\mathrm{group}}\) measures variability of 50-environment group means. " + r"KDM/KDM-F use one candidate without selection; their daggered worst-case entries are minima of seed means. " + r"Other worst-case entries are minima of group means. Bold entries denote the best displayed value; " + r"protocol differences limit direct comparisons." + ) + if old_caption not in table: + raise ValueError("Original table caption changed unexpectedly") + table = table.replace(old_caption, new_caption) + block_pattern = r"(\\multirow\{8\}\{\*\}\{(Walker2d|HalfCheetah)\})(.*?)(?=\\midrule|\\bottomrule)" + row_pattern = re.compile(r"&(?P