#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

import logging
import os
import random
import unittest
from typing import Optional

import numpy as np
import pytorch_lightning as pl
import torch
import torch.nn as nn
from parameterized import parameterized
from reagent.core import types as rlt
from reagent.core.parameters import (
    NormalizationData,
    NormalizationParameters,
    ProblemDomain,
    Seq2RewardTrainerParameters,
)
from reagent.gym.envs import Gym
from reagent.gym.utils import create_df_from_replay_buffer
from reagent.models.seq2reward_model import Seq2RewardNetwork
from reagent.net_builder.value.fully_connected import FullyConnected
from reagent.prediction.predictor_wrapper import (
    Seq2RewardWithPreprocessor,
    Seq2RewardPlanShortSeqWithPreprocessor,
    FAKE_STATE_ID_LIST_FEATURES,
    FAKE_STATE_ID_SCORE_LIST_FEATURES,
)
from reagent.preprocessing.identify_types import DO_NOT_PREPROCESS
from reagent.preprocessing.preprocessor import Preprocessor
from reagent.training.utils import gen_permutations
from reagent.training.world_model.compress_model_trainer import CompressModelTrainer
from reagent.training.world_model.seq2reward_trainer import get_Q, Seq2RewardTrainer
from torch.utils.data import DataLoader

logger = logging.getLogger(__name__)

SEED = 0
STRING_GAME_TESTS = [(False,), (True,)]


class FakeStepPredictionNetwork(nn.Module):
    def __init__(self, look_ahead_steps):
        super().__init__()
        self.look_ahead_steps = look_ahead_steps

    def forward(self, state: torch.Tensor):
        """
        Given the current state, predict the probability of
        experiencing next n steps (1 <=n <= look_ahead_steps)

        For the test purpose, it outputs fixed fake numbers
        """
        batch_size, _ = state.shape
        return torch.ones(batch_size, self.look_ahead_steps).float()


class FakeSeq2RewardNetwork(nn.Module):
    def forward(
        self,
        state: rlt.FeatureData,
        action: rlt.FeatureData,
        valid_reward_len: Optional[torch.Tensor] = None,
    ):
        """
        Mimic I/O of Seq2RewardNetwork but return fake reward
        Reward is the concatenation of action indices, independent
        of state.

        For example, when seq_len = 3, batch_size = 1, action_num = 2,
        acc_reward = tensor(
            [[  0.],
            [  1.],
            [ 10.],
            [ 11.],
            [100.],
            [101.],
            [110.],
            [111.]]
        )

        Input action shape: seq_len, batch_size, num_action
        Output acc_reward shape: batch_size, 1
        """
        # pyre-fixme[9]: action has type `FeatureData`; used as `Tensor`.
        action = action.float_features.transpose(0, 1)
        action_indices = torch.argmax(action, dim=2).tolist()
        acc_reward = torch.tensor(
            list(map(lambda x: float("".join(map(str, x))), action_indices))
        ).reshape(-1, 1)
        logger.info(f"acc_reward: {acc_reward}")
        return rlt.Seq2RewardOutput(acc_reward=acc_reward)


def create_string_game_data(
    dataset_size=10000, training_data_ratio=0.9, filter_short_sequence=False
):
    SEQ_LEN = 6
    NUM_ACTION = 2
    NUM_MDP_PER_BATCH = 5

    env = Gym(env_name="StringGame-v0", set_max_steps=SEQ_LEN)
    df = create_df_from_replay_buffer(
        env=env,
        problem_domain=ProblemDomain.DISCRETE_ACTION,
        desired_size=dataset_size,
        multi_steps=None,
        ds="2020-10-10",
    )

    if filter_short_sequence:
        batch_size = NUM_MDP_PER_BATCH
        time_diff = torch.ones(SEQ_LEN, batch_size)
        valid_step = SEQ_LEN * torch.ones(batch_size, dtype=torch.int64)[:, None]
        not_terminal = torch.Tensor(
            [0 if i == SEQ_LEN - 1 else 1 for i in range(SEQ_LEN)]
        )
        not_terminal = torch.transpose(not_terminal.tile(NUM_MDP_PER_BATCH, 1), 0, 1)
    else:
        batch_size = NUM_MDP_PER_BATCH * SEQ_LEN
        time_diff = torch.ones(SEQ_LEN, batch_size)
        valid_step = torch.arange(SEQ_LEN, 0, -1).tile(NUM_MDP_PER_BATCH)[:, None]
        not_terminal = torch.transpose(
            torch.tril(torch.ones(SEQ_LEN, SEQ_LEN), diagonal=-1).tile(
                NUM_MDP_PER_BATCH, 1
            ),
            0,
            1,
        )

    num_batches = int(dataset_size / SEQ_LEN / NUM_MDP_PER_BATCH)
    batches = [None for _ in range(num_batches)]
    batch_count, batch_seq_count = 0, 0
    batch_reward = torch.zeros(SEQ_LEN, batch_size)
    batch_action = torch.zeros(SEQ_LEN, batch_size, NUM_ACTION)
    batch_state = torch.zeros(SEQ_LEN, batch_size, NUM_ACTION)
    for mdp_id in sorted(set(df.mdp_id)):
        mdp = df[df["mdp_id"] == mdp_id].sort_values("sequence_number", ascending=True)
        if len(mdp) != SEQ_LEN:
            continue

        all_step_reward = torch.Tensor(list(mdp["reward"]))
        all_step_state = torch.Tensor([list(s.values()) for s in mdp["state_features"]])
        all_step_action = torch.zeros_like(all_step_state)
        all_step_action[torch.arange(SEQ_LEN), [int(a) for a in mdp["action"]]] = 1.0

        for j in range(SEQ_LEN):
            if filter_short_sequence and j > 0:
                break

            reward = torch.zeros_like(all_step_reward)
            reward[: SEQ_LEN - j] = all_step_reward[-(SEQ_LEN - j) :]
            batch_reward[:, batch_seq_count] = reward

            state = torch.zeros_like(all_step_state)
            state[: SEQ_LEN - j] = all_step_state[-(SEQ_LEN - j) :]
            batch_state[:, batch_seq_count] = state

            action = torch.zeros_like(all_step_action)
            action[: SEQ_LEN - j] = all_step_action[-(SEQ_LEN - j) :]
            batch_action[:, batch_seq_count] = action

            batch_seq_count += 1

        if batch_seq_count == batch_size:
            batches[batch_count] = rlt.MemoryNetworkInput(
                reward=batch_reward,
                action=rlt.FeatureData(float_features=batch_action),
                state=rlt.FeatureData(float_features=batch_state),
                next_state=rlt.FeatureData(
                    float_features=torch.zeros_like(batch_state)
                ),  # fake, not used anyway
                not_terminal=not_terminal,
                time_diff=time_diff,
                valid_step=valid_step,
                step=None,
            )
            batch_count += 1
            batch_seq_count = 0
            batch_reward = torch.zeros_like(batch_reward)
            batch_action = torch.zeros_like(batch_action)
            batch_state = torch.zeros_like(batch_state)
    assert batch_count == num_batches

    num_training_batches = int(training_data_ratio * num_batches)
    training_data = DataLoader(
        batches[:num_training_batches], collate_fn=lambda x: x[0]
    )
    eval_data = DataLoader(batches[num_training_batches:], collate_fn=lambda x: x[0])
    return training_data, eval_data


def train_seq2reward_model(training_data, learning_rate=0.01, num_epochs=5):
    SEQ_LEN, batch_size, NUM_ACTION = next(
        iter(training_data)
    ).action.float_features.shape
    assert SEQ_LEN == 6 and NUM_ACTION == 2

    seq2reward_network = Seq2RewardNetwork(
        state_dim=NUM_ACTION,
        action_dim=NUM_ACTION,
        num_hiddens=64,
        num_hidden_layers=2,
    )

    trainer_param = Seq2RewardTrainerParameters(
        learning_rate=learning_rate,
        multi_steps=SEQ_LEN,
        action_names=["0", "1"],
        gamma=1.0,
        view_q_value=True,
    )

    trainer = Seq2RewardTrainer(
        seq2reward_network=seq2reward_network, params=trainer_param
    )

    pl.seed_everything(SEED)
    pl_trainer = pl.Trainer(max_epochs=num_epochs, deterministic=True)
    pl_trainer.fit(trainer, training_data)

    return trainer


def eval_seq2reward_model(eval_data, seq2reward_trainer):
    SEQ_LEN, batch_size, NUM_ACTION = next(iter(eval_data)).action.float_features.shape

    initial_state = torch.Tensor([[0, 0]])
    initial_state_q_values = torch.squeeze(
        get_Q(
            seq2reward_trainer.seq2reward_network,
            initial_state,
            seq2reward_trainer.all_permut,
        )
    )

    total_mse_loss = 0
    total_q_values = torch.zeros(NUM_ACTION)
    total_action_distribution = torch.zeros(NUM_ACTION)
    for idx, batch in enumerate(eval_data):
        (
            mse_loss,
            _,
            q_values,
            action_distribution,
        ) = seq2reward_trainer.validation_step(batch, idx)
        total_mse_loss += mse_loss
        total_q_values += torch.tensor(q_values)
        total_action_distribution += torch.tensor(action_distribution)

    N_eval = len(eval_data)
    eval_mse_loss = total_mse_loss / N_eval
    eval_q_values = total_q_values / N_eval
    eval_action_distribution = total_action_distribution / N_eval

    return (
        initial_state_q_values,
        eval_mse_loss,
        eval_q_values,
        eval_action_distribution,
    )


def train_seq2reward_compress_model(
    training_data, seq2reward_network, learning_rate=0.1, num_epochs=5
):
    SEQ_LEN, batch_size, NUM_ACTION = next(
        iter(training_data)
    ).action.float_features.shape
    assert SEQ_LEN == 6 and NUM_ACTION == 2

    compress_net_builder = FullyConnected(sizes=[8, 8])
    state_normalization_data = NormalizationData(
        dense_normalization_parameters={
            0: NormalizationParameters(feature_type=DO_NOT_PREPROCESS),
            1: NormalizationParameters(feature_type=DO_NOT_PREPROCESS),
        }
    )
    compress_model_network = compress_net_builder.build_value_network(
        state_normalization_data,
        output_dim=NUM_ACTION,
    )

    trainer_param = Seq2RewardTrainerParameters(
        learning_rate=0.0,
        multi_steps=SEQ_LEN,
        action_names=["0", "1"],
        compress_model_learning_rate=learning_rate,
        gamma=1.0,
        view_q_value=True,
    )

    trainer = CompressModelTrainer(
        compress_model_network=compress_model_network,
        seq2reward_network=seq2reward_network,
        params=trainer_param,
    )

    pl.seed_everything(SEED)
    pl_trainer = pl.Trainer(max_epochs=num_epochs, deterministic=True)
    pl_trainer.fit(trainer, training_data)

    return trainer


def eval_seq2reward_compress_model(eval_data, compress_model_trainer):
    SEQ_LEN, batch_size, NUM_ACTION = next(iter(eval_data)).action.float_features.shape
    total_mse_loss = 0
    total_q_values = torch.zeros(NUM_ACTION)
    total_action_distribution = torch.zeros(NUM_ACTION)
    for idx, batch in enumerate(eval_data):
        (
            mse_loss,
            q_values,
            action_distribution,
            _,
        ) = compress_model_trainer.validation_step(batch, idx)
        total_mse_loss += mse_loss
        total_q_values += torch.tensor(q_values)
        total_action_distribution += torch.tensor(action_distribution)

    N_eval = len(eval_data)
    eval_mse_loss = total_mse_loss / N_eval
    eval_q_values = total_q_values / N_eval
    eval_action_distribution = total_action_distribution / N_eval

    return eval_mse_loss, eval_q_values, eval_action_distribution


class TestSeq2Reward(unittest.TestCase):
    def test_seq2reward_with_preprocessor_plan_short_sequence(self):
        self._test_seq2reward_with_preprocessor(plan_short_sequence=True)

    def test_seq2reward_with_preprocessor_plan_full_sequence(self):
        self._test_seq2reward_with_preprocessor(plan_short_sequence=False)

    def _test_seq2reward_with_preprocessor(self, plan_short_sequence):
        state_dim = 4
        action_dim = 2
        seq_len = 3
        model = FakeSeq2RewardNetwork()
        state_normalization_parameters = {
            i: NormalizationParameters(
                feature_type=DO_NOT_PREPROCESS, mean=0.0, stddev=1.0
            )
            for i in range(1, state_dim)
        }
        state_preprocessor = Preprocessor(state_normalization_parameters, False)

        if plan_short_sequence:
            step_prediction_model = FakeStepPredictionNetwork(seq_len)
            model_with_preprocessor = Seq2RewardPlanShortSeqWithPreprocessor(
                model,
                step_prediction_model,
                state_preprocessor,
                seq_len,
                action_dim,
            )
        else:
            model_with_preprocessor = Seq2RewardWithPreprocessor(
                model,
                state_preprocessor,
                seq_len,
                action_dim,
            )
        input_prototype = rlt.ServingFeatureData(
            float_features_with_presence=state_preprocessor.input_prototype(),
            id_list_features=FAKE_STATE_ID_LIST_FEATURES,
            id_score_list_features=FAKE_STATE_ID_SCORE_LIST_FEATURES,
        )
        q_values = model_with_preprocessor(input_prototype)
        if plan_short_sequence:
            # When planning for 1, 2, and 3 steps ahead,
            # the expected q values are respectively:
            # [0, 1], [1, 11], [11, 111]
            # Weighting the expected q values by predicted step
            # probabilities [0.33, 0.33, 0.33], we have [4, 41]
            expected_q_values = torch.tensor([[4.0, 41.0]])
        else:
            expected_q_values = torch.tensor([[11.0, 111.0]])
        assert torch.all(expected_q_values == q_values)

    def test_get_Q(self):
        NUM_ACTION = 2
        MULTI_STEPS = 3
        BATCH_SIZE = 2
        STATE_DIM = 4
        all_permut = gen_permutations(MULTI_STEPS, NUM_ACTION)
        seq2reward_network = FakeSeq2RewardNetwork()
        state = torch.zeros(BATCH_SIZE, STATE_DIM)
        q_values = get_Q(seq2reward_network, state, all_permut)
        expected_q_values = torch.tensor([[11.0, 111.0], [11.0, 111.0]])
        logger.info(f"q_values: {q_values}")
        assert torch.all(expected_q_values == q_values)

    def test_gen_permutations_seq_len_1_action_6(self):
        SEQ_LEN = 1
        NUM_ACTION = 6
        expected_outcome = torch.tensor([[0], [1], [2], [3], [4], [5]])
        self._test_gen_permutations(SEQ_LEN, NUM_ACTION, expected_outcome)

    def test_gen_permutations_seq_len_3_num_action_2(self):
        SEQ_LEN = 3
        NUM_ACTION = 2
        expected_outcome = torch.tensor(
            [
                [0, 0, 0],
                [0, 0, 1],
                [0, 1, 0],
                [0, 1, 1],
                [1, 0, 0],
                [1, 0, 1],
                [1, 1, 0],
                [1, 1, 1],
            ]
        )
        self._test_gen_permutations(SEQ_LEN, NUM_ACTION, expected_outcome)

    def _test_gen_permutations(self, SEQ_LEN, NUM_ACTION, expected_outcome):
        # expected shape: SEQ_LEN, PERM_NUM, ACTION_DIM
        result = gen_permutations(SEQ_LEN, NUM_ACTION)
        assert result.shape == (SEQ_LEN, NUM_ACTION ** SEQ_LEN, NUM_ACTION)
        outcome = torch.argmax(result.transpose(0, 1), dim=-1)
        assert torch.all(outcome == expected_outcome)

    @parameterized.expand(STRING_GAME_TESTS)
    @unittest.skipIf("SANDCASTLE" in os.environ, "Skipping long test on sandcastle.")
    def test_seq2reward_on_string_game_v0(self, filter_short_sequence):
        np.random.seed(SEED)
        random.seed(SEED)
        torch.manual_seed(SEED)
        training_data, eval_data = create_string_game_data(
            filter_short_sequence=filter_short_sequence
        )
        seq2reward_trainer = train_seq2reward_model(training_data)
        (
            initial_state_q_values,
            eval_mse_loss,
            eval_q_values,
            eval_action_distribution,
        ) = eval_seq2reward_model(eval_data, seq2reward_trainer)

        assert abs(initial_state_q_values[0].item() - 10) < 1.0
        assert abs(initial_state_q_values[1].item() - 5) < 1.0

        if filter_short_sequence:
            assert eval_mse_loss < 0.1
        else:
            # Same short sequences may have different total rewards due to the missing
            # states and actions in previous steps, so the trained network is not able
            # to reduce the mse loss to values close to zero.
            assert eval_mse_loss < 10

        compress_model_trainer = train_seq2reward_compress_model(
            training_data, seq2reward_trainer.seq2reward_network
        )
        (
            compress_eval_mse_loss,
            compress_eval_q_values,
            compress_eval_action_distribution,
        ) = eval_seq2reward_compress_model(eval_data, compress_model_trainer)

        assert compress_eval_mse_loss < 1e-5
        assert torch.all(eval_q_values - compress_eval_q_values < 1e-5)
        assert torch.all(
            eval_action_distribution - compress_eval_action_distribution < 1e-5
        )

[(39, 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), False, 'import logging\n'), (105, 'reagent.gym.envs.Gym', 'Gym', ([], {'env_name': '"""StringGame-v0"""', 'set_max_steps': 'SEQ_LEN'}), False, 'from reagent.gym.envs import Gym\n'), (106, 'reagent.gym.utils.create_df_from_replay_buffer', 'create_df_from_replay_buffer', ([], {'env': 'env', 'problem_domain': 'ProblemDomain.DISCRETE_ACTION', 'desired_size': 'dataset_size', 'multi_steps': 'None', 'ds': '"""2020-10-10"""'}), False, 'from reagent.gym.utils import create_df_from_replay_buffer\n'), (137, 'torch.zeros', 'torch.zeros', (['SEQ_LEN', 'batch_size'], {}), False, 'import torch\n'), (138, 'torch.zeros', 'torch.zeros', (['SEQ_LEN', 'batch_size', 'NUM_ACTION'], {}), False, 'import torch\n'), (139, 'torch.zeros', 'torch.zeros', (['SEQ_LEN', 'batch_size', 'NUM_ACTION'], {}), False, 'import torch\n'), (189, 'torch.utils.data.DataLoader', 'DataLoader', (['batches[:num_training_batches]'], {'collate_fn': '(lambda x: x[0])'}), False, 'from torch.utils.data import DataLoader\n'), (192, 'torch.utils.data.DataLoader', 'DataLoader', (['batches[num_training_batches:]'], {'collate_fn': '(lambda x: x[0])'}), False, 'from torch.utils.data import DataLoader\n'), (202, 'reagent.models.seq2reward_model.Seq2RewardNetwork', 'Seq2RewardNetwork', ([], {'state_dim': 'NUM_ACTION', 'action_dim': 'NUM_ACTION', 'num_hiddens': '(64)', 'num_hidden_layers': '(2)'}), False, 'from reagent.models.seq2reward_model import Seq2RewardNetwork\n'), (209, 'reagent.core.parameters.Seq2RewardTrainerParameters', 'Seq2RewardTrainerParameters', ([], {'learning_rate': 'learning_rate', 'multi_steps': 'SEQ_LEN', 'action_names': "['0', '1']", 'gamma': '(1.0)', 'view_q_value': '(True)'}), False, 'from reagent.core.parameters import NormalizationData, NormalizationParameters, ProblemDomain, Seq2RewardTrainerParameters\n'), (217, 'reagent.training.world_model.seq2reward_trainer.Seq2RewardTrainer', 'Seq2RewardTrainer', ([], {'seq2reward_network': 'seq2reward_network', 'params': 'trainer_param'}), False, 'from reagent.training.world_model.seq2reward_trainer import get_Q, Seq2RewardTrainer\n'), (221, 'pytorch_lightning.seed_everything', 'pl.seed_everything', (['SEED'], {}), True, 'import pytorch_lightning as pl\n'), (222, 'pytorch_lightning.Trainer', 'pl.Trainer', ([], {'max_epochs': 'num_epochs', 'deterministic': '(True)'}), True, 'import pytorch_lightning as pl\n'), (231, 'torch.Tensor', 'torch.Tensor', (['[[0, 0]]'], {}), False, 'import torch\n'), (241, 'torch.zeros', 'torch.zeros', (['NUM_ACTION'], {}), False, 'import torch\n'), (242, 'torch.zeros', 'torch.zeros', (['NUM_ACTION'], {}), False, 'import torch\n'), (275, 'reagent.net_builder.value.fully_connected.FullyConnected', 'FullyConnected', ([], {'sizes': '[8, 8]'}), False, 'from reagent.net_builder.value.fully_connected import FullyConnected\n'), (287, 'reagent.core.parameters.Seq2RewardTrainerParameters', 'Seq2RewardTrainerParameters', ([], {'learning_rate': '(0.0)', 'multi_steps': 'SEQ_LEN', 'action_names': "['0', '1']", 'compress_model_learning_rate': 'learning_rate', 'gamma': '(1.0)', 'view_q_value': '(True)'}), False, 'from reagent.core.parameters import NormalizationData, NormalizationParameters, ProblemDomain, Seq2RewardTrainerParameters\n'), (296, 'reagent.training.world_model.compress_model_trainer.CompressModelTrainer', 'CompressModelTrainer', ([], {'compress_model_network': 'compress_model_network', 'seq2reward_network': 'seq2reward_network', 'params': 'trainer_param'}), False, 'from reagent.training.world_model.compress_model_trainer import CompressModelTrainer\n'), (302, 'pytorch_lightning.seed_everything', 'pl.seed_everything', (['SEED'], {}), True, 'import pytorch_lightning as pl\n'), (303, 'pytorch_lightning.Trainer', 'pl.Trainer', ([], {'max_epochs': 'num_epochs', 'deterministic': '(True)'}), True, 'import pytorch_lightning as pl\n'), (312, 'torch.zeros', 'torch.zeros', (['NUM_ACTION'], {}), False, 'import torch\n'), (313, 'torch.zeros', 'torch.zeros', (['NUM_ACTION'], {}), False, 'import torch\n'), (429, 'parameterized.parameterized.expand', 'parameterized.expand', (['STRING_GAME_TESTS'], {}), False, 'from parameterized import parameterized\n'), (430, 'unittest.skipIf', 'unittest.skipIf', (["('SANDCASTLE' in os.environ)", '"""Skipping long test on sandcastle."""'], {}), False, 'import unittest\n'), (95, 'reagent.core.types.Seq2RewardOutput', 'rlt.Seq2RewardOutput', ([], {'acc_reward': 'acc_reward'}), True, 'from reagent.core import types as rlt\n'), (116, 'torch.ones', 'torch.ones', (['SEQ_LEN', 'batch_size'], {}), False, 'import torch\n'), (124, 'torch.ones', 'torch.ones', (['SEQ_LEN', 'batch_size'], {}), False, 'import torch\n'), (147, 'torch.zeros_like', 'torch.zeros_like', (['all_step_state'], {}), False, 'import torch\n'), (233, 'reagent.training.world_model.seq2reward_trainer.get_Q', 'get_Q', (['seq2reward_trainer.seq2reward_network', 'initial_state', 'seq2reward_trainer.all_permut'], {}), False, 'from reagent.training.world_model.seq2reward_trainer import get_Q, Seq2RewardTrainer\n'), (251, 'torch.tensor', 'torch.tensor', (['q_values'], {}), False, 'import torch\n'), (252, 'torch.tensor', 'torch.tensor', (['action_distribution'], {}), False, 'import torch\n'), (322, 'torch.tensor', 'torch.tensor', (['q_values'], {}), False, 'import torch\n'), (323, 'torch.tensor', 'torch.tensor', (['action_distribution'], {}), False, 'import torch\n'), (351, 'reagent.preprocessing.preprocessor.Preprocessor', 'Preprocessor', (['state_normalization_parameters', '(False)'], {}), False, 'from reagent.preprocessing.preprocessor import Preprocessor\n'), (384, 'torch.all', 'torch.all', (['(expected_q_values == q_values)'], {}), False, 'import torch\n'), (391, 'reagent.training.utils.gen_permutations', 'gen_permutations', (['MULTI_STEPS', 'NUM_ACTION'], {}), False, 'from reagent.training.utils import gen_permutations\n'), (393, 'torch.zeros', 'torch.zeros', (['BATCH_SIZE', 'STATE_DIM'], {}), False, 'import torch\n'), (394, 'reagent.training.world_model.seq2reward_trainer.get_Q', 'get_Q', (['seq2reward_network', 'state', 'all_permut'], {}), False, 'from reagent.training.world_model.seq2reward_trainer import get_Q, Seq2RewardTrainer\n'), (395, 'torch.tensor', 'torch.tensor', (['[[11.0, 111.0], [11.0, 111.0]]'], {}), False, 'import torch\n'), (397, 'torch.all', 'torch.all', (['(expected_q_values == q_values)'], {}), False, 'import torch\n'), (402, 'torch.tensor', 'torch.tensor', (['[[0], [1], [2], [3], [4], [5]]'], {}), False, 'import torch\n'), (408, 'torch.tensor', 'torch.tensor', (['[[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0\n    ], [1, 1, 1]]'], {}), False, 'import torch\n'), (424, 'reagent.training.utils.gen_permutations', 'gen_permutations', (['SEQ_LEN', 'NUM_ACTION'], {}), False, 'from reagent.training.utils import gen_permutations\n'), (427, 'torch.all', 'torch.all', (['(outcome == expected_outcome)'], {}), False, 'import torch\n'), (432, 'numpy.random.seed', 'np.random.seed', (['SEED'], {}), True, 'import numpy as np\n'), (433, 'random.seed', 'random.seed', (['SEED'], {}), False, 'import random\n'), (434, 'torch.manual_seed', 'torch.manual_seed', (['SEED'], {}), False, 'import torch\n'), (467, 'torch.all', 'torch.all', (['(eval_q_values - compress_eval_q_values < 1e-05)'], {}), False, 'import torch\n'), (468, 'torch.all', 'torch.all', (['(eval_action_distribution - compress_eval_action_distribution < 1e-05)'], {}), False, 'import torch\n'), (154, 'torch.zeros_like', 'torch.zeros_like', (['all_step_reward'], {}), False, 'import torch\n'), (158, 'torch.zeros_like', 'torch.zeros_like', (['all_step_state'], {}), False, 'import torch\n'), (162, 'torch.zeros_like', 'torch.zeros_like', (['all_step_action'], {}), False, 'import torch\n'), (183, 'torch.zeros_like', 'torch.zeros_like', (['batch_reward'], {}), False, 'import torch\n'), (184, 'torch.zeros_like', 'torch.zeros_like', (['batch_action'], {}), False, 'import torch\n'), (185, 'torch.zeros_like', 'torch.zeros_like', (['batch_state'], {}), False, 'import torch\n'), (346, 'reagent.core.parameters.NormalizationParameters', 'NormalizationParameters', ([], {'feature_type': 'DO_NOT_PREPROCESS', 'mean': '(0.0)', 'stddev': '(1.0)'}), False, 'from reagent.core.parameters import NormalizationData, NormalizationParameters, ProblemDomain, Seq2RewardTrainerParameters\n'), (355, 'reagent.prediction.predictor_wrapper.Seq2RewardPlanShortSeqWithPreprocessor', 'Seq2RewardPlanShortSeqWithPreprocessor', (['model', 'step_prediction_model', 'state_preprocessor', 'seq_len', 'action_dim'], {}), False, 'from reagent.prediction.predictor_wrapper import Seq2RewardWithPreprocessor, Seq2RewardPlanShortSeqWithPreprocessor, FAKE_STATE_ID_LIST_FEATURES, FAKE_STATE_ID_SCORE_LIST_FEATURES\n'), (363, 'reagent.prediction.predictor_wrapper.Seq2RewardWithPreprocessor', 'Seq2RewardWithPreprocessor', (['model', 'state_preprocessor', 'seq_len', 'action_dim'], {}), False, 'from reagent.prediction.predictor_wrapper import Seq2RewardWithPreprocessor, Seq2RewardPlanShortSeqWithPreprocessor, FAKE_STATE_ID_LIST_FEATURES, FAKE_STATE_ID_SCORE_LIST_FEATURES\n'), (381, 'torch.tensor', 'torch.tensor', (['[[4.0, 41.0]]'], {}), False, 'import torch\n'), (383, 'torch.tensor', 'torch.tensor', (['[[11.0, 111.0]]'], {}), False, 'import torch\n'), (58, 'torch.ones', 'torch.ones', (['batch_size', 'self.look_ahead_steps'], {}), False, 'import torch\n'), (90, 'torch.argmax', 'torch.argmax', (['action'], {'dim': '(2)'}), False, 'import torch\n'), (117, 'torch.ones', 'torch.ones', (['batch_size'], {'dtype': 'torch.int64'}), False, 'import torch\n'), (278, 'reagent.core.parameters.NormalizationParameters', 'NormalizationParameters', ([], {'feature_type': 'DO_NOT_PREPROCESS'}), False, 'from reagent.core.parameters import NormalizationData, NormalizationParameters, ProblemDomain, Seq2RewardTrainerParameters\n'), (279, 'reagent.core.parameters.NormalizationParameters', 'NormalizationParameters', ([], {'feature_type': 'DO_NOT_PREPROCESS'}), False, 'from reagent.core.parameters import NormalizationData, NormalizationParameters, ProblemDomain, Seq2RewardTrainerParameters\n'), (125, 'torch.arange', 'torch.arange', (['SEQ_LEN', '(0)', '(-1)'], {}), False, 'import torch\n'), (148, 'torch.arange', 'torch.arange', (['SEQ_LEN'], {}), False, 'import torch\n'), (171, 'reagent.core.types.FeatureData', 'rlt.FeatureData', ([], {'float_features': 'batch_action'}), True, 'from reagent.core import types as rlt\n'), (172, 'reagent.core.types.FeatureData', 'rlt.FeatureData', ([], {'float_features': 'batch_state'}), True, 'from reagent.core import types as rlt\n'), (127, 'torch.ones', 'torch.ones', (['SEQ_LEN', 'SEQ_LEN'], {}), False, 'import torch\n'), (174, 'torch.zeros_like', 'torch.zeros_like', (['batch_state'], {}), False, 'import torch\n')]