def fit_model(_run, ray_server, activation_glob, output_root, max_timesteps, data_type,
model_class, model_kwargs, train_opponent, train_percentage):
"""Fits density models for each environment and victim type in activation_dir,
saving resulting models to output_root. Works by repeatedly calling `density_fitter`,
running in parallel via Ray."""
ray.init(redis_address=ray_server)
# Find activation paths for each environment & victim-path tuple
stem_pattern = re.compile(r'(.*)_opponent_.*\.npz')
opponent_pattern = re.compile(r'.*_opponent_([^\s]+)+\.npz')
# activation_paths is indexed by [env_victim][opponent_type] where env_victim is
# e.g. 'SumoHumans-v0_victim_zoo_1' and opponent_type is e.g. 'ppo2_1'.
activation_paths = {}
for activation_path in glob.glob(activation_glob):
activation_dir = os.path.basename(activation_path)
stem_match = stem_pattern.match(activation_dir)
if stem_match is None:
logger.debug(f"Skipping {activation_path}")
continue
stem = stem_match.groups()[0]
opponent_match = opponent_pattern.match(activation_dir)
opponent_type = opponent_match.groups()[0]
activation_paths.setdefault(stem, {})[opponent_type] = activation_path
# Create temporary output directory (if needed)
tmp_dir = None
if output_root is None:
tmp_dir = tempfile.TemporaryDirectory()
output_root = tmp_dir.name
else:
exp_name = gen_exp_name(model_class, model_kwargs)
output_root = os.path.join(output_root, exp_name)
# Fit density model and save weights
results = []
for stem, paths in activation_paths.items():
output_dir = osp.join(output_root, stem)
os.makedirs(output_dir)
future = density_fitter.remote(paths, output_dir, model_class, model_kwargs,
max_timesteps, data_type, train_opponent,
train_percentage)
results.append(future)
ray.get(results) # block until all jobs have finished
utils.add_artifacts(_run, output_root, ingredient=fit_model_ex)
# Clean up temporary directory (if needed)
if tmp_dir is not None:
tmp_dir.cleanup()
ray.shutdown()
def visualize(_run, model_glob, output_root):
# Output directory
tmp_dir = None
if output_root is None:
tmp_dir = tempfile.TemporaryDirectory()
output_root = tmp_dir.name
for model_dir in glob.glob(model_glob):
model_name = os.path.basename(model_dir)
output_dir = osp.join(output_root, model_name)
os.makedirs(output_dir)
_visualize_helper(model_dir, output_dir)
utils.add_artifacts(_run, output_root, ingredient=visualize_ex)
if tmp_dir is not None:
tmp_dir.cleanup()
def generate_activations(_run, out_dir, score_configs, score_update,
adversary_path, ray_upload_dir):
"""Uses multi.score to generate activations, then extracts them into a convenient
directory structure."""
logger.info("Generating activations")
activation_dirs = run_external(score_configs,
post_named_configs=['save_activations'],
config_updates=score_update,
adversary_path=adversary_path)
os.makedirs(out_dir)
extract_data(_activations_path_generator, out_dir, activation_dirs,
ray_upload_dir)
logger.info("Activations saved")
utils.add_artifacts(_run, out_dir)
def generate_activations(_run, out_dir, score_configs, score_update,
adversary_path, ray_upload_dir):
"""Uses multi.score to generate activations, then extracts them into a convenient
directory structure."""
logger.info("Generating activations")
activation_dirs = run_external(score_configs,
post_named_configs=['save_activations'],
config_updates=score_update,
adversary_path=adversary_path)
def path_generator(trial_root, env_name, victim_index, victim_type,
victim_path, opponent_type, opponent_path):
src_path = osp.join(trial_root, 'data', 'trajectories',
f'agent_{victim_index}.npz')
new_name = (f'{env_name}_victim_{victim_type}_{victim_path}'
f'_opponent_{opponent_type}_{opponent_path}')
return src_path, new_name, 'npz'
os.makedirs(out_dir)
extract_data(path_generator, out_dir, activation_dirs, ray_upload_dir)
logger.info("Activations saved")
utils.add_artifacts(_run, out_dir)