HBBC部署到代码中

2026-03-02 10:58:20 +08:00
parent 8a75f0db0d
commit be35650533
23 changed files with 1293 additions and 83 deletions
--- a/scripts/README.md
+++ b/scripts/README.md
@@ -16,9 +16,19 @@

 | 脚本 | 用途 | 用法示例 |
 |------|------|----------|
-| [generate_expert_data.py](generate_expert_data.py) | 从 Waymo 数据生成专家 (obs, act) 的 pkl | `python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100` |
+| [generate_expert_data.py](generate_expert_data.py) | 从 Waymo 数据生成专家 (obs, act) 的 pkl | 见下方 |

-**常用参数**：`--data_dir`（默认 `data/exp_filtered`）、`--output_dir`（默认 `data/training_data`）、`--start_index`、`--num_scenarios`。
+**多智能体**（输出 `expert_data_{start_index}_{num_scenarios}.pkl`）：
+```bash
+python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0
+```
+
+**单智能体**（仅采集 ego 车轨迹，输出 `expert_data_ego_{start_index}_{num_scenarios}.pkl`，用于单智能体 BC）：
+```bash
+python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0 --ego_only
+```
+
+**常用参数**：`--data_dir`（默认 `data/exp_filtered`）、`--output_dir`（默认 `data/training_data`）、`--start_index`、`--num_scenarios`、`--ego_only`（仅保存 default_agent 轨迹，输出使用 `expert_data_ego_*.pkl` 前缀）。

 ---

@@ -40,13 +50,32 @@ python scripts/visualize.py replay --data_dir data/exp_filtered --num_scenarios
 python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt --data_dir data/exp_filtered --num_scenarios 1
 python scripts/visualize.py policy --policy_type magail --model_path models/magail/model_50_actor.pth --num_scenarios 1 --deterministic
 ```
+- **policy + 仅自车策略、其他车回放**（BC 单智能体模型）：加 `--ego_only`，自车由策略控制，其余车辆按专家轨迹回放。
+```bash
+python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt --data_dir data/exp_filtered --num_scenarios 1 --ego_only
+```
+
+- **policy + HBBC 动态背景车**（仅动态背景车启用，静态背景车保持原样）：
+```bash
+python scripts/visualize.py policy \
+  --policy_type bc \
+  --model_path models/bc/policy_best.pt \
+  --data_dir data/exp_filtered \
+  --num_scenarios 1 \
+  --ego_only \
+  --enable_hbbc_background \
+  --hbbc_model_path models/hbbc/hbbc.pt \
+  --hbbc_inference_device cpu \
+  --hbbc_latent_mode per_vehicle_fixed \
+  --hbbc_latent_json_path docs/examples/hbbc_latent_example.json
+```

 - **trajectory**（专家轨迹 matplotlib 俯视图动画）：
 ```bash
 python scripts/visualize.py trajectory --data_dir data/exp_filtered --scenario_idx 0
 ```

-**公共参数**：`--data_dir`（默认 `data/exp_filtered`）、`--start_index`、`--num_scenarios`、`--horizon`。policy 模式另有 `--policy_type`（auto/bc/magail）、`--model_path`、`--deterministic`（仅 MAGAIL）。
+**公共参数**：`--data_dir`（默认 `data/exp_filtered`）、`--start_index`、`--num_scenarios`、`--horizon`。policy 模式另有 `--policy_type`（auto/bc/magail）、`--model_path`、`--deterministic`（仅 MAGAIL）、`--ego_only`（仅 BC：自车用策略，其他车专家回放）、`--enable_hbbc_background`、`--hbbc_model_path`、`--hbbc_inference_device`、`--hbbc_latent_mode`、`--hbbc_latent_json_path`。

 ---

@@ -70,7 +99,7 @@ python scripts/visualize.py trajectory --data_dir data/exp_filtered --scenario_i

 ## 与训练流程的对应关系

-1. **数据准备**：`generate_expert_data.py` → 输出到 `data/training_data/*.pkl`
-2. **BC 训练**：根目录 `train_bc.py` → 模型保存到 `models/bc/`，日志到 `logs/bc/`
+1. **数据准备**：`generate_expert_data.py` → 输出到 `data/training_data/*.pkl`（多智能体 `expert_data_*.pkl`，单智能体 `expert_data_ego_*.pkl`）
+2. **BC 训练**：根目录 `train_bc.py` → 模型保存到 `models/bc/`，日志到 `logs/bc/`。单智能体模式加 `--single_agent` 并指定 ego-only 的 pkl。
 3. **MAGAIL 训练**：根目录 `train_magail.py` → 模型保存到 `models/magail/`，日志到 `logs/magail/`
 4. **可视化**：`scripts/visualize.py`（子命令 replay / policy / trajectory）→ 数据目录默认 `data/exp_filtered`
--- a/scripts/generate_expert_data.py
+++ b/scripts/generate_expert_data.py
@@ -102,7 +102,9 @@ def generate_data(args):
            
            # Post-process episode data
            for agent_id, data in episode_data.items():
-                if len(data['obs']) > 10: # Minimum length filter
+                if args.ego_only and agent_id != "default_agent":
+                    continue
+                if len(data['obs']) > 10:  # Minimum length filter
                    expert_trajectories.append({
                        'obs': np.array(data['obs']),
                        'acts': np.array(data['acts']),
@@ -120,9 +122,14 @@ def generate_data(args):
        pass
        
    # Save data
-    output_file = os.path.join(args.output_dir, f"expert_data_{args.start_index}_{args.num_scenarios}.pkl")
+    if args.ego_only:
+        output_file = os.path.join(args.output_dir, f"expert_data_ego_{args.start_index}_{args.num_scenarios}.pkl")
+    else:
+        output_file = os.path.join(args.output_dir, f"expert_data_{args.start_index}_{args.num_scenarios}.pkl")
    os.makedirs(args.output_dir, exist_ok=True)
-    
+
+    if args.ego_only:
+        print("Ego-only mode: saved trajectories are SDC (default_agent) only.")
    print(f"Saving {len(expert_trajectories)} trajectories to {output_file}")
    with open(output_file, 'wb') as f:
        pickle.dump(expert_trajectories, f)
@@ -157,6 +164,6 @@ if __name__ == "__main__":
    parser.add_argument("--output_dir", type=str, default="data/training_data", help="Output directory")
    parser.add_argument("--start_index", type=int, default=0)
    parser.add_argument("--num_scenarios", type=int, default=10)
-    
+    parser.add_argument("--ego_only", action="store_true", help="Only collect and save ego (default_agent) trajectories; output uses expert_data_ego_*.pkl prefix")
    args = parser.parse_args()
    generate_data(args)
--- a/scripts/visualize.py
+++ b/scripts/visualize.py
@@ -111,36 +111,49 @@ def _resolve_model_path(model_path, policy_type):

 def _run_policy(args):
    from Env.bc_env import BCScenarioEnv
+    from Env.bc_ego_replay_env import BCEgoReplayEnv
    from metadrive.engine.engine_utils import close_engine

    policy_type = (args.policy_type or "auto").lower()
    if policy_type == "auto":
        policy_type = "bc" if args.model_path.endswith(".pt") else "magail"
+    ego_only = getattr(args, "ego_only", False)
+    if ego_only and policy_type != "bc":
+        print("[WARN] --ego_only is supported for BC policy only; MAGAIL will run in multi-agent mode.")

    data_dir = _resolve_data_dir(args.data_dir)
    data_path = os.path.abspath(data_dir)
    env_config = {
        "data_directory": data_path,
        "is_multi_agent": True,
-        "num_controlled_agents": 3,
+        "num_controlled_agents": 100 if ego_only else 3,
        "horizon": args.horizon,
        "use_render": True,
        "sequential_seed": True,
        "start_scenario_index": args.start_index,
        "num_scenarios": args.num_scenarios,
        "log_level": 40,
+        "enable_hbbc_background": bool(getattr(args, "enable_hbbc_background", False)),
+        "hbbc_model_path": getattr(args, "hbbc_model_path", "models/hbbc/hbbc.pt"),
+        "hbbc_inference_device": getattr(args, "hbbc_inference_device", "cpu"),
+        "hbbc_latent_mode": getattr(args, "hbbc_latent_mode", "per_vehicle_fixed"),
+        "hbbc_latent_json_path": getattr(args, "hbbc_latent_json_path", None),
    }

-    print(f"Initializing BCScenarioEnv (policy_type={policy_type})...")
+    if ego_only and policy_type == "bc":
+        print("Initializing BCEgoReplayEnv (ego-only: policy on self, others replayed)...")
+    else:
+        print(f"Initializing BCScenarioEnv (policy_type={policy_type})...")
+
    try:
-        env = BCScenarioEnv(env_config, agent2policy={})
+        env = BCEgoReplayEnv(config=env_config) if (ego_only and policy_type == "bc") else BCScenarioEnv(env_config, agent2policy={})
    except Exception as e:
        print(f"Error init env: {e}. Trying to close lingering engine...")
        try:
            close_engine()
        except Exception:
            pass
-        env = BCScenarioEnv(env_config, agent2policy={})
+        env = BCEgoReplayEnv(config=env_config) if (ego_only and policy_type == "bc") else BCScenarioEnv(env_config, agent2policy={})

    state_dim = 45
    action_dim = 2
@@ -156,7 +169,11 @@ def _run_policy(args):
            hidden_units=(256, 256),
            hidden_activation=torch.nn.Tanh(),
        ).to(device)
-        policy.load_state_dict(torch.load(model_path, map_location=device))
+        try:
+            state = torch.load(model_path, map_location=device, weights_only=True)
+        except TypeError:
+            state = torch.load(model_path, map_location=device)
+        policy.load_state_dict(state)
        policy.eval()
    else:
        from train_magail import Actor
@@ -178,7 +195,16 @@ def _run_policy(args):
                    pass
                continue

-            print(f"Scenario loaded. Controlled agents (current): {len(obs_dict)}, total in scenario: {env.num_controlled_in_scenario}")
+            n_total = getattr(env, "num_controlled_in_scenario", len(obs_dict))
+            mode_note = " (ego only, others replayed)" if (ego_only and policy_type == "bc") else ""
+            if ego_only and policy_type == "bc" and bool(env_config.get("enable_hbbc_background", False)):
+                mode_note = " (ego only, dynamic background via HBBC)"
+            print(f"Scenario loaded. Controlled agents (current): {len(obs_dict)}, total in scenario: {n_total}{mode_note}")
+            if ego_only and policy_type == "bc" and len(obs_dict) == 1:
+                if bool(env_config.get("enable_hbbc_background", False)):
+                    print("  [Ego control: policy injected — dynamic background vehicles use HBBC; static background stays static.]")
+                else:
+                    print("  [Ego control: policy injected — ego uses model output each step; other vehicles expert replay.]")
            if len(obs_dict) == 0:
                print(f"Scenario {i} has no controlled agents (all filtered out). Skipping.")
                continue
@@ -370,6 +396,12 @@ def main():
    pp.add_argument("--policy_type", type=str, default="auto", choices=["auto", "bc", "magail"])
    pp.add_argument("--model_path", type=str, default="models/bc/policy_best.pt")
    pp.add_argument("--deterministic", action="store_true", help="MAGAIL: use mean action")
+    pp.add_argument("--ego_only", action="store_true", help="BC only: inject policy into ego only; other vehicles use expert replay")
+    pp.add_argument("--enable_hbbc_background", action="store_true", help="Enable HBBC policy for dynamic background vehicles")
+    pp.add_argument("--hbbc_model_path", type=str, default="models/hbbc/hbbc.pt")
+    pp.add_argument("--hbbc_inference_device", type=str, default="cpu")
+    pp.add_argument("--hbbc_latent_mode", type=str, default="per_vehicle_fixed", choices=["per_vehicle_fixed", "per_episode_reset"])
+    pp.add_argument("--hbbc_latent_json_path", type=str, default=None, help="Optional JSON for per-vehicle latent override")

    # trajectory
    pt = subparsers.add_parser("trajectory", help="2D matplotlib animation of expert trajectories")