This commit is contained in:
2026-02-10 22:44:02 +08:00
parent 0f9f080e77
commit 8a75f0db0d
12 changed files with 115 additions and 17 deletions

View File

@@ -88,10 +88,48 @@ class BCScenarioEnv(MultiAgentScenarioEnv):
self._spawn_controlled_agents() self._spawn_controlled_agents()
self._update_background_vehicles() self._update_background_vehicles()
obs = self._get_all_obs() obs = self._get_all_obs()
rewards = {aid: 0.0 for aid in self.controlled_agents}
# Reward shaping for evaluation/rollout monitoring (BC training itself doesn't use env reward).
speed_coef = float(self.config.get("reward_speed_coef", 0.05))
collision_distance = float(self.config.get("collision_distance", 6.0))
collision_penalty = float(self.config.get("collision_penalty", 100.0))
# Pre-collect all active vehicles (includes background vehicles).
active_agents = list(self.engine.agent_manager.active_agents.items())
rewards = {}
infos = {}
for aid, vehicle in self.controlled_agents.items():
# Speed reward
speed = getattr(vehicle, "speed", None)
if speed is None:
speed = float(np.linalg.norm(vehicle.velocity))
r_speed = speed_coef * float(speed)
# Near-collision penalty (distance-based, simulator-agnostic)
min_dist = float("inf")
for other_id, other_vehicle in active_agents:
if other_id == aid:
continue
try:
dist = float(np.linalg.norm(vehicle.position - other_vehicle.position))
except Exception:
continue
if dist < min_dist:
min_dist = dist
near_collision = bool(min_dist < collision_distance)
r_collision = -collision_penalty if near_collision else 0.0
rewards[aid] = float(r_speed + r_collision)
infos[aid] = {
"near_collision": near_collision,
"min_dist": (min_dist if np.isfinite(min_dist) else None),
"r_speed": float(r_speed),
"r_collision": float(r_collision),
}
dones = {aid: False for aid in self.controlled_agents} dones = {aid: False for aid in self.controlled_agents}
dones["__all__"] = self.episode_step >= self.config["horizon"] dones["__all__"] = self.episode_step >= self.config["horizon"]
infos = {aid: {} for aid in self.controlled_agents}
return obs, rewards, dones, infos return obs, rewards, dones, infos
def _get_all_obs(self): def _get_all_obs(self):

View File

@@ -2,7 +2,7 @@ import numpy as np
import math import math
class InverseDynamics: class InverseDynamics:
def __init__(self, max_steering=0.7, max_acc=15.0, length=4.5): def __init__(self, max_steering=0.7, max_acc=8.0, length=4.5):
""" """
:param max_steering: Max steering angle in radians (approx 40 degrees) :param max_steering: Max steering angle in radians (approx 40 degrees)
:param max_acc: Max acceleration in m/s^2 :param max_acc: Max acceleration in m/s^2

View File

@@ -93,6 +93,10 @@ python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir
### 2. 行为克隆 (BC) ### 2. 行为克隆 (BC)
- **训练**`python train_bc.py`(模型保存到 `models/bc/`,日志到 `logs/bc/` - **训练**`python train_bc.py`(模型保存到 `models/bc/`,日志到 `logs/bc/`
```
# 注意替换文件名
python train_bc.py --expert_data_path ./data/training/expert_data_0_50.pkl --epochs 100
```
- **可视化**`python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt` - **可视化**`python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt`
### 3. 多智能体对抗模仿学习 (MAGAIL) ### 3. 多智能体对抗模仿学习 (MAGAIL)

View File

@@ -10,8 +10,15 @@ import torch
from torch.utils.data import Dataset from torch.utils.data import Dataset
def load_expert_pkl(expert_data_path): def load_expert_pkl(expert_data_path, *, filter_terminal_last_step: bool = False):
"""从目录或单个 pkl 加载专家 (obs, acts),返回 concat 后的 obs_data, act_data。""" """从目录或单个 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): if os.path.isdir(expert_data_path):
pkl_files = glob.glob(os.path.join(expert_data_path, "*.pkl")) pkl_files = glob.glob(os.path.join(expert_data_path, "*.pkl"))
if not pkl_files: if not pkl_files:
@@ -30,12 +37,27 @@ def load_expert_pkl(expert_data_path):
if isinstance(data, list): if isinstance(data, list):
for traj in data: for traj in data:
if "obs" in traj and "acts" in traj: if "obs" in traj and "acts" in traj:
obs_data.append(traj["obs"]) obs = traj["obs"]
act_data.append(traj["acts"]) 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): elif isinstance(data, dict):
if "observations" in data and "actions" in data: if "observations" in data and "actions" in data:
obs_data.append(data["observations"]) obs = data["observations"]
act_data.append(data["actions"]) 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: else:
print(f"Skipping {pkl_file}: Unknown data format {type(data)}") print(f"Skipping {pkl_file}: Unknown data format {type(data)}")
except Exception as e: except Exception as e:
@@ -79,7 +101,7 @@ def get_expert_scenario_ids(expert_data_path, max_ids=10):
class MAGAILExpertDataset(Dataset): class MAGAILExpertDataset(Dataset):
def __init__(self, data_dir, transform=None): def __init__(self, data_dir, transform=None, *, filter_terminal_last_step: bool = False):
""" """
Args: Args:
data_dir (str): Directory containing .pkl files from generate_expert_data.py data_dir (str): Directory containing .pkl files from generate_expert_data.py
@@ -110,7 +132,10 @@ class MAGAILExpertDataset(Dataset):
acts = traj["acts"] acts = traj["acts"]
# obs: (T, 45), acts: (T, 2) # 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])) self.flat_data.append((obs[i], acts[i]))
print(f"Total samples: {len(self.flat_data)}") print(f"Total samples: {len(self.flat_data)}")

View File

@@ -35,6 +35,8 @@ def evaluate_policy(policy, args, device):
scenario_ids = [0, 1, 2] scenario_ids = [0, 1, 2]
total_rewards = [] total_rewards = []
total_steps = []
collision_episodes = 0
horizon = 200 horizon = 200
for idx, scenario_id in enumerate(scenario_ids): for idx, scenario_id in enumerate(scenario_ids):
@@ -67,6 +69,7 @@ def evaluate_policy(policy, args, device):
episode_reward = 0.0 episode_reward = 0.0
step_count = 0 step_count = 0
had_near_collision = False
dones = {"__all__": False} dones = {"__all__": False}
while not dones["__all__"] and step_count < horizon: while not dones["__all__"] and step_count < horizon:
step_count += 1 step_count += 1
@@ -80,10 +83,18 @@ def evaluate_policy(policy, args, device):
actions, _ = policy.sample(obs_tensor) actions, _ = policy.sample(obs_tensor)
actions = actions.cpu().numpy() actions = actions.cpu().numpy()
action_dict = {aid: act for aid, act in zip(agent_ids, actions)} action_dict = {aid: act for aid, act in zip(agent_ids, actions)}
obs_dict, rewards, dones, _ = env.step(action_dict) obs_dict, rewards, dones, infos = env.step(action_dict)
episode_reward += sum(rewards.values()) 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_rewards.append(episode_reward)
total_steps.append(step_count)
if had_near_collision:
collision_episodes += 1
print( print(
f" Eval Episode {idx} (scenario {scenario_id}): Total Reward {episode_reward:.2f}, steps {step_count}, " f" Eval Episode {idx} (scenario {scenario_id}): Total Reward {episode_reward:.2f}, steps {step_count}, "
f"agents (current): {n_controlled}, total in scenario: {n_total_in_scenario}" f"agents (current): {n_controlled}, total in scenario: {n_total_in_scenario}"
@@ -92,10 +103,15 @@ def evaluate_policy(policy, args, device):
if not total_rewards: if not total_rewards:
print(" No valid eval episodes (all skipped or failed).") print(" No valid eval episodes (all skipped or failed).")
return 0.0 return 0.0, 0.0, 0.0
avg_reward = float(np.mean(total_rewards)) avg_reward = float(np.mean(total_rewards))
print(f" Average Evaluation Reward: {avg_reward:.2f}") avg_steps = float(np.mean(total_steps)) if total_steps else 0.0
return avg_reward 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): def main(args):
@@ -108,7 +124,10 @@ def main(args):
print(f"TensorBoard logging to: {log_dir}") print(f"TensorBoard logging to: {log_dir}")
os.makedirs(args.save_dir, exist_ok=True) os.makedirs(args.save_dir, exist_ok=True)
obs_data, act_data = load_expert_pkl(args.expert_data_path) obs_data, act_data = load_expert_pkl(
args.expert_data_path,
filter_terminal_last_step=args.filter_terminal_last_step,
)
obs_tensor = torch.FloatTensor(obs_data) obs_tensor = torch.FloatTensor(obs_data)
act_tensor = torch.FloatTensor(act_data) act_tensor = torch.FloatTensor(act_data)
dataset = TensorDataset(obs_tensor, act_tensor) dataset = TensorDataset(obs_tensor, act_tensor)
@@ -147,9 +166,15 @@ def main(args):
best_val_loss = avg_val_loss best_val_loss = avg_val_loss
torch.save(policy.state_dict(), os.path.join(args.save_dir, "policy_best.pt")) 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: if (epoch + 1) % args.eval_freq == 0:
eval_reward = evaluate_policy(policy, args, device) eval_reward, eval_collision_rate, eval_mean_steps = evaluate_policy(policy, args, device)
writer.add_scalar("Reward/eval", eval_reward, epoch) 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")) torch.save(policy.state_dict(), os.path.join(args.save_dir, "policy_final.pt"))
writer.close() writer.close()
@@ -164,5 +189,11 @@ if __name__ == "__main__":
parser.add_argument("--batch_size", type=int, default=64) parser.add_argument("--batch_size", type=int, default=64)
parser.add_argument("--lr", type=float, default=3e-4) parser.add_argument("--lr", type=float, default=3e-4)
parser.add_argument("--eval_freq", type=int, default=10) 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).",
)
args = parser.parse_args() args = parser.parse_args()
main(args) main(args)