Initial upload of BFN-hybrid + DDPM PushT-xarm policies

README.md

@@ -0,0 +1,71 @@
+---
+license: mit
+tags:
+- robotics
+- bfn
+- bayesian-flow-networks
+- diffusion-policy
+- pusht
+- xarm
+- hybrid-action
+---
+
+# PushT-xarm Real-Robot Policies (BFN-Hybrid vs DDPM-OneHot)
+
+Real-robot push-T policies trained on [borueihuang/pusht_xarm_merged](https://huggingface.co/datasets/borueihuang/pusht_xarm_merged).
+Both policies were trained on 144 episodes / 7835 frames at 30 Hz, top camera only, 200 epochs.
+
+## Action space
+
+- **Discrete:** 8 push directions (`action.direction` in {0..7})
+- **Continuous:** push distance (`action.distance` in `[0, 50]`)
+
+## Observation space
+
+- `camera_0`: top-view RGB image, 3x224x224, two-step history (`n_obs_steps=2`)
+
+## Policies
+
+| File | Method | Action treatment | Inference steps |
+|------|--------|------------------|-----------------|
+| `bfn/latest.ckpt`  | **BFN-Hybrid** (categorical + continuous Bayesian flow) | true hybrid | 20 |
+| `ddpm/latest.ckpt` | **DDPM** | one-hot continuous (9D) | 100 |
+
+## Quick start (BFN)
+
+```bash
+pip install -r requirements.txt
+python inference.py --ckpt bfn/latest.ckpt --config bfn/policy_config.yaml
+```
+
+Programmatic use:
+
+```python
+from inference import load_bfn_policy, infer_step
+
+policy = load_bfn_policy("bfn/latest.ckpt", "bfn/policy_config.yaml", "cuda")
+actions = infer_step(policy, cam0_now, cam0_prev, "cuda")
+# actions: List[{"direction": int 0..7, "distance": float 0..50}], len = n_action_steps (8)
+```
+
+## DDPM checkpoint
+
+The DDPM policy uses `diffusion_policy.policy.diffusion_unet_hybrid_image_policy.DiffusionUnetHybridImagePolicy`
+from the `diffusion-policy` library. Action is a 9D continuous vector: `[one_hot(8), distance]`.
+At inference time, take `argmax` of the first 8 dims for the direction, and the 9th dim for distance.
+
+## Files
+
+```
+bfn/
+  latest.ckpt
+  policy_config.yaml
+ddpm/
+  latest.ckpt
+  policy_config.yaml
+bfn_hybrid_image_policy.py   # standalone BFN policy class
+policies/base.py             # BasePolicy abstract class
+networks/base.py             # BFNetwork wrapper
+inference.py                 # example loader + inference
+requirements.txt
+```

bfn/latest.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd3e457f54d235da8d379d0e1771ffeb1a2d2e420b5fdd86212c1cdfb62777b7
+size 697951714

bfn/policy_config.yaml

@@ -0,0 +1,127 @@
+_target_: workspaces.train_bfn_workspace.TrainBFNWorkspace
+name: train_bfn_pusht_xarm_top
+task_name: pusht_xarm
+exp_name: bfn_hybrid_top
+shape_meta:
+  obs:
+    camera_0:
+      shape:
+      - 3
+      - 224
+      - 224
+      type: rgb
+  action:
+    shape:
+    - 2
+horizon: 16
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+dataset_obs_steps: 2
+past_action_visible: false
+obs_as_global_cond: true
+policy:
+  _target_: policies.bfn_hybrid_image_policy.BFNHybridImagePolicy
+  shape_meta: ${shape_meta}
+  horizon: ${horizon}
+  n_action_steps: ${n_action_steps}
+  n_obs_steps: ${n_obs_steps}
+  num_discrete_actions: 8
+  continuous_param_dim: 1
+  sigma_1: 0.001
+  beta_1: 0.2
+  n_timesteps: 20
+  crop_shape:
+  - 216
+  - 216
+  obs_encoder_group_norm: true
+  eval_fixed_crop: true
+  diffusion_step_embed_dim: 128
+  down_dims:
+  - 256
+  - 512
+  - 1024
+  kernel_size: 5
+  n_groups: 8
+  cond_predict_scale: true
+task:
+  dataset:
+    _target_: dataset.pusht_xarm_dataset.PushTXArmDataset
+    zarr_path: data/pusht_xarm_merged/replay.zarr
+    horizon: ${horizon}
+    pad_before: 1
+    pad_after: 7
+    n_obs_steps: ${n_obs_steps}
+    seed: 42
+    val_ratio: 0.1
+    cameras:
+    - camera_0
+    action_mode: hybrid
+  env_runner: null
+ema:
+  _target_: diffusion_policy.model.diffusion.ema_model.EMAModel
+  update_after_step: 0
+  inv_gamma: 1.0
+  power: 0.75
+  min_value: 0.0
+  max_value: 0.9999
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 0.0001
+  betas:
+  - 0.95
+  - 0.999
+  eps: 1.0e-08
+  weight_decay: 1.0e-06
+training:
+  device: cuda:0
+  seed: 42
+  debug: false
+  resume: false
+  lr_scheduler: cosine
+  lr_warmup_steps: 500
+  num_epochs: 200
+  gradient_accumulate_every: 1
+  use_ema: true
+  rollout_every: 50
+  checkpoint_every: 50
+  val_every: 1
+  sample_every: 5
+  max_train_steps: null
+  max_val_steps: null
+  tqdm_interval_sec: 1.0
+dataloader:
+  batch_size: 32
+  num_workers: 4
+  shuffle: true
+  pin_memory: true
+  persistent_workers: false
+val_dataloader:
+  batch_size: 32
+  num_workers: 2
+  shuffle: false
+  pin_memory: true
+  persistent_workers: false
+checkpoint:
+  topk:
+    monitor_key: train_loss
+    mode: min
+    k: 1
+    format_str: epoch={epoch:04d}-train_loss={train_loss:.4f}.ckpt
+  save_last_ckpt: true
+  save_last_snapshot: false
+logging:
+  project: pusht_xarm_bfn
+  resume: true
+  mode: offline
+  name: ${now:%Y.%m.%d-%H.%M.%S}_${name}
+  tags:
+  - bfn
+  - pusht_xarm
+  - hybrid
+  - top
+  id: null
+  group: null
+multi_run:
+  run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}
+  wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}

bfn_hybrid_image_policy.py

@@ -0,0 +1,356 @@
+"""BFN Hybrid Image Policy: image observations + categorical-discrete + continuous action heads.
+
+This is the policy for real-robot PushT with hybrid action space:
+- Discrete: 8 push directions
+- Continuous: push distance
+- Observation: one or more RGB cameras (cam0 top, cam1 side)
+"""
+from __future__ import annotations
+
+import math
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusion_policy.model.common.normalizer import LinearNormalizer
+from diffusion_policy.model.diffusion.conditional_unet1d import ConditionalUnet1D
+from diffusion_policy.common.pytorch_util import dict_apply
+
+from policies.base import BasePolicy
+from networks.base import BFNetwork
+
+try:
+    import robomimic.models.obs_core as rmbn
+    import diffusion_policy.model.vision.crop_randomizer as dmvc
+    from diffusion_policy.common.pytorch_util import replace_submodules
+    import robomimic.utils.obs_utils as ObsUtils
+    from robomimic.config import config_factory
+    from robomimic.algo import algo_factory, PolicyAlgo
+    from diffusion_policy.common.robomimic_config_util import get_robomimic_config
+    HAS_ROBOMIMIC = True
+except ImportError:
+    HAS_ROBOMIMIC = False
+
+
+__all__ = ["BFNHybridImagePolicy"]
+
+
+class HybridUnetWrapper(BFNetwork):
+    def __init__(self, model, horizon, continuous_dim, discrete_configs, cond_dim):
+        super().__init__(is_conditional_model=True)
+        self.model = model
+        self.horizon = horizon
+        self.continuous_dim = continuous_dim
+        self.discrete_configs = discrete_configs
+        self.cond_dim = cond_dim
+        self.cond_is_discrete = False
+        total_disc = sum(n for _, n in discrete_configs)
+        self.input_dim = continuous_dim + total_disc
+
+    def forward(self, x, t, cond=None):
+        B = x.shape[0]
+        x = x.view(B, self.horizon, self.input_dim)
+        out = self.model(x, t, global_cond=cond)
+        return out.reshape(B, -1)
+
+
+class BFNHybridImagePolicy(BasePolicy):
+    """BFN policy with image observations + hybrid (categorical + continuous) action head."""
+
+    def __init__(
+        self,
+        shape_meta: dict,
+        horizon: int = 16,
+        n_action_steps: int = 8,
+        n_obs_steps: int = 2,
+        num_discrete_actions: int = 8,
+        continuous_param_dim: int = 1,
+        sigma_1: float = 0.001,
+        beta_1: float = 0.2,
+        n_timesteps: int = 20,
+        crop_shape: tuple = (216, 216),
+        obs_encoder_group_norm: bool = True,
+        eval_fixed_crop: bool = True,
+        diffusion_step_embed_dim: int = 128,
+        down_dims: tuple = (256, 512, 1024),
+        kernel_size: int = 5,
+        n_groups: int = 8,
+        cond_predict_scale: bool = True,
+        device: str = "cpu",
+        dtype: str = "float32",
+        clip_actions: bool = True,
+        **kwargs,
+    ):
+        super().__init__(action_space=None, device=device, dtype=dtype, clip_actions=clip_actions)
+
+        self.horizon = horizon
+        self.n_action_steps = n_action_steps
+        self.n_obs_steps = n_obs_steps
+        self.num_discrete_actions = num_discrete_actions
+        self.continuous_dim = continuous_param_dim
+        self.discrete_configs = [(0, num_discrete_actions)]
+        self.discrete_action_indices = {0}
+        self.total_action_dim = 1 + continuous_param_dim
+        self.sigma_1 = sigma_1
+        self.beta_1 = beta_1
+        self.n_timesteps = n_timesteps
+
+        # Parse shape_meta for image obs keys
+        obs_shape_meta = shape_meta["obs"]
+        obs_config = {"low_dim": [], "rgb": [], "depth": [], "scan": []}
+        obs_key_shapes = {}
+        self.rgb_keys: List[str] = []
+        for key, attr in obs_shape_meta.items():
+            obs_key_shapes[key] = list(attr["shape"])
+            t = attr.get("type", "low_dim")
+            if t == "rgb":
+                obs_config["rgb"].append(key)
+                self.rgb_keys.append(key)
+            elif t == "low_dim":
+                obs_config["low_dim"].append(key)
+            else:
+                raise ValueError(f"Unsupported obs type: {t}")
+        assert HAS_ROBOMIMIC, "robomimic required for image policy"
+
+        self.obs_encoder = self._build_robomimic_encoder(
+            obs_config, obs_key_shapes, crop_shape, obs_encoder_group_norm, eval_fixed_crop
+        )
+        obs_feature_dim = self.obs_encoder.output_shape()[0]
+        global_cond_dim = obs_feature_dim * n_obs_steps
+
+        # U-Net input/output dim = continuous + discrete-logits
+        unet_dim = continuous_param_dim + num_discrete_actions
+
+        self.model = ConditionalUnet1D(
+            input_dim=unet_dim,
+            local_cond_dim=None,
+            global_cond_dim=global_cond_dim,
+            diffusion_step_embed_dim=diffusion_step_embed_dim,
+            down_dims=list(down_dims),
+            kernel_size=kernel_size,
+            n_groups=n_groups,
+            cond_predict_scale=cond_predict_scale,
+        )
+        self.unet_wrapper = HybridUnetWrapper(
+            model=self.model,
+            horizon=horizon,
+            continuous_dim=continuous_param_dim,
+            discrete_configs=self.discrete_configs,
+            cond_dim=global_cond_dim,
+        )
+        self.normalizer = LinearNormalizer()
+        self.global_cond_dim = global_cond_dim
+
+        print(f"BFN Hybrid Image Policy:")
+        print(f"  cameras: {self.rgb_keys}")
+        print(f"  discrete: {num_discrete_actions}, continuous: {continuous_param_dim}")
+        print(f"  obs_feature_dim: {obs_feature_dim}, global_cond_dim: {global_cond_dim}")
+        print(f"  U-Net params: {sum(p.numel() for p in self.model.parameters()):.2e}")
+        print(f"  Vision params: {sum(p.numel() for p in self.obs_encoder.parameters()):.2e}")
+
+    def _build_robomimic_encoder(self, obs_config, obs_key_shapes, crop_shape, group_norm, eval_fixed_crop):
+        config = get_robomimic_config(algo_name="bc_rnn", hdf5_type="image", task_name="square", dataset_type="ph")
+        with config.unlocked():
+            config.observation.modalities.obs = obs_config
+            if crop_shape is None:
+                for key, modality in config.observation.encoder.items():
+                    if modality.obs_randomizer_class == "CropRandomizer":
+                        modality["obs_randomizer_class"] = None
+            else:
+                ch, cw = crop_shape
+                for key, modality in config.observation.encoder.items():
+                    if modality.obs_randomizer_class == "CropRandomizer":
+                        modality.obs_randomizer_kwargs.crop_height = ch
+                        modality.obs_randomizer_kwargs.crop_width = cw
+        ObsUtils.initialize_obs_utils_with_config(config)
+        policy: PolicyAlgo = algo_factory(
+            algo_name=config.algo_name, config=config,
+            obs_key_shapes=obs_key_shapes, ac_dim=1 + self.num_discrete_actions, device="cpu",
+        )
+        obs_encoder = policy.nets["policy"].nets["encoder"].nets["obs"]
+        if group_norm:
+            replace_submodules(
+                root_module=obs_encoder,
+                predicate=lambda x: isinstance(x, nn.BatchNorm2d),
+                func=lambda x: nn.GroupNorm(num_groups=x.num_features // 16, num_channels=x.num_features),
+            )
+        if eval_fixed_crop:
+            replace_submodules(
+                root_module=obs_encoder,
+                predicate=lambda x: isinstance(x, rmbn.CropRandomizer),
+                func=lambda x: dmvc.CropRandomizer(
+                    input_shape=x.input_shape, crop_height=x.crop_height,
+                    crop_width=x.crop_width, num_crops=x.num_crops, pos_enc=x.pos_enc,
+                ),
+            )
+        return obs_encoder
+
+    def set_normalizer(self, normalizer: LinearNormalizer):
+        self.normalizer.load_state_dict(normalizer.state_dict())
+
+    def _encode_obs(self, nobs: Dict[str, torch.Tensor]) -> torch.Tensor:
+        """Encode obs dict to [B, global_cond_dim]."""
+        B = nobs[self.rgb_keys[0]].shape[0]
+        To = self.n_obs_steps
+        # Stack across time: build dict of [B*To, C, H, W]
+        flat = {}
+        for k, v in nobs.items():
+            v_t = v[:, :To]  # [B, To, C, H, W]
+            flat[k] = v_t.reshape(B * To, *v_t.shape[2:])
+        feats = self.obs_encoder(flat)  # [B*To, feat_dim]
+        feats = feats.reshape(B, To, -1)
+        return feats.reshape(B, -1)
+
+    def forward(self, obs, *, deterministic: bool = False, **kwargs):
+        if isinstance(obs, torch.Tensor):
+            obs = {"obs": obs}
+        return self.predict_action(obs)["action"]
+
+    def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        nobs = self.normalizer.normalize(obs_dict)
+        cond = self._encode_obs(nobs)
+        B = cond.shape[0]
+        device = cond.device
+        dtype = cond.dtype
+        naction = self._sample_hybrid_bfn(B, self.horizon, cond, device, dtype)
+        start = self.n_obs_steps - 1
+        end = start + self.n_action_steps
+        action = naction[:, start:end]
+        action_unnorm = action.clone()
+        if action.shape[-1] > 1:
+            full_unnorm = self.normalizer["action"].unnormalize(action.clone())
+            action_unnorm[:, :, 1:] = full_unnorm[:, :, 1:]
+        return {"action": action_unnorm, "action_pred": naction}
+
+    @torch.no_grad()
+    def _sample_hybrid_bfn(self, B, T, cond, device, dtype):
+        n_steps = self.n_timesteps
+        cont_dim = self.continuous_dim
+        disc_configs = self.discrete_configs
+
+        mu_cont = torch.zeros(B, T, cont_dim, device=device, dtype=dtype)
+        rho_cont = 1.0
+        theta_list = [
+            torch.full((B, T, n), 1.0 / n, device=device, dtype=dtype) for _, n in disc_configs
+        ]
+
+        for i in range(1, n_steps + 1):
+            t_val = (i - 1) / n_steps
+            t_batch = torch.full((B,), t_val, device=device, dtype=dtype)
+            net_input = torch.cat([mu_cont, *theta_list], dim=-1) if theta_list else mu_cont
+            out_flat = self.unet_wrapper(net_input.reshape(B, -1), t_batch, cond=cond)
+            out = out_flat.reshape(B, T, -1)
+
+            x_cont_pred = out[:, :, :cont_dim]
+            alpha_cont = (self.sigma_1 ** (-2.0 * i / n_steps)) * (1.0 - self.sigma_1 ** (2.0 / n_steps))
+            sender_std = 1.0 / (alpha_cont ** 0.5 + 1e-8)
+            y_cont = x_cont_pred + sender_std * torch.randn_like(x_cont_pred)
+            new_rho = rho_cont + alpha_cont
+            mu_cont = (rho_cont * mu_cont + alpha_cont * y_cont) / new_rho
+            rho_cont = new_rho
+
+            alpha_disc = self.beta_1 * (2 * i - 1) / (n_steps ** 2)
+            offset = cont_dim
+            new_theta_list = []
+            for j, (_, n_classes) in enumerate(disc_configs):
+                logits = out[:, :, offset:offset + n_classes]
+                probs = torch.softmax(logits, dim=-1)
+                probs_flat = probs.reshape(-1, n_classes)
+                k_samples = torch.multinomial(probs_flat, num_samples=1).squeeze(-1).reshape(B, T)
+                e_k = F.one_hot(k_samples, num_classes=n_classes).float()
+                y_mean = alpha_disc * (n_classes * e_k - 1)
+                y_std = (alpha_disc * n_classes + 1e-8) ** 0.5
+                y_disc = y_mean + y_std * torch.randn_like(y_mean)
+                log_theta = torch.log(theta_list[j] + 1e-8)
+                theta_new = torch.softmax(log_theta + y_disc, dim=-1)
+                new_theta_list.append(theta_new)
+                offset += n_classes
+            theta_list = new_theta_list
+
+        # Final
+        t_final = torch.ones(B, device=device, dtype=dtype)
+        net_input = torch.cat([mu_cont, *theta_list], dim=-1) if theta_list else mu_cont
+        out_final = self.unet_wrapper(net_input.reshape(B, -1), t_final, cond=cond).reshape(B, T, -1)
+        x_cont_final = out_final[:, :, :cont_dim].clamp(-1.0, 1.0)
+
+        disc_values = []
+        offset = cont_dim
+        for j, (_, n_classes) in enumerate(disc_configs):
+            logits = out_final[:, :, offset:offset + n_classes]
+            disc_values.append(logits.argmax(dim=-1).float().unsqueeze(-1))
+            offset += n_classes
+
+        if disc_values:
+            return torch.cat([torch.cat(disc_values, dim=-1), x_cont_final], dim=-1)
+        return x_cont_final
+
+    def compute_loss(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
+        nobs = self.normalizer.normalize(batch["obs"])
+        cond = self._encode_obs(nobs)
+        raw_action = batch["action"]
+        discrete_k = raw_action[:, :, 0].long()
+        naction = self.normalizer["action"].normalize(raw_action)
+        continuous_x = naction[:, :, 1:]
+
+        B = raw_action.shape[0]
+        T = self.horizon
+        device = raw_action.device
+        dtype = raw_action.dtype
+
+        t = torch.rand(B, device=device, dtype=dtype).clamp(min=1e-5, max=1.0 - 1e-5)
+        t_exp = t.view(B, 1, 1)
+        gamma = 1.0 - (self.sigma_1 ** (2.0 * t_exp))
+        var = gamma * (1.0 - gamma)
+        std = (var + 1e-8).sqrt()
+        mu_cont = gamma * continuous_x + std * torch.randn_like(continuous_x)
+
+        beta = self.beta_1 * t_exp.pow(2.0)
+        theta_list = []
+        disc_targets = []
+        for j, (_, n) in enumerate(self.discrete_configs):
+            d = discrete_k.clamp(0, n - 1)
+            disc_targets.append(d)
+            e_x = F.one_hot(d, num_classes=n).float()
+            mean = beta * (n * e_x - 1)
+            std_disc = (beta * n + 1e-8).sqrt()
+            y = mean + std_disc * torch.randn_like(mean)
+            theta_list.append(torch.softmax(y, dim=-1))
+
+        net_input = torch.cat([mu_cont, *theta_list], dim=-1) if theta_list else mu_cont
+        out_flat = self.unet_wrapper(net_input.reshape(B, -1), t, cond=cond)
+        out = out_flat.reshape(B, T, -1)
+
+        x_cont_pred = out[:, :, :self.continuous_dim]
+        cont_loss = (gamma * (continuous_x - x_cont_pred).pow(2.0)).mean()
+
+        disc_loss = 0.0
+        offset = self.continuous_dim
+        for j, (_, n) in enumerate(self.discrete_configs):
+            logits = out[:, :, offset:offset + n]
+            disc_loss = disc_loss + F.cross_entropy(
+                logits.reshape(-1, n), disc_targets[j].reshape(-1)
+            )
+            offset += n
+
+        return cont_loss + disc_loss
+
+    def state_dict(self):
+        return {
+            "obs_encoder": self.obs_encoder.state_dict(),
+            "model": self.model.state_dict(),
+            "normalizer": self.normalizer.state_dict(),
+        }
+
+    def load_state_dict(self, state_dict):
+        self.obs_encoder.load_state_dict(state_dict["obs_encoder"])
+        self.model.load_state_dict(state_dict["model"])
+        if "normalizer" in state_dict:
+            self.normalizer.load_state_dict(state_dict["normalizer"])
+
+    def set_actions(self, action: torch.Tensor):
+        pass
+
+    def reset(self):
+        pass

ddpm/latest.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af0076703b823c77145c56da3234220b9aa4336fac725ebccab39e6c559b3e73
+size 697958870

ddpm/policy_config.yaml

@@ -0,0 +1,133 @@
+_target_: workspaces.train_bfn_workspace.TrainBFNWorkspace
+name: train_ddpm_pusht_xarm_top
+task_name: pusht_xarm
+exp_name: ddpm_top
+shape_meta:
+  obs:
+    camera_0:
+      shape:
+      - 3
+      - 224
+      - 224
+      type: rgb
+  action:
+    shape:
+    - 9
+horizon: 16
+n_obs_steps: 2
+n_action_steps: 8
+n_latency_steps: 0
+dataset_obs_steps: 2
+past_action_visible: false
+obs_as_global_cond: true
+policy:
+  _target_: diffusion_policy.policy.diffusion_unet_hybrid_image_policy.DiffusionUnetHybridImagePolicy
+  shape_meta: ${shape_meta}
+  horizon: ${horizon}
+  n_action_steps: ${n_action_steps}
+  n_obs_steps: ${n_obs_steps}
+  noise_scheduler:
+    _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler
+    num_train_timesteps: 100
+    beta_start: 0.0001
+    beta_end: 0.02
+    beta_schedule: squaredcos_cap_v2
+    variance_type: fixed_small
+    clip_sample: true
+    prediction_type: epsilon
+  num_inference_steps: 100
+  obs_as_global_cond: ${obs_as_global_cond}
+  obs_encoder_group_norm: true
+  eval_fixed_crop: true
+  crop_shape:
+  - 216
+  - 216
+  diffusion_step_embed_dim: 128
+  down_dims:
+  - 256
+  - 512
+  - 1024
+  kernel_size: 5
+  n_groups: 8
+  cond_predict_scale: true
+task:
+  dataset:
+    _target_: dataset.pusht_xarm_dataset.PushTXArmDataset
+    zarr_path: data/pusht_xarm_merged/replay.zarr
+    horizon: ${horizon}
+    pad_before: 1
+    pad_after: 7
+    n_obs_steps: ${n_obs_steps}
+    seed: 42
+    val_ratio: 0.1
+    cameras:
+    - camera_0
+    action_mode: onehot
+  env_runner: null
+ema:
+  _target_: diffusion_policy.model.diffusion.ema_model.EMAModel
+  update_after_step: 0
+  inv_gamma: 1.0
+  power: 0.75
+  min_value: 0.0
+  max_value: 0.9999
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 0.0001
+  betas:
+  - 0.95
+  - 0.999
+  eps: 1.0e-08
+  weight_decay: 1.0e-06
+training:
+  device: cuda:0
+  seed: 42
+  debug: false
+  resume: false
+  lr_scheduler: cosine
+  lr_warmup_steps: 500
+  num_epochs: 200
+  gradient_accumulate_every: 1
+  use_ema: true
+  rollout_every: 50
+  checkpoint_every: 50
+  val_every: 1
+  sample_every: 5
+  max_train_steps: null
+  max_val_steps: null
+  tqdm_interval_sec: 1.0
+dataloader:
+  batch_size: 32
+  num_workers: 4
+  shuffle: true
+  pin_memory: true
+  persistent_workers: false
+val_dataloader:
+  batch_size: 32
+  num_workers: 2
+  shuffle: false
+  pin_memory: true
+  persistent_workers: false
+checkpoint:
+  topk:
+    monitor_key: train_loss
+    mode: min
+    k: 1
+    format_str: epoch={epoch:04d}-train_loss={train_loss:.4f}.ckpt
+  save_last_ckpt: true
+  save_last_snapshot: false
+logging:
+  project: pusht_xarm_ddpm
+  resume: true
+  mode: offline
+  name: ${now:%Y.%m.%d-%H.%M.%S}_${name}
+  tags:
+  - ddpm
+  - pusht_xarm
+  - onehot
+  - top
+  id: null
+  group: null
+multi_run:
+  run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}
+  wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}

inference.py

@@ -0,0 +1,100 @@
+"""Standalone real-robot inference for PushT-xarm BFN-hybrid policy.
+
+Inputs each step:
+    cam0: HxWx3 uint8 RGB image from the top camera (any size; will be resized to 224x224)
+    cam0_prev: same, one step earlier (n_obs_steps=2)
+
+Output per step (predicts horizon=16, returns next n_action_steps=8):
+    direction: int in {0..7}
+    distance:  float in [0, 50]
+"""
+import argparse
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from PIL import Image
+
+THIS_DIR = Path(__file__).resolve().parent
+sys.path.insert(0, str(THIS_DIR))
+
+from bfn_hybrid_image_policy import BFNHybridImagePolicy  # noqa: E402
+
+
+def load_bfn_policy(ckpt_path: str, config_path: str, device: str = "cuda"):
+    with open(config_path) as f:
+        cfg = yaml.safe_load(f)
+    pcfg = cfg["policy"]
+    policy = BFNHybridImagePolicy(
+        shape_meta=cfg["shape_meta"],
+        horizon=cfg["horizon"],
+        n_action_steps=cfg["n_action_steps"],
+        n_obs_steps=cfg["n_obs_steps"],
+        num_discrete_actions=pcfg.get("num_discrete_actions", 8),
+        continuous_param_dim=pcfg.get("continuous_param_dim", 1),
+        sigma_1=pcfg.get("sigma_1", 0.001),
+        beta_1=pcfg.get("beta_1", 0.2),
+        n_timesteps=pcfg.get("n_timesteps", 20),
+        crop_shape=tuple(pcfg.get("crop_shape", [216, 216])),
+        obs_encoder_group_norm=pcfg.get("obs_encoder_group_norm", True),
+        eval_fixed_crop=pcfg.get("eval_fixed_crop", True),
+        diffusion_step_embed_dim=pcfg.get("diffusion_step_embed_dim", 128),
+        down_dims=tuple(pcfg.get("down_dims", [256, 512, 1024])),
+        kernel_size=pcfg.get("kernel_size", 5),
+        n_groups=pcfg.get("n_groups", 8),
+        cond_predict_scale=pcfg.get("cond_predict_scale", True),
+    )
+    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+    state = ckpt["state_dicts"]["model"] if "state_dicts" in ckpt else ckpt
+    policy.load_state_dict(state)
+    policy.to(device).eval()
+    return policy
+
+
+def preprocess_image(img: np.ndarray) -> np.ndarray:
+    """Resize HxWx3 uint8 -> 3x224x224 float32 in [0,1]."""
+    if img.shape[:2] != (224, 224):
+        img = np.array(Image.fromarray(img).resize((224, 224), Image.BILINEAR))
+    img = img.astype(np.float32) / 255.0
+    return img.transpose(2, 0, 1)
+
+
+def infer_step(policy, cam0_now: np.ndarray, cam0_prev: np.ndarray, device: str = "cuda"):
+    """One inference call. Returns a list of dicts [{direction, distance}, ...] of length n_action_steps."""
+    a = preprocess_image(cam0_prev)
+    b = preprocess_image(cam0_now)
+    obs = torch.from_numpy(np.stack([a, b])).unsqueeze(0).to(device)  # [1, 2, 3, 224, 224]
+    with torch.no_grad():
+        out = policy.predict_action({"camera_0": obs})
+    actions = out["action"][0].cpu().numpy()  # [n_action_steps, 2] = [direction, distance]
+    return [
+        {"direction": int(round(a[0])) % 8, "distance": float(np.clip(a[1], 0, 50))}
+        for a in actions
+    ]
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--ckpt", required=True)
+    p.add_argument("--config", required=True)
+    p.add_argument("--device", default="cuda")
+    args = p.parse_args()
+
+    print(f"Loading policy from {args.ckpt}...")
+    policy = load_bfn_policy(args.ckpt, args.config, args.device)
+    print("Policy loaded.")
+
+    # Dummy roundtrip test
+    dummy = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
+    t0 = time.time()
+    actions = infer_step(policy, dummy, dummy, args.device)
+    dt = (time.time() - t0) * 1000
+    print(f"Smoke test: {len(actions)} actions in {dt:.1f} ms")
+    print(f"First action: {actions[0]}")
+
+
+if __name__ == "__main__":
+    main()

networks/base.py

@@ -0,0 +1,146 @@
+# Copyright (C) 2023 Maxime Robeyns <dev@maximerobeyns.com>
+#
+# 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.
+"""Base network classes for torch_bfn, as well as some utilities"""
+
+
+import math
+import torch as t
+import torch.nn as nn
+
+from abc import abstractmethod
+from typing import Optional
+from functools import partial
+
+try:
+    from torchtyping import TensorType as Tensor
+except ImportError:  # pragma: no cover - optional dependency
+
+    class _TensorAlias:
+        def __class_getitem__(cls, key):
+            return t.Tensor
+
+    Tensor = _TensorAlias
+
+__all__ = [
+    "BFNetwork",
+    "DiscreteBFNetwork",
+    "SinusoidalPosEmb",
+    "RandomOrLearnedSinusoidalPosEmb",
+]
+
+
+class BFNetwork(nn.Module):
+    """
+    Abstraact base class for neural networks (nn.Module) for use with
+    torch_bfn.
+    """
+
+    def __init__(self, is_conditional_model: bool = False):
+        super().__init__()
+        self.is_conditional_model = is_conditional_model
+
+    @abstractmethod
+    def forward(
+        self,
+        x: Tensor["B", "D"],
+        time: Tensor["B"],
+        cond: Optional[Tensor["B", "C"]] = None,
+        cond_drop_prob: Optional[float] = None,
+    ) -> Tensor["B", "D"]:
+        """Returns a value of the same shape as x (e.g. predicts the noise
+        applied to x) at time t with optional conditioning information.
+
+        Args:
+            x: the current parameter vector
+            time: current timestep
+            cond: conditioning information
+            cond_drop_prob: probability of dropping conditioning info out for
+                classifier-free guidance.
+
+        Returns:
+            Tensor["B", "D"]: updated parameter vector
+        """
+        raise NotImplementedError
+
+    def forward_with_cond_scale(
+        self, *args, cond_scale=1.0, rescaled_phi=0.0, **kwargs
+    ) -> Tensor["B", "D"]:
+        """For conditional sampling, this additionally scales the conditional
+        guidance, and sharpens phi.
+
+        This abstract class just invokes the forward method as a fallback.
+        """
+        logits = self.forward(*args, cond_drop_prob=0.0, **kwargs)
+        if cond_scale == 1.0:
+            return logits
+
+        null_logits = self.forward(*args, cond_drop_prob=1.0, **kwargs)
+        scaled_logits = null_logits + (logits - null_logits) * cond_scale
+
+        if rescaled_phi == 0.0:
+            return scaled_logits
+
+        std_fn = partial(t.std, dim=tuple(range(1, scaled_logits.ndim)), keepdim=True)
+        rescaled_logits = scaled_logits * (std_fn(logits) / std_fn(scaled_logits))
+        return rescaled_logits * rescaled_phi + scaled_logits * (1.0 - rescaled_phi)
+
+
+class DiscreteBFNetwork(BFNetwork):
+    """
+    For discrete variants of BFNs, we require networks to use the last tensor
+    dimemsion as the class label dimension, as is conventional in transformer
+    models for language.
+    """
+
+    @abstractmethod
+    def forward(
+        self,
+        x: Tensor["B", "D", "K"],
+        time: Tensor["B"],
+        cond: Optional[Tensor["B", "C"]],
+        cond_drop_prob: Optional[float] = None,
+    ) -> Tensor["B", "D", "K"]:
+        raise NotImplementedError
+
+
+class SinusoidalPosEmb(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x: t.Tensor) -> t.Tensor:
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = t.exp(t.arange(half_dim, device=x.device) * -emb)
+        emb = x[:, None] * emb[None, :]
+        emb = t.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class RandomOrLearnedSinusoidalPosEmb(nn.Module):
+    """
+    https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/models/danbooru_128.py#L8
+    """
+
+    def __init__(self, dim: int, is_random: bool = False):
+        super().__init__()
+        assert dim % 2 == 0, "Sinusoidal positional embedding dim must be even"
+        half_dim = dim // 2
+        self.weights = nn.Parameter(t.randn(half_dim), requires_grad=not is_random)
+
+    def forward(self, x: Tensor["B", 1]) -> Tensor["B", "dim+1"]:
+        freqs = x * self.weights[None, :] * 2 * math.pi
+        fouriered = t.cat((freqs.sin(), freqs.cos()), -1)
+        fouriered = t.cat((x, fouriered), -1)
+        return fouriered

policies/base.py

@@ -0,0 +1,242 @@
+"""Base class for all robotics policies.
+
+This module defines the abstract base class `BasePolicy`, which standardizes
+the interface for environment interaction (`act`), model inference (`forward`),
+training (`compute_loss`), and device/data management.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from utils.bfn_utils import str_to_torch_dtype
+
+__all__ = ["BasePolicy"]
+
+
+class BasePolicy(nn.Module, ABC):
+    """Abstract base class for robotics policies.
+
+    Provides shared functionality for:
+    1. Device and dtype management.
+    2. Data conversion (Numpy <-> Torch).
+    3. Automatic batch dimension handling.
+    4. Action clipping and normalization hooks.
+
+    Subclasses must implement:
+    - `forward(obs, ...)`: The core PyTorch inference logic.
+    - `compute_loss(batch)`: The training logic (optional, raises NotImplementedError by default).
+    """
+
+    def __init__(
+        self,
+        action_space: Any,
+        *,
+        device: str = "cpu",
+        dtype: str = "float32",
+        clip_actions: bool = True,
+        normalizer: Optional[Any] = None,
+    ):
+        """Initializes the BasePolicy.
+
+        Args:
+            action_space: The Gym action space (used for clipping).
+            device: Default device to place tensors on ('cpu', 'cuda').
+            dtype: Default dtype for tensor creation ('float32', 'float16').
+            clip_actions: Whether to clamp output actions to the action_space bounds.
+            normalizer: Optional normalization module (e.g., LinearNormalizer).
+        """
+        super().__init__()
+        self.action_space = action_space
+        self._device = torch.device(device)
+        self._dtype = str_to_torch_dtype(dtype)
+        self.clip_actions = clip_actions
+        self.normalizer = normalizer
+
+    # --- Properties ---
+
+    @property
+    def device(self) -> torch.device:
+        """Returns the current device of the policy.
+
+        Infers device from the first parameter if available (robust to .to() calls),
+        otherwise falls back to the initialization value.
+        """
+        try:
+            return next(self.parameters()).device
+        except StopIteration:
+            return self._device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """Returns the current dtype of the policy."""
+        try:
+            return next(self.parameters()).dtype
+        except StopIteration:
+            return self._dtype
+
+    # --- Public Interface ---
+
+    def set_normalizer(self, normalizer: Any) -> None:
+        """Updates the normalizer used by the policy."""
+        self.normalizer = normalizer
+
+    @abstractmethod
+    def forward(
+        self,
+        obs: Union[torch.Tensor, Dict[str, torch.Tensor]],
+        *,
+        deterministic: bool = False,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        """Core inference method returning a batch of actions.
+
+        Args:
+            obs: Observations with leading batch dimension [B, ...].
+            deterministic: Whether to sample deterministically (policy dependent).
+            **kwargs: Additional arguments (e.g., conditioning info).
+
+        Returns:
+            Action tensor of shape [B, ActionDim].
+        """
+        raise NotImplementedError
+
+    def compute_loss(self, batch: Any) -> torch.Tensor:
+        """Computes training loss for the policy.
+
+        Args:
+            batch: A batch of data (usually containing 'obs', 'action').
+
+        Returns:
+            Scalar loss tensor.
+
+        Raises:
+            NotImplementedError: If the policy does not support internal training logic.
+        """
+        raise NotImplementedError(
+            f"compute_loss is not implemented for {self.__class__.__name__}."
+        )
+
+    @torch.inference_mode()
+    def act(
+        self,
+        obs: Any,
+        *,
+        deterministic: bool = False,
+        return_torch: bool = False,
+        **kwargs: Any,
+    ) -> Union[np.ndarray, torch.Tensor]:
+        """Convenience wrapper for environment interaction.
+
+        Handles:
+        1. conversion of numpy obs -> torch tensors.
+        2. adding batch dimension if missing.
+        3. inference via `forward`.
+        4. clipping actions.
+        5. conversion of torch output -> numpy action (optional).
+
+        Args:
+            obs: Observation from the environment (Numpy array, Dict, or Tensor).
+            deterministic: Whether to use deterministic mode.
+            return_torch: If True, returns a Tensor on device; else returns Numpy array.
+            **kwargs: Passed to `forward`.
+
+        Returns:
+            Action (Numpy array or Tensor).
+        """
+        # 1. Convert to Tensor
+        obs_t = self._to_tensor(obs)
+
+        # 2. Auto-Batching
+        # We assume environment interaction usually provides a single unbatched observation
+        obs_t, batch_added = self._maybe_add_batch_dim(obs_t)
+
+        # 3. Inference
+        action = self.forward(obs_t, deterministic=deterministic, **kwargs)
+
+        # 4. Clipping
+        action = self._clip_actions(action)
+
+        # 5. Return
+        if return_torch:
+            return action
+
+        # Remove batch dim if we added it
+        if batch_added:
+            action = action.squeeze(0)
+
+        return action.detach().cpu().numpy()
+
+    # --- Internal Helpers ---
+
+    def _to_tensor(self, data: Any) -> Any:
+        """Recursively converts input data to tensors on the correct device/dtype."""
+        if isinstance(data, torch.Tensor):
+            return data.to(device=self.device, dtype=self.dtype)
+
+        if isinstance(data, Mapping):
+            return {k: self._to_tensor(v) for k, v in data.items()}
+
+        if isinstance(data, (list, tuple)):
+            return type(data)(self._to_tensor(v) for v in data)
+
+        # Fallback for numpy arrays / scalars
+        return torch.as_tensor(data, device=self.device, dtype=self.dtype)
+
+    def _maybe_add_batch_dim(self, obs: Any) -> Tuple[Any, bool]:
+        """Adds a leading batch dimension if the input appears to be unbatched.
+
+        Note: This uses a heuristic. If the input is a Tensor, we assume it is
+        unbatched if it matches the observation space shape (not implemented here generic enough)
+        OR we rely on the caller context (usually `act` is single-step).
+
+        Here, we unconditionally unsqueeze dim 0 for `act` convenience.
+        """
+        batch_added = False
+
+        if isinstance(obs, torch.Tensor):
+            # Heuristic: We assume `act` is called with single observations.
+            # For robust batch detection, one would check obs_space.shape.
+            # Here we simply unsqueeze to ensure [1, ...] shape.
+            obs = obs.unsqueeze(0)
+            batch_added = True
+
+        elif isinstance(obs, Mapping):
+            # Handle Dict inputs (e.g. {'image': ..., 'state': ...})
+            # Only unsqueeze tensors.
+            new_obs = {}
+            for k, v in obs.items():
+                if isinstance(v, torch.Tensor):
+                    new_obs[k] = v.unsqueeze(0)
+                    batch_added = True  # Mark true if ANY tensor was unsqueezed
+                else:
+                    new_obs[k] = v
+            obs = new_obs
+
+        return obs, batch_added
+
+    def _clip_actions(self, action: torch.Tensor) -> torch.Tensor:
+        """Clips actions to the environment bounds if `clip_actions` is True."""
+        if not self.clip_actions:
+            return action
+
+        # Check if action space has bounds
+        if not hasattr(self.action_space, "low") or not hasattr(
+            self.action_space, "high"
+        ):
+            return action
+
+        # Create tensor bounds on the fly (caching could be an optimization)
+        low = torch.as_tensor(
+            self.action_space.low, device=action.device, dtype=action.dtype
+        )
+        high = torch.as_tensor(
+            self.action_space.high, device=action.device, dtype=action.dtype
+        )
+
+        return torch.clamp(action, low, high)

requirements.txt

@@ -0,0 +1,12 @@
+torch>=2.0
+torchvision
+diffusers>=0.18
+einops
+numpy
+zarr
+pyyaml
+pillow
+hydra-core
+robomimic
+# diffusion-policy library is needed for DDPM/DDIM/EDM/Consistency policies.
+# Install from source: pip install git+https://github.com/real-stanford/diffusion_policy.git