zyworkflow 0.0.1__py3-none-any.whl

zyworkflow/policy/train_pick_policy.py

@@ -0,0 +1,834 @@
+import os
+import time
+import glob
+import torch
+import traceback
+import numpy as np
+import pandas as pd
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.multiprocessing as mp
+from PIL import Image
+from tqdm import tqdm
+from bnn.simulate import Simulate
+from torch.utils.data import Dataset, DataLoader
+from concurrent.futures import ThreadPoolExecutor
+from zyworkflow.utils.logger_config import setup_train_pick_policy_logger
+
+
+logger = setup_train_pick_policy_logger()
+GRIPPER_CLOSE_THRESH = 500.0
+SUCCESS_THRESH = 0.5
+COL_JOINTS = ["j1", "j2", "j3", "j4", "j5", "j6"]
+COL_GRIPPER = "Gripper_Set_Position(‰)"
+COL_SUCCESS = "success_flag"
+
+
+def build_success_targets_and_mask(
+    raw_success_flags: torch.Tensor,
+    t: int,
+    traj_len: int,
+    chunk_len: int,
+    pad_len: int,
+    mode: str,
+    thr: float,
+):
+    device = raw_success_flags.device
+    succ = (raw_success_flags > thr).float()
+
+    if mode == "within_horizon":
+        rev = torch.flip(succ, dims=[0])
+        rev_cum = torch.cumsum(rev, dim=0)
+        suffix_any = (torch.flip(rev_cum, dims=[0]) > 0).float()
+        c_s = suffix_any
+        c_sm = torch.ones_like(c_s)
+
+    elif mode == "terminal_only":
+        c_s = torch.zeros((chunk_len, 1), dtype=torch.float32, device=device)
+        c_sm = torch.zeros((chunk_len, 1), dtype=torch.float32, device=device)
+        final_idx = traj_len - 1
+        if t <= final_idx < (t + chunk_len):
+            off = final_idx - t
+            c_s[off, 0] = succ[off, 0]
+            c_sm[off, 0] = 1.0
+    else:
+        c_s = succ
+        c_sm = torch.ones_like(c_s)
+
+    if pad_len > 0:
+        if chunk_len > 0:
+            last_s = c_s[-1:].repeat(pad_len, 1)
+        else:
+            last_s = torch.zeros((pad_len, 1), device=device)
+
+        c_s = torch.cat([c_s, last_s], dim=0)
+        pad_zeros = torch.zeros((pad_len, 1), dtype=torch.float32, device=device)
+        c_sm = torch.cat([c_sm, pad_zeros], dim=0)
+
+    return c_s, c_sm
+
+
+class BNNWorker(mp.Process):
+    def __init__(self, pipe):
+        super().__init__()
+        self.pipe = pipe
+        self.bnn_instance = None
+
+    def run(self):
+        try:
+            os.environ["OMP_NUM_THREADS"] = "1"
+            torch.set_num_threads(1)
+
+            self.bnn_instance = Simulate()
+
+            while True:
+                cmd, data = self.pipe.recv()
+                if cmd == "STEP":
+                    inp = torch.tensor(data, dtype=torch.float32)
+                    with torch.no_grad():
+                        out = self.bnn_instance.run_simulation(inp)
+                    if isinstance(out, torch.Tensor):
+                        out = out.detach().cpu().numpy()
+                    self.pipe.send(out)
+                elif cmd == "RESET":
+                    if self.bnn_instance:
+                        self.bnn_instance.reset_state()
+                    self.pipe.send("OK")
+                elif cmd == "CLOSE":
+                    break
+        except Exception as e:
+            logger.error(f"BNN Worker Error: {e}")
+            logger.error(traceback.format_exc())
+            try:
+                self.pipe.send(None)
+            except:
+                pass
+
+
+class PersistentBNNPool:
+    def __init__(self, num_workers):
+        self.num_workers = num_workers
+        self.workers = []
+        self.pipes = []
+        logger.info(f"启动 {num_workers} 个 BNN 常驻进程...")
+        for _ in range(num_workers):
+            parent_conn, child_conn = mp.Pipe()
+            p = BNNWorker(child_conn)
+            p.daemon = True
+            p.start()
+            self.workers.append(p)
+            self.pipes.append(parent_conn)
+
+    def reset_all(self, n_used=None):
+        if n_used is None:
+            n_used = len(self.pipes)
+        for p in self.pipes[:n_used]:
+            p.send(("RESET", None))
+        for p in self.pipes[:n_used]:
+            p.recv()
+
+    def step_batch(self, batch_inputs_np):
+        n = len(batch_inputs_np)
+        for i, inp in enumerate(batch_inputs_np):
+            self.pipes[i].send(("STEP", inp))
+        return [self.pipes[i].recv() for i in range(n)]
+
+    def close(self, timeout: float = 2.0):
+        for p in self.pipes:
+            try:
+                p.send(("CLOSE", None))
+            except Exception:
+                pass
+
+        deadline = time.time() + timeout
+        for w in self.workers:
+            remaining = max(0.0, deadline - time.time())
+            try:
+                w.join(timeout=remaining)
+            except Exception:
+                pass
+
+        for w in self.workers:
+            if w.is_alive():
+                try:
+                    logger.warning(f"强制终止残留BNN进程: pid={w.pid}")
+                    w.terminate()
+                    w.join(timeout=1.0)
+                except Exception as e:
+                    logger.error(f"强制终止BNN进程失败: pid={w.pid}, err={e}")
+
+        for p in self.pipes:
+            try:
+                p.close()
+            except Exception:
+                pass
+
+
+class SingleViewRobotTrajectoryDataset(Dataset):
+    def __init__(self, root_dir, min_frames_per_traj=6, time_round=3, time_tol=1e-3, debug_max_bad=3):
+        self.root_dir = root_dir
+        self.trajectories = []
+
+        self.min_frames_per_traj = min_frames_per_traj
+        self.time_round = time_round
+        self.time_tol = time_tol
+
+        self.debug_stats = {
+            "root_dir_exists": os.path.exists(root_dir),
+            "traj_dirs_found": 0,
+            "traj_used": 0,
+            "skip_missing_csv": 0,
+            "skip_missing_imgdir": 0,
+            "skip_csv_read_error": 0,
+            "skip_missing_cols": 0,
+            "skip_no_images": 0,
+            "skip_no_parsable_images": 0,
+            "skip_no_matched": 0,
+            "skip_too_few_images": 0,
+            "total_frames_csv": 0,
+            "total_images_in_dir": 0,
+            "matched_frames": 0,
+        }
+        self.bad_examples = []
+
+        traj_dirs = sorted([os.path.join(root_dir, d) for d in os.listdir(root_dir) if d.startswith("traj_")])
+        self.debug_stats["traj_dirs_found"] = len(traj_dirs)
+        logger.info(f"Dataset: 正在扫描数据集: {root_dir} | traj_dirs={len(traj_dirs)}")
+
+        num_pos = 0
+        num_neg = 0
+        all_joints_list = []
+
+        for traj_path in tqdm(traj_dirs):
+            csv_path = os.path.join(traj_path, "actions.csv")
+            if not os.path.exists(csv_path):
+                self.debug_stats["skip_missing_csv"] += 1
+                continue
+
+            img_dir = None
+            for cand in ["images", "image"]:
+                p = os.path.join(traj_path, cand)
+                if os.path.isdir(p):
+                    img_dir = p
+                    break
+            if img_dir is None:
+                self.debug_stats["skip_missing_imgdir"] += 1
+                continue
+
+            try:
+                df = pd.read_csv(csv_path, header=0)
+            except Exception:
+                self.debug_stats["skip_csv_read_error"] += 1
+                continue
+
+            if "Time(s)" not in df.columns:
+                self.debug_stats["skip_missing_cols"] += 1
+                if len(self.bad_examples) < debug_max_bad:
+                    self.bad_examples.append({
+                        "traj": traj_path,
+                        "reason": "missing Time(s) col",
+                        "df_cols": list(df.columns)[:30],
+                    })
+                continue
+
+            need_cols = [COL_GRIPPER, COL_SUCCESS] + COL_JOINTS
+            if any(c not in df.columns for c in need_cols):
+                self.debug_stats["skip_missing_cols"] += 1
+                if len(self.bad_examples) < debug_max_bad:
+                    self.bad_examples.append({
+                        "traj": traj_path,
+                        "reason": "missing required cols",
+                        "need_cols": need_cols,
+                        "df_cols": list(df.columns),
+                    })
+                continue
+
+            df = df.sort_values("Time(s)").reset_index(drop=True)
+            T = len(df)
+            self.debug_stats["total_frames_csv"] += T
+
+            img_files = glob.glob(os.path.join(img_dir, "*.png"))
+            self.debug_stats["total_images_in_dir"] += len(img_files)
+            if len(img_files) == 0:
+                self.debug_stats["skip_no_images"] += 1
+                continue
+
+            img_entries = []
+            img_name_examples = []
+            for fp in img_files:
+                stem = os.path.splitext(os.path.basename(fp))[0]
+                if len(img_name_examples) < 8:
+                    img_name_examples.append(stem)
+                try:
+                    tf = float(stem)
+                    img_entries.append((tf, fp))
+                except Exception:
+                    pass
+
+            if len(img_entries) == 0:
+                self.debug_stats["skip_no_parsable_images"] += 1
+                if len(self.bad_examples) < debug_max_bad:
+                    self.bad_examples.append({
+                        "traj": traj_path,
+                        "reason": "no parsable image filenames (stem->float failed)",
+                        "img_dir": img_dir,
+                        "img_stems_sample": img_name_examples,
+                    })
+                continue
+
+            img_entries.sort(key=lambda x: x[0])
+            img_times = np.array([x[0] for x in img_entries], dtype=np.float64)
+            img_paths = [x[1] for x in img_entries]
+
+            img_map = {}
+            for tf, fp in img_entries:
+                k = round(float(tf), self.time_round)
+                if k not in img_map:
+                    img_map[k] = fp
+
+            joints_np = df[COL_JOINTS].to_numpy(dtype=np.float32)
+            gripper_np = df[COL_GRIPPER].to_numpy(dtype=np.float32)
+            success_np = df[COL_SUCCESS].to_numpy(dtype=np.float32)
+            times_np = df["Time(s)"].to_numpy(dtype=np.float64)
+
+            valid_img_paths = []
+            valid_targets = []
+
+            for i in range(T):
+                t_csv = float(times_np[i])
+                key = round(t_csv, self.time_round)
+
+                fp = img_map.get(key, None)
+
+                if fp is None:
+                    idx = int(np.searchsorted(img_times, t_csv))
+                    cand = []
+                    if 0 <= idx < len(img_times):
+                        cand.append(idx)
+                    if 0 <= idx - 1 < len(img_times):
+                        cand.append(idx - 1)
+                    best_fp = None
+                    best_dt = 1e9
+                    for ci in cand:
+                        dt = abs(float(img_times[ci]) - t_csv)
+                        if dt < best_dt:
+                            best_dt = dt
+                            best_fp = img_paths[ci]
+                    if best_fp is not None and best_dt <= self.time_tol:
+                        fp = best_fp
+
+                if fp is None or (not os.path.exists(fp)):
+                    continue
+
+                joints_val = joints_np[i]
+                gripper_val = float(gripper_np[i])
+                success_val = float(success_np[i])
+
+                target_vec = np.concatenate([joints_val, [gripper_val], [success_val]]).astype(np.float32)
+                valid_img_paths.append(fp)
+                valid_targets.append(target_vec)
+
+                if gripper_val < GRIPPER_CLOSE_THRESH:
+                    num_pos += 1
+                else:
+                    num_neg += 1
+
+            self.debug_stats["matched_frames"] += len(valid_img_paths)
+
+            if len(valid_img_paths) == 0:
+                self.debug_stats["skip_no_matched"] += 1
+                if len(self.bad_examples) < debug_max_bad:
+                    self.bad_examples.append({
+                        "traj": traj_path,
+                        "reason": "matched_frames=0 (Time(s) vs img filename mismatch)",
+                        "csv_times_sample": [float(x) for x in times_np[:8]],
+                        "img_stems_sample": img_name_examples,
+                        "time_round": self.time_round,
+                        "time_tol": self.time_tol,
+                    })
+                continue
+
+            if len(valid_img_paths) < self.min_frames_per_traj:
+                self.debug_stats["skip_too_few_images"] += 1
+                if len(self.bad_examples) < debug_max_bad:
+                    self.bad_examples.append({
+                        "traj": traj_path,
+                        "reason": f"too few matched frames (<{self.min_frames_per_traj})",
+                        "matched": len(valid_img_paths),
+                        "csv_len": T,
+                        "images_in_dir": len(img_files),
+                    })
+                continue
+
+            try:
+                with Image.open(valid_img_paths[0]) as im:
+                    exp_hw = (im.height, im.width)
+            except Exception:
+                exp_hw = (480, 640)
+
+            self.trajectories.append({
+                "traj_id": traj_path,
+                "view_paths": valid_img_paths,
+                "targets": np.array(valid_targets, dtype=np.float32),
+                "length": len(valid_img_paths),
+                "exp_hw": exp_hw
+            })
+            self.debug_stats["traj_used"] += 1
+
+            all_joints_list.append(np.array(valid_targets, dtype=np.float32)[:, :6])
+
+        if len(all_joints_list) > 0:
+            all_joints_np = np.concatenate(all_joints_list, axis=0)
+            self.joint_mean = torch.tensor(np.mean(all_joints_np, axis=0), dtype=torch.float32)
+            self.joint_std = torch.tensor(np.std(all_joints_np, axis=0), dtype=torch.float32)
+            self.joint_std = torch.where(self.joint_std < 1e-6, torch.ones_like(self.joint_std), self.joint_std)
+        else:
+            self.joint_mean = torch.zeros(6)
+            self.joint_std = torch.ones(6)
+
+        self.pos_weight = (num_neg / max(num_pos, 1)) if (num_pos + num_neg) > 0 else 1.0
+
+        logger.success(f"数据加载完毕: {len(self.trajectories)} 条轨迹.")
+        logger.info(f"PosWeight={self.pos_weight:.2f}")
+        logger.info(f"Joint Mean: {self.joint_mean.numpy().round(3)}")
+
+        logger.debug("Dataset build summary:")
+        for k, v in self.debug_stats.items():
+            logger.debug(f"  - {k}: {v}")
+        if len(self.bad_examples) > 0:
+            logger.warning(f"发现 {len(self.bad_examples)} 个错误示例:")
+            for ex in self.bad_examples:
+                for kk, vv in ex.items():
+                    logger.debug(f"{kk}: {vv}")
+
+    def __len__(self):
+        return len(self.trajectories)
+
+    def __getitem__(self, idx):
+        return self.trajectories[idx]
+
+
+def traj_collate_fn(batch):
+    return batch
+
+
+class FiLM(nn.Module):
+    def __init__(self, dim_in, dim_out, num_layers=2, hidden_dim=128):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(dim_in, hidden_dim), nn.ReLU(),
+            nn.Linear(hidden_dim, dim_out * 2)
+        )
+
+    def forward(self, x, cond):
+        cond_flat = cond.reshape(-1, cond.size(-1))
+        params = self.mlp(cond_flat)
+        gamma, beta = params.chunk(2, dim=-1)
+        B, S, D = x.shape
+        gamma = gamma.view(B, S, D)
+        beta = beta.view(B, S, D)
+        return gamma * x + beta
+
+
+class SingleViewBNNActionPolicy(nn.Module):
+    def __init__(self, seq_len=4, action_chunk=8, dim_to_bnn=15, dim_bnn_output=80):
+        super().__init__()
+        self.seq_len = seq_len
+        self.action_chunk = action_chunk
+
+        self.conv_layers = nn.Sequential(
+            nn.Conv2d(3, 24, kernel_size=5, stride=2, padding=2), nn.ReLU(),
+            nn.Conv2d(24, 36, kernel_size=5, stride=2, padding=2), nn.ReLU(),
+            nn.Conv2d(36, 48, kernel_size=3, stride=2, padding=1), nn.ReLU(),
+            nn.Conv2d(48, 64, kernel_size=3, stride=1, padding=1), nn.ReLU(),
+            nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1), nn.ReLU()
+        )
+
+        self.before_bnn_mlp = nn.Sequential(
+            nn.Linear(80 * 60, 512), nn.ReLU(),
+            nn.Linear(512, dim_to_bnn)
+        )
+        self.bnn_adapter = nn.Sequential(nn.LayerNorm(dim_to_bnn), nn.Tanh())
+
+        self.film = FiLM(dim_in=dim_to_bnn, dim_out=dim_bnn_output)
+
+        feature_dim = (dim_to_bnn + dim_bnn_output) * seq_len
+        self.shared_backbone = nn.Sequential(
+            nn.Linear(feature_dim, 512), nn.ReLU(),
+            nn.Linear(512, 256)
+        )
+
+        self.head_joints = nn.Sequential(
+            nn.Linear(256, 128), nn.ReLU(),
+            nn.Linear(128, 64), nn.ReLU(),
+            nn.Linear(64, 6 * action_chunk)
+        )
+        self.head_gripper = nn.Sequential(
+            nn.Linear(256, 64), nn.ReLU(),
+            nn.Linear(64, 1 * action_chunk)
+        )
+        self.head_success = nn.Sequential(
+            nn.Linear(256, 64), nn.ReLU(),
+            nn.Linear(64, 1 * action_chunk)
+        )
+
+    def encode_visual(self, x):
+        B, S, C, H, W = x.shape
+        x = x.contiguous().view(B * S, C, H, W)
+        x = self.conv_layers(x)
+        x = x.view(B, S, -1)
+        feat = self.before_bnn_mlp(x)
+        return self.bnn_adapter(feat)
+
+    def decode_action(self, bnn_in, bnn_out):
+        bnn_out_film = self.film(bnn_out, bnn_in)
+        bnn_out_film = F.layer_norm(bnn_out_film, (bnn_out_film.size(-1),))
+        feats = torch.cat([bnn_in, bnn_out_film], dim=-1).view(bnn_in.size(0), -1)
+        shared = self.shared_backbone(feats)
+
+        j = self.head_joints(shared).view(-1, self.action_chunk, 6)
+        g = self.head_gripper(shared).view(-1, self.action_chunk, 1)
+        s = self.head_success(shared).view(-1, self.action_chunk, 1)
+        return j, g, s
+
+
+def _load_single_window(args):
+    v_paths, t, seq_len, exp_hw = args
+    target_w, target_h = 640, 480
+
+    v_list = []
+    start_idx = t - seq_len + 1
+
+    for i in range(start_idx, t + 1):
+        idx = i if i >= 0 else 0
+        try:
+            with Image.open(v_paths[idx]) as img:
+                arr = np.array(img.resize((target_w, target_h)), dtype=np.float32) / 255.0
+                v_list.append(arr.transpose(2, 0, 1))
+        except Exception:
+            v_list.append(np.zeros((3, target_h, target_w), dtype=np.float32))
+
+    return np.stack(v_list)
+
+
+def train_single_view_parallel_chunk(
+    task_name,
+    root_dir,
+    batch_size=50,
+    seq_len=4,
+    action_chunk=8,
+    lr=1e-4,
+    num_epochs=500,
+    start_epoch=0,
+    lambda_joints=10.0, lambda_grip=5.0, lambda_success=2.0,
+    log_path=None, ckpt_dir=None, success_mode="within_horizon",
+    report_url=None,
+):
+    logger.info(f"启动单视角训练 | BS={batch_size} | Chunk={action_chunk} | Mode={success_mode}")
+
+    bnn_pool = None
+    io_pool = None
+    device = None
+
+    last_saved_ckpt_path = None
+    try:
+        if log_path:
+            os.makedirs(os.path.dirname(log_path), exist_ok=True)
+        os.makedirs(ckpt_dir, exist_ok=True)
+
+        dataset = SingleViewRobotTrajectoryDataset(
+            root_dir=root_dir,
+            min_frames_per_traj=6,
+            time_round=3,
+            time_tol=1e-3,
+            debug_max_bad=3
+        )
+
+        if len(dataset) == 0:
+            raise RuntimeError(
+                "Dataset size = 0，训练无法开始。\n"
+                "请看上面 [DEBUG] 的 bad examples：基本都是 Time(s) 与图片文件名不匹配导致 matched_frames=0。\n"
+            )
+
+        loader = DataLoader(dataset, batch_size=batch_size, shuffle=True,
+                            collate_fn=traj_collate_fn, drop_last=True)
+
+        bnn_pool = PersistentBNNPool(num_workers=batch_size)
+        io_pool = ThreadPoolExecutor(max_workers=min(64, batch_size * 2))
+
+        device = torch.device("cuda")
+
+        joint_mean_gpu = dataset.joint_mean.to(device)
+        joint_std_gpu = dataset.joint_std.to(device)
+
+        model = SingleViewBNNActionPolicy(seq_len, action_chunk).to(device)
+        optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
+
+        crit_mse = nn.MSELoss(reduction='none')
+        crit_bce_w = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([dataset.pos_weight], device=device), reduction='none')
+        crit_bce = nn.BCEWithLogitsLoss(reduction='none')
+
+        if start_epoch > 0:
+            p = os.path.join(ckpt_dir, f"epoch_{start_epoch}.pth")
+            if os.path.exists(p):
+                checkpoint = torch.load(p, map_location=device)
+                if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+                    model.load_state_dict(checkpoint['model_state_dict'])
+                    logger.info(f"Loaded Checkpoint Epoch {checkpoint.get('epoch')}")
+                else:
+                    model.load_state_dict(checkpoint)
+                    logger.info("Loaded Legacy State Dict")
+
+        update_freq = 5
+
+        for epoch in range(start_epoch, num_epochs):
+            epoch_start_time = time.time()
+            model.train()
+            total_loss, steps = 0.0, 0
+            real_joint_err_accum = 0.0
+
+            # pbar = tqdm(loader, desc=f"Ep {epoch+1}")
+            for batch_trajs in loader:
+                curr_bs = len(batch_trajs)
+                bnn_pool.reset_all(curr_bs)
+
+                optimizer.zero_grad(set_to_none=True)
+                accum_steps = 0
+
+                lengths = [t["length"] for t in batch_trajs]
+                max_len = max(lengths)
+                exp_hw = batch_trajs[0]["exp_hw"]
+
+                bnn_hist = torch.zeros(curr_bs, max_len, 80, device=device)
+
+                for t in range(max_len):
+
+                    futures = [
+                        io_pool.submit(_load_single_window, (tr["view_paths"], t, seq_len, exp_hw))
+                        if t < tr["length"] else None
+                        for tr in batch_trajs
+                    ]
+
+                    imgs, masks = [], []
+                    for f in futures:
+                        if f:
+                            v = f.result()
+                            imgs.append(torch.from_numpy(v))
+                            masks.append(True)
+                        else:
+                            z = torch.zeros(seq_len, 3, 480, 640)
+                            imgs.append(z)
+                            masks.append(False)
+
+                    b_imgs = torch.stack(imgs).to(device, non_blocking=True)
+                    bnn_in = model.encode_visual(b_imgs)
+
+                    curr_feat = bnn_in[:, -1, :].detach().cpu().numpy() / 10.0
+                    bnn_outs = bnn_pool.step_batch([curr_feat[b].reshape(1, -1) for b in range(curr_bs)])
+
+                    clean_outs = [
+                        o.T if o is not None and hasattr(o, "shape") and o.shape == (1, 80)
+                        else (o if o is not None else np.zeros(80))
+                        for o in bnn_outs
+                    ]
+                    clean_outs = [np.array(o).squeeze() for o in clean_outs]
+                    bnn_curr = torch.tensor(np.stack(clean_outs), device=device, dtype=torch.float32)
+                    bnn_hist[:, t, :] = bnn_curr
+
+                    s_idx = t - seq_len + 1
+                    if s_idx >= 0:
+                        bnn_seq = bnn_hist[:, s_idx:t+1]
+                    else:
+                        first = bnn_hist[:, 0:1, :].repeat(1, -s_idx, 1)
+                        bnn_seq = torch.cat([first, bnn_hist[:, 0:t+1]], dim=1)
+
+                    p_j, p_g, p_s = model.decode_action(bnn_in, bnn_seq)
+
+                    t_j_list, t_g_list, t_s_list, m_list = [], [], [], []
+                    for b_idx in range(curr_bs):
+                        if not masks[b_idx]:
+                            z = torch.zeros(action_chunk, 1, device=device)
+                            t_j_list.append(torch.zeros(action_chunk, 6, device=device))
+                            t_g_list.append(z)
+                            t_s_list.append(z)
+                            m_list.append(z)
+                            continue
+
+                        traj = batch_trajs[b_idx]
+                        real_end = min(t + action_chunk, traj["length"])
+                        chunk_sz = real_end - t
+                        pad_sz = action_chunk - chunk_sz
+
+                        raw = torch.from_numpy(traj["targets"][t:real_end]).to(device)
+                        c_j = raw[:, 0:6]
+                        c_g = (raw[:, 6:7] < GRIPPER_CLOSE_THRESH).float()
+                        c_s, _ = build_success_targets_and_mask(
+                            raw[:, 7:8], t, traj["length"], chunk_sz, pad_sz, success_mode, SUCCESS_THRESH
+                        )
+                        c_s = c_s.to(device)
+
+                        if pad_sz > 0:
+                            c_j = torch.cat([c_j, c_j[-1:].repeat(pad_sz, 1)], 0)
+                            c_g = torch.cat([c_g, c_g[-1:].repeat(pad_sz, 1)], 0)
+
+                        mask = torch.cat([torch.ones(chunk_sz, 1), torch.zeros(pad_sz, 1)], 0).to(device)
+                        t_j_list.append(c_j)
+                        t_g_list.append(c_g)
+                        t_s_list.append(c_s)
+                        m_list.append(mask)
+
+                    t_j = torch.stack(t_j_list)
+                    t_g = torch.stack(t_g_list)
+                    t_s = torch.stack(t_s_list)
+                    loss_mask = torch.stack(m_list)
+
+                    t_j_norm = (t_j - joint_mean_gpu) / joint_std_gpu
+
+                    valid = loss_mask.sum()
+                    if valid > 0:
+                        l_j = (crit_mse(p_j, t_j_norm).mean(-1, keepdim=True) * loss_mask).sum() / valid
+                        l_g = (crit_bce_w(p_g, t_g) * loss_mask).sum() / valid
+                        l_s = (crit_bce(p_s, t_s) * loss_mask).sum() / valid
+
+                        loss = lambda_joints * l_j + lambda_grip * l_g + lambda_success * l_s
+                        (loss / update_freq).backward()
+
+                        accum_steps += 1
+                        total_loss += loss.item()
+                        steps += 1
+
+                        if accum_steps % update_freq == 0:
+                            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=20.0)
+                            optimizer.step()
+                            optimizer.zero_grad(set_to_none=True)
+
+                        with torch.no_grad():
+                            real_p = p_j[:, 0, :] * joint_std_gpu + joint_mean_gpu
+                            real_t = t_j[:, 0, :]
+                            v0 = loss_mask[:, 0, :]
+                            if v0.sum() > 0:
+                                err = torch.abs(real_p - real_t).mean(dim=1)
+                                real_joint_err_accum += (err * v0.squeeze()).sum().item() / (v0.sum().item() + 1e-6)
+
+                if accum_steps % update_freq != 0:
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=20.0)
+                    optimizer.step()
+                    optimizer.zero_grad(set_to_none=True)
+
+                # pbar.set_postfix({
+                #     "Loss": f"{total_loss/steps:.3f}" if steps > 0 else "0",
+                #     "J_Err": f"{real_joint_err_accum/steps:.3f}" if steps > 0 else "0"
+                # })
+
+            epoch_duration = time.time() - epoch_start_time
+
+            checkpoint = {
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'joint_mean': dataset.joint_mean,
+                'joint_std': dataset.joint_std,
+                'epoch': epoch + 1
+            }
+            last_saved_ckpt_path = os.path.join(ckpt_dir, f"epoch_{epoch+1}.pth")
+            torch.save(checkpoint, last_saved_ckpt_path)
+
+            avg_loss = total_loss / steps if steps > 0 else 0.0
+            avg_j_err = real_joint_err_accum / steps if steps > 0 else 0.0
+            msg_core = f"Ep {epoch+1} Saved. Time: {epoch_duration:.2f}s, Avg Loss: {avg_loss:.4f}, J_Err: {avg_j_err:.4f}"
+            msg = f"[{task_name}] {msg_core}" if task_name else msg_core
+
+            logger.info(msg)
+            # if log_path:
+            #     with open(log_path, "a") as f:
+            #         f.write(msg + "\n")
+
+            if report_url and task_name:
+                try:
+                    import requests
+                    payload = {
+                        "task_name": task_name,
+                        "epoch": epoch + 1,
+                        "duration_sec": epoch_duration,
+                        "avg_loss": avg_loss,
+                        "j_err": avg_j_err,
+                        "msg": msg_core,
+                        "is_finished": False if epoch < num_epochs - 1 else True,
+                        "model_path": last_saved_ckpt_path,
+                    }
+                    requests.post(report_url, json=payload, timeout=3)
+                except Exception as e:
+                    logger.warning(f"上报失败: {e}")
+    except KeyboardInterrupt:
+        logger.warning("训练被用户中断")
+        raise
+    except Exception as e:
+        logger.error(f"训练过程中发生错误: {str(e)}\n{traceback.format_exc()}")
+        raise
+    finally:
+        logger.info("正在清理训练资源...")
+        if bnn_pool is not None:
+            try:
+                bnn_pool.close()
+                logger.info("BNN 进程池已关闭")
+            except Exception as e:
+                logger.error(f"关闭 BNN 进程池时出错: {str(e)}")
+
+        if io_pool is not None:
+            try:
+                io_pool.shutdown(wait=True, cancel_futures=True)
+                logger.info("IO 线程池已关闭")
+            except Exception as e:
+                logger.error(f"关闭 IO 线程池时出错: {str(e)}")
+
+        if device is not None and device.type == 'cuda':
+            try:
+                torch.cuda.empty_cache()
+                torch.cuda.ipc_collect()
+                logger.info("已清空 GPU 缓存")
+            except Exception as e:
+                logger.error(f"清空 GPU 缓存时出错: {str(e)}")
+
+        logger.info("资源清理完成")
+
+
+if __name__ == "__main__":
+    mp.set_start_method('spawn', force=True)
+
+    import argparse
+
+    parser = argparse.ArgumentParser(description="单视角训练脚本")
+    parser.add_argument("--task_name", type=str, required=False, default=None)
+    parser.add_argument("--report_url", type=str, required=False, default=None)
+    parser.add_argument("--root_dir", type=str, required=True)
+    parser.add_argument("--batch_size", type=int, default=48)
+    parser.add_argument("--seq_len", type=int, default=4)
+    parser.add_argument("--action_chunk", type=int, default=8)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--num_epochs", type=int, default=500)
+    parser.add_argument("--start_epoch", type=int, default=0)
+    parser.add_argument("--lambda_joints", type=float, default=10.0)
+    parser.add_argument("--lambda_grip", type=float, default=5.0)
+    parser.add_argument("--lambda_success", type=float, default=2.0)
+    parser.add_argument("--log_path", type=str, default=None)
+    parser.add_argument("--ckpt_dir", type=str, required=True)
+    parser.add_argument("--success_mode", type=str, default="within_horizon")
+
+    args = parser.parse_args()
+
+    train_single_view_parallel_chunk(
+        task_name=args.task_name,
+        root_dir=args.root_dir,
+        batch_size=args.batch_size,
+        seq_len=args.seq_len,
+        action_chunk=args.action_chunk,
+        lr=args.lr,
+        num_epochs=args.num_epochs,
+        start_epoch=args.start_epoch,
+        log_path=args.log_path,
+        ckpt_dir=args.ckpt_dir,
+        lambda_joints=args.lambda_joints,
+        lambda_grip=args.lambda_grip,
+        lambda_success=args.lambda_success,
+        success_mode=args.success_mode,
+        report_url=args.report_url,
+    )