BC训练

dataset/loader.py

@@ -10,8 +10,15 @@ import torch
 from torch.utils.data import Dataset
 
 
-def load_expert_pkl(expert_data_path):
-    """从目录或单个 pkl 加载专家 (obs, acts)，返回 concat 后的 obs_data, act_data。"""
+def load_expert_pkl(expert_data_path, *, filter_terminal_last_step: bool = False):
+    """从目录或单个 pkl 加载专家 (obs, acts)，返回 concat 后的 obs_data, act_data。
+
+    Args:
+        expert_data_path: Directory containing pkl files or a single pkl file.
+        filter_terminal_last_step: If True, drop the last (obs, act) pair of each trajectory.
+            This approximates II's \"train only on non-terminal steps\" when the dataset doesn't
+            explicitly store dones.
+    """
     if os.path.isdir(expert_data_path):
         pkl_files = glob.glob(os.path.join(expert_data_path, "*.pkl"))
         if not pkl_files:
@@ -30,12 +37,27 @@ def load_expert_pkl(expert_data_path):
             if isinstance(data, list):
                 for traj in data:
                     if "obs" in traj and "acts" in traj:
-                        obs_data.append(traj["obs"])
-                        act_data.append(traj["acts"])
+                        obs = traj["obs"]
+                        acts = traj["acts"]
+                        if filter_terminal_last_step and len(obs) > 0 and len(acts) > 0:
+                            # Drop last step of each trajectory
+                            obs = obs[:-1]
+                            acts = acts[:-1]
+                        if len(obs) == 0 or len(acts) == 0:
+                            continue
+                        obs_data.append(obs)
+                        act_data.append(acts)
             elif isinstance(data, dict):
                 if "observations" in data and "actions" in data:
-                    obs_data.append(data["observations"])
-                    act_data.append(data["actions"])
+                    obs = data["observations"]
+                    acts = data["actions"]
+                    if filter_terminal_last_step and len(obs) > 0 and len(acts) > 0:
+                        obs = obs[:-1]
+                        acts = acts[:-1]
+                    if len(obs) == 0 or len(acts) == 0:
+                        continue
+                    obs_data.append(obs)
+                    act_data.append(acts)
             else:
                 print(f"Skipping {pkl_file}: Unknown data format {type(data)}")
         except Exception as e:
@@ -79,7 +101,7 @@ def get_expert_scenario_ids(expert_data_path, max_ids=10):
 
 
 class MAGAILExpertDataset(Dataset):
-    def __init__(self, data_dir, transform=None):
+    def __init__(self, data_dir, transform=None, *, filter_terminal_last_step: bool = False):
         """
         Args:
             data_dir (str): Directory containing .pkl files from generate_expert_data.py
@@ -110,7 +132,10 @@ class MAGAILExpertDataset(Dataset):
             acts = traj["acts"]
 
             # obs: (T, 45), acts: (T, 2)
-            for i in range(len(obs)):
+            max_i = len(obs)
+            if filter_terminal_last_step and max_i > 0:
+                max_i -= 1
+            for i in range(max_i):
                 self.flat_data.append((obs[i], acts[i]))
 
         print(f"Total samples: {len(self.flat_data)}")