def test_get_onnx_deterministic_tensors():
inp_size = 4
act_size = 2
action_model, masks = create_action_model(inp_size, act_size)
sample_inp = torch.ones((1, inp_size))
out_tensors = action_model.get_action_out(sample_inp, masks=masks)
(
continuous_out,
discrete_out,
action_out_deprecated,
deterministic_continuous_out,
deterministic_discrete_out,
) = out_tensors
assert continuous_out.shape == (1, 2)
assert discrete_out.shape == (1, 2)
assert deterministic_discrete_out.shape == (1, 2)
assert deterministic_continuous_out.shape == (1, 2)
# Second sampling from same distribution
out_tensors2 = action_model.get_action_out(sample_inp, masks=masks)
(
continuous_out_2,
discrete_out_2,
action_out_2_deprecated,
deterministic_continuous_out_2,
deterministic_discrete_out_2,
) = out_tensors2
assert ~torch.all(torch.eq(continuous_out, continuous_out_2))
assert torch.all(
torch.eq(deterministic_continuous_out, deterministic_continuous_out_2)
)
def test_initialization_layer():
torch.manual_seed(0)
# Test Zero
layer = linear_layer(
3, 4, kernel_init=Initialization.Zero, bias_init=Initialization.Zero
)
assert torch.all(torch.eq(layer.weight.data, torch.zeros_like(layer.weight.data)))
assert torch.all(torch.eq(layer.bias.data, torch.zeros_like(layer.bias.data)))
def test_lstm_layer():
torch.manual_seed(0)
# Test zero for LSTM
layer = lstm_layer(
4, 4, kernel_init=Initialization.Zero, bias_init=Initialization.Zero
)
for name, param in layer.named_parameters():
if "weight" in name:
assert torch.all(torch.eq(param.data, torch.zeros_like(param.data)))
elif "bias" in name:
assert torch.all(
torch.eq(param.data[4:8], torch.ones_like(param.data[4:8]))
)
def test_layer_norm():
torch.manual_seed(0)
torch_ln = torch.nn.LayerNorm(10, elementwise_affine=False)
cust_ln = LayerNorm()
sample_input = torch.rand(10)
assert torch.all(
torch.isclose(
torch_ln(sample_input), cust_ln(sample_input), atol=1e-5, rtol=0.0
)
)
sample_input = torch.rand((4, 10))
assert torch.all(
torch.isclose(
torch_ln(sample_input), cust_ln(sample_input), atol=1e-5, rtol=0.0
)
)
sample_input = torch.rand((7, 6, 10))
assert torch.all(
torch.isclose(
torch_ln(sample_input), cust_ln(sample_input), atol=1e-5, rtol=0.0
)
)
def test_swish():
layer = Swish()
input_tensor = torch.Tensor([[1, 2, 3], [4, 5, 6]])
target_tensor = torch.mul(input_tensor, torch.sigmoid(input_tensor))
assert torch.all(torch.eq(layer(input_tensor), target_tensor))