Repository: ray-project/ray
File: rllib/algorithms/dqn/tests/test_dqn_rl_module.py (Python)

Code before:


class TestDQNRLModule:
    @pytest.mark.parametrize("obs_space_name", ["box", "image", "dict"])
    @pytest.mark.parametrize("forward_method", ["train", "exploration", "inference"])
    @pytest.mark.parametrize("double_q", [True, False])
    @pytest.mark.parametrize("dueling", [True, False])
    def test_forward(self, obs_space_name, forward_method, double_q, dueling):
        """Test forward methods with different obs spaces and config settings."""
        obs_space = OBS_SPACES[obs_space_name]
        action_space = Discrete(4)

        module = _get_dqn_module(
            obs_space, action_space, double_q=double_q, dueling=dueling
        )

        if (
            forward_method == "train"
        ):  # forward train needs batching, exploration and inference don't
            module.train()
            # Create a batch first
            batch_size = 4
            obs_list = [obs_space.sample() for _ in range(batch_size)]
            next_obs_list = [obs_space.sample() for _ in range(batch_size)]

            obs_batch = tree.map_structure(lambda *x: np.stack(x, axis=0), *obs_list)
            next_obs_batch = tree.map_structure(
                lambda *x: np.stack(x, axis=0), *next_obs_list
            )

            obs_batch = tree.map_structure(lambda x: x.astype(np.float32), obs_batch)
            next_obs_batch = tree.map_structure(
                lambda x: x.astype(np.float32), next_obs_batch
            )

            batch = {
                Columns.OBS: convert_to_torch_tensor(obs_batch),
                Columns.NEXT_OBS: convert_to_torch_tensor(next_obs_batch),
                Columns.ACTIONS: convert_to_torch_tensor(
                    np.array([0] * batch_size, dtype=np.int64)
                ),
                Columns.REWARDS: convert_to_torch_tensor(
                    np.array([1.0] * batch_size, dtype=np.float32)
                ),
                Columns.TERMINATEDS: convert_to_torch_tensor(
                    np.array([False] * batch_size, dtype=np.bool_)
                ),
                Columns.TRUNCATEDS: convert_to_torch_tensor(
                    np.array([False] * batch_size, dtype=np.bool_)
                ),
            }

Review comment (suggestion):
We can add `np.stack(x, axis=0, dtype=np.float32)` to remove the following two lines

Code after:
    return module


class TestDQNRLModule:
    @pytest.mark.parametrize("obs_space_name", ["box", "image", "dict"])
    @pytest.mark.parametrize("forward_method", ["train", "exploration", "inference"])
    @pytest.mark.parametrize("double_q", [True, False])
    @pytest.mark.parametrize("dueling", [True, False])
    def test_forward(self, obs_space_name, forward_method, double_q, dueling):
        """Test forward methods with different obs spaces and config settings."""
        obs_space = OBS_SPACES[obs_space_name]
        action_space = Discrete(4)

        module = _get_dqn_module(
            obs_space, action_space, double_q=double_q, dueling=dueling
        )

        if (
            forward_method == "train"
        ):  # forward train needs batching, exploration and inference don't
            module.train()
            # Create a batch first
            batch_size = 4
            obs_list = [obs_space.sample() for _ in range(batch_size)]
            next_obs_list = [obs_space.sample() for _ in range(batch_size)]

            obs_batch = tree.map_structure(
                lambda *x: np.stack(x, axis=0, dtype=np.float32), *obs_list
            )
            next_obs_batch = tree.map_structure(
                lambda *x: np.stack(x, axis=0, dtype=np.float32), *next_obs_list
            )

            batch = {
                Columns.OBS: convert_to_torch_tensor(obs_batch),
                Columns.NEXT_OBS: convert_to_torch_tensor(next_obs_batch),
                Columns.ACTIONS: convert_to_torch_tensor(
                    np.array([0] * batch_size, dtype=np.int64)
                ),
                Columns.REWARDS: convert_to_torch_tensor(
                    np.array([1.0] * batch_size, dtype=np.float32)
                ),
                Columns.TERMINATEDS: convert_to_torch_tensor(
                    np.array([False] * batch_size, dtype=np.bool_)
                ),
                Columns.TRUNCATEDS: convert_to_torch_tensor(
                    np.array([False] * batch_size, dtype=np.bool_)
                ),
            }

            # Forward pass and check outputs