def get_td3_trainer(env, parameters, use_gpu):
state_dim = get_num_output_features(env.normalization)
action_dim = get_num_output_features(env.normalization_action)
q1_network = FullyConnectedParametricDQN(
state_dim,
action_dim,
parameters.q_network.layers,
parameters.q_network.activations,
)
q2_network = None
if parameters.training.use_2_q_functions:
q2_network = FullyConnectedParametricDQN(
state_dim,
action_dim,
parameters.q_network.layers,
parameters.q_network.activations,
)
actor_network = FullyConnectedActor(
state_dim,
action_dim,
parameters.actor_network.layers,
parameters.actor_network.activations,
)
min_action_range_tensor_training = torch.full((1, action_dim), -1)
max_action_range_tensor_training = torch.full((1, action_dim), 1)
min_action_range_tensor_serving = torch.FloatTensor(
env.action_space.low).unsqueeze(dim=0)
max_action_range_tensor_serving = torch.FloatTensor(
env.action_space.high).unsqueeze(dim=0)
if use_gpu:
q1_network.cuda()
if q2_network:
q2_network.cuda()
actor_network.cuda()
min_action_range_tensor_training = min_action_range_tensor_training.cuda(
)
max_action_range_tensor_training = max_action_range_tensor_training.cuda(
)
min_action_range_tensor_serving = min_action_range_tensor_serving.cuda(
)
max_action_range_tensor_serving = max_action_range_tensor_serving.cuda(
)
trainer_args = [q1_network, actor_network, parameters]
trainer_kwargs = {
"q2_network": q2_network,
"min_action_range_tensor_training": min_action_range_tensor_training,
"max_action_range_tensor_training": max_action_range_tensor_training,
"min_action_range_tensor_serving": min_action_range_tensor_serving,
"max_action_range_tensor_serving": max_action_range_tensor_serving,
}
return TD3Trainer(*trainer_args, use_gpu=use_gpu, **trainer_kwargs)
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_actor_wrapper(self):
state_normalization_parameters = {i: _cont_norm() for i in range(1, 5)}
action_normalization_parameters = {
i: _cont_action_norm() for i in range(101, 105)
}
state_preprocessor = Preprocessor(state_normalization_parameters, False)
postprocessor = Postprocessor(action_normalization_parameters, False)
# Test with FullyConnectedActor to make behavior deterministic
actor = FullyConnectedActor(
state_dim=len(state_normalization_parameters),
action_dim=len(action_normalization_parameters),
sizes=[16],
activations=["relu"],
)
actor_with_preprocessor = ActorWithPreprocessor(
actor, state_preprocessor, postprocessor
)
wrapper = ActorPredictorWrapper(actor_with_preprocessor)
input_prototype = actor_with_preprocessor.input_prototype()
action = wrapper(*input_prototype)
self.assertEqual(action.shape, (1, len(action_normalization_parameters)))
expected_output = postprocessor(
actor(
rlt.PreprocessedState.from_tensor(
state_preprocessor(*input_prototype[0])
)
).action
)
self.assertTrue((expected_output == action).all())
def test_basic(self):
state_dim = 8
action_dim = 4
model = FullyConnectedActor(
state_dim,
action_dim,
sizes=[7, 6],
activations=["relu", "relu"],
use_batch_norm=True,
)
input = model.input_prototype()
self.assertEqual((1, state_dim), input.state.float_features.shape)
# Using batch norm requires more than 1 example in training, avoid that
model.eval()
action = model(input)
self.assertEqual((1, action_dim), action.action.shape)
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_basic(self):
state_dim = 8
action_dim = 4
model = FullyConnectedActor(
state_dim,
action_dim,
sizes=[7, 6],
activations=["relu", "relu"],
use_batch_norm=True,
)
input = model.input_prototype()
self.assertEqual((1, state_dim), input.state.float_features.shape)
# Using batch norm requires more than 1 example in training, avoid that
model.eval()
action = model(input)
self.assertEqual((1, action_dim), action.action.shape)
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 get_td3_trainer(self, env, parameters, use_gpu):
state_dim = get_num_output_features(env.normalization)
action_dim = get_num_output_features(env.normalization_action)
q1_network = FullyConnectedParametricDQN(
state_dim,
action_dim,
parameters.q_network.layers,
parameters.q_network.activations,
)
q2_network = None
if parameters.training.use_2_q_functions:
q2_network = FullyConnectedParametricDQN(
state_dim,
action_dim,
parameters.q_network.layers,
parameters.q_network.activations,
)
actor_network = FullyConnectedActor(
state_dim,
action_dim,
parameters.actor_network.layers,
parameters.actor_network.activations,
)
if use_gpu:
q1_network.cuda()
if q2_network:
q2_network.cuda()
actor_network.cuda()
return TD3Trainer(
q1_network,
actor_network,
parameters,
q2_network=q2_network,
use_gpu=use_gpu,
)