您最多选择25个主题
主题必须以中文或者字母或数字开头，可以包含连字符 (-)，并且长度不得超过35个字符
36 行
1.3 KiB
36 行
1.3 KiB
import pytest
|
|
|
|
from mlagents.trainers.tf.models import ModelUtils
|
|
from mlagents.tf_utils import tf
|
|
from mlagents_envs.base_env import BehaviorSpec, ActionSpec, SensorType
|
|
|
|
|
|
def create_behavior_spec(num_visual, num_vector, vector_size):
|
|
obs_shapes = [(84, 84, 3)] * int(num_visual) + [(vector_size,)] * int(num_vector)
|
|
sensor_types = [SensorType.OBSERVATION for _ in range(len(obs_shapes))]
|
|
behavior_spec = BehaviorSpec(
|
|
obs_shapes, sensor_types, ActionSpec.create_discrete((1,))
|
|
)
|
|
return behavior_spec
|
|
|
|
|
|
@pytest.mark.parametrize("num_visual", [1, 2, 4])
|
|
@pytest.mark.parametrize("num_vector", [1, 2, 4])
|
|
def test_create_input_placeholders(num_vector, num_visual):
|
|
vec_size = 8
|
|
name_prefix = "test123"
|
|
bspec = create_behavior_spec(num_visual, num_vector, vec_size)
|
|
vec_in, vis_in = ModelUtils.create_input_placeholders(
|
|
bspec.observation_shapes, name_prefix=name_prefix
|
|
)
|
|
|
|
assert isinstance(vis_in, list)
|
|
assert len(vis_in) == num_visual
|
|
assert isinstance(vec_in, tf.Tensor)
|
|
assert vec_in.get_shape().as_list()[1] == num_vector * 8
|
|
|
|
# Check names contain prefix and vis shapes are correct
|
|
for _vis in vis_in:
|
|
assert _vis.get_shape().as_list() == [None, 84, 84, 3]
|
|
assert _vis.name.startswith(name_prefix)
|
|
assert vec_in.name.startswith(name_prefix)
|