213 lines
9.2 KiB
Python
213 lines
9.2 KiB
Python
"""
|
||
BC 训练脚本：负责数据加载、环境评估、日志与保存；BC 算法由 Algorithm.bc 提供。
|
||
使用方式不变：python train_bc.py [--expert_data_path data/training_data] [--save_dir models/bc] ...
|
||
"""
|
||
import os
|
||
import numpy as np
|
||
import torch
|
||
import argparse
|
||
from torch.utils.data import DataLoader, TensorDataset
|
||
from torch.optim import Adam
|
||
from torch.optim.lr_scheduler import ExponentialLR
|
||
from datetime import datetime
|
||
from torch.utils.tensorboard import SummaryWriter
|
||
|
||
from Algorithm.policy import StateIndependentPolicy
|
||
from Algorithm.bc import train_bc_epoch, eval_bc_epoch
|
||
from Env.bc_env import BCScenarioEnv
|
||
from Env.bc_ego_replay_env import BCEgoReplayEnv
|
||
from dataset.loader import load_expert_pkl, get_expert_scenario_ids
|
||
|
||
|
||
def evaluate_policy(policy, args, device):
|
||
"""在 BCScenarioEnv（多智能体）或 BCEgoReplayEnv（单智能体）中评估策略。
|
||
仅使用专家数据中出现过的 scenario_id。单智能体模式下仅 ego 受策略控制，其他车专家回放。"""
|
||
waymo_data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
|
||
data_dir = os.path.join(waymo_data_dir, "exp_filtered")
|
||
if not os.path.exists(data_dir):
|
||
data_dir = os.path.join(waymo_data_dir, "exp_converted")
|
||
if not os.path.exists(data_dir):
|
||
print(f"[ERROR] Could not find scenario data in {waymo_data_dir}. Evaluation skipped.")
|
||
return 0.0, 0.0, 0.0
|
||
|
||
scenario_ids = get_expert_scenario_ids(args.expert_data_path, max_ids=5)
|
||
if not scenario_ids:
|
||
print("[WARN] No scenario_id in expert pkl, falling back to scenarios [0,1,2]. Eval may have 0 controlled agents.")
|
||
scenario_ids = [0, 1, 2]
|
||
|
||
total_rewards = []
|
||
total_steps = []
|
||
collision_episodes = 0
|
||
horizon = 200
|
||
single_agent = getattr(args, "single_agent", False)
|
||
|
||
for idx, scenario_id in enumerate(scenario_ids):
|
||
env_config = {
|
||
"data_directory": data_dir,
|
||
"is_multi_agent": True,
|
||
"num_controlled_agents": 100,
|
||
"use_render": False,
|
||
"sequential_seed": True,
|
||
"horizon": horizon,
|
||
"start_scenario_index": scenario_id,
|
||
"num_scenarios": 1,
|
||
"log_level": 50,
|
||
}
|
||
if single_agent:
|
||
env = BCEgoReplayEnv(config=env_config)
|
||
else:
|
||
env = BCScenarioEnv(env_config, agent2policy=None)
|
||
try:
|
||
obs_dict = env.reset(seed=scenario_id)
|
||
except Exception as e:
|
||
print(f" Eval Episode {idx} (scenario {scenario_id}): reset failed: {e}")
|
||
env.close()
|
||
continue
|
||
|
||
n_controlled = len(env.controlled_agents)
|
||
n_total_in_scenario = getattr(env, "num_controlled_in_scenario", n_controlled) if not single_agent else 1
|
||
if n_controlled == 0:
|
||
print(
|
||
f" Eval Episode {idx} (scenario {scenario_id}): 0 controlled agents (total in scenario: {n_total_in_scenario}), skip."
|
||
)
|
||
env.close()
|
||
continue
|
||
|
||
episode_reward = 0.0
|
||
step_count = 0
|
||
had_near_collision = False
|
||
dones = {"__all__": False}
|
||
while not dones["__all__"] and step_count < horizon:
|
||
step_count += 1
|
||
if not obs_dict:
|
||
obs_dict, _, dones, _ = env.step({})
|
||
continue
|
||
agent_ids = list(obs_dict.keys())
|
||
obs_list = [obs_dict[aid] for aid in agent_ids]
|
||
obs_tensor = torch.FloatTensor(np.array(obs_list)).to(device)
|
||
with torch.no_grad():
|
||
actions, _ = policy.sample(obs_tensor)
|
||
actions = actions.cpu().numpy()
|
||
action_dict = {aid: act for aid, act in zip(agent_ids, actions)}
|
||
obs_dict, rewards, dones, infos = env.step(action_dict)
|
||
episode_reward += sum(rewards.values())
|
||
if infos:
|
||
for _aid, info in infos.items():
|
||
if isinstance(info, dict) and info.get("near_collision", False):
|
||
had_near_collision = True
|
||
break
|
||
|
||
total_rewards.append(episode_reward)
|
||
total_steps.append(step_count)
|
||
if had_near_collision:
|
||
collision_episodes += 1
|
||
mode_str = "single-agent (ego)" if single_agent else f"agents (current): {n_controlled}, total in scenario: {n_total_in_scenario}"
|
||
print(
|
||
f" Eval Episode {idx} (scenario {scenario_id}): Total Reward {episode_reward:.2f}, steps {step_count}, {mode_str}"
|
||
)
|
||
env.close()
|
||
|
||
if not total_rewards:
|
||
print(" No valid eval episodes (all skipped or failed).")
|
||
return 0.0, 0.0, 0.0
|
||
avg_reward = float(np.mean(total_rewards))
|
||
avg_steps = float(np.mean(total_steps)) if total_steps else 0.0
|
||
collision_rate = float(collision_episodes / max(1, len(total_rewards)))
|
||
print(
|
||
f" Average Evaluation Reward: {avg_reward:.2f} | Mean Episode Length: {avg_steps:.1f} | "
|
||
f"Collision Rate (near): {collision_rate:.3f}"
|
||
)
|
||
return avg_reward, collision_rate, avg_steps
|
||
|
||
|
||
def main(args):
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
print(f"Using device: {device}")
|
||
|
||
os.makedirs("logs/bc", exist_ok=True)
|
||
log_dir = os.path.join("logs", "bc", datetime.now().strftime("%Y%m%d-%H%M%S"))
|
||
writer = SummaryWriter(log_dir)
|
||
print(f"TensorBoard logging to: {log_dir}")
|
||
os.makedirs(args.save_dir, exist_ok=True)
|
||
|
||
agent_id_filter = "default_agent" if getattr(args, "single_agent", False) else None
|
||
obs_data, act_data = load_expert_pkl(
|
||
args.expert_data_path,
|
||
filter_terminal_last_step=args.filter_terminal_last_step,
|
||
agent_id_filter=agent_id_filter,
|
||
)
|
||
obs_tensor = torch.FloatTensor(obs_data)
|
||
act_tensor = torch.FloatTensor(act_data)
|
||
dataset = TensorDataset(obs_tensor, act_tensor)
|
||
train_size = int(0.8 * len(dataset))
|
||
val_size = len(dataset) - train_size
|
||
train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
|
||
train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
|
||
val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)
|
||
print(f"Dataset loaded. Train size: {len(train_dataset)}, Val size: {len(val_dataset)}")
|
||
|
||
state_dim = obs_data.shape[1]
|
||
action_dim = act_data.shape[1]
|
||
print(f"State Dim: {state_dim}, Action Dim: {action_dim}")
|
||
|
||
policy = StateIndependentPolicy(
|
||
state_shape=(state_dim,),
|
||
action_shape=(action_dim,),
|
||
hidden_units=(256, 256),
|
||
hidden_activation=torch.nn.Tanh(),
|
||
).to(device)
|
||
optimizer = Adam(policy.parameters(), lr=args.lr)
|
||
scheduler = ExponentialLR(optimizer, gamma=0.99)
|
||
|
||
best_val_loss = float("inf")
|
||
for epoch in range(args.epochs):
|
||
avg_train_loss = train_bc_epoch(policy, train_loader, optimizer, device)
|
||
scheduler.step()
|
||
avg_val_loss = eval_bc_epoch(policy, val_loader, device)
|
||
|
||
print(f"Epoch {epoch+1}/{args.epochs} | Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")
|
||
writer.add_scalar("Loss/train", avg_train_loss, epoch)
|
||
writer.add_scalar("Loss/val", avg_val_loss, epoch)
|
||
writer.add_scalar("Learning_rate", scheduler.get_last_lr()[0], epoch)
|
||
|
||
if avg_val_loss < best_val_loss:
|
||
best_val_loss = avg_val_loss
|
||
torch.save(policy.state_dict(), os.path.join(args.save_dir, "policy_best.pt"))
|
||
|
||
# Periodic checkpointing (II-style)
|
||
if args.checkpoint_freq > 0 and (epoch + 1) % args.checkpoint_freq == 0:
|
||
torch.save(policy.state_dict(), os.path.join(args.save_dir, f"policy_epoch{epoch+1}.pt"))
|
||
|
||
if (epoch + 1) % args.eval_freq == 0:
|
||
eval_reward, eval_collision_rate, eval_mean_steps = evaluate_policy(policy, args, device)
|
||
writer.add_scalar("Reward/eval", eval_reward, epoch)
|
||
writer.add_scalar("Eval/collision_rate_near", eval_collision_rate, epoch)
|
||
writer.add_scalar("Eval/mean_episode_length", eval_mean_steps, epoch)
|
||
|
||
torch.save(policy.state_dict(), os.path.join(args.save_dir, "policy_final.pt"))
|
||
writer.close()
|
||
print("Training finished.")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--expert_data_path", type=str, default="data/training_data", help="Path to expert data pickle or directory")
|
||
parser.add_argument("--save_dir", type=str, default="models/bc", help="Directory to save models")
|
||
parser.add_argument("--epochs", type=int, default=100)
|
||
parser.add_argument("--batch_size", type=int, default=64)
|
||
parser.add_argument("--lr", type=float, default=3e-4)
|
||
parser.add_argument("--eval_freq", type=int, default=10)
|
||
parser.add_argument("--checkpoint_freq", type=int, default=50, help="Save policy_epochN.pt every N epochs. Set <=0 to disable.")
|
||
parser.add_argument(
|
||
"--filter_terminal_last_step",
|
||
action="store_true",
|
||
help="Drop the last (obs, act) pair of each trajectory to approximate training on non-terminal steps (II-style).",
|
||
)
|
||
parser.add_argument(
|
||
"--single_agent",
|
||
action="store_true",
|
||
help="Use single-agent (ego) expert data and evaluation; load only default_agent trajectories and evaluate with BCEgoReplayEnv.",
|
||
)
|
||
args = parser.parse_args()
|
||
main(args)
|