import pytest from mlagents.trainers.tf.models import ModelUtils from mlagents.tf_utils import tf from mlagents_envs.base_env import BehaviorSpec, ActionSpec def create_behavior_spec(num_visual, num_vector, vector_size): behavior_spec = BehaviorSpec( [(84, 84, 3)] * int(num_visual) + [(vector_size,)] * int(num_vector), 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)