def save(self, path): logger.log(f'\nSaving weights to {path}') self.model.save_weights(path)
def load(self, path): logger.log(f'\nLoading weights from {path}') self.model.load_weights(path)
def plot(paths, x_axis, y_axis, x_label, y_label, window, interval, show_seeds, columns, x_min, x_max, baselines, baselines_source, name, save_formats, cmap, legend_columns, legend_marker_size, dpi, fig): '''Plots results from experiments and benchmark data.''' # Extract the data. logger.log('Loading data...') data = get_data(paths, baselines, baselines_source, x_axis, y_axis, x_min, x_max, window) # List the environments. envs = sorted(data.keys(), key=str.casefold) num_envs = len(envs) if num_envs == 0: logger.error('No logs found.') return # List the agents. agents = set() for env in data: for agent in data[env]: agents.add(agent) agents = sorted(agents, key=str.casefold) num_agents = len(agents) if not cmap: if num_agents <= 10: cmap = 'tab10' elif num_agents <= 20: cmap = 'tab20' else: cmap = 'rainbow' cmap = plt.get_cmap(cmap) if isinstance(cmap, mpl.colors.ListedColormap): colors = cmap(range(num_agents)) else: colors = list(cmap(np.linspace(0, 1, num_agents))) agent_colors = {a: c for a, c in zip(agents, colors)} if columns is None: columns = int(np.ceil(np.sqrt(num_envs))) else: columns = min(columns, num_envs) rows = int(np.ceil(num_envs / columns)) if fig is None: plt.ion() fig = plt.figure(figsize=(columns * 6, rows * 5)) else: fig.clear() grid = gridspec.GridSpec(rows + 1, 1 + columns, height_ratios=[1] * rows + [0.1], width_ratios=[0] + [1] * columns) axes = [] for i in range(num_envs): ax = fig.add_subplot(grid[i // columns, 1 + i % columns]) axes.append(ax) logger.log('Plotting...') for env, ax in zip(envs, axes): # Plot intervals first. if interval in ['std', 'bounds']: for agent in sorted(data[env], key=str.casefold): color = agent_colors[agent] x, mean, min_mean, max_mean, std = data[env][agent]['stats'] if interval == 'std': if std is None: logger.error('No std found in the data.') else: ax.fill_between(x, mean - std, mean + std, color=color, alpha=0.1, lw=0) elif interval == 'bounds': ax.fill_between(x, min_mean, max_mean, color=color, alpha=0.1, lw=0) # Plot means after. for agent in sorted(data[env], key=str.casefold): color = agent_colors[agent] # Plot each run if needed. xs, means = data[env][agent]['seeds'] if show_seeds and len(xs) > 1: for x, mean in zip(xs, means): ax.plot(x, mean, c=color, lw=1, alpha=0.5) # Plot the overall mean. x, mean = data[env][agent]['stats'][:2] ax.plot(x, mean, c=color, lw=2, alpha=1) # Finalize the figures. ax.locator_params(axis='x', nbins=6) ax.locator_params(axis='y', tight=True, nbins=6) ax.get_yaxis().set_major_formatter( mpl.ticker.FuncFormatter(lambda x, p: f'{x:,g}')) low, high = ax.get_xlim() if max(abs(low), abs(high)) >= 1e3: ax.ticklabel_format(style='sci', axis='x', scilimits=(0, 0)) ax.xaxis.grid(linewidth=0.5, alpha=0.5) ax.yaxis.grid(linewidth=0.5) ax.spines['top'].set_visible(False) ax.spines['right'].set_visible(False) ax.tick_params(axis='both', which='both', length=0) if x_label is None: x_label = 'Steps' if x_axis == 'train/steps' else x_axis if y_label is None: y_label = 'Score' if y_axis == 'test/episode_score' else y_axis ax.set_xlabel(x_label) ax.set_ylabel(y_label) ax.set_title(env) # Add a legend bellow. legend_ax = fig.add_subplot(grid[-1:, :]) legend_ax.set_axis_off() handles = [] for color in colors: marker = lines.Line2D(range(1), range(1), marker='o', markerfacecolor=color, markersize=legend_marker_size, linewidth=0, markeredgewidth=0) handles.append(marker) # Find a nice number of legend columns. Labels should not overlap. if legend_columns is None: legend_columns = range(num_agents, 0, -1) else: legend_columns = [legend_columns] for ncol in legend_columns: legend = legend_ax.legend(flip(handles, ncol), flip(agents, ncol), loc='center', mode='expand', borderaxespad=0, borderpad=0, handlelength=0.9, ncol=ncol, numpoints=1) legend_frame = legend.get_frame() legend_frame.set_linewidth(0) fig.tight_layout(pad=0, w_pad=0, h_pad=1.0) fig.canvas.draw() renderer = legend_ax.get_renderer_cache() h_packer = legend.get_children()[0].get_children()[1] target_width = h_packer.get_extent(renderer)[0] current_width = sum( [ch.get_extent(renderer)[0] for ch in h_packer.get_children()]) if target_width > 1.3 * current_width: break # Save the plot in different formats if needed. if save_formats: logger.log('Saving...') if name is None: if len(envs) > 1: name = 'results' else: name = envs[0] for save_format in save_formats: file_name = name + '.' + save_format fig.savefig(file_name, facecolor=fig.get_facecolor(), dpi=dpi) print(file_name) print('to', os.getcwd()) return fig
def get_data(paths, baselines, baselines_source, x_axis, y_axis, x_min, x_max, window): '''Gets data from the paths and benchmark.''' data = {} # List the log files in paths. log_paths = [] for path in paths: if os.path.isdir(path): log_paths.extend(pathlib.Path(path).rglob('log.*')) elif path[-7:-3] == 'log.': log_paths.append(path) # Load the data from the log files. for path in log_paths: sub_path, file = os.path.split(path) dfs = {} if file == 'log.csv': # Extract the environment, agent and seed from the paths. env, agent, seed = sub_path.split(os.sep)[-3:] # Load the data. df_seed = pd.read_csv(path, sep=',', engine='python') x = df_seed[x_axis].values # The x axis should be sorted. if not np.all(np.diff(x) > 0): logger.warning(f'Skipping unsorted {env} {agent} {seed}') continue if x_min and x[-1] < x_min: logger.warning( f'Skipping {env} {agent} {seed} ({x[-1]} steps)') continue dfs[seed] = df_seed elif file == 'log.pkl': # Extract the environment, agent and seed from the paths. env, agent = sub_path.split(os.sep)[-2:] # Load the data. df = pd.read_pickle(path, compression='zip') for seed, df_seed in df.groupby('seed'): x = df_seed[x_axis].values # The x axis should be sorted. if not np.all(np.diff(x) > 0): logger.warning(f'Skipping unsorted {env} {agent} {seed}') continue if x_min and x[-1] < x_min: logger.warning( f'Skipping {env} {agent} {seed} ({x[-1]} steps)') continue dfs[seed] = df_seed for seed, df in dfs.items(): if env not in data: data[env] = {} if agent not in data[env]: data[env][agent] = {} assert seed not in data[env][agent] # Extract the data. x = df[x_axis].values if y_axis in df: mean = df[y_axis].values if y_axis[-5:] == '/mean' and y_axis[:-5] + '/std' in df: std = df[y_axis[:-5] + '/std'].values else: std = None elif y_axis + '/mean' in df: mean = df[y_axis + '/mean'].values if y_axis + '/std' in df: std = df[y_axis + '/std'].values else: std = None else: raise KeyError(f'Key {y_axis} not found.') # Limit the range of x values and smooth. if x_max: max_index = np.argmax(x > x_max) or len(x) else: max_index = len(x) x = x[:max_index] mean = smooth(mean, window) mean = mean[:max_index] if std is not None: std = smooth(std, window) std = std[:max_index] data[env][agent][seed] = x, mean, std # Aggregate the runs and compute the statistics. for env, env_data in data.items(): for agent, agent_data in env_data.items(): xs, means, stds = [], [], [] for seed, (x, mean, std) in agent_data.items(): xs.append(x) means.append(mean) stds.append(std) if stds[0] is None: stds = None print(env, agent) env_data[agent] = dict(seeds=(xs, means), stats=stats(xs, means, stds)) # Load the data from the benchmark if needed. if baselines: full_benchmark = len(paths) == 0 path = __file__ tonic_path = os.path.split(os.path.split(path)[0])[0] # Find the benchmark data to use. if baselines_source == 'tensorflow': logger.log('Loading TensorFlow baselines...') path = os.path.join(tonic_path, 'data/logs/tensorflow_logs.pkl') print('Path:', path) else: logger.log(f'Loading baselines from {baselines_source}...') path = baselines_source # Load the compressed benchmark data. df = pd.read_pickle(path, compression='zip') # Search for data on the corresponding environments and agents. for env, df_env in df.groupby('environment'): if full_benchmark: data[env] = {} elif env not in data: continue for agent, df_agent in df_env.groupby('agent'): # All agents or a specific list of agents can be used. if baselines != ['all'] and agent not in baselines: continue # Rename agents with the same name. if agent in data[env]: logger.warning(f'Renaming baseline {agent}') agent = agent + ' (baseline)' # Extract the data. xs, means, stds = [], [], [] for seed, df_seed in df_agent.groupby('seed'): x = df_seed[x_axis].values if x_min and x[-1] < x_min: logger.warning( f'Skipping {env} {agent} {seed} ({x[-1]} steps)') continue if y_axis in df_seed: mean = df_seed[y_axis].values if (y_axis[-5:] == '/mean' and y_axis[:-5] + '/std' in df_seed): std = df_seed[y_axis[:-5] + '/std'].values else: std = None elif y_axis + '/mean' in df_seed: mean = df_seed[y_axis + '/mean'].values if y_axis + '/std' in df_seed: std = df_seed[y_axis + '/std'].values else: std = None else: raise KeyError(f'Key {y_axis} not found.') # Limit the range of x values and smooth. if x_max: max_index = np.argmax(x > x_max) or len(x) else: max_index = len(x) x = x[:max_index] mean = smooth(mean, window) mean = mean[:max_index] if std is not None: std = smooth(std, window) std = std[:max_index] xs.append(x) means.append(mean) stds.append(std) if stds[0] is None: stds = None print(env, agent) data[env][agent] = dict(seeds=(xs, means), stats=stats(xs, means, stds)) return data
def load(self, path): logger.log('Loading weights from {}'.format(path)) path = path + '.pt' self.model.load_state_dict(torch.load(path))
def save(self, path): path = path + '.pt' logger.log('Saving weights to {}'.format(path)) os.makedirs(os.path.dirname(path), exist_ok=True) torch.save(self.model.state_dict(), path)
def load(self, path): path = path + '.pt' logger.log(f'Loading weights from {path}') self.model.load_state_dict(torch.load(path))
def load(self, path): logger.log('Loading weights from {}'.format(path)) self.model.load_weights(path)
def save(self, path): logger.log('Saving weights to {}'.format(path)) self.model.save_weights(path)