text: """
  Tests for environment wrappers.
  """
  
  import os
  import tempfile
  
  import gym
  import numpy as np
  import pandas
  import pytest
  import tensorflow as tf
  
  from anyrl.envs import batched_gym_env
  from anyrl.envs.wrappers import (RL2Env, DownsampleEnv, GrayscaleEnv, FrameStackEnv,
                                   MaxEnv, ResizeImageEnv, BatchedFrameStack, ObservationPadEnv,
                                   LoggedEnv)
  from anyrl.tests import SimpleEnv
  
  
  def test_rl2_num_eps():
      """
      Test that RL^2 meta-episodes contain the right number
      of sub-episodes.
      """
      real_env = SimpleEnv(1, (3, 4), 'uint8')
      env = RL2Env(real_env, real_env.action_space.sample(),
                   num_eps=3, warmup_eps=1)
      done_eps = 0
      env.reset()
      while done_eps < 3:
          obs, _, done, _ = env.step(env.action_space.sample())
          if obs[3]:
              done_eps += 1
          if done_eps == 3:
              assert done
              break
          else:
              assert not done
  
  
  def test_rl2_rewards():
      """
      Test that rewards in RL^2 are masked properly.
      """
      real_env = SimpleEnv(1, (3, 4), 'uint8')
      env = RL2Env(real_env, real_env.action_space.sample(), num_eps=5, warmup_eps=-2)
      env.reset()
      done_eps = 0
      while done_eps < 3:
          obs, rew, _, _ = env.step(env.action_space.sample())
          assert rew == 0
          if obs[3]:
              done_eps += 1
      while True:
          _, rew, done, _ = env.step(env.action_space.sample())
          assert rew != 0
          if done:
              break
  
  
  def test_downsample_rate_1():
      """
      Test DownsampleEnv with rate=1.
      """
      low = np.array([[1, 2], [3, 4]])
      high = np.array([[3, 4], [5, 6]])
      env = DownsampleEnv(ShapeEnv(low, high), 1)
      assert (env.observation_space.low == low).all()
      assert (env.observation_space.high == high).all()
  
  
  def test_downsample_rate_2():
      """
      Test DownsampleEnv with rate=2.
      """
      low = np.array([[1, 2, 3, 4], [5, 6, 7, 8]])
      high = np.array([[9, 10, 11, 12], [13, 14, 15, 16]])
      env = DownsampleEnv(ShapeEnv(low, high), 2)
      assert env.observation_space.shape == (1, 2)
      assert (env.observation_space.low == np.array([[1, 3]])).all()
      assert (env.observation_space.high == np.array([[9, 11]])).all()
  
      low = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 6, 7, 8]])
      high = np.array([[9, 10, 11, 12], [13, 14, 15, 16], [16, 14, 15, 16]])
      env = DownsampleEnv(ShapeEnv(low, high), 2)
      assert env.observation_space.shape == (2, 2)
      assert (env.observation_space.low == np.array([[1, 3], [9, 7]])).all()
      assert (env.observation_space.high == np.array([[9, 11], [16, 15]])).all()
  
  
  def test_grayscale_integers():
      """
      Tests GrayscaleEnv for integer observations.
      """
      low = np.array([[[0, 0, 0], [16, 16, 50]]])
      high = np.array([[[19, 255, 13], [255, 255, 255]]])
      env = GrayscaleEnv(ShapeEnv(low, high))
      assert env.observation_space.shape == (1, 2, 1)
      assert (env.observation_space.low == np.array([[[0], [26]]])).all()
      assert (env.observation_space.high == np.array([[[95], [255]]])).all()
  
  
  def test_stack_3():
      """
      Test FrameStackEnv for 3 frames.
      """
      low = np.zeros((4, 5, 2))
      high = np.zeros((4, 5, 2)) + 0xff
      env = FrameStackEnv(ShapeEnv(low, high), 3)
      assert env.observation_space.shape == (4, 5, 6)
      obs1 = env.reset()
      obs2, _, _, _ = env.step(0)
      obs3, _, _, _ = env.step(0)
      assert not (obs1 == obs2).all()
      assert not (obs1 == obs3).all()
      assert not (obs2 == obs3).all()
      assert (obs1[:, :, 2:] == obs2[:, :, :4]).all()
      assert (obs2[:, :, 2:] == obs3[:, :, :4]).all()
  
  
  def test_stack_3_no_concat():
      """
      Test FrameStackEnv for 3 frames with no concatenation.
      """
      low = np.zeros((4, 5, 2))
      high = np.zeros((4, 5, 2)) + 0xff
      env = FrameStackEnv(ShapeEnv(low, high), 3, concat=False)
      assert env.observation_space.box.shape == (4, 5, 2)
      assert env.observation_space.count == 3
      obs1 = env.reset()
      obs2, _, _, _, = env.step(0)
      assert np.allclose(obs2[0], obs1[0])
      assert not np.allclose(obs2[-1], obs1[0])
  
  
  def test_max_2():
      """
      Test MaxEnv for 2 frames.
      """
      env = SimpleEnv(5, (3, 2, 5), 'float32')
      actions = [env.action_space.sample() for _ in range(4)]
      frame1 = env.reset()
      frame2 = env.step(actions[0])[0]
      frame3 = env.step(actions[1])[0]
      frame4 = env.step(actions[2])[0]
      frame5 = env.step(actions[3])[0]
      wrapped = MaxEnv(env, num_images=2)
      max1 = wrapped.reset()
      max2 = wrapped.step(actions[0])[0]
      max3 = wrapped.step(actions[1])[0]
      max4 = wrapped.step(actions[2])[0]
      max5 = wrapped.step(actions[3])[0]
      assert (max1 == frame1).all()
      assert (max2 == np.max([frame1, frame2], axis=0)).all()
      assert (max3 == np.max([frame2, frame3], axis=0)).all()
      assert (max4 == np.max([frame3, frame4], axis=0)).all()
      assert (max5 == np.max([frame4, frame5], axis=0)).all()
  
  
  def test_resize_even():
      """
      Test ResizeImageEnv for an even number of pixels.
      """
      env = SimpleEnv(5, (13, 5, 3), 'float32')
      frame = env.reset()
      actual = ResizeImageEnv(env, size=(5, 4)).reset()
      expected = tf.Session().run(
          tf.image.resize_images(frame, [5, 4], method=tf.image.ResizeMethod.AREA))
      assert actual.shape == (5, 4, 3)
      assert np.allclose(actual, expected)
  
  
  @pytest.mark.parametrize('concat', [False, True])
  def test_batched_stack(concat):
      """
      Test that BatchedFrameStack is equivalent to a regular
      batched FrameStackEnv.
      """
      envs = [lambda idx=i: SimpleEnv(idx+2, (3, 2, 5), 'float32') for i in range(6)]
      env1 = BatchedFrameStack(batched_gym_env(envs, num_sub_batches=3, sync=True),
                               concat=concat)
      env2 = batched_gym_env([lambda env=e: FrameStackEnv(env(), concat=concat) for e in envs],
                             num_sub_batches=3, sync=True)
      for j in range(50):
          for i in range(3):
              if j == 0 or (j + i) % 17 == 0:
                  env1.reset_start(sub_batch=i)
                  env2.reset_start(sub_batch=i)
                  obs1 = env1.reset_wait(sub_batch=i)
                  obs2 = env2.reset_wait(sub_batch=i)
                  assert np.allclose(obs1, obs2)
              actions = [env1.action_space.sample() for _ in range(2)]
              env1.step_start(actions, sub_batch=i)
              env2.step_start(actions, sub_batch=i)
              obs1, rews1, dones1, _ = env1.step_wait(sub_batch=i)
              obs2, rews2, dones2, _ = env2.step_wait(sub_batch=i)
              assert np.allclose(obs1, obs2)
              assert np.array(rews1 == rews2).all()
              assert np.array(dones1 == dones2).all()
  
  
  def test_obs_pad_centered():
      """
      Test ObservationPadEnv with centered padding.
      """
      env = ShapeEnv(np.array([[1, 2, 5], [3, 4, 0]]),
                     np.array([[12, 21, 15], [31, 14, 10]]))
      padded = ObservationPadEnv(env, (3, 5))
      assert padded.observation_space.shape == (3, 5)
      assert np.allclose(padded.observation_space.low,
                         np.array([[0, 1, 2, 5, 0], [0, 3, 4, 0, 0], [0]*5]))
      assert np.allclose(padded.observation_space.high,
                         np.array([[0, 12, 21, 15, 0], [0, 31, 14, 10, 0], [0]*5]))
  
  
  def test_obs_pad_uncentered():
      """
      Test ObservationPadEnv with uncentered padding.
      """
      env = ShapeEnv(np.array([[1, 2, 5], [3, 4, 0]]),
                     np.array([[12, 21, 15], [31, 14, 10]]))
      padded = ObservationPadEnv(env, (3, 5), center=False)
      assert padded.observation_space.shape == (3, 5)
      assert np.allclose(padded.observation_space.low,
                         np.array([[1, 2, 5, 0, 0], [3, 4, 0, 0, 0], [0]*5]))
      assert np.allclose(padded.observation_space.high,
                         np.array([[12, 21, 15, 0, 0], [31, 14, 10, 0, 0], [0]*5]))
  
  
  def test_logged_single_env():
      """
      Test LoggedEnv for a single environment.
      """
      with tempfile.TemporaryDirectory() as dirpath:
          log_file = os.path.join(dirpath, 'monitor.csv')
          env = LoggedEnv(SimpleEnv(2, (3,), 'float32'), log_file)
          for _ in range(4):
              env.reset()
              while not env.step(env.action_space.sample())[2]:
                  pass
          env.close()
          with open(log_file, 'rt'):
              log_contents = pandas.read_csv(log_file)
              assert list(log_contents['r']) == [2] * 4
              assert list(log_contents['l']) == [3] * 4
  
  
  def test_multi_env():
      """
      Test monitoring for concurrent environments.
      """
      with tempfile.TemporaryDirectory() as dirpath:
          log_file = os.path.join(dirpath, 'monitor.csv')
          env1 = LoggedEnv(SimpleEnv(2, (3,), 'float32'), log_file, use_locking=True)
          env2 = LoggedEnv(SimpleEnv(3, (3,), 'float32'), log_file, use_locking=True)
  
          env1.reset()
          env2.reset()
          for _ in range(13):
              for env in [env1, env2]:
                  if env.step(env.action_space.sample())[2]:
                      env.reset()
          env1.close()
          env2.close()
  
          with open(log_file, 'rt'):
              log_contents = pandas.read_csv(log_file)
              assert list(log_contents['r']) == [2, 2.5, 2, 2.5, 2, 2, 2.5]
              assert list(log_contents['l']) == [3, 4, 3, 4, 3, 3, 4]
  
  
  class ShapeEnv(gym.Env):
      """
      An environment with a pre-defined observation shape.
      """
  
      def __init__(self, low, high):
          super(ShapeEnv, self).__init__()
          self.observation_space = gym.spaces.Box(low, high, dtype=low.dtype)
  
      def reset(self):
          return self.observation_space.sample()
  
      def step(self, action):
          return self.observation_space.sample(), 0, False, {}
  
      def render(self, mode='human'):
          pass
  