def test_get_predictor_export_meta_and_workspace(self):
model = Model()
# 2 params + 1 const
expected_num_params, expected_num_inputs, expected_num_outputs = 3, 2, 4
check_save_load(self, model, expected_num_params, expected_num_inputs,
expected_num_outputs)
def test_save_load(self):
state_dim = 8
model = ExampleSequenceModel(state_dim)
# ONNX sure exports a lot of parameters...
expected_num_params, expected_num_inputs, expected_num_outputs = 133, 3, 1
check_save_load(
self, model, expected_num_params, expected_num_inputs, expected_num_outputs
)
def test_save_load_parametric_action(self):
state_dim = 8
action_dim = 4
model = ParametricDuelingQNetwork.make_fully_connected(
state_dim, action_dim, [8, 4], ["relu", "relu"]
)
expected_num_params, expected_num_inputs, expected_num_outputs = 22, 2, 1
check_save_load(
self, model, expected_num_params, expected_num_inputs, expected_num_outputs
)
def test_save_load_discrete_action(self):
state_dim = 8
action_dim = 4
model = DuelingQNetwork.make_fully_connected(
state_dim, action_dim, layers=[8, 4], activations=["relu", "relu"]
)
expected_num_params, expected_num_inputs, expected_num_outputs = 22, 1, 1
check_save_load(
self, model, expected_num_params, expected_num_inputs, expected_num_outputs
)
def test_save_load_discrete_action(self):
state_dim = 8
action_dim = 4
model = DuelingQNetwork(
layers=[state_dim, 8, 4, action_dim],
activations=["relu", "relu", "linear"],
use_batch_norm=False,
)
expected_num_params, expected_num_inputs, expected_num_outputs = 22, 1, 1
check_save_load(self, model, expected_num_params, expected_num_inputs,
expected_num_outputs)
def test_save_load(self):
state_dim = 8
action_dim = 4
model = FullyConnectedActor(
state_dim,
action_dim,
sizes=[7, 6],
activations=["relu", "relu"],
use_batch_norm=False,
)
expected_num_params, expected_num_inputs, expected_num_outputs = 6, 1, 1
check_save_load(self, model, expected_num_params, expected_num_inputs,
expected_num_outputs)
def test_save_load_parametric_action_batch_norm(self):
state_dim = 8
action_dim = 4
model = ParametricDuelingQNetwork.make_fully_connected(
state_dim, action_dim, [8, 4], ["relu", "relu"]
)
# Freezing batch_norm
model.eval()
# Number of expected params is the same because DuelingQNetwork always
# initialize batch norm layer even if it doesn't use it.
expected_num_params, expected_num_inputs, expected_num_outputs = 22, 2, 1
check_save_load(
self, model, expected_num_params, expected_num_inputs, expected_num_outputs
)
def test_save_load_batch_norm(self):
state_dim = 8
action_dim = 4
model = FullyConnectedActor(
state_dim,
action_dim,
sizes=[7, 6],
activations=["relu", "relu"],
use_batch_norm=True,
)
# Freezing batch_norm
model.eval()
expected_num_params, expected_num_inputs, expected_num_outputs = 21, 1, 1
check_save_load(self, model, expected_num_params, expected_num_inputs,
expected_num_outputs)
def test_save_load_discrete_action_batch_norm(self):
state_dim = 8
action_dim = 4
model = DuelingQNetwork(
layers=[state_dim, 8, 4, action_dim],
activations=["relu", "relu", "linear"],
use_batch_norm=False,
)
# Freezing batch_norm
model.eval()
# Number of expected params is the same because DuelingQNetwork always
# initialize batch norm layer even if it doesn't use it.
expected_num_params, expected_num_inputs, expected_num_outputs = 22, 1, 1
check_save_load(self, model, expected_num_params, expected_num_inputs,
expected_num_outputs)
def test_save_load(self):
state_dim = 8
action_dim = 4
q_network = FullyConnectedDQN(state_dim,
action_dim,
sizes=[8, 4],
activations=["relu", "relu"])
imitator_network = FullyConnectedNetwork(
layers=[state_dim, 8, 4, action_dim],
activations=["relu", "relu", "linear"])
model = BatchConstrainedDQN(
state_dim=state_dim,
q_network=q_network,
imitator_network=imitator_network,
bcq_drop_threshold=0.05,
)
# 6 for DQN + 6 for Imitator Network + 2 for BCQ constants
expected_num_params, expected_num_inputs, expected_num_outputs = 14, 1, 1
check_save_load(self, model, expected_num_params, expected_num_inputs,
expected_num_outputs)
def test_save_load(self):
state_dim = 8
action_dim = 4
model = GaussianFullyConnectedActor(
state_dim,
action_dim,
sizes=[7, 6],
activations=["relu", "relu"],
use_batch_norm=False,
)
expected_num_params, expected_num_inputs, expected_num_outputs = 6, 1, 1
# Actor output is stochastic and won't match between PyTorch & Caffe2
check_save_load(
self,
model,
expected_num_params,
expected_num_inputs,
expected_num_outputs,
check_equality=False,
)
