#!/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', False, 'import logging\n'), (105, 'reagent.gym.envs.Gym', 'Gym', False, 'from reagent.gym.envs import Gym\n'), (106, 'reagent.gym.utils.create_df_from_replay_buffer', 'create_df_from_replay_buffer', False, 'from reagent.gym.utils import create_df_from_replay_buffer\n'), (137, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (138, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (139, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (189, 'torch.utils.data.DataLoader', 'DataLoader', False, 'from torch.utils.data import DataLoader\n'), (192, 'torch.utils.data.DataLoader', 'DataLoader', False, 'from torch.utils.data import DataLoader\n'), (202, 'reagent.models.seq2reward_model.Seq2RewardNetwork', 'Seq2RewardNetwork', False, 'from reagent.models.seq2reward_model import Seq2RewardNetwork\n'), (209, 'reagent.core.parameters.Seq2RewardTrainerParameters', 'Seq2RewardTrainerParameters', False, 'from reagent.core.parameters import NormalizationData, NormalizationParameters, ProblemDomain, Seq2RewardTrainerParameters\n'), (217, 'reagent.training.world_model.seq2reward_trainer.Seq2RewardTrainer', 'Seq2RewardTrainer', False, 'from reagent.training.world_model.seq2reward_trainer import get_Q, Seq2RewardTrainer\n'), (221, 'pytorch_lightning.seed_everything', 'pl.seed_everything', True, 'import pytorch_lightning as pl\n'), (222, 'pytorch_lightning.Trainer', 'pl.Trainer', True, 'import pytorch_lightning as pl\n'), (231, 'torch.Tensor', 'torch.Tensor', False, 'import torch\n'), (241, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (242, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (275, 'reagent.net_builder.value.fully_connected.FullyConnected', 'FullyConnected', False, 'from reagent.net_builder.value.fully_connected import FullyConnected\n'), (287, 'reagent.core.parameters.Seq2RewardTrainerParameters', 'Seq2RewardTrainerParameters', False, 'from reagent.core.parameters import NormalizationData, NormalizationParameters, ProblemDomain, Seq2RewardTrainerParameters\n'), (296, 'reagent.training.world_model.compress_model_trainer.CompressModelTrainer', 'CompressModelTrainer', False, 'from reagent.training.world_model.compress_model_trainer import CompressModelTrainer\n'), (302, 'pytorch_lightning.seed_everything', 'pl.seed_everything', True, 'import pytorch_lightning as pl\n'), (303, 'pytorch_lightning.Trainer', 'pl.Trainer', True, 'import pytorch_lightning as pl\n'), (312, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (313, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (429, 'parameterized.parameterized.expand', 'parameterized.expand', False, 'from parameterized import parameterized\n'), (430, 'unittest.skipIf', 'unittest.skipIf', False, 'import unittest\n'), (95, 'reagent.core.types.Seq2RewardOutput', 'rlt.Seq2RewardOutput', True, 'from reagent.core import types as rlt\n'), (116, 'torch.ones', 'torch.ones', False, 'import torch\n'), (124, 'torch.ones', 'torch.ones', False, 'import torch\n'), (147, 'torch.zeros_like', 'torch.zeros_like', False, 'import torch\n'), (233, 'reagent.training.world_model.seq2reward_trainer.get_Q', 'get_Q', False, 'from reagent.training.world_model.seq2reward_trainer import get_Q, Seq2RewardTrainer\n'), (251, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (252, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (322, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (323, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (351, 'reagent.preprocessing.preprocessor.Preprocessor', 'Preprocessor', False, 'from reagent.preprocessing.preprocessor import Preprocessor\n'), (384, 'torch.all', 'torch.all', False, 'import torch\n'), (391, 'reagent.training.utils.gen_permutations', 'gen_permutations', False, 'from reagent.training.utils import gen_permutations\n'), (393, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (394, 'reagent.training.world_model.seq2reward_trainer.get_Q', 'get_Q', False, 'from reagent.training.world_model.seq2reward_trainer import get_Q, Seq2RewardTrainer\n'), (395, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (397, 'torch.all', 'torch.all', False, 'import torch\n'), (402, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (408, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (424, 'reagent.training.utils.gen_permutations', 'gen_permutations', False, 'from reagent.training.utils import gen_permutations\n'), (427, 'torch.all', 'torch.all', False, 'import torch\n'), (432, 'numpy.random.seed', 'np.random.seed', True, 'import numpy as np\n'), (433, 'random.seed', 'random.seed', False, 'import random\n'), (434, 'torch.manual_seed', 'torch.manual_seed', False, 'import torch\n'), (467, 'torch.all', 'torch.all', False, 'import torch\n'), (468, 'torch.all', 'torch.all', False, 'import torch\n'), (154, 'torch.zeros_like', 'torch.zeros_like', False, 'import torch\n'), (158, 'torch.zeros_like', 'torch.zeros_like', False, 'import torch\n'), (162, 'torch.zeros_like', 'torch.zeros_like', False, 'import torch\n'), (183, 'torch.zeros_like', 'torch.zeros_like', False, 'import torch\n'), (184, 'torch.zeros_like', 'torch.zeros_like', False, 'import torch\n'), (185, 'torch.zeros_like', 'torch.zeros_like', False, 'import torch\n'), (346, 'reagent.core.parameters.NormalizationParameters', 'NormalizationParameters', False, 'from reagent.core.parameters import NormalizationData, NormalizationParameters, ProblemDomain, Seq2RewardTrainerParameters\n'), (355, 'reagent.prediction.predictor_wrapper.Seq2RewardPlanShortSeqWithPreprocessor', 'Seq2RewardPlanShortSeqWithPreprocessor', 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', 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', False, 'import torch\n'), (383, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (58, 'torch.ones', 'torch.ones', False, 'import torch\n'), (90, 'torch.argmax', 'torch.argmax', False, 'import torch\n'), (117, 'torch.ones', 'torch.ones', False, 'import torch\n'), (278, 'reagent.core.parameters.NormalizationParameters', 'NormalizationParameters', False, 'from reagent.core.parameters import NormalizationData, NormalizationParameters, ProblemDomain, Seq2RewardTrainerParameters\n'), (279, 'reagent.core.parameters.NormalizationParameters', 'NormalizationParameters', False, 'from reagent.core.parameters import NormalizationData, NormalizationParameters, ProblemDomain, Seq2RewardTrainerParameters\n'), (125, 'torch.arange', 'torch.arange', False, 'import torch\n'), (148, 'torch.arange', 'torch.arange', False, 'import torch\n'), (171, 'reagent.core.types.FeatureData', 'rlt.FeatureData', True, 'from reagent.core import types as rlt\n'), (172, 'reagent.core.types.FeatureData', 'rlt.FeatureData', True, 'from reagent.core import types as rlt\n'), (127, 'torch.ones', 'torch.ones', False, 'import torch\n'), (174, 'torch.zeros_like', 'torch.zeros_like', False, 'import torch\n')]