def _load_models(self, name, model, mlp_model):
# Loads a network based on the provided config.
model_class = U.import_obj('Model', model.class_path)
mlp_model_class = U.import_obj('Model', mlp_model.class_path)
del model['class_path'], mlp_model['class_path']
with tf.variable_scope(name):
self._mlp_model = mlp_model_class(seed=self.seed, **mlp_model)
with tf.variable_scope('gcn'):
self._model = model_class(seed=self.seed, **model)
def _load_model(self, name, class_path, seed=None, **kwargs):
# Loads a network based on the provided config.
model_class = U.import_obj('Model', class_path)
if seed is None:
seed = self.seed
with tf.variable_scope(name):
self._model = model_class(seed=seed, **kwargs)
return self._model
def make_env():
env_config = get_env_config()
return ParallelBatchedEnv(args.bs,
U.import_obj(env_config.class_name,
env_config.class_path),
[env_config] * args.bs,
args.seed,
use_threads=not args.use_procs)
def run_evaluator(self, id: str):
env_config, sess_config, agent_config = (self.env_config,
self.sess_config,
self.agent_config)
eval_config = self.eval_config
agent_class = U.import_obj(agent_config.class_name,
agent_config.class_path)
shell_class = U.import_obj(sess_config.shell.class_name,
sess_config.shell.class_path)
env_class = U.import_obj(env_config.class_name, env_config.class_path)
agent_config = copy.deepcopy(agent_config)
agent_config.update(evaluation_mode=True)
shell_config = dict(agent_class=agent_class,
agent_config=agent_config,
**self.sess_config.shell)
env_configs = []
for i in range(eval_config.batch_size):
env_config = ConfigDict(**self.env_config)
env_config.update({
eval_config.dataset_type_field: id,
'graph_start_idx': i,
**eval_config.env_config
})
env_configs.append(env_config)
evaluator_config = dict(
shell_class=shell_class,
shell_config=shell_config,
env_class=env_class,
env_configs=env_configs,
loggers=self._setup_evaluator_loggers(id),
heuristic_loggers=self._setup_evaluator_loggers(f'heuristic-{id}'),
seed=self.seed,
**eval_config)
from liaison.distributed import Evaluator
evaluator = Evaluator(**evaluator_config)
t = evaluator.get_heuristic_thread()
t.start()
evaluator.run_loop(int(1e9))
t.join()
def run_actor(self, actor_id):
"""
Launches an actor process with actor_id
Args:
actor_id (int): actor's id
"""
agent_config, env_config, sess_config = (self.agent_config,
self.env_config,
self.sess_config)
agent_class = U.import_obj(agent_config.class_name,
agent_config.class_path)
shell_class = U.import_obj(sess_config.shell.class_name,
sess_config.shell.class_path)
env_class = U.import_obj(env_config.class_name, env_config.class_path)
shell_config = dict(agent_class=agent_class,
agent_config=agent_config,
**self.sess_config.shell)
actor_config = dict(
actor_id=actor_id,
shell_class=shell_class,
shell_config=shell_config,
agent_class=agent_class,
agent_config=agent_config,
env_class=env_class,
env_configs=[self.env_config] * self.batch_size,
traj_length=self.traj_length,
seed=self.seed + actor_id * self.batch_size,
batch_size=self.batch_size,
system_loggers=self._setup_actor_system_loggers(actor_id)
if actor_id == 0 else [],
**self.sess_config.actor)
actor_class = U.import_obj(sess_config.actor.class_name,
sess_config.actor.class_path)
actor_class(**actor_config) # blocking constructor.
def run_debug():
acts = read_pkl(args.debug_actions_fname)['actions']
idx = args.debug_actions_index
env_config = get_env_config()
env = U.import_obj(env_config.class_name,
env_config.class_path)(**env_config,
seed=args.seed,
id=idx)
ts = env.reset()
for act in acts[idx]:
env.step(act)
def _create_shell(env):
agent_config = get_agent_config()
agent_class = U.import_obj(agent_config.class_name,
agent_config.class_path)
action_spec = env.action_spec()
obs_spec = env.observation_spec()
return Shell(action_spec,
obs_spec,
seed=SEED,
agent_class=agent_class,
agent_config=agent_config,
batch_size=B,
use_gpu=FLAGS.enable_gpu)
def get_shell_config():
config = ConfigDict()
agent_config = get_agent_config()
# shell class path is default to the distributed folder.
config.class_path = 'liaison.distributed.shell_for_test'
config.class_name = 'Shell'
config.agent_scope = 'shell'
config.use_gpu = True
config.agent_class = U.import_obj(agent_config.class_name,
agent_config.class_path)
config.agent_config = agent_config
config.agent_config.update(evaluation_mode=True)
return config
def make_env(k_val):
env_config = ConfigDict(to_nested_dicts(args.env_config))
env_config.lp_features = False
env_config.k = k_val
env_config.n_local_moves = args.n_local_moves
env_config.primal_gap_reward = True
env_config.delta_reward = False
env_config.disable_maxcuts = args.disable_maxcuts
assert env_config.n_graphs == 1
env_class = U.import_obj(env_config.class_name, env_config.class_path)
env = env_class(id=0, seed=args.seed, **env_config)
return env
def _create_learner(env, shell):
agent_config = get_agent_config()
agent_class = U.import_obj(agent_config.class_name,
agent_config.class_path)
traj_spec = Trajectory(env.observation_spec(),
step_output_spec=shell.step_output_spec()).spec
traj_spec = Trajectory.format_traj_spec(traj_spec, B, T)
action_spec = env.action_spec()
action_spec.set_shape((B, ) + action_spec.shape[1:])
return Learner(seed=SEED,
traj_spec=traj_spec,
action_spec=action_spec,
agent_class=agent_class,
agent_config=dict(vis_loggers=[], **agent_config),
use_gpu=FLAGS.enable_gpu)
def run_replay_worker(self, replay_id):
"""
Launches a single replay server
Args:
replay_id: The id of the replay server
"""
replay_class = U.import_obj(self.sess_config.replay.class_name,
self.sess_config.replay.class_path)
replay = replay_class(seed=self.seed,
index=replay_id,
tensorplex_config=self.sess_config.tensorplex,
**self.sess_config.replay)
replay.start_threads()
replay.join()
def run_learner(self, iterations=None):
"""
Launches the learner process.
Learner consumes experience from replay
and publishes experience to parameter server
"""
agent_class = U.import_obj(self.agent_config.class_name,
self.agent_config.class_path)
loggers, var_loggers, vis_loggers = self._setup_learner_loggers()
learner = Learner(agent_class=agent_class,
agent_config=dict(vis_loggers=vis_loggers,
**self.agent_config),
traj_length=self.traj_length,
seed=self.seed,
loggers=loggers,
var_loggers=var_loggers,
system_loggers=self._setup_learner_system_loggers(),
**self.sess_config.learner)
learner.main()
def __init__(self, seed, model_kwargs, **kwargs):
self.seed = seed
self.config = ConfigDict(kwargs)
with tf.variable_scope('gcn_model'):
self._gcn_model = U.import_obj(
'Model', model_kwargs['class_path'])(seed=seed, **model_kwargs)
with tf.variable_scope('transformer'):
self._trans = Transformer(**self.config)
with tf.variable_scope('enc_node_embeds'):
self._node_emb_mlp = make_mlp([self.config.d_model] * 2,
'relu',
activate_final=False,
seed=seed,
layer_norm=False)
with tf.variable_scope('out_projection'):
self._out_projection_mlp = snt.nets.MLP(
[self.config.d_model],
initializers=dict(w=glorot_uniform(seed),
b=initializers.init_ops.Constant(0)),
activate_final=False,
)
def make_env(): env_config = get_env_config() Env = U.import_obj(env_config.class_name, env_config.class_path) return Env(0, 42, **env_config)
def get_env_class():
env_config = get_env_config()
return U.import_obj(env_config.class_name, env_config.class_path)
def main(argv):
global args
args = parser.parse_args(argv[1:])
if args.gpu_ids:
os.environ['CUDA_VISIBLE_DEVICES'] = '_'.join(map(str, args.gpu_ids))
else:
os.environ['CUDA_VISIBLE_DEVICES'] = ''
sess_config = ConfigDict(to_nested_dicts(args.sess_config))
env_config = ConfigDict(to_nested_dicts(args.env_config))
agent_config = ConfigDict(to_nested_dicts(args.agent_config))
shell_class = U.import_obj(sess_config.shell.class_name,
sess_config.shell.class_path)
env_class = U.import_obj(env_config.class_name, env_config.class_path)
agent_class = U.import_obj(agent_config.class_name,
agent_config.class_path)
if args.standalone:
if args.heuristic:
results_dir = Path(
f'/data/nms/tfp/evaluation/standalone/{args.heuristic}/{args.name}/{env_config.graph_start_idx}/'
)
else:
results_dir = Path(
f'/data/nms/tfp/evaluation/standalone/agent/{args.name}/{env_config.graph_start_idx}/'
)
elif args.without_agent:
results_dir = Path(
f'/data/nms/tfp/evaluation/without_agent/{args.name}/{env_config.graph_start_idx}/'
)
elif args.heuristic:
results_dir = Path(
f'/data/nms/tfp/evaluation/{args.heuristic}/{args.name}/{env_config.graph_start_idx}/'
)
else:
results_dir = Path(
f'/data/nms/tfp/evaluation/scip/{args.name}/{env_config.graph_start_idx}'
)
results_dir.mkdir(parents=True, exist_ok=True)
evaluator = Evaluator(shell_class=shell_class,
shell_config=sess_config.shell,
agent_class=agent_class,
agent_config=agent_config,
env_class=env_class,
env_config=env_config,
seed=args.seed,
dataset=env_config.dataset,
dataset_type=env_config.dataset_type,
graph_start_idx=env_config.graph_start_idx,
gap=args.gap,
max_nodes=args.max_nodes,
batch_size=args.batch_size,
n_local_moves=args.n_local_moves,
results_dir=results_dir,
use_parallel_envs=args.use_parallel_envs,
use_threaded_envs=args.use_threaded_envs,
heur_frequency=args.heur_frequency,
create_shell=(args.heuristic is None),
**sess_config)
evaluator.run(standalone=args.standalone,
without_agent=args.without_agent,
heuristic=args.heuristic)
print('Done!')
def get_shell_class():
shell_config = get_shell_config()
return U.import_obj(shell_config.class_name, shell_config.class_path)
def _get_env():
env_config = get_env_config()
env_class = U.import_obj(env_config.class_name, env_config.class_path)
return SerialBatchedEnv(B, env_class, [env_config] * B, SEED)