text: # coding=utf-8
  # Copyright 2018 The TF-Agents Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Tests for the Stationary Stochastic Structured Bandit environment."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import
  from tf_agents.bandits.environments import stationary_stochastic_structured_py_environment as ssspe
  from tf_agents.policies import random_py_policy
  from tf_agents.specs import array_spec
  
  
  def check_unbatched_time_step_spec(time_step, time_step_spec, batch_size):
    """Checks if time step conforms array spec, even if batched."""
    if batch_size is None:
      return array_spec.check_arrays_nest(time_step, time_step_spec)
  
    return array_spec.check_arrays_nest(
        time_step, array_spec.add_outer_dims_nest(time_step_spec, (batch_size,)))
  
  
  class StationaryStochasticStructuredBanditPyEnvironmentTest(tf.test.TestCase):
  
    def test_with_random_policy(self):
  
      def _global_context_sampling_fn():
        abc = np.array(['a', 'b', 'c'])
        return {'global1': np.random.randint(-2, 3, [3, 4]),
                'global2': abc[np.random.randint(0, 2, [1])]}
  
      def _arm_context_sampling_fn():
        aabbcc = np.array(['aa', 'bb', 'cc'])
        return {'arm1': np.random.randint(-3, 4, [5]),
                'arm2': np.random.randint(-3, 4, [3, 1]),
                'arm3': aabbcc[np.random.randint(0, 2, [1])]}
  
      def _reward_fn(global_obs, arm_obs):
        return global_obs['global1'][2, 1] + arm_obs['arm1'][4]
  
      env = ssspe.StationaryStochasticStructuredPyEnvironment(
          _global_context_sampling_fn,
          _arm_context_sampling_fn,
          6,
          _reward_fn,
          batch_size=2)
      time_step_spec = env.time_step_spec()
      action_spec = array_spec.BoundedArraySpec(
          shape=(), minimum=0, maximum=5, dtype=np.int32)
  
      random_policy = random_py_policy.RandomPyPolicy(
          time_step_spec=time_step_spec, action_spec=action_spec)
  
      for _ in range(5):
        time_step = env.reset()
        self.assertTrue(
            check_unbatched_time_step_spec(
                time_step=time_step,
                time_step_spec=time_step_spec,
                batch_size=env.batch_size))
  
        action = random_policy.action(time_step).action
        self.assertAllEqual(action.shape, [2])
        self.assertAllGreaterEqual(action, 0)
        self.assertAllLess(action, 6)
        time_step = env.step(action)
        self.assertEqual(time_step.reward.shape, (2,))
  
  
  if __name__ == '__main__':
    tf.test.main()
  