class DQNPolicy(Agent): lane_actions = [ "keep_lane", "slow_down", "change_lane_left", "change_lane_right" ] def __init__( self, policy_params=None, checkpoint_dir=None, ): self.policy_params = policy_params network_class = DQNWithSocialEncoder self.epsilon_obj = EpsilonExplore(1.0, 0.05, 100000) action_space_type = policy_params["action_space_type"] if action_space_type == "continuous": discrete_action_spaces = [ np.asarray([-0.25, 0.0, 0.5, 0.75, 1.0]), np.asarray([ -1.0, -0.75, -0.5, -0.25, -0.1, 0.0, 0.1, 0.25, 0.5, 0.75, 1.0 ]), ] else: discrete_action_spaces = [[0], [1]] action_size = discrete_action_spaces self.merge_action_spaces = 0 if action_space_type == "continuous" else -1 self.step_count = 0 self.update_count = 0 self.num_updates = 0 self.current_sticky = 0 self.current_iteration = 0 lr = float(policy_params["lr"]) seed = int(policy_params["seed"]) self.train_step = int(policy_params["train_step"]) self.target_update = float(policy_params["target_update"]) self.device_name = "cuda:0" if torch.cuda.is_available() else "cpu" self.device = torch.device(self.device_name) self.warmup = int(policy_params["warmup"]) self.gamma = float(policy_params["gamma"]) self.batch_size = int(policy_params["batch_size"]) self.use_ddqn = policy_params["use_ddqn"] self.sticky_actions = int(policy_params["sticky_actions"]) prev_action_size = int(policy_params["prev_action_size"]) self.prev_action = np.zeros(prev_action_size) if self.merge_action_spaces == 1: index2action, action2index = merge_discrete_action_spaces( *action_size) self.index2actions = [index2action] self.action2indexs = [action2index] self.num_actions = [len(self.index2actions)] elif self.merge_action_spaces == 0: self.index2actions = [ merge_discrete_action_spaces([each])[0] for each in action_size ] self.action2indexs = [ merge_discrete_action_spaces([each])[1] for each in action_size ] self.num_actions = [len(e) for e in action_size] else: index_to_actions = [ e.tolist() if not isinstance(e, list) else e for e in action_size ] action_to_indexs = { str(k): v for k, v in zip( index_to_actions, np.arange(len(index_to_actions)).astype(np.int)) } self.index2actions, self.action2indexs = ( [index_to_actions], [action_to_indexs], ) self.num_actions = [len(index_to_actions)] # state preprocessing self.social_policy_hidden_units = int( policy_params["social_vehicles"].get("social_policy_hidden_units", 0)) self.social_capacity = int(policy_params["social_vehicles"].get( "social_capacity", 0)) self.observation_num_lookahead = int( policy_params.get("observation_num_lookahead", 0)) self.social_polciy_init_std = int(policy_params["social_vehicles"].get( "social_polciy_init_std", 0)) self.num_social_features = int(policy_params["social_vehicles"].get( "num_social_features", 0)) self.social_vehicle_config = get_social_vehicle_configs( **policy_params["social_vehicles"]) self.social_vehicle_encoder = self.social_vehicle_config["encoder"] self.state_description = get_state_description( policy_params["social_vehicles"], policy_params["observation_num_lookahead"], prev_action_size, ) self.social_feature_encoder_class = self.social_vehicle_encoder[ "social_feature_encoder_class"] self.social_feature_encoder_params = self.social_vehicle_encoder[ "social_feature_encoder_params"] self.checkpoint_dir = checkpoint_dir self.reset() torch.manual_seed(seed) network_params = { "state_size": self.state_size, "social_feature_encoder_class": self.social_feature_encoder_class, "social_feature_encoder_params": self.social_feature_encoder_params, } self.online_q_network = network_class( num_actions=self.num_actions, **(network_params if network_params else {}), ).to(self.device) self.target_q_network = network_class( num_actions=self.num_actions, **(network_params if network_params else {}), ).to(self.device) self.update_target_network() self.optimizers = torch.optim.Adam( params=self.online_q_network.parameters(), lr=lr) self.loss_func = nn.MSELoss(reduction="none") if self.checkpoint_dir: self.load(self.checkpoint_dir) self.action_space_type = "continuous" self.to_real_action = to_3d_action self.state_preprocessor = StatePreprocessor(preprocess_state, to_2d_action, self.state_description) self.replay = ReplayBuffer( buffer_size=int(policy_params["replay_buffer"]["buffer_size"]), batch_size=int(policy_params["replay_buffer"]["batch_size"]), state_preprocessor=self.state_preprocessor, device_name=self.device_name, ) def lane_action_to_index(self, state): state = state.copy() if (len(state["action"]) == 3 and (state["action"] == np.asarray( [0, 0, 0])).all()): # initial action state["action"] = np.asarray([0]) else: state["action"] = self.lane_actions.index(state["action"]) return state @property def state_size(self): # Adjusting state_size based on number of features (ego+social) size = sum(self.state_description["low_dim_states"].values()) if self.social_feature_encoder_class: size += self.social_feature_encoder_class( **self.social_feature_encoder_params).output_dim else: size += self.social_capacity * self.num_social_features return size def reset(self): self.eps_throttles = [] self.eps_steers = [] self.eps_step = 0 self.current_sticky = 0 def soft_update(self, target, src, tau): for target_param, param in zip(target.parameters(), src.parameters()): target_param.detach_() target_param.copy_(target_param * (1.0 - tau) + param * tau) def update_target_network(self): self.target_q_network.load_state_dict( self.online_q_network.state_dict().copy()) def act(self, *args, **kwargs): if self.current_sticky == 0: self.action = self._act(*args, **kwargs) self.current_sticky = (self.current_sticky + 1) % self.sticky_actions self.current_iteration += 1 return self.to_real_action(self.action) def _act(self, state, explore=True): epsilon = self.epsilon_obj.get_epsilon() if not explore or np.random.rand() > epsilon: state = self.state_preprocessor( state, normalize=True, unsqueeze=True, device=self.device, social_capacity=self.social_capacity, observation_num_lookahead=self.observation_num_lookahead, social_vehicle_config=self.social_vehicle_config, prev_action=self.prev_action, ) self.online_q_network.eval() with torch.no_grad(): qs = self.online_q_network(state) qs = [q.data.cpu().numpy().flatten() for q in qs] # out_str = " || ".join( # [ # " ".join( # [ # "{}: {:.4f}".format(index2action[j], q[j]) # for j in range(num_action) # ] # ) # for index2action, q, num_action in zip( # self.index2actions, qs, self.num_actions # ) # ] # ) # print(out_str) inds = [np.argmax(q) for q in qs] else: inds = [ np.random.randint(num_action) for num_action in self.num_actions ] action = [] for j, ind in enumerate(inds): action.extend(self.index2actions[j][ind]) self.epsilon_obj.step() self.eps_step += 1 action = np.asarray(action) return action def save(self, model_dir): model_dir = pathlib.Path(model_dir) torch.save(self.online_q_network.state_dict(), model_dir / "online.pth") torch.save(self.target_q_network.state_dict(), model_dir / "target.pth") def load(self, model_dir, cpu=False): model_dir = pathlib.Path(model_dir) print("loading from :", model_dir) map_location = None if cpu: map_location = torch.device("cpu") self.online_q_network.load_state_dict( torch.load(model_dir / "online.pth", map_location=map_location)) self.target_q_network.load_state_dict( torch.load(model_dir / "target.pth", map_location=map_location)) print("Model loaded") def step(self, state, action, reward, next_state, done, others=None): # dont treat timeout as done equal to True max_steps_reached = state["events"].reached_max_episode_steps if max_steps_reached: done = False if self.action_space_type == "continuous": action = to_2d_action(action) _action = ([[e] for e in action] if not self.merge_action_spaces else [action.tolist()]) action_index = np.asarray([ action2index[str(e)] for action2index, e in zip(self.action2indexs, _action) ]) else: action_index = self.lane_actions.index(action) action = action_index self.replay.add( state=state, action=action_index, reward=reward, next_state=next_state, done=done, others=others, social_capacity=self.social_capacity, observation_num_lookahead=self.observation_num_lookahead, social_vehicle_config=self.social_vehicle_config, prev_action=self.prev_action, ) if (self.step_count % self.train_step == 0 and len(self.replay) >= self.batch_size and (self.warmup is None or len(self.replay) >= self.warmup)): out = self.learn() self.update_count += 1 else: out = {} if self.target_update > 1 and self.step_count % self.target_update == 0: self.update_target_network() elif self.target_update < 1.0: self.soft_update(self.target_q_network, self.online_q_network, self.target_update) self.step_count += 1 self.prev_action = action return out def learn(self): states, actions, rewards, next_states, dones, others = self.replay.sample( device=self.device) if not self.merge_action_spaces: actions = torch.chunk(actions, len(self.num_actions), -1) else: actions = [actions] self.target_q_network.eval() with torch.no_grad(): qs_next_target = self.target_q_network(next_states) if self.use_ddqn: self.online_q_network.eval() with torch.no_grad(): qs_next_online = self.online_q_network(next_states) next_actions = [ torch.argmax(q_next_online, dim=1, keepdim=True) for q_next_online in qs_next_online ] else: next_actions = [ torch.argmax(q_next_target, dim=1, keepdim=True) for q_next_target in qs_next_target ] qs_next_target = [ torch.gather(q_next_target, 1, next_action) for q_next_target, next_action in zip(qs_next_target, next_actions) ] self.online_q_network.train() qs, aux_losses = self.online_q_network(states, training=True) qs = [ torch.gather(q, 1, action.long()) for q, action in zip(qs, actions) ] qs_target_value = [ rewards + self.gamma * (1 - dones) * q_next_target for q_next_target in qs_next_target ] td_loss = [ self.loss_func(q, q_target_value).mean() for q, q_target_value in zip(qs, qs_target_value) ] mean_td_loss = sum(td_loss) / len(td_loss) loss = mean_td_loss + sum( [e["value"] * e["weight"] for e in aux_losses.values()]) self.optimizers.zero_grad() loss.backward() self.optimizers.step() out = {} out.update({ "loss/td{}".format(j): { "type": "scalar", "data": td_loss[j].data.cpu().numpy(), "freq": 10, } for j in range(len(td_loss)) }) out.update({ "loss/{}".format(k): { "type": "scalar", "data": v["value"], # .detach().cpu().numpy(), "freq": 10, } for k, v in aux_losses.items() }) out.update({"loss/all": {"type": "scalar", "data": loss, "freq": 10}}) self.num_updates += 1 return out
class DQNPolicy(Agent): lane_actions = [ "keep_lane", "slow_down", "change_lane_left", "change_lane_right" ] def __init__( self, policy_params=None, checkpoint_dir=None, ): self.policy_params = policy_params self.lr = float(policy_params["lr"]) self.seed = int(policy_params["seed"]) self.train_step = int(policy_params["train_step"]) self.target_update = float(policy_params["target_update"]) self.warmup = int(policy_params["warmup"]) self.gamma = float(policy_params["gamma"]) self.batch_size = int(policy_params["batch_size"]) self.use_ddqn = policy_params["use_ddqn"] self.sticky_actions = int(policy_params["sticky_actions"]) self.epsilon_obj = EpsilonExplore(1.0, 0.05, 100000) self.step_count = 0 self.update_count = 0 self.num_updates = 0 self.current_sticky = 0 self.current_iteration = 0 self.action_type = adapters.type_from_string( policy_params["action_type"]) self.observation_type = adapters.type_from_string( policy_params["observation_type"]) self.reward_type = adapters.type_from_string( policy_params["reward_type"]) if self.action_type == adapters.AdapterType.DefaultActionContinuous: discrete_action_spaces = [ np.asarray([-0.25, 0.0, 0.5, 0.75, 1.0]), np.asarray([ -1.0, -0.75, -0.5, -0.25, -0.1, 0.0, 0.1, 0.25, 0.5, 0.75, 1.0 ]), ] self.index2actions = [ merge_discrete_action_spaces([discrete_action_space])[0] for discrete_action_space in discrete_action_spaces ] self.action2indexs = [ merge_discrete_action_spaces([discrete_action_space])[1] for discrete_action_space in discrete_action_spaces ] self.merge_action_spaces = 0 self.num_actions = [ len(discrete_action_space) for discrete_action_space in discrete_action_spaces ] self.action_size = 2 self.to_real_action = to_3d_action elif self.action_type == adapters.AdapterType.DefaultActionDiscrete: discrete_action_spaces = [[0], [1], [2], [3]] index_to_actions = [ discrete_action_space.tolist() if not isinstance(discrete_action_space, list) else discrete_action_space for discrete_action_space in discrete_action_spaces ] action_to_indexs = { str(discrete_action): index for discrete_action, index in zip( index_to_actions, np.arange(len(index_to_actions)).astype(np.int)) } self.index2actions = [index_to_actions] self.action2indexs = [action_to_indexs] self.merge_action_spaces = -1 self.num_actions = [len(index_to_actions)] self.action_size = 1 self.to_real_action = lambda action: self.lane_actions[action[0]] else: raise Exception( f"DQN baseline does not support the '{self.action_type}' action type." ) if self.observation_type == adapters.AdapterType.DefaultObservationVector: observation_space = adapters.space_from_type(self.observation_type) low_dim_states_size = observation_space["low_dim_states"].shape[0] social_capacity = observation_space["social_vehicles"].shape[0] num_social_features = observation_space["social_vehicles"].shape[1] # Get information to build the encoder. encoder_key = policy_params["social_vehicles"]["encoder_key"] social_policy_hidden_units = int( policy_params["social_vehicles"].get( "social_policy_hidden_units", 0)) social_policy_init_std = int(policy_params["social_vehicles"].get( "social_policy_init_std", 0)) social_vehicle_config = get_social_vehicle_configs( encoder_key=encoder_key, num_social_features=num_social_features, social_capacity=social_capacity, seed=self.seed, social_policy_hidden_units=social_policy_hidden_units, social_policy_init_std=social_policy_init_std, ) social_vehicle_encoder = social_vehicle_config["encoder"] social_feature_encoder_class = social_vehicle_encoder[ "social_feature_encoder_class"] social_feature_encoder_params = social_vehicle_encoder[ "social_feature_encoder_params"] # Calculate the state size based on the number of features (ego + social). state_size = low_dim_states_size if social_feature_encoder_class: state_size += social_feature_encoder_class( **social_feature_encoder_params).output_dim else: state_size += social_capacity * num_social_features # Add the action size to account for the previous action. state_size += self.action_size network_class = DQNWithSocialEncoder network_params = { "num_actions": self.num_actions, "state_size": state_size, "social_feature_encoder_class": social_feature_encoder_class, "social_feature_encoder_params": social_feature_encoder_params, } elif self.observation_type == adapters.AdapterType.DefaultObservationImage: observation_space = adapters.space_from_type(self.observation_type) stack_size = observation_space.shape[0] image_shape = (observation_space.shape[1], observation_space.shape[2]) network_class = DQNCNN network_params = { "n_in_channels": stack_size, "image_dim": image_shape, "num_actions": self.num_actions, } else: raise Exception( f"DQN baseline does not support the '{self.observation_type}' " f"observation type.") self.prev_action = np.zeros(self.action_size) self.checkpoint_dir = checkpoint_dir torch.manual_seed(self.seed) self.device_name = "cuda:0" if torch.cuda.is_available() else "cpu" self.device = torch.device(self.device_name) self.online_q_network = network_class(**network_params).to(self.device) self.target_q_network = network_class(**network_params).to(self.device) self.update_target_network() self.optimizers = torch.optim.Adam( params=self.online_q_network.parameters(), lr=self.lr) self.loss_func = nn.MSELoss(reduction="none") self.replay = ReplayBuffer( buffer_size=int(policy_params["replay_buffer"]["buffer_size"]), batch_size=int(policy_params["replay_buffer"]["batch_size"]), observation_type=self.observation_type, device_name=self.device_name, ) self.reset() if self.checkpoint_dir: self.load(self.checkpoint_dir) def lane_action_to_index(self, state): state = state.copy() if (len(state["action"]) == 3 and (state["action"] == np.asarray( [0, 0, 0])).all()): # initial action state["action"] = np.asarray([0]) else: state["action"] = self.lane_actions.index(state["action"]) return state def reset(self): self.eps_throttles = [] self.eps_steers = [] self.eps_step = 0 self.current_sticky = 0 def soft_update(self, target, src, tau): for target_param, param in zip(target.parameters(), src.parameters()): target_param.detach_() target_param.copy_(target_param * (1.0 - tau) + param * tau) def update_target_network(self): self.target_q_network.load_state_dict( self.online_q_network.state_dict().copy()) def act(self, *args, **kwargs): if self.current_sticky == 0: self.action = self._act(*args, **kwargs) self.current_sticky = (self.current_sticky + 1) % self.sticky_actions self.current_iteration += 1 return self.to_real_action(self.action) def _act(self, state, explore=True): epsilon = self.epsilon_obj.get_epsilon() if not explore or np.random.rand() > epsilon: state = copy.deepcopy(state) if self.observation_type == adapters.AdapterType.DefaultObservationVector: # Default vector observation type. state["low_dim_states"] = np.float32( np.append(state["low_dim_states"], self.prev_action)) state["social_vehicles"] = (torch.from_numpy( state["social_vehicles"]).unsqueeze(0).to(self.device)) state["low_dim_states"] = (torch.from_numpy( state["low_dim_states"]).unsqueeze(0).to(self.device)) else: # Default image observation type. state = torch.from_numpy(state).unsqueeze(0).to(self.device) self.online_q_network.eval() with torch.no_grad(): qs = self.online_q_network(state) qs = [q.data.cpu().numpy().flatten() for q in qs] # out_str = " || ".join( # [ # " ".join( # [ # "{}: {:.4f}".format(index2action[j], q[j]) # for j in range(num_action) # ] # ) # for index2action, q, num_action in zip( # self.index2actions, qs, self.num_actions # ) # ] # ) # print(out_str) inds = [np.argmax(q) for q in qs] else: inds = [ np.random.randint(num_action) for num_action in self.num_actions ] action = [] for j, ind in enumerate(inds): action.extend(self.index2actions[j][ind]) self.epsilon_obj.step() self.eps_step += 1 action = np.asarray(action) return action def save(self, model_dir): model_dir = pathlib.Path(model_dir) torch.save(self.online_q_network.state_dict(), model_dir / "online.pth") torch.save(self.target_q_network.state_dict(), model_dir / "target.pth") def load(self, model_dir, cpu=False): model_dir = pathlib.Path(model_dir) print("loading from :", model_dir) map_location = None if cpu: map_location = torch.device("cpu") self.online_q_network.load_state_dict( torch.load(model_dir / "online.pth", map_location=map_location)) self.target_q_network.load_state_dict( torch.load(model_dir / "target.pth", map_location=map_location)) print("Model loaded") def step(self, state, action, reward, next_state, done, info, others=None): # dont treat timeout as done equal to True max_steps_reached = info["logs"]["events"].reached_max_episode_steps if max_steps_reached: done = False if self.action_type == adapters.AdapterType.DefaultActionContinuous: action = to_2d_action(action) _action = ([[e] for e in action] if not self.merge_action_spaces else [action.tolist()]) action_index = np.asarray([ action2index[str(e)] for action2index, e in zip(self.action2indexs, _action) ]) else: action_index = self.lane_actions.index(action) action = action_index self.replay.add( state=state, action=action_index, reward=reward, next_state=next_state, done=done, others=others, prev_action=self.prev_action, ) if (self.step_count % self.train_step == 0 and len(self.replay) >= self.batch_size and (self.warmup is None or len(self.replay) >= self.warmup)): out = self.learn() self.update_count += 1 else: out = {} if self.target_update > 1 and self.step_count % self.target_update == 0: self.update_target_network() elif self.target_update < 1.0: self.soft_update(self.target_q_network, self.online_q_network, self.target_update) self.step_count += 1 self.prev_action = action return out def learn(self): states, actions, rewards, next_states, dones, others = self.replay.sample( device=self.device) if not self.merge_action_spaces: actions = torch.chunk(actions, len(self.num_actions), -1) else: actions = [actions] self.target_q_network.eval() with torch.no_grad(): qs_next_target = self.target_q_network(next_states) if self.use_ddqn: self.online_q_network.eval() with torch.no_grad(): qs_next_online = self.online_q_network(next_states) next_actions = [ torch.argmax(q_next_online, dim=1, keepdim=True) for q_next_online in qs_next_online ] else: next_actions = [ torch.argmax(q_next_target, dim=1, keepdim=True) for q_next_target in qs_next_target ] qs_next_target = [ torch.gather(q_next_target, 1, next_action) for q_next_target, next_action in zip(qs_next_target, next_actions) ] self.online_q_network.train() qs, aux_losses = self.online_q_network(states, training=True) qs = [ torch.gather(q, 1, action.long()) for q, action in zip(qs, actions) ] qs_target_value = [ rewards + self.gamma * (1 - dones) * q_next_target for q_next_target in qs_next_target ] td_loss = [ self.loss_func(q, q_target_value).mean() for q, q_target_value in zip(qs, qs_target_value) ] mean_td_loss = sum(td_loss) / len(td_loss) loss = mean_td_loss + sum( [e["value"] * e["weight"] for e in aux_losses.values()]) self.optimizers.zero_grad() loss.backward() self.optimizers.step() out = {} out.update({ "loss/td{}".format(j): { "type": "scalar", "data": td_loss[j].data.cpu().numpy(), "freq": 10, } for j in range(len(td_loss)) }) out.update({ "loss/{}".format(k): { "type": "scalar", "data": v["value"], # .detach().cpu().numpy(), "freq": 10, } for k, v in aux_losses.items() }) out.update({"loss/all": {"type": "scalar", "data": loss, "freq": 10}}) self.num_updates += 1 return out