599 lines
24 KiB
Python
599 lines
24 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import torch.optim as optim
|
|
from torch.distributions import Normal
|
|
import numpy as np
|
|
import os
|
|
import argparse
|
|
import signal
|
|
import sys
|
|
from torch.utils.data import DataLoader
|
|
from dataset.magail_dataset import MAGAILExpertDataset
|
|
|
|
# --- Networks ---
|
|
|
|
class Actor(nn.Module):
|
|
def __init__(self, state_dim, action_dim, hidden_dim=256):
|
|
super(Actor, self).__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Linear(state_dim, hidden_dim),
|
|
nn.Tanh(),
|
|
nn.Linear(hidden_dim, hidden_dim),
|
|
nn.Tanh(),
|
|
)
|
|
self.mu_head = nn.Linear(hidden_dim, action_dim)
|
|
self.log_std_head = nn.Parameter(torch.zeros(1, action_dim))
|
|
|
|
def forward(self, state):
|
|
x = self.net(state)
|
|
mu = torch.tanh(self.mu_head(x)) # Action range [-1, 1]
|
|
if mu.dim() == 1:
|
|
mu = mu.unsqueeze(0) # Handle single sample
|
|
log_std = self.log_std_head.expand_as(mu)
|
|
std = torch.exp(log_std)
|
|
dist = Normal(mu, std)
|
|
return dist
|
|
|
|
class Critic(nn.Module):
|
|
def __init__(self, state_dim, hidden_dim=256):
|
|
super(Critic, self).__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Linear(state_dim, hidden_dim),
|
|
nn.Tanh(),
|
|
nn.Linear(hidden_dim, hidden_dim),
|
|
nn.Tanh(),
|
|
nn.Linear(hidden_dim, 1)
|
|
)
|
|
|
|
def forward(self, state):
|
|
return self.net(state)
|
|
|
|
class Discriminator(nn.Module):
|
|
def __init__(self, state_dim, action_dim, hidden_dim=256):
|
|
super(Discriminator, self).__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Linear(state_dim + action_dim, hidden_dim),
|
|
nn.Tanh(),
|
|
nn.Linear(hidden_dim, hidden_dim),
|
|
nn.Tanh(),
|
|
nn.Linear(hidden_dim, 1),
|
|
nn.Sigmoid()
|
|
)
|
|
|
|
def forward(self, state, action):
|
|
x = torch.cat([state, action], dim=-1)
|
|
return self.net(x)
|
|
|
|
# --- PPO Algorithm ---
|
|
|
|
class PPO:
|
|
def __init__(self, state_dim, action_dim, lr=3e-4, gamma=0.99, eps_clip=0.2, K_epochs=10):
|
|
self.actor = Actor(state_dim, action_dim).cuda()
|
|
self.critic = Critic(state_dim).cuda()
|
|
self.optimizer_actor = optim.Adam(self.actor.parameters(), lr=lr)
|
|
self.optimizer_critic = optim.Adam(self.critic.parameters(), lr=lr)
|
|
|
|
self.gamma = gamma
|
|
self.eps_clip = eps_clip
|
|
self.K_epochs = K_epochs
|
|
self.mse_loss = nn.MSELoss()
|
|
|
|
def _log_prob_from_dist(self, dist, pre_tanh_action):
|
|
# Tanh-squashed Gaussian log-prob with correction term.
|
|
log_prob = dist.log_prob(pre_tanh_action)
|
|
correction = torch.log(1 - torch.tanh(pre_tanh_action) ** 2 + 1e-6)
|
|
return (log_prob - correction).sum(dim=-1)
|
|
|
|
def select_action(self, state):
|
|
with torch.no_grad():
|
|
state = torch.FloatTensor(state).cuda()
|
|
dist = self.actor(state)
|
|
pre_tanh_action = dist.sample()
|
|
action = torch.tanh(pre_tanh_action)
|
|
action_logprob = self._log_prob_from_dist(dist, pre_tanh_action)
|
|
return (
|
|
action.cpu().numpy(),
|
|
action_logprob.cpu().numpy(),
|
|
pre_tanh_action.cpu().numpy()
|
|
)
|
|
|
|
def update(self, memory):
|
|
# Convert memory to tensors
|
|
states = torch.FloatTensor(np.array(memory['states'])).cuda()
|
|
actions = torch.FloatTensor(np.array(memory['actions'])).cuda()
|
|
pre_tanh_actions = torch.FloatTensor(np.array(memory['pre_tanh_actions'])).cuda()
|
|
logprobs = torch.FloatTensor(np.array(memory['logprobs'])).cuda()
|
|
rewards = torch.FloatTensor(np.array(memory['rewards'])).cuda()
|
|
next_states = torch.FloatTensor(np.array(memory['next_states'])).cuda()
|
|
dones = torch.FloatTensor(np.array(memory['dones'])).cuda()
|
|
|
|
# Monte Carlo estimate of state rewards (or GAE if implemented, simplistic here)
|
|
# Usually for PPO we use GAE. Let's do a simple discounted return for now or bootstrapping.
|
|
# Let's use bootstrapping from critic for returns.
|
|
|
|
returns = []
|
|
discounted_reward = 0
|
|
# This simple loop assumes full episode or consistent batch.
|
|
# For multi-agent disjoint steps, bootstrapping is better.
|
|
# But let's calculate advantage using GAE for stability.
|
|
|
|
values = self.critic(states).detach()
|
|
next_values = self.critic(next_states).detach()
|
|
|
|
# GAE
|
|
advantages = []
|
|
gae = 0
|
|
for i in reversed(range(len(rewards))):
|
|
delta = rewards[i] + self.gamma * next_values[i] * (1 - dones[i]) - values[i]
|
|
gae = delta + self.gamma * 0.95 * (1 - dones[i]) * gae
|
|
advantages.insert(0, gae)
|
|
|
|
advantages = torch.FloatTensor(advantages).cuda()
|
|
returns = advantages + values.squeeze()
|
|
|
|
# Optimize policy for K epochs:
|
|
for _ in range(self.K_epochs):
|
|
# Evaluating old actions and values :
|
|
dist = self.actor(states)
|
|
action_logprobs = self._log_prob_from_dist(dist, pre_tanh_actions)
|
|
dist_entropy = dist.entropy().sum(dim=-1)
|
|
state_values = self.critic(states).squeeze()
|
|
|
|
# Finding the ratio (pi_theta / pi_theta__old):
|
|
ratios = torch.exp(action_logprobs - logprobs)
|
|
|
|
# Finding Surrogate Loss:
|
|
surr1 = ratios * advantages
|
|
surr2 = torch.clamp(ratios, 1-self.eps_clip, 1+self.eps_clip) * advantages
|
|
loss = -torch.min(surr1, surr2) + 0.5*self.mse_loss(state_values, returns) - 0.01*dist_entropy
|
|
|
|
# take gradient step
|
|
self.optimizer_actor.zero_grad()
|
|
self.optimizer_critic.zero_grad()
|
|
loss.mean().backward()
|
|
self.optimizer_actor.step()
|
|
self.optimizer_critic.step()
|
|
|
|
return loss.mean().item()
|
|
|
|
def save(self, checkpoint_path):
|
|
torch.save(self.actor.state_dict(), checkpoint_path + "_actor.pth")
|
|
torch.save(self.critic.state_dict(), checkpoint_path + "_critic.pth")
|
|
|
|
from Env.scenario_env import MultiAgentScenarioEnv
|
|
|
|
class MAGAILScenarioEnv(MultiAgentScenarioEnv):
|
|
def _get_all_obs(self):
|
|
# Same logic as ExpertReplayEnv to ensure compatibility
|
|
obs_dict = {}
|
|
for agent_id, vehicle in self.controlled_agents.items():
|
|
# 1. Ego State
|
|
ego_state = [
|
|
vehicle.position[0], vehicle.position[1],
|
|
vehicle.velocity[0], vehicle.velocity[1],
|
|
vehicle.heading_theta
|
|
]
|
|
|
|
# 2. Neighbors
|
|
candidates = []
|
|
for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
|
|
if other_id == agent_id:
|
|
continue
|
|
dist = np.linalg.norm(vehicle.position - other_vehicle.position)
|
|
if dist < 30.0:
|
|
candidates.append((dist, other_vehicle))
|
|
|
|
candidates.sort(key=lambda x: x[0])
|
|
top_10 = candidates[:10]
|
|
|
|
neighbor_feats = []
|
|
for _, neighbor in top_10:
|
|
neighbor_feats.extend([
|
|
neighbor.position[0] - vehicle.position[0],
|
|
neighbor.position[1] - vehicle.position[1],
|
|
neighbor.velocity[0],
|
|
neighbor.velocity[1]
|
|
])
|
|
|
|
missing = 10 - len(top_10)
|
|
if missing > 0:
|
|
neighbor_feats.extend([0.0] * (4 * missing))
|
|
|
|
obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
|
|
obs_dict[agent_id] = obs
|
|
return obs_dict
|
|
|
|
# --- Training Loop ---
|
|
|
|
def train(args):
|
|
# 1. Setup Environment (Dummy for now, usually you run simulation here)
|
|
# But for MAGAIL we need to collect generated trajectories.
|
|
# We need the Env class to be importable.
|
|
from Env.scenario_env import MultiAgentScenarioEnv
|
|
from Env.simple_idm_policy import ConstantVelocityPolicy # Just for init
|
|
|
|
# Config for Env
|
|
env_config = {
|
|
"data_directory": args.data_dir,
|
|
"is_multi_agent": True,
|
|
"num_controlled_agents": 3, # Dynamic
|
|
"horizon": 200,
|
|
"use_render": False,
|
|
"sequential_seed": True,
|
|
"start_scenario_index": 0,
|
|
"num_scenarios": args.num_scenarios # Use argument
|
|
}
|
|
|
|
# Ideally we use a wrapper for RL
|
|
# env = MultiAgentScenarioEnv(config=env_config) # This requires Waymo data loader setup
|
|
|
|
# 2. Setup Models
|
|
state_dim = 45
|
|
action_dim = 2
|
|
|
|
ppo_agent = PPO(state_dim, action_dim)
|
|
discriminator = Discriminator(state_dim, action_dim).cuda()
|
|
disc_optimizer = optim.Adam(discriminator.parameters(), lr=3e-4)
|
|
disc_criterion = nn.BCELoss()
|
|
|
|
# 3. Load Expert Data
|
|
expert_dataset = MAGAILExpertDataset(args.expert_data_dir)
|
|
# Ensure batch_size is not larger than dataset
|
|
if len(expert_dataset) < args.batch_size:
|
|
print(f"Warning: Expert dataset size {len(expert_dataset)} < batch_size {args.batch_size}. Adjusting batch_size.")
|
|
args.batch_size = len(expert_dataset)
|
|
if args.batch_size == 0:
|
|
raise ValueError("Expert dataset is empty!")
|
|
|
|
expert_loader = DataLoader(expert_dataset, batch_size=args.batch_size, shuffle=True, drop_last=True)
|
|
|
|
# Create an infinite iterator
|
|
def cycle(loader):
|
|
while True:
|
|
for batch in loader:
|
|
yield batch
|
|
expert_iter = cycle(expert_loader)
|
|
|
|
# 4. Initialize Env
|
|
from Env.expert_replay_env import ExpertReplayEnv # Using ReplayEnv for config, but we need ScenarioEnv for simulation?
|
|
# Actually we need MultiAgentScenarioEnv for interactive training, not Replay.
|
|
from Env.scenario_env import MultiAgentScenarioEnv
|
|
from Env.simple_idm_policy import ConstantVelocityPolicy # Placeholder policy for init
|
|
|
|
# 2. Setup Models
|
|
# Determine state dim from environment if possible, or use fixed
|
|
# Expert data has 45 dim?
|
|
# But Env might return something else if we are using default ScenarioEnv settings.
|
|
# ScenarioEnv returns list of obs.
|
|
# The error says: "mat1 and mat2 shapes cannot be multiplied (1x108 and 45x256)"
|
|
# This means the Env is returning 108-dim observation (MetaDrive default + Lidar),
|
|
# but our Actor expects 45 (which is what we saved in expert data).
|
|
|
|
# We must align the environment observation space with our expert data format.
|
|
# Our ExpertReplayEnv used a custom _get_all_obs.
|
|
# We need to inject that same logic into the training env, OR
|
|
# subclass MultiAgentScenarioEnv in the training script to override observation.
|
|
|
|
# class MAGAILScenarioEnv(MultiAgentScenarioEnv):
|
|
# def _get_all_obs(self):
|
|
# # Same logic as ExpertReplayEnv to ensure compatibility
|
|
# obs_dict = {}
|
|
# for agent_id, vehicle in self.controlled_agents.items():
|
|
# # 1. Ego State
|
|
# ego_state = [
|
|
# vehicle.position[0], vehicle.position[1],
|
|
# vehicle.velocity[0], vehicle.velocity[1],
|
|
# vehicle.heading_theta
|
|
# ]
|
|
#
|
|
# # 2. Neighbors
|
|
# candidates = []
|
|
# for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
|
|
# if other_id == agent_id:
|
|
# continue
|
|
# dist = np.linalg.norm(vehicle.position - other_vehicle.position)
|
|
# if dist < 30.0:
|
|
# candidates.append((dist, other_vehicle))
|
|
#
|
|
# candidates.sort(key=lambda x: x[0])
|
|
# top_10 = candidates[:10]
|
|
#
|
|
# neighbor_feats = []
|
|
# for _, neighbor in top_10:
|
|
# neighbor_feats.extend([
|
|
# neighbor.position[0] - vehicle.position[0],
|
|
# neighbor.position[1] - vehicle.position[1],
|
|
# neighbor.velocity[0],
|
|
# neighbor.velocity[1]
|
|
# ])
|
|
#
|
|
# missing = 10 - len(top_10)
|
|
# if missing > 0:
|
|
# neighbor_feats.extend([0.0] * (4 * missing))
|
|
#
|
|
# obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
|
|
# obs_dict[agent_id] = obs
|
|
# return obs_dict
|
|
|
|
env = MAGAILScenarioEnv(config=env_config, agent2policy={}) # Pass empty dict if we control all externally
|
|
|
|
print("Starting training...")
|
|
|
|
# Tensorboard
|
|
try:
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
writer = SummaryWriter(log_dir=args.log_dir)
|
|
except ImportError:
|
|
print("TensorBoard not installed. Logging to console only.")
|
|
writer = None
|
|
|
|
global_step = 0
|
|
|
|
for i_episode in range(args.max_episodes):
|
|
# --- 1. Collect Rollouts (Interaction) ---
|
|
memory = {
|
|
'states': [],
|
|
'actions': [],
|
|
'pre_tanh_actions': [],
|
|
'logprobs': [],
|
|
'rewards': [],
|
|
'next_states': [],
|
|
'dones': []
|
|
}
|
|
|
|
# Prepare seed
|
|
available_scenarios = env.config["num_scenarios"]
|
|
start_index = env.config["start_scenario_index"]
|
|
seed = np.random.randint(start_index, start_index + available_scenarios)
|
|
|
|
# Reset Env
|
|
try:
|
|
# MetaDrive sometimes complains about uncleared objects if reset happens too fast or with lingering objs
|
|
# We can try to force clear before reset or handle exception
|
|
# But standard reset should handle it.
|
|
# The error "You should clear all generated objects..." means some manager didn't clear its objects.
|
|
# This is likely due to TrafficManager or AgentManager holding refs.
|
|
|
|
# Re-creating env is safer but slower.
|
|
# Let's try closing and re-creating if reset fails frequently.
|
|
# Or just ignore this error and try reset again? No, reset failing is fatal usually.
|
|
|
|
# Hack: Manually clear objects if we can access engine
|
|
if env.engine is not None:
|
|
env.engine.clear_objects(list(env.engine.get_objects().keys()))
|
|
|
|
obs_dict = env.reset(seed=seed)
|
|
except Exception as e:
|
|
# print(f"Env reset failed: {e}. Recreating environment...")
|
|
try:
|
|
env.close()
|
|
except:
|
|
pass
|
|
|
|
# Ensure engine is closed properly
|
|
from metadrive.engine.engine_utils import close_engine
|
|
try:
|
|
close_engine()
|
|
except Exception as e2:
|
|
# Force cleanup of singleton if close failed
|
|
from metadrive.engine.base_engine import BaseEngine
|
|
if BaseEngine.singleton is not None:
|
|
BaseEngine.singleton = None
|
|
|
|
# Also need to clear ShowBase
|
|
try:
|
|
from direct.showbase.ShowBase import ShowBase
|
|
if hasattr(base, 'destroy'):
|
|
base.destroy()
|
|
except:
|
|
pass
|
|
|
|
# Brutal force: delete base from builtins if it exists
|
|
import builtins
|
|
if hasattr(builtins, 'base'):
|
|
del builtins.base
|
|
|
|
# print(f"Error closing engine: {e2}")
|
|
|
|
# Explicitly delete old env object to free memory
|
|
del env
|
|
import gc
|
|
gc.collect()
|
|
|
|
env = MAGAILScenarioEnv(config=env_config, agent2policy={})
|
|
obs_dict = env.reset(seed=seed)
|
|
|
|
episode_reward = 0
|
|
steps = 0
|
|
|
|
# Rollout loop
|
|
while True:
|
|
# Select actions for all agents
|
|
actions = {}
|
|
action_logprobs = {}
|
|
pre_tanh_actions = {}
|
|
|
|
# obs_dict: {agent_id: obs}
|
|
# MultiAgentScenarioEnv usually returns a dict {agent_id: obs}
|
|
# BUT wait, check scenario_env.py implementation
|
|
|
|
if isinstance(obs_dict, list):
|
|
# This happens if the environment returns a list instead of a dict
|
|
# MultiAgentScenarioEnv._get_all_obs returns a list in original implementation?
|
|
# Let's check scenario_env.py
|
|
# If it returns list, we need to map it to agent ids or just iterate
|
|
pass
|
|
|
|
# Temporary fix if it returns list (which means my previous edit to Env/expert_replay_env.py
|
|
# changed it there, but maybe not in Env/scenario_env.py which we are using here!)
|
|
|
|
if isinstance(obs_dict, list):
|
|
# We need agent IDs to step
|
|
# In MultiAgentScenarioEnv, controlled_agents is a dict.
|
|
# If obs is a list, it probably corresponds to controlled_agents.values() order?
|
|
# This is risky.
|
|
# Let's assume obs_dict is actually just observations.
|
|
# We need to keys to create action dict.
|
|
|
|
current_agent_ids = list(env.controlled_agents.keys())
|
|
# Ensure length matches
|
|
if len(obs_dict) != len(current_agent_ids):
|
|
# print(f"Warning: Obs list len {len(obs_dict)} != agents {len(current_agent_ids)}")
|
|
pass
|
|
|
|
# Reconstruct dict
|
|
new_obs_dict = {}
|
|
for i, agent_id in enumerate(current_agent_ids):
|
|
if i < len(obs_dict):
|
|
new_obs_dict[agent_id] = obs_dict[i]
|
|
obs_dict = new_obs_dict
|
|
|
|
for agent_id, obs in obs_dict.items():
|
|
act, logprob, pre_tanh = ppo_agent.select_action(obs) # Select action returns numpy
|
|
actions[agent_id] = act.flatten() # (2,)
|
|
action_logprobs[agent_id] = logprob # scalar
|
|
pre_tanh_actions[agent_id] = pre_tanh.flatten()
|
|
|
|
# Step Env
|
|
next_obs_dict, rewards, dones, infos = env.step(actions)
|
|
|
|
# Store in memory
|
|
for agent_id, obs in obs_dict.items():
|
|
if agent_id in actions:
|
|
memory['states'].append(obs)
|
|
memory['actions'].append(actions[agent_id])
|
|
memory['pre_tanh_actions'].append(pre_tanh_actions[agent_id])
|
|
memory['logprobs'].append(action_logprobs[agent_id])
|
|
|
|
# Store standard environmental reward for logging (not used for update in GAIL)
|
|
# For GAIL update we use Discriminator reward later
|
|
memory['rewards'].append(0) # Placeholder
|
|
|
|
# Next state
|
|
if agent_id in next_obs_dict:
|
|
memory['next_states'].append(next_obs_dict[agent_id])
|
|
memory['dones'].append(dones.get("__all__", False))
|
|
else:
|
|
# Agent finished/vanished
|
|
# We need a dummy next state or handle done correctly
|
|
# Just duplicate current state and mark done?
|
|
memory['next_states'].append(obs)
|
|
memory['dones'].append(True)
|
|
|
|
obs_dict = next_obs_dict
|
|
steps += 1
|
|
|
|
if dones["__all__"] or steps >= 200: # Limit horizon
|
|
break
|
|
|
|
# Initialize losses to 0/None before potential loop skip
|
|
disc_loss = torch.tensor(0.0)
|
|
ppo_loss = 0.0
|
|
all_gail_rewards = [0.0]
|
|
|
|
# --- 2. Train Discriminator ---
|
|
# Convert policy memory to tensors
|
|
policy_states = torch.FloatTensor(np.array(memory['states'])).cuda()
|
|
policy_actions = torch.FloatTensor(np.array(memory['actions'])).cuda()
|
|
|
|
# Sample expert batch
|
|
expert_batch = next(expert_iter)
|
|
|
|
expert_states = expert_batch['state'].cuda()
|
|
expert_actions = expert_batch['action'].cuda()
|
|
|
|
# Minibatch size matching
|
|
batch_size = min(policy_states.size(0), expert_states.size(0))
|
|
|
|
if batch_size > 0: # Only train if we have data
|
|
policy_states = policy_states[:batch_size]
|
|
policy_actions = policy_actions[:batch_size]
|
|
expert_states = expert_states[:batch_size]
|
|
expert_actions = expert_actions[:batch_size]
|
|
|
|
# Update Discriminator
|
|
# Label 1 for Expert, 0 for Policy
|
|
# Train Expert
|
|
disc_optimizer.zero_grad()
|
|
|
|
exp_preds = discriminator(expert_states, expert_actions)
|
|
exp_loss = disc_criterion(exp_preds, torch.ones_like(exp_preds))
|
|
|
|
pol_preds = discriminator(policy_states.detach(), policy_actions.detach()) # Detach policy data
|
|
pol_loss = disc_criterion(pol_preds, torch.zeros_like(pol_preds))
|
|
|
|
disc_loss = exp_loss + pol_loss
|
|
disc_loss.backward()
|
|
disc_optimizer.step()
|
|
|
|
with torch.no_grad():
|
|
disc_acc_exp = (exp_preds > 0.5).float().mean().item()
|
|
disc_acc_pol = (pol_preds < 0.5).float().mean().item()
|
|
|
|
# --- 3. Update Policy with GAIL Rewards ---
|
|
# Reward = -log(1 - D(s, a))
|
|
# Or more stable: log(D(s, a)) ? Original GAIL uses -log(1-D) which is log(D) roughly.
|
|
# Let's use -log(1 - D(s, a) + eps)
|
|
|
|
# Actually PPO needs the full trajectory for GAE.
|
|
# So we should compute rewards for ALL policy samples in memory.
|
|
|
|
all_policy_states = torch.FloatTensor(np.array(memory['states'])).cuda()
|
|
all_policy_actions = torch.FloatTensor(np.array(memory['actions'])).cuda()
|
|
|
|
with torch.no_grad():
|
|
all_d_val = discriminator(all_policy_states, all_policy_actions)
|
|
all_gail_rewards = -torch.log(1 - all_d_val + 1e-8).cpu().numpy().flatten()
|
|
|
|
# Replace placeholders
|
|
memory['rewards'] = all_gail_rewards.tolist()
|
|
|
|
# Update PPO
|
|
ppo_loss = ppo_agent.update(memory)
|
|
|
|
# Clean up memory
|
|
del policy_states, policy_actions, expert_states, expert_actions, exp_preds, exp_loss, pol_preds, pol_loss
|
|
del all_policy_states, all_policy_actions, all_d_val
|
|
torch.cuda.empty_cache()
|
|
else:
|
|
print(f"Episode {i_episode}: No data collected (Env might have crashed or no agents). Skipping update.")
|
|
|
|
# --- 4. Logging ---
|
|
if writer:
|
|
writer.add_scalar('Loss/Discriminator', disc_loss.item(), i_episode)
|
|
writer.add_scalar('Loss/Policy', ppo_loss, i_episode)
|
|
writer.add_scalar('Reward/Mean_GAIL', np.mean(all_gail_rewards), i_episode)
|
|
if batch_size > 0:
|
|
writer.add_scalar('Acc/Disc_Expert', disc_acc_exp, i_episode)
|
|
writer.add_scalar('Acc/Disc_Policy', disc_acc_pol, i_episode)
|
|
if len(memory['actions']) > 0:
|
|
action_arr = np.array(memory['actions'])
|
|
action_clip_ratio = (np.abs(action_arr) > 0.98).mean()
|
|
writer.add_scalar('Policy/ActionClipRatio', action_clip_ratio, i_episode)
|
|
|
|
print(f"Episode {i_episode}: Disc Loss {disc_loss.item():.4f} | PPO Loss {ppo_loss:.4f} | Mean Reward {np.mean(all_gail_rewards):.4f}")
|
|
|
|
if i_episode % 50 == 0:
|
|
ppo_agent.save(os.path.join(args.log_dir, f"model_{i_episode}"))
|
|
|
|
env.close()
|
|
if writer:
|
|
writer.close()
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--expert_data_dir", type=str, default="data/training_data", help="Directory with .pkl expert data")
|
|
parser.add_argument("--data_dir", type=str, default="data/exp_filtered", help="Waymo data dir for Env")
|
|
parser.add_argument("--batch_size", type=int, default=1024)
|
|
parser.add_argument("--max_episodes", type=int, default=1000)
|
|
parser.add_argument("--num_scenarios", type=int, default=100)
|
|
parser.add_argument("--log_dir", type=str, default="runs/magail_exp")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Create log dir
|
|
os.makedirs(args.log_dir, exist_ok=True)
|
|
|
|
train(args)
|