Add Relay-Kitchen FDK student + DDPM/EDM teachers (fw0 step80k + comparison archive) (part 4)

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)

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)

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()

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

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

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

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

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

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

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)

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

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

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)

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

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

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,))

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

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"])
+        
+
+

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

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
+

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

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

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

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.")

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)

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]

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()

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

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

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()

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()

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

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()

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,
+)

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,
+)

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)),
+)

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)),
+)

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))

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))

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)
+        )

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

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))

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)

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

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
+            )

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,
+)

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

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)

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

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