您最多选择25个主题
主题必须以中文或者字母或数字开头，可以包含连字符 (-)，并且长度不得超过35个字符
123 行
3.7 KiB
123 行
3.7 KiB
import unittest.mock as mock
|
|
import pytest
|
|
|
|
from mlagents.trainers import tf
|
|
import yaml
|
|
|
|
from mlagents.trainers.ppo.multi_gpu_policy import MultiGpuPPOPolicy
|
|
from mlagents.trainers.tests.mock_brain import create_mock_brainparams
|
|
|
|
|
|
@pytest.fixture
|
|
def dummy_config():
|
|
return yaml.safe_load(
|
|
"""
|
|
trainer: ppo
|
|
batch_size: 32
|
|
beta: 5.0e-3
|
|
buffer_size: 512
|
|
epsilon: 0.2
|
|
hidden_units: 128
|
|
lambd: 0.95
|
|
learning_rate: 3.0e-4
|
|
max_steps: 5.0e4
|
|
normalize: true
|
|
num_epoch: 5
|
|
num_layers: 2
|
|
time_horizon: 64
|
|
sequence_length: 64
|
|
summary_freq: 1000
|
|
use_recurrent: false
|
|
memory_size: 8
|
|
curiosity_strength: 0.0
|
|
curiosity_enc_size: 1
|
|
reward_signals:
|
|
extrinsic:
|
|
strength: 1.0
|
|
gamma: 0.99
|
|
"""
|
|
)
|
|
|
|
|
|
@mock.patch("mlagents.trainers.ppo.multi_gpu_policy.get_devices")
|
|
def test_create_model(mock_get_devices, dummy_config):
|
|
tf.reset_default_graph()
|
|
mock_get_devices.return_value = [
|
|
"/device:GPU:0",
|
|
"/device:GPU:1",
|
|
"/device:GPU:2",
|
|
"/device:GPU:3",
|
|
]
|
|
|
|
trainer_parameters = dummy_config
|
|
trainer_parameters["model_path"] = ""
|
|
trainer_parameters["keep_checkpoints"] = 3
|
|
brain = create_mock_brainparams()
|
|
|
|
policy = MultiGpuPPOPolicy(0, brain, trainer_parameters, False, False)
|
|
assert len(policy.towers) == len(mock_get_devices.return_value)
|
|
|
|
|
|
@mock.patch("mlagents.trainers.ppo.multi_gpu_policy.get_devices")
|
|
def test_average_gradients(mock_get_devices, dummy_config):
|
|
tf.reset_default_graph()
|
|
mock_get_devices.return_value = [
|
|
"/device:GPU:0",
|
|
"/device:GPU:1",
|
|
"/device:GPU:2",
|
|
"/device:GPU:3",
|
|
]
|
|
|
|
trainer_parameters = dummy_config
|
|
trainer_parameters["model_path"] = ""
|
|
trainer_parameters["keep_checkpoints"] = 3
|
|
brain = create_mock_brainparams()
|
|
with tf.Session() as sess:
|
|
policy = MultiGpuPPOPolicy(0, brain, trainer_parameters, False, False)
|
|
var = tf.Variable(0)
|
|
tower_grads = [
|
|
[(tf.constant(0.1), var)],
|
|
[(tf.constant(0.2), var)],
|
|
[(tf.constant(0.3), var)],
|
|
[(tf.constant(0.4), var)],
|
|
]
|
|
avg_grads = policy.average_gradients(tower_grads)
|
|
|
|
init = tf.global_variables_initializer()
|
|
sess.run(init)
|
|
run_out = sess.run(avg_grads)
|
|
assert run_out == [(0.25, 0)]
|
|
|
|
|
|
@mock.patch("mlagents.trainers.tf_policy.TFPolicy._execute_model")
|
|
@mock.patch("mlagents.trainers.ppo.policy.PPOPolicy.construct_feed_dict")
|
|
@mock.patch("mlagents.trainers.ppo.multi_gpu_policy.get_devices")
|
|
def test_update(
|
|
mock_get_devices, mock_construct_feed_dict, mock_execute_model, dummy_config
|
|
):
|
|
tf.reset_default_graph()
|
|
mock_get_devices.return_value = ["/device:GPU:0", "/device:GPU:1"]
|
|
mock_construct_feed_dict.return_value = {}
|
|
mock_execute_model.return_value = {
|
|
"value_loss": 0.1,
|
|
"policy_loss": 0.3,
|
|
"update_batch": None,
|
|
}
|
|
|
|
trainer_parameters = dummy_config
|
|
trainer_parameters["model_path"] = ""
|
|
trainer_parameters["keep_checkpoints"] = 3
|
|
brain = create_mock_brainparams()
|
|
policy = MultiGpuPPOPolicy(0, brain, trainer_parameters, False, False)
|
|
mock_mini_batch = mock.Mock()
|
|
mock_mini_batch.items.return_value = [("action", [1, 2]), ("value", [3, 4])]
|
|
run_out = policy.update(mock_mini_batch, 1)
|
|
|
|
assert mock_mini_batch.items.call_count == len(mock_get_devices.return_value)
|
|
assert mock_construct_feed_dict.call_count == len(mock_get_devices.return_value)
|
|
assert run_out["Losses/Value Loss"] == 0.1
|
|
assert run_out["Losses/Policy Loss"] == 0.3
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main()
|