BC训练
This commit is contained in:
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -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):
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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)}")
|
||||||
|
|||||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
43
train_bc.py
43
train_bc.py
@@ -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)
|
||||||
|
|||||||
Reference in New Issue
Block a user