@ex.automain
def train(dataset, test_size, max_episode_steps, reward_type, input_meas_type, meas_transform, target_transform, node_history, gamma, target_update_freq, cat_features, feature_range, replay_capacity, min_horizon, epsilon_start, epsilon_end, exploration_frac, n_train_steps, train_freq, loss, batch_size, lr, n_test_episodes, init_eval, n_eval_artifacts, test_freq, log_freq, device, seed, data_seed, save_model, _log, _run, _config):
    np.set_printoptions(precision=2, suppress=True)
    if device.startswith('cuda'):
        assert torch.cuda.is_available()
    logger = _log
    device = torch.device(device)
    (train_set, test_set) = get_datasets(dataset, seed=data_seed, test_size=test_size)
    max_nodes = max(train_set.max_nodes, test_set.max_nodes)
    max_edges = (2 * max(train_set.max_edges, test_set.max_edges))
    test_loader = get_test_loader(test_set, seed=seed, num_samples=n_test_episodes)
    train_gen = get_train_generator(train_set, seed=seed)
    preprocessor = Preprocessor(input_meas_type=input_meas_type, output_meas_type=input_meas_type, feature_range=feature_range, meas_transform=meas_transform, target_transform=target_transform, temporal_offsets=[1.0], max_nodes=max_nodes, device=device)
    train_env_config = dict(max_episode_steps=max_episode_steps, reward_type=reward_type, max_nodes=max_nodes, max_edges=max_edges, nn_feat=('N' in cat_features))
    train_env = make_env(**train_env_config, data_generator=train_gen, seed=seed)
    test_env_config = copy.deepcopy(train_env_config)
    test_env_config.update(sample_goals=False, data_generator=None)
    test_env = make_env(**test_env_config, seed=seed)
    (neg_label, pos_label) = feature_range
    mem_features = dict(cat=cat_features)
    graph_mem_config = dict(max_episode_steps=max_episode_steps, max_nodes=max_nodes, max_edges=max_edges, history=node_history, memory_type='cat', features=mem_features, neg_label=neg_label, pos_label=pos_label)
    eval_memory = make_memory(online=True, **graph_mem_config)
    acting_memory = make_memory(online=True, **graph_mem_config)
    model_config = dict(dim_node=eval_memory.dim_node, dim_meas=preprocessor.dim_input_meas, dim_goal=1, max_edges=max_edges, **_config['model'])
    network = make_network(**model_config).to(device)
    eval_policy = GraphDQNPolicy(network, eval_memory, preprocessor=preprocessor, device=device)
    evaluator = Evaluator(test_loader, test_env, eval_policy)
    exploration_steps = int((exploration_frac * n_train_steps))
    exploration_schedule = LinearSchedule(epsilon_start, epsilon_end, exploration_steps)
    acting_policy = GraphDQNPolicy(network, graph_memory=acting_memory, preprocessor=preprocessor, exploration_schedule=exploration_schedule, device=device)
    replay_buffer = Replay(capacity=replay_capacity, ob_space=train_env.observation_space, graph_mem_config=graph_mem_config, min_horizon=min_horizon)
    actor = Actor(train_env, acting_policy, replay_buffer)
    optimizer = torch.optim.Adam(network.parameters(), lr=lr)
    if (loss == 'mse'):
        criterion = torch.nn.MSELoss()
    else:
        raise ValueError(f'Unsupported loss: {loss}')
    trainer = DQNTrainer(gamma=gamma, target_update_freq=target_update_freq, replay_buffer=replay_buffer, batch_size=batch_size, network=network, preprocessor=preprocessor, criterion=criterion, optimizer=optimizer, device=device)
    network.eval()
    logger.info(f'Filling up the replay buffer...')
    actor.step(n=replay_capacity, use_tqdm=True)
    logger.info(f'Replay buffer filled: [{len(replay_buffer)} / {replay_capacity}]')
    preprocessor.fit(replay_buffer._measurements)
    best_perf = main_loop(actor, trainer, evaluator, network, exploration_schedule, init_eval, n_eval_artifacts, n_train_steps, train_freq, log_freq, test_freq, save_model)
    train_env.close()
    evaluator.close()
    return best_perf
@ex.automain
def train(dataset, test_size, max_episode_steps, reward_type, input_meas_type, meas_transform, target_transform, node_history, gamma, target_update_freq, cat_features, feature_range, replay_capacity, min_horizon, epsilon_start, epsilon_end, exploration_frac, n_train_steps, train_freq, loss, batch_size, lr, n_test_episodes, init_eval, n_eval_artifacts, test_freq, log_freq, device, seed, data_seed, save_model, _log, _run, _config):
    np.set_printoptions(precision=2, suppress=True)
    if device.startswith('cuda'):
        assert torch.cuda.is_available()
    logger = _log
    device = torch.device(device)
    (train_set, test_set) = get_datasets(dataset, seed=data_seed, test_size=test_size)
    max_nodes = max(train_set.max_nodes, test_set.max_nodes)
    max_edges = (2 * max(train_set.max_edges, test_set.max_edges))
    test_loader = get_test_loader(test_set, seed=seed, num_samples=n_test_episodes)
    train_gen = get_train_generator(train_set, seed=seed)
    preprocessor = Preprocessor(input_meas_type=input_meas_type, output_meas_type=input_meas_type, feature_range=feature_range, meas_transform=meas_transform, target_transform=target_transform, temporal_offsets=[1.0], max_nodes=max_nodes, device=device)
    train_env_config = dict(max_episode_steps=max_episode_steps, reward_type=reward_type, max_nodes=max_nodes, max_edges=max_edges, nn_feat=('N' in cat_features))
    train_env = make_env(**train_env_config, data_generator=train_gen, seed=seed)
    test_env_config = copy.deepcopy(train_env_config)
    test_env_config.update(sample_goals=False, data_generator=None)
    test_env = make_env(**test_env_config, seed=seed)
    (neg_label, pos_label) = feature_range
    mem_features = dict(cat=cat_features)
    graph_mem_config = dict(max_episode_steps=max_episode_steps, max_nodes=max_nodes, max_edges=max_edges, history=node_history, memory_type='cat', features=mem_features, neg_label=neg_label, pos_label=pos_label)
    eval_memory = make_memory(online=True, **graph_mem_config)
    acting_memory = make_memory(online=True, **graph_mem_config)
    model_config = dict(dim_node=eval_memory.dim_node, dim_meas=preprocessor.dim_input_meas, dim_goal=1, max_edges=max_edges, **_config['model'])
    network = make_network(**model_config).to(device)
    eval_policy = GraphDQNPolicy(network, eval_memory, preprocessor=preprocessor, device=device)
    evaluator = Evaluator(test_loader, test_env, eval_policy)
    exploration_steps = int((exploration_frac * n_train_steps))
    exploration_schedule = LinearSchedule(epsilon_start, epsilon_end, exploration_steps)
    acting_policy = GraphDQNPolicy(network, graph_memory=acting_memory, preprocessor=preprocessor, exploration_schedule=exploration_schedule, device=device)
    replay_buffer = Replay(capacity=replay_capacity, ob_space=train_env.observation_space, graph_mem_config=graph_mem_config, min_horizon=min_horizon)
    actor = Actor(train_env, acting_policy, replay_buffer)
    optimizer = torch.optim.Adam(network.parameters(), lr=lr)
    if (loss == 'mse'):
        criterion = torch.nn.MSELoss()
    else:
        raise ValueError(f'Unsupported loss: {loss}')
    trainer = DQNTrainer(gamma=gamma, target_update_freq=target_update_freq, replay_buffer=replay_buffer, batch_size=batch_size, network=network, preprocessor=preprocessor, criterion=criterion, optimizer=optimizer, device=device)
    network.eval()
    logger.info(f'Filling up the replay buffer...')
    actor.step(n=replay_capacity, use_tqdm=True)
    logger.info(f'Replay buffer filled: [{len(replay_buffer)} / {replay_capacity}]')
    preprocessor.fit(replay_buffer._measurements)
    best_perf = main_loop(actor, trainer, evaluator, network, exploration_schedule, init_eval, n_eval_artifacts, n_train_steps, train_freq, log_freq, test_freq, save_model)
    train_env.close()
    evaluator.close()
    return best_perf