Spaces:
Runtime error
Runtime error
| import io | |
| import warnings | |
| from typing import Tuple, Dict, Optional, List, Text | |
| import gym | |
| import math | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import pickle, os | |
| from numpy import ndarray | |
| from facility_location.utils.config import Config | |
| from facility_location.env.facility_location_client import FacilityLocationClient | |
| from facility_location.env.obs_extractor import ObsExtractor | |
| from stable_baselines3 import PPO | |
| from stable_baselines3.common.vec_env import DummyVecEnv | |
| from facility_location.agent import MaskedFacilityLocationActorCriticPolicy | |
| from facility_location.utils.policy import get_policy_kwargs | |
| class PMPEnv(gym.Env): | |
| EPSILON = 1e-6 | |
| def __init__(self, | |
| cfg: Config): | |
| self.cfg = cfg | |
| self._train_region = cfg.env_specs['region'] | |
| self._eval_region = cfg.eval_specs['region'] | |
| self._min_n = cfg.env_specs['min_n'] | |
| self._max_n = cfg.env_specs['max_n'] | |
| self._min_p_ratio = cfg.env_specs['min_p_ratio'] | |
| self._max_p_ratio = cfg.env_specs['max_p_ratio'] | |
| self._max_steps_scale = cfg.env_specs['max_steps_scale'] | |
| self._tabu_stable_steps_scale = cfg.env_specs['tabu_stable_steps_scale'] | |
| self._popstar = cfg.env_specs['popstar'] | |
| self._seed(cfg.seed) | |
| self._done = False | |
| self._set_node_edge_range() | |
| self._flc = FacilityLocationClient(cfg, self._np_random) | |
| self._obs_extractor = ObsExtractor(cfg, self._flc, self._node_range, self._edge_range) | |
| self._declare_spaces() | |
| def _declare_spaces(self) -> None: | |
| self.observation_space = gym.spaces.Dict({ | |
| 'node_features': gym.spaces.Box(low=0, high=1, shape=(self._node_range, self.get_node_feature_dim())), | |
| 'static_adjacency_list': gym.spaces.Box(low=0, high=self._node_range, shape=(self._edge_range, 2), dtype=np.int64), | |
| 'dynamic_adjacency_list': gym.spaces.Box(low=0, high=self._node_range, shape=(self._edge_range, 2), dtype=np.int64), | |
| 'node_mask': gym.spaces.Box(low=0, high=1, shape=(self._node_range,), dtype=np.bool), | |
| 'static_edge_mask': gym.spaces.Box(low=0, high=1, shape=(self._edge_range,), dtype=np.bool), | |
| 'dynamic_edge_mask': gym.spaces.Box(low=0, high=1, shape=(self._edge_range,), dtype=np.bool), | |
| }) | |
| if not self._popstar: | |
| self.action_space = gym.spaces.Discrete(self._node_range ** 2) | |
| else: | |
| self.action_space = gym.spaces.Discrete(self._node_range ** 2) | |
| def _set_node_edge_range(self) -> None: | |
| self._node_range = self._max_n + 2 | |
| self._edge_range = int(self._max_n ** 2 * self._max_p_ratio) | |
| def get_node_feature_dim(self) -> int: | |
| return self._obs_extractor.get_node_dim() | |
| def _seed(self, seed: int) -> None: | |
| self._np_random = np.random.default_rng(seed) | |
| def get_reward(self) -> float: | |
| reward = self._obj_value[self._t - 1] - self._obj_value[self._t] | |
| return reward | |
| def _transform_action(self, action: np.ndarray) -> np.ndarray: | |
| if self._popstar: | |
| action = np.array(np.unravel_index(action, (self._node_range, self._node_range))) | |
| action = action - 1 | |
| return action | |
| def step(self, action: np.ndarray): | |
| if self._done: | |
| raise RuntimeError('Action taken after episode is done.') | |
| obj_value, solution, info = self._flc.swap(action, self._t) | |
| self._t += 1 | |
| self._done = (self._t == self._max_steps) | |
| self._obj_value[self._t] = obj_value | |
| self._solution[self._t] = solution | |
| reward = self.get_reward() | |
| if obj_value < self._best_obj_value - self.EPSILON: | |
| self._best_obj_value = obj_value | |
| self._best_solution = solution | |
| self._last_best_t = self._t | |
| elif (self._t - self._last_best_t) % self._tabu_stable_steps == 0: | |
| self._flc.reset_tabu_time() | |
| # if self._done: | |
| # print('done') | |
| # for i in range(self._t): | |
| # print(f'{i}:',np.where(self._solution[i])) | |
| return self._get_obs(self._t), reward, self._done, False, info | |
| def reset(self, seed = 0) -> Optional[Dict]: | |
| if self._train_region is None: | |
| points, demands, n, p = self._generate_new_instance() | |
| self._flc.set_instance(points, demands, n, p, False) | |
| else: | |
| points, demands, n, p = self._use_real_instance() | |
| self._flc.set_instance(points, demands, n, p, True) | |
| return self.prepare(n, p), {} | |
| def prepare(self, n: int, p: int) -> Dict: | |
| initial_obj_value, initial_solution = self._flc.compute_initial_solution() | |
| self._obs_extractor.reset() | |
| self._done = False | |
| self._t = 0 | |
| self._max_steps = max(int(p * self._max_steps_scale), 5) | |
| self._obj_value = np.zeros(self._max_steps + 1) | |
| self._obj_value[0] = initial_obj_value | |
| self._solution = np.zeros((self._max_steps + 1, n), dtype=bool) | |
| self._solution[0] = initial_solution | |
| self._best_solution = initial_solution | |
| self._best_obj_value = initial_obj_value | |
| self._last_best_t = 0 | |
| self._tabu_stable_steps = max(1, round(self._max_steps * self._tabu_stable_steps_scale)) | |
| return self._get_obs(self._t) | |
| def render(self, mode='human', dpi=300) -> Optional[np.ndarray]: | |
| gdf, facilities = self._flc.get_gdf_facilities() | |
| if len(facilities) > 10: | |
| warnings.warn('Too many facilities to render. Only rendering the first 10.') | |
| facilities = facilities[:10] | |
| cm = plt.get_cmap('tab10') | |
| fig, axs = plt.subplots(1, 2, figsize=(12, 6), dpi=dpi) | |
| for i, f in enumerate(facilities): | |
| gdf.loc[gdf['assignment'] == f].plot(ax=axs[0], | |
| zorder=2, | |
| alpha=0.7, | |
| edgecolor="k", | |
| color=cm(i)) | |
| gdf.loc[[f]].plot(ax=axs[0], | |
| marker='*', | |
| markersize=300, | |
| zorder=3, | |
| alpha=0.7, | |
| edgecolor="k", | |
| color=cm(i)) | |
| axs[0].set_title("Facility Location", fontweight="bold") | |
| plot_obj_value = self._obj_value[:self._t + 1] | |
| axs[1].plot(plot_obj_value, marker='.', markersize=10, color='k') | |
| axs[1].set_title("Objective Value", fontweight="bold") | |
| axs[1].set_xticks(np.arange(self._max_steps + 1, step=math.ceil((self._max_steps + 1) / 10))) | |
| fig.tight_layout() | |
| if mode == 'human': | |
| plt.show() | |
| else: | |
| io_buf = io.BytesIO() | |
| fig.savefig(io_buf, format='raw', dpi=dpi) | |
| io_buf.seek(0) | |
| img_arr = np.reshape(np.frombuffer(io_buf.getvalue(), dtype=np.uint8), | |
| newshape=(int(fig.bbox.bounds[3]), int(fig.bbox.bounds[2]), -1)) | |
| io_buf.close() | |
| return img_arr | |
| def close(self): | |
| plt.close() | |
| def _generate_new_instance(self) -> Tuple[np.ndarray, np.ndarray, int, int]: | |
| n = self._np_random.integers(self._min_n, self._max_n, endpoint=True) | |
| p_ratio = self._np_random.uniform(self._min_p_ratio, self._max_p_ratio) | |
| p = int(max(n * p_ratio, 4)) | |
| points = self._np_random.uniform(size=(n, 2)) | |
| while np.unique(points, axis=0).shape[0] != n: | |
| points = self._np_random.uniform(size=(n, 2)) | |
| demands = self._np_random.random(size=(n,)) | |
| return points, demands, n, p | |
| def _use_real_instance(self) -> Tuple[np.ndarray, np.ndarray, int, int]: | |
| data_path = './data/{}/pkl'.format(self.cfg.eval_specs['region']) | |
| files = os.listdir(data_path) | |
| files = [f for f in files if f.endswith('.pkl')] | |
| sample_data_path = os.path.join(data_path, files[self._np_random.integers(len(files))]) | |
| with open(sample_data_path, 'rb') as f: | |
| np_data = pickle.load(f) | |
| n = self._np_random.integers(self._min_n, self._max_n, endpoint=True) | |
| p = max(int(n * self._np_random.uniform(self._min_p_ratio, self._max_p_ratio)), 4) | |
| sample_cbgs = self._np_random.choice(list(np_data[1].keys()), n, replace=False) | |
| points = [] | |
| demands = [] | |
| for cbg in sample_cbgs: | |
| points.append(np_data[1][cbg]['pos']) | |
| demands.append(np_data[1][cbg]['demand']) | |
| points = np.array(points) | |
| demands = np.array(demands) | |
| return points, demands, n, p | |
| def _get_obs(self, t: int) -> Dict: | |
| return self._obs_extractor.get_obs(t) | |
| def get_initial_solution(self) -> np.ndarray: | |
| return self._solution[0] | |
| class EvalPMPEnv(PMPEnv): | |
| def __init__(self, | |
| cfg: Config, | |
| positions, demands, n, p, boost=False): | |
| self._eval_np = (n,p) | |
| self._eval_seed = cfg.eval_specs['seed'] | |
| self._boost = boost | |
| self.points = positions | |
| self.demands = demands | |
| self._n = n | |
| self._p = p | |
| super().__init__(cfg) | |
| def _set_node_edge_range(self) -> None: | |
| n, p = self._eval_np | |
| self._node_range = n + 2 | |
| self._edge_range = n * p | |
| def get_eval_num_cases(self) -> int: | |
| return self._eval_num_cases | |
| def get_eval_np(self) -> Tuple[int, int]: | |
| return self._eval_np | |
| def reset_instance_id(self) -> None: | |
| self._instance_id = 0 | |
| def step(self, action: np.ndarray): | |
| if self._done: | |
| raise RuntimeError('Action taken after episode is done.') | |
| obj_value, solution, info = self._flc.swap(action, self._t) | |
| self._t += 1 | |
| self._done = (self._t == self._max_steps) | |
| self._obj_value[self._t] = obj_value | |
| self._solution[self._t] = solution | |
| reward = self.get_reward() | |
| if obj_value < self._best_obj_value - self.EPSILON: | |
| self._best_obj_value = obj_value | |
| self._best_solution = solution | |
| self._last_best_t = self._t | |
| elif (self._t - self._last_best_t) % self._tabu_stable_steps == 0: | |
| self._flc.reset_tabu_time() | |
| print(self._t, self._max_steps) | |
| return self._get_obs(self._t), reward, self._done, False, info | |
| def get_reward(self) -> float: | |
| if self._done: | |
| reward = -np.min(self._obj_value) | |
| else: | |
| reward = 0.0 | |
| return reward | |
| def get_best_solution(self) -> np.ndarray: | |
| return self._best_solution | |
| def reset(self, seed = 0) -> Dict: | |
| self._flc.set_instance(self.points, self.demands, self._n, self._p, False) | |
| return self.prepare(self._n, self._p, self._boost), {} | |
| def prepare(self, n: int, p: int, boost: bool) -> Dict: | |
| initial_obj_value, initial_solution = self._flc.compute_initial_solution() | |
| self._obs_extractor.reset() | |
| self._done = False | |
| self._t = 0 | |
| self._max_steps = max(int(p * self._max_steps_scale), 5) | |
| if boost: | |
| self._max_steps = max(int(self._max_steps_scale / 10), 5) | |
| self._obj_value = np.zeros(self._max_steps + 1) | |
| self._obj_value[0] = initial_obj_value | |
| self._solution = np.zeros((self._max_steps + 1, n), dtype=bool) | |
| self._solution[0] = initial_solution | |
| self._best_solution = initial_solution | |
| self._best_obj_value = initial_obj_value | |
| self._last_best_t = 0 | |
| self._tabu_stable_steps = max(1, round(self._max_steps * self._tabu_stable_steps_scale)) | |
| return self._get_obs(self._t) | |
| def get_instance(self) -> Tuple[np.ndarray, np.ndarray, int, int]: | |
| points, demands, n, p = self._flc.get_instance() | |
| return points, demands, n, p | |
| def get_distance_and_cost(self) -> Tuple[np.ndarray, np.ndarray]: | |
| return self._flc.get_distance_and_cost_matrix() | |
| def evaluate(self, solution: np.ndarray) -> float: | |
| self._flc.set_solution(solution) | |
| obj_value = self._flc.compute_obj_value() | |
| return obj_value | |
| class MULTIPMP(PMPEnv): | |
| EPSILON = 1e-6 | |
| def __init__(self, | |
| cfg, | |
| data_npy, | |
| boost = False): | |
| self.cfg = cfg | |
| self.data_npy = data_npy | |
| self._boost = boost | |
| self._all_points, self._all_demands, self._n, self._all_p = self._load_multi_facility_data(data_npy) | |
| self.boost = boost | |
| self._all_solutions = self._load_multi_facility_solutions(boost) | |
| print('all_solutions:', self._all_solutions) | |
| self._final_solutions = list(self._all_solutions) | |
| self._num_types = len(self._all_p) | |
| self._current_type = 0 | |
| self._all_max_steps, self._old_mask, self._new_mask = self._get_max_steps() | |
| super().__init__(cfg) | |
| def _set_node_edge_range(self) -> None: | |
| self._node_range = self._n + 2 | |
| self._edge_range = self._n * max(self._all_p) | |
| def step(self, action: np.ndarray): | |
| if self._num_types == 1: | |
| reward = self.get_reward() | |
| self._done = True | |
| pickle.dump(self._final_solutions, open('./facility_location/solutions.pkl', 'wb')) | |
| return self._get_obs(self._t), reward, self._done, False, {} | |
| if self._done: | |
| raise RuntimeError('Action taken after episode is done.') | |
| obj_value, solution, info = self._flc.swap(action, self._t, stage=2) | |
| self._t += 1 | |
| self._done = (self._t == self._all_max_steps[-1] and self._current_type == len(self._all_max_steps) - 1) | |
| self._obj_value[self._t] = obj_value | |
| self._solution[self._t] = solution | |
| reward = self.get_reward() | |
| if obj_value < self._best_obj_value - self.EPSILON: | |
| self._best_obj_value = obj_value | |
| self._best_solution = solution | |
| self._last_best_t = self._t | |
| elif (self._t - self._last_best_t) % self._tabu_stable_steps == 0: | |
| self._flc.reset_tabu_time() | |
| print(self._t, self._all_max_steps[self._current_type]) | |
| if self._t == self._all_max_steps[self._current_type] and not self._done: | |
| self._t = 0 | |
| self._multi_obj += obj_value | |
| self._final_solutions[self._current_type] = solution | |
| self._update_type() | |
| if self._done: | |
| pickle.dump(self._final_solutions, open('./facility_location/solutions.pkl', 'wb')) | |
| return self._get_obs(self._t), reward, self._done, False, info | |
| def reset(self, seed = 0) -> Optional[Dict]: | |
| self._current_type = 0 | |
| points = self._all_points | |
| demands = self._all_demands[:,0] | |
| n = self._n | |
| p = self._all_p[0] | |
| solution = self._all_solutions[0] | |
| self._multi_obj = 0 | |
| self._flc.set_instance(points, demands, n, p, True) | |
| return self.prepare(n, p, solution), {} | |
| def _update_type(self): | |
| if self._current_type >= self._num_types: | |
| raise RuntimeError('Action taken after episode is done.') | |
| if self._current_type < self._num_types - 1: | |
| self._current_type += 1 | |
| print(f'current type: {self._current_type}') | |
| print(self._num_types) | |
| points = self._all_points | |
| demands = self._all_demands[:,self._current_type] | |
| n = self._n | |
| p = self._all_p[self._current_type] | |
| solution = self._all_solutions[self._current_type] | |
| self._flc.set_instance(points, demands, n, p, True) | |
| self.prepare(n, p, solution) | |
| def prepare(self, n: int, p: int, solution: list) -> Dict: | |
| self._obs_extractor.reset() | |
| self._done = False | |
| self._t = 0 | |
| if len(self._all_p) > 1: | |
| self._max_steps = self._all_max_steps[self._current_type] | |
| self._flc.init_facility_mask(self._old_mask[self._current_type], self._new_mask[self._current_type]) | |
| else: | |
| self._max_steps = 0 | |
| initial_solution = solution | |
| initial_obj_value = self._flc.compute_obj_value_from_solution(initial_solution,stage=2) | |
| self._obj_value = np.zeros(self._max_steps + 1) | |
| self._obj_value[0] = initial_obj_value | |
| self._solution = np.zeros((self._max_steps + 1, n), dtype=bool) | |
| self._solution[0] = initial_solution | |
| self._best_solution = initial_solution | |
| self._best_obj_value = initial_obj_value | |
| self._last_best_t = 0 | |
| self._tabu_stable_steps = max(1, round(self._max_steps * self._tabu_stable_steps_scale)) | |
| return self._get_obs(self._t) | |
| def _get_max_steps(self) -> list: | |
| # print(self._all_solutions) | |
| tmp_all_solitions = list(self._all_solutions) | |
| max_steps = [] | |
| old_idx = [] | |
| new_idx = [] | |
| for t in range(self._num_types): | |
| count_true = [sum(s) for s in zip(*(tmp_all_solitions[t:]))] | |
| # print(count_true) | |
| old = [i for i in range(len(count_true)) if count_true[i] > 1 and tmp_all_solitions[t][i]] | |
| new = [i for i in range(len(count_true)) if count_true[i] == 0] | |
| if len(old): | |
| old_idx.append(old) | |
| new_idx.append(new) | |
| max_steps.append(len(old)) | |
| for i in old: | |
| count_true[i] = count_true[i] - 1 | |
| # print(max_steps, old_idx, new_idx) | |
| return max_steps, old_idx, new_idx | |
| def _generate_new_instance(self) -> Tuple[np.ndarray, np.ndarray, int, int]: | |
| n = self._np_random.integers(self._min_n, self._max_n, endpoint=True) | |
| p_ratio = self._np_random.uniform(self._min_p_ratio, self._max_p_ratio) | |
| p = int(max(n * p_ratio, 4)) | |
| points = self._np_random.uniform(size=(n, 2)) | |
| while np.unique(points, axis=0).shape[0] != n: | |
| points = self._np_random.uniform(size=(n, 2)) | |
| demands = self._np_random.random(size=(n,)) | |
| return points, demands, n, p | |
| def _load_multi_facility_data(self, data_npy) -> Tuple[np.ndarray, np.ndarray]: | |
| data = data_npy.split('\n') | |
| n = len(data) | |
| p = int((len(data[0].split(' '))-2) / 2) | |
| positions = [] | |
| demands = [] | |
| actual_facilities = [] | |
| ps = [] | |
| for row in data: | |
| row = row.split(' ') | |
| row = [x for x in row if len(x)] | |
| positions.append([float(row[0]), float(row[1])]) | |
| demand = [] | |
| for i in range(2, 2+p): | |
| demand.append(float(row[i])) | |
| demands.append(demand) | |
| actual_facility = [] | |
| for i in range(2+p, 2+2*p): | |
| actual_facility.append(bool(int(float(row[i])))) | |
| actual_facilities.append(actual_facility) | |
| positions = np.array(positions) | |
| positions = np.deg2rad(positions) | |
| demands = np.array(demands) | |
| actual_facilities = np.array(actual_facilities) | |
| ps = actual_facilities.sum(axis=0) | |
| return positions, demands, n, ps | |
| def _load_multi_facility_solutions(self, boost) -> list: | |
| def load_model(positions, demands, n, p, boost): | |
| eval_env = EvalPMPEnv(self.cfg, positions, demands, n, p, boost) | |
| eval_env = DummyVecEnv([lambda: eval_env]) | |
| policy_kwargs = get_policy_kwargs(self.cfg) | |
| test_model = PPO(MaskedFacilityLocationActorCriticPolicy, | |
| eval_env, | |
| verbose=1, | |
| policy_kwargs=policy_kwargs, | |
| device='cuda:1') | |
| train_model = PPO.load(self.cfg.load_model_path) | |
| test_model.set_parameters(train_model.get_parameters()) | |
| return test_model, eval_env | |
| def get_optimal_solution(model, eval_env): | |
| obs = eval_env.reset() | |
| done = False | |
| while not done: | |
| action, _ = model.predict(obs, deterministic=True) | |
| obs, _, done, info = eval_env.step(action) | |
| return eval_env.get_attr('_best_solution')[0] | |
| multi_solutions = [] | |
| for i in range(len(self._all_p)): | |
| positions = self._all_points | |
| demands = self._all_demands[:,i] | |
| n = self._n | |
| p = self._all_p[i] | |
| model, env = load_model(positions,demands,n,p,boost) | |
| multi_solutions.append(get_optimal_solution(model, env)) | |
| return multi_solutions | |
| def get_reward(self) -> float: | |
| if self._done: | |
| reward = np.min(self._obj_value) | |
| else: | |
| reward = 0.0 | |
| return reward | |