wehub-resource-sync
98 lines
3.4 KiB
Python
98 lines
3.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Author: SilverWings
|
|
GitHub: https://github.com/silverwingsbot
|
|
|
|
Simple example: Run a trained Diffusion_QL model in easycarla
|
|
"""
|
|
|
|
import gym
|
|
import easycarla
|
|
import numpy as np
|
|
import torch
|
|
import os
|
|
from agents.ql_diffusion import Diffusion_QL
|
|
|
|
# ===================== Helper Functions =====================
|
|
def convert_obs_dict_to_vector(obs_dict):
|
|
"""Convert observation dictionary to a flattened state vector."""
|
|
return np.concatenate([
|
|
obs_dict['ego_state'], # 9 dimensions
|
|
obs_dict['lane_info'], # 2 dimensions
|
|
obs_dict['lidar'], # 240 dimensions
|
|
obs_dict['nearby_vehicles'], # 20 dimensions
|
|
obs_dict['waypoints'] # 36 dimensions
|
|
]).astype(np.float32)
|
|
|
|
# ===================== Environment Configuration =====================
|
|
carla_params = {
|
|
'number_of_vehicles': 100,
|
|
'number_of_walkers': 0,
|
|
'dt': 0.1, # time interval between two frames
|
|
'ego_vehicle_filter': 'vehicle.tesla.model3', # filter for defining ego vehicle
|
|
'surrounding_vehicle_spawned_randomly': True, # Whether surrounding vehicles are spawned randomly (True) or set manually (False)
|
|
'port': 2000, # connection port
|
|
'town': 'Town03', # which town to simulate
|
|
'max_time_episode': 1000, # maximum timesteps per episode
|
|
'max_waypoints': 12, # maximum number of waypoints
|
|
'visualize_waypoints': True, # Whether to visualize waypoints (default: True)
|
|
'desired_speed': 8, # desired speed (m/s)
|
|
'max_ego_spawn_times': 200, # maximum times to spawn ego vehicle
|
|
'view_mode' : 'top', # 'top' for bird's-eye view, 'follow' for third-person view
|
|
'traffic': 'off', # 'on' for normal traffic lights, 'off' for always green and frozen
|
|
'lidar_max_range': 50.0, # Maximum LIDAR perception range (meters)
|
|
'max_nearby_vehicles': 5, # Maximum number of nearby vehicles to observe
|
|
}
|
|
|
|
# ===================== Initialize Environment =====================
|
|
env = gym.make('carla-v0', params=carla_params)
|
|
|
|
# ===================== Initialize Model =====================
|
|
state_dim = 307
|
|
action_dim = 3
|
|
max_action = 1.0
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
model = Diffusion_QL(
|
|
state_dim=state_dim,
|
|
action_dim=action_dim,
|
|
max_action=max_action,
|
|
device=device,
|
|
discount=0.99,
|
|
tau=0.005,
|
|
eta=0.01,
|
|
beta_schedule='vp',
|
|
n_timesteps=5
|
|
)
|
|
|
|
# ===================== Load Pretrained Model =====================
|
|
model_id = 200 # Model checkpoint ID to load
|
|
save_path = './params_dql' # Model checkpoint directory
|
|
model.load_model(save_path, id=model_id)
|
|
print(f"Successfully loaded model ID {model_id}")
|
|
|
|
# ===================== Run One Episode =====================
|
|
obs = env.reset()
|
|
done = False
|
|
step = 0
|
|
episode_reward = 0.0
|
|
|
|
while not done:
|
|
obs_vec = convert_obs_dict_to_vector(obs)
|
|
action = model.sample_action(obs_vec)
|
|
|
|
try:
|
|
next_obs, reward, cost, done, info = env.step(action)
|
|
except Exception as e:
|
|
print(f"[Error] Carla step failed: {e}")
|
|
obs = env.reset()
|
|
continue
|
|
|
|
obs = next_obs
|
|
episode_reward += reward
|
|
step += 1
|
|
|
|
# Optional: add a delay for better visualization
|
|
# time.sleep(0.05)
|
|
|
|
print(f"Episode finished. Total reward: {episode_reward:.2f}, Total steps: {step}") |