
"""simple docstring"""
import inspect
import unittest

from transformers import DecisionTransformerConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import DecisionTransformerModel
      from transformers.models.decision_transformer.modeling_decision_transformer import (
          DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
      )
class 			UpperCamelCase      :
     def __init__(     self  ,     snake_case__  ,     snake_case__=13  ,     snake_case__=7  ,     snake_case__=6  ,     snake_case__=17  ,     snake_case__=23  ,     snake_case__=11  ,     snake_case__=True  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Any							    = parent
       _SCREAMING_SNAKE_CASE   :				Tuple							    = batch_size
       _SCREAMING_SNAKE_CASE   :				int							    = seq_length
       _SCREAMING_SNAKE_CASE   :				List[str]							    = act_dim
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = state_dim
       _SCREAMING_SNAKE_CASE   :				int							    = hidden_size
       _SCREAMING_SNAKE_CASE   :				int							    = max_length
       _SCREAMING_SNAKE_CASE   :				str							    = is_training
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = floats_tensor((self.batch_size, self.seq_length, self.state_dim)				)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = floats_tensor((self.batch_size, self.seq_length, self.act_dim)				)
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = floats_tensor((self.batch_size, self.seq_length, 1)				)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = floats_tensor((self.batch_size, self.seq_length, 1)				)
       _SCREAMING_SNAKE_CASE   :				str							    = ids_tensor((self.batch_size, self.seq_length)  ,     vocab_size=1000				)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = random_attention_mask((self.batch_size, self.seq_length)				)

       _SCREAMING_SNAKE_CASE   :				str							    = self.get_config()

       return (
           config,
           states,
           actions,
           rewards,
           returns_to_go,
           timesteps,
           attention_mask,
       )
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       return DecisionTransformerConfig(
           batch_size=self.batch_size  ,     seq_length=self.seq_length  ,     act_dim=self.act_dim  ,     state_dim=self.state_dim  ,     hidden_size=self.hidden_size  ,     max_length=self.max_length  ,     )
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Any							    = DecisionTransformerModel(config=snake_case__				)
       model.to(snake_case__				)
       model.eval()
       _SCREAMING_SNAKE_CASE   :				int							    = model(snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__				)

       self.parent.assertEqual(result.state_preds.shape  ,     states.shape				)
       self.parent.assertEqual(result.action_preds.shape  ,     actions.shape				)
       self.parent.assertEqual(result.return_preds.shape  ,     returns_to_go.shape				)
       self.parent.assertEqual(
           result.last_hidden_state.shape  ,     (self.batch_size, self.seq_length * 3, self.hidden_size)				)  # seq length *3 as there are 3 modelities: states, returns and actions




     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.prepare_config_and_inputs()
       (
           (
           _SCREAMING_SNAKE_CASE
       )    ,		(
           _SCREAMING_SNAKE_CASE
       )    ,		(
           _SCREAMING_SNAKE_CASE
       )    ,		(
           _SCREAMING_SNAKE_CASE
       )    ,		(
           _SCREAMING_SNAKE_CASE
       )    ,		(
           _SCREAMING_SNAKE_CASE
       )    ,		(
           _SCREAMING_SNAKE_CASE
       )    ,		
       )   :				Any							    = config_and_inputs
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = {
           "states": states,
           "actions": actions,
           "rewards": rewards,
           "returns_to_go": returns_to_go,
           "timesteps": timesteps,
           "attention_mask": attention_mask,
       }
       return config, inputs_dict
@require_torch
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE      ,       __SCREAMING_SNAKE_CASE      ,       __SCREAMING_SNAKE_CASE      ,       unittest.TestCase  ):
     A__      =			(DecisionTransformerModel,) if is_torch_available() else ()
     A__      =			()
     A__      =			{"""feature-extraction""": DecisionTransformerModel} if is_torch_available() else {}

     # Ignoring of a failing test from GenerationTesterMixin, as the model does not use inputs_ids
     A__      =			False

     # Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
     A__      =			False
     A__      =			False
     A__      =			False
     A__      =			False
     A__      =			False
     A__      =			False
     A__      =			False
     A__      =			False
     A__      =			False
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Dict							    = DecisionTransformerModelTester(self				)
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = ConfigTester(self  ,     config_class=snake_case__  ,     hidden_size=37				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       self.config_tester.run_common_tests()
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[str]							    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*snake_case__				)
     @slow
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         _SCREAMING_SNAKE_CASE   :				Tuple							    = DecisionTransformerModel.from_pretrained(snake_case__				)
         self.assertIsNotNone(snake_case__				)




     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				str							    = self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
         _SCREAMING_SNAKE_CASE   :				List[Any]							    = model_class(snake_case__				)
         _SCREAMING_SNAKE_CASE   :				int							    = inspect.signature(model.forward				)
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         _SCREAMING_SNAKE_CASE   :				str							    = [*signature.parameters.keys()]

         _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = [
             "states",
             "actions",
             "rewards",
             "returns_to_go",
             "timesteps",
             "attention_mask",
         ]

         self.assertListEqual(arg_names[: len(snake_case__				)]  ,     snake_case__				)
@require_torch
class 			UpperCamelCase      (							unittest.TestCase  ):
     @slow
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				int							    = 2  # number of steps of autoregressive prediction we will perform
       _SCREAMING_SNAKE_CASE   :				Dict							    = 10  # defined by the RL environment, may be normalized
       _SCREAMING_SNAKE_CASE   :				Tuple							    = DecisionTransformerModel.from_pretrained("edbeeching/decision-transformer-gym-hopper-expert"				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = model.to(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				str							    = model.config
       torch.manual_seed(0				)
       _SCREAMING_SNAKE_CASE   :				str							    = torch.randn(1  ,     1  ,     config.state_dim				).to(device=snake_case__  ,     dtype=torch.floataa				)  # env.reset()

       _SCREAMING_SNAKE_CASE   :				Tuple							    = torch.tensor(
           [[0.242_793, -0.28_693_074, 0.8_742_613], [0.67_815_274, -0.08_101_085, -0.12_952_147]]  ,     device=snake_case__				)

       _SCREAMING_SNAKE_CASE   :				List[Any]							    = torch.tensor(snake_case__  ,     device=snake_case__  ,     dtype=torch.floataa				).reshape(1  ,     1  ,     1				)
       _SCREAMING_SNAKE_CASE   :				int							    = state
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = torch.zeros(1  ,     0  ,     config.act_dim  ,     device=snake_case__  ,     dtype=torch.floataa				)
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = torch.zeros(1  ,     0  ,     device=snake_case__  ,     dtype=torch.floataa				)
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = torch.tensor(0  ,     device=snake_case__  ,     dtype=torch.long				).reshape(1  ,     1				)

       for step in range(snake_case__				):
         _SCREAMING_SNAKE_CASE   :				int							    = torch.cat([actions, torch.zeros(1  ,     1  ,     config.act_dim  ,     device=snake_case__				)]  ,     dim=1				)
         _SCREAMING_SNAKE_CASE   :				Tuple							    = torch.cat([rewards, torch.zeros(1  ,     1  ,     device=snake_case__				)]  ,     dim=1				)

         _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = torch.ones(1  ,     states.shape[1]				).to(dtype=torch.long  ,     device=states.device				)

         with torch.no_grad():
           _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				List[Any]							    = model(
               states=snake_case__  ,     actions=snake_case__  ,     rewards=snake_case__  ,     returns_to_go=snake_case__  ,     timesteps=snake_case__  ,     attention_mask=snake_case__  ,     return_dict=snake_case__  ,     )

         self.assertEqual(action_pred.shape  ,     actions.shape				)
         self.assertTrue(torch.allclose(action_pred[0, -1]  ,     expected_outputs[step]  ,     atol=1E-4				)				)
         _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				Tuple							    = (  # env.step(action)
             torch.randn(1  ,     1  ,     config.state_dim				).to(device=snake_case__  ,     dtype=torch.floataa				),
             1.0,
             False,
             {},
         )

         _SCREAMING_SNAKE_CASE   :				Any							    = action_pred[0, -1]
         _SCREAMING_SNAKE_CASE   :				Tuple							    = torch.cat([states, state]  ,     dim=1				)
         _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = returns_to_go[0, -1] - reward
         _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = torch.cat([returns_to_go, pred_return.reshape(1  ,     1  ,     1				)]  ,     dim=1				)
         _SCREAMING_SNAKE_CASE   :				Tuple							    = torch.cat(
             [timesteps, torch.ones((1, 1)  ,     device=snake_case__  ,     dtype=torch.long				) * (step + 1)]  ,     dim=1				)





"""simple docstring"""
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, Encoder, VectorQuantizer
@dataclass
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE  ):
     A__      =			42
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE      ,       __SCREAMING_SNAKE_CASE  ):
     @register_to_config
     def __init__(     self  ,     snake_case__ = 3  ,     snake_case__ = 3  ,     snake_case__ = ("DownEncoderBlock2D",)  ,     snake_case__ = ("UpDecoderBlock2D",)  ,     snake_case__ = (64,)  ,     snake_case__ = 1  ,     snake_case__ = "silu"  ,     snake_case__ = 3  ,     snake_case__ = 32  ,     snake_case__ = 256  ,     snake_case__ = 32  ,     snake_case__ = None  ,     snake_case__ = 0.18_215  ,     snake_case__ = "group"  ,     ):




       """simple docstring"""


       super().__init__()

       # pass init params to Encoder
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = Encoder(
           in_channels=snake_case__  ,     out_channels=snake_case__  ,     down_block_types=snake_case__  ,     block_out_channels=snake_case__  ,     layers_per_block=snake_case__  ,     act_fn=snake_case__  ,     norm_num_groups=snake_case__  ,     double_z=snake_case__  ,     )

       _SCREAMING_SNAKE_CASE   :				Dict							    = vq_embed_dim if vq_embed_dim is not None else latent_channels

       _SCREAMING_SNAKE_CASE   :				Tuple							    = nn.Convad(snake_case__  ,     snake_case__  ,     1				)
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = VectorQuantizer(snake_case__  ,     snake_case__  ,     beta=0.25  ,     remap=snake_case__  ,     sane_index_shape=snake_case__				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = nn.Convad(snake_case__  ,     snake_case__  ,     1				)

       # pass init params to Decoder
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = Decoder(
           in_channels=snake_case__  ,     out_channels=snake_case__  ,     up_block_types=snake_case__  ,     block_out_channels=snake_case__  ,     layers_per_block=snake_case__  ,     act_fn=snake_case__  ,     norm_num_groups=snake_case__  ,     norm_type=snake_case__  ,     )
     @apply_forward_hook
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__ = True				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[str]							    = self.encoder(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				Dict							    = self.quant_conv(snake_case__				)

       if not return_dict:
         return (h,)

       return VQEncoderOutput(latents=snake_case__				)
     @apply_forward_hook
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__ = False  ,     snake_case__ = True				):




       """simple docstring"""


       if not force_not_quantize:
         _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				int							    = self.quantize(snake_case__				)
       else:
         _SCREAMING_SNAKE_CASE   :				int							    = h
       _SCREAMING_SNAKE_CASE   :				Dict							    = self.post_quant_conv(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				int							    = self.decoder(snake_case__  ,     quant if self.config.norm_type == "spatial" else None				)

       if not return_dict:
         return (dec,)

       return DecoderOutput(sample=snake_case__				)




     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__ = True				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Tuple							    = sample
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.encode(snake_case__				).latents
       _SCREAMING_SNAKE_CASE   :				Any							    = self.decode(snake_case__				).sample

       if not return_dict:
         return (dec,)

       return DecoderOutput(sample=snake_case__				)



