"""Tests for the replay module

Basically copied from:
https://github.com/deepmind/dqn_zoo/blob/master/dqn_zoo/replay_test.py
"""

# pylint: disable=missing-class-docstring
# pylint: disable=missing-function-docstring
# pylint: disable=protected-access

import collections

import chex
import numpy as np
from absl.testing import absltest

from rljaxff import parts
from rljaxff import replay as replay_lib


Pair = collections.namedtuple("Pair", ['a', 'b'])


class TransitionReplayTest(absltest.TestCase):

    def setUp(self):
        super().setUp()
        self.capacity = 10
        self.replay = replay_lib.TransitionReplay(
            capacity=self.capacity,
            structure=Pair(a=None, b=None),
            random_state=np.random.RandomState(1)
        )

        self.items = [
            Pair(a=1, b=2),
            Pair(a=11, b=22),
            Pair(a=111, b=222),
            Pair(a=1111, b=2222),
        ]
        for item in self.items:
            self.replay.add(item)

    def test_size(self):
        self.assertLen(self.items, self.replay.size)

    def test_capacity(self):
        self.assertEqual(self.capacity, self.replay.capacity)

    def test_sample(self):
        num_samples = 2
        samples = self.replay.sample(num_samples)
        chex.assert_shape(samples.a, (num_samples, ))

    def test_sample2(self):
        num_samples = 2
        samples = self.replay.sample(num_samples)

        valid_a_values = [p.a for p in self.items]
        valid_b_values = [p.b for p in self.items]

        valid_a = [a for a in samples.a if a in valid_a_values]
        valid_b = [b for b in samples.b if b in valid_b_values]
        self.assertEqual(num_samples, len(valid_a))
        self.assertEqual(num_samples, len(valid_b))


class TransitionAccumulatorTest(absltest.TestCase):

    def setUp(self):
        super().setUp()
        self.transition_accumulator = replay_lib.TransitionAccumulator()

        self.num_timesteps = 10
        self.observation_dims = (10, 5)
        self.discount = 0.99
        self.timesteps = []
        self.actions = []
        for i in range(self.num_timesteps):
            if i == 0:
                step_type = parts.StepType.FIRST
            elif i == self.num_timesteps - 1:
                step_type = parts.StepType.LAST
            else:
                step_type = parts.StepType.MID
            timestep = parts.Timestep(
                step_type=step_type,
                reward=parts.Reward(i),
                discount=self.discount,
                observation=np.array([i])
            )
            self.timesteps.append(timestep)
            self.actions.append(parts.Action(i))

    def test_reset(self):
        self.transition_accumulator.reset()
        self.assertEqual(self.transition_accumulator._timestep_tm1, None)
        self.assertEqual(self.transition_accumulator._a_tm1, None)

    def test_first_step(self):
        transitions = self.transition_accumulator.step(
            self.timesteps[0], self.actions[0]
        )
        self.assertEqual([], list(transitions))

    def test_mid(self):
        n = self.num_timesteps // 2
        for i in range(n + 1):
            transitions = list(self.transition_accumulator.step(
                self.timesteps[i], self.actions[i]
            ))

        transition = transitions[0]
        t_t = self.timesteps[n]
        t_tm1 = self.timesteps[n-1]
        a_tm1 = self.actions[n-1]

        self.assertEqual(t_tm1.observation, transition.s_tm1)
        self.assertEqual(a_tm1, transition.a_tm1)
        self.assertEqual(t_t.reward, transition.r_t)
        self.assertEqual(t_t.discount, transition.discount_t)
        self.assertEqual(t_t.observation, transition.s_t)


if __name__ == "__main__":
    absltest.main()

"""Tests for the replay module

Basically copied from:
https://github.com/deepmind/dqn_zoo/blob/master/dqn_zoo/replay_test.py
"""

# pylint: disable=missing-class-docstring
# pylint: disable=missing-function-docstring
# pylint: disable=protected-access

import collections

import chex
import numpy as np
from absl.testing import absltest

from rljaxff import parts
from rljaxff import replay as replay_lib


Pair = collections.namedtuple("Pair", ['a', 'b'])


class TransitionReplayTest(absltest.TestCase):

    def setUp(self):
        super().setUp()
        self.capacity = 10
        self.replay = replay_lib.TransitionReplay(
            capacity=self.capacity,
            structure=Pair(a=None, b=None),
            random_state=np.random.RandomState(1)
        )

        self.items = [
            Pair(a=1, b=2),
            Pair(a=11, b=22),
            Pair(a=111, b=222),
            Pair(a=1111, b=2222),
        ]
        for item in self.items:
            self.replay.add(item)

    def test_size(self):
        self.assertLen(self.items, self.replay.size)

    def test_capacity(self):
        self.assertEqual(self.capacity, self.replay.capacity)

    def test_sample(self):
        num_samples = 2
        samples = self.replay.sample(num_samples)
        chex.assert_shape(samples.a, (num_samples, ))

    def test_sample2(self):
        num_samples = 2
        samples = self.replay.sample(num_samples)

        valid_a_values = [p.a for p in self.items]
        valid_b_values = [p.b for p in self.items]

        valid_a = [a for a in samples.a if a in valid_a_values]
        valid_b = [b for b in samples.b if b in valid_b_values]
        self.assertEqual(num_samples, len(valid_a))
        self.assertEqual(num_samples, len(valid_b))


class TransitionAccumulatorTest(absltest.TestCase):

    def setUp(self):
        super().setUp()
        self.transition_accumulator = replay_lib.TransitionAccumulator()

        self.num_timesteps = 10
        self.observation_dims = (10, 5)
        self.discount = 0.99
        self.timesteps = []
        self.actions = []
        for i in range(self.num_timesteps):
            if i == 0:
                step_type = parts.StepType.FIRST
            elif i == self.num_timesteps - 1:
                step_type = parts.StepType.LAST
            else:
                step_type = parts.StepType.MID
            timestep = parts.Timestep(
                step_type=step_type,
                reward=parts.Reward(i),
                discount=self.discount,
                observation=np.array([i])
            )
            self.timesteps.append(timestep)
            self.actions.append(parts.Action(i))

    def test_reset(self):
        self.transition_accumulator.reset()
        self.assertEqual(self.transition_accumulator._timestep_tm1, None)
        self.assertEqual(self.transition_accumulator._a_tm1, None)

    def test_first_step(self):
        transitions = self.transition_accumulator.step(
            self.timesteps[0], self.actions[0]
        )
        self.assertEqual([], list(transitions))

    def test_mid(self):
        n = self.num_timesteps // 2
        for i in range(n + 1):
            transitions = list(self.transition_accumulator.step(
                self.timesteps[i], self.actions[i]
            ))

        transition = transitions[0]
        t_t = self.timesteps[n]
        t_tm1 = self.timesteps[n-1]
        a_tm1 = self.actions[n-1]

        self.assertEqual(t_tm1.observation, transition.s_tm1)
        self.assertEqual(a_tm1, transition.a_tm1)
        self.assertEqual(t_t.reward, transition.r_t)
        self.assertEqual(t_t.discount, transition.discount_t)
        self.assertEqual(t_t.observation, transition.s_t)


if __name__ == "__main__":
    absltest.main()

Tests for the replay module  Basically copied from: https://github.com/deepmind/dqn_zoo/blob/master/dqn_zoo/replay_test.py # pylint: disable=missing-class-docstring # pylint: disable=missing-function-docstring # pylint: disable=protected-access