'''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 					__SCREAMING_SNAKE_CASE						:







		'''simple docstring'''






		def __init__(			self				,		snake_case_				,		snake_case_=1_3				,		snake_case_=7				,		snake_case_=6				,		snake_case_=1_7				,		snake_case_=2_3				,		snake_case_=1_1				,		snake_case_=True				,		):




							'''simple docstring'''

							UpperCAmelCase_							:       Optional[int]       =				parent
							UpperCAmelCase_							:       List[str]       =				batch_size
							UpperCAmelCase_							:       Any       =				seq_length
							UpperCAmelCase_							:       Any       =				act_dim
							UpperCAmelCase_							:       Union[str, Any]       =				state_dim
							UpperCAmelCase_							:       List[Any]       =				hidden_size
							UpperCAmelCase_							:       Union[str, Any]       =				max_length
							UpperCAmelCase_							:       Any       =				is_training


		def _UpperCamelCase					(			self						):




							'''simple docstring'''

							UpperCAmelCase_							:       Tuple       =				floats_tensor((self.batch_size, self.seq_length, self.state_dim)						)
							UpperCAmelCase_							:       Dict       =				floats_tensor((self.batch_size, self.seq_length, self.act_dim)						)
							UpperCAmelCase_							:       Union[str, Any]       =				floats_tensor((self.batch_size, self.seq_length, 1)						)
							UpperCAmelCase_							:       Any       =				floats_tensor((self.batch_size, self.seq_length, 1)						)
							UpperCAmelCase_							:       int       =				ids_tensor((self.batch_size, self.seq_length)				,		vocab_size=1_0_0_0						)
							UpperCAmelCase_							:       int       =				random_attention_mask((self.batch_size, self.seq_length)						)

							UpperCAmelCase_							:       Optional[int]       =				self.get_config()

							return (
							    config,
							    states,
							    actions,
							    rewards,
							    returns_to_go,
							    timesteps,
							    attention_mask,
							)


		def _UpperCamelCase					(			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 _UpperCamelCase					(			self				,		snake_case_				,		snake_case_				,		snake_case_				,		snake_case_				,		snake_case_				,		snake_case_				,		snake_case_				,		):




							'''simple docstring'''

							UpperCAmelCase_							:       int       =				DecisionTransformerModel(config=snake_case_						)
							model.to(snake_case_						)
							model.eval()
							UpperCAmelCase_							:       List[Any]       =				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 _UpperCamelCase					(			self						):




							'''simple docstring'''

							UpperCAmelCase_							:       List[str]       =				self.prepare_config_and_inputs()
							(
							    UpperCAmelCase_      
							)							:       Optional[Any]       =				config_and_inputs
							UpperCAmelCase_							:       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 					__SCREAMING_SNAKE_CASE						(	lowerCamelCase_						,     lowerCamelCase_						,     lowerCamelCase_						,     unittest.TestCase	):







		'''simple docstring'''






		lowerCamelCase_		:List[Any]						   =		(DecisionTransformerModel,) if is_torch_available() else ()
		lowerCamelCase_		:Union[str, Any]						   =		()
		lowerCamelCase_		:Tuple						   =		{'''feature-extraction''': DecisionTransformerModel} if is_torch_available() else {}

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

		# Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
		lowerCamelCase_		:Union[str, Any]						   =		False
		lowerCamelCase_		:List[Any]						   =		False
		lowerCamelCase_		:Dict						   =		False
		lowerCamelCase_		:List[Any]						   =		False
		lowerCamelCase_		:List[Any]						   =		False
		lowerCamelCase_		:int						   =		False
		lowerCamelCase_		:Dict						   =		False
		lowerCamelCase_		:List[Any]						   =		False
		lowerCamelCase_		:List[Any]						   =		False


		def _UpperCamelCase					(			self						):




							'''simple docstring'''

							UpperCAmelCase_							:       str       =				DecisionTransformerModelTester(self						)
							UpperCAmelCase_							:       Optional[Any]       =				ConfigTester(self				,		config_class=snake_case_				,		hidden_size=3_7						)


		def _UpperCamelCase					(			self						):




							'''simple docstring'''

							self.config_tester.run_common_tests()


		def _UpperCamelCase					(			self						):




							'''simple docstring'''

							UpperCAmelCase_							:       Optional[Any]       =				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*snake_case_						)


		@slow
		def _UpperCamelCase					(			self						):




							'''simple docstring'''

							for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												UpperCAmelCase_							:       int       =				DecisionTransformerModel.from_pretrained(snake_case_						)
												self.assertIsNotNone(snake_case_						)






		def _UpperCamelCase					(			self						):




							'''simple docstring'''

							UpperCAmelCase_							:       Dict       =				self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
												UpperCAmelCase_							:       Optional[Any]       =				model_class(snake_case_						)
												UpperCAmelCase_							:       List[Any]       =				inspect.signature(model.forward						)
												# signature.parameters is an OrderedDict => so arg_names order is deterministic
												UpperCAmelCase_							:       Tuple       =				[*signature.parameters.keys()]

												UpperCAmelCase_							:       List[str]       =				[
												    'states',
												    'actions',
												    'rewards',
												    'returns_to_go',
												    'timesteps',
												    'attention_mask',
												]

												self.assertListEqual(arg_names[: len(snake_case_						)]				,		snake_case_						)







@require_torch
class 					__SCREAMING_SNAKE_CASE						(	unittest.TestCase	):







		'''simple docstring'''






		@slow
		def _UpperCamelCase					(			self						):




							'''simple docstring'''

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

							UpperCAmelCase_							:       Optional[int]       =				torch.tensor(
							    [[0.24_27_93, -0.28_69_30_74, 0.8_74_26_13], [0.67_81_52_74, -0.08_10_10_85, -0.12_95_21_47]]				,		device=snake_case_						)

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

							for step in range(snake_case_						):
												UpperCAmelCase_							:       Optional[int]       =				torch.cat([actions, torch.zeros(1				,		1				,		config.act_dim				,		device=snake_case_						)]				,		dim=1						)
												UpperCAmelCase_							:       str       =				torch.cat([rewards, torch.zeros(1				,		1				,		device=snake_case_						)]				,		dim=1						)

												UpperCAmelCase_							:       List[Any]       =				torch.ones(1				,		states.shape[1]						).to(dtype=torch.long				,		device=states.device						)

												with torch.no_grad():
																	UpperCAmelCase_							:       List[str]       =				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						)						)
												UpperCAmelCase_							:       Dict       =				(  # env.step(action)
												    torch.randn(1				,		1				,		config.state_dim						).to(device=snake_case_				,		dtype=torch.floataa						),
												    1.0,
												    False,
												    {},
												)

												UpperCAmelCase_							:       Any       =				action_pred[0, -1]
												UpperCAmelCase_							:       str       =				torch.cat([states, state]				,		dim=1						)
												UpperCAmelCase_							:       str       =				returns_to_go[0, -1] - reward
												UpperCAmelCase_							:       List[Any]       =				torch.cat([returns_to_go, pred_return.reshape(1				,		1				,		1						)]				,		dim=1						)
												UpperCAmelCase_							:       int       =				torch.cat(
												    [timesteps, torch.ones((1, 1)				,		device=snake_case_				,		dtype=torch.long						) * (step + 1)]				,		dim=1						)







'''simple docstring'''


import unittest

from transformers.utils.backbone_utils import (
    BackboneMixin,
    get_aligned_output_features_output_indices,
    verify_out_features_out_indices,
)






class 					__SCREAMING_SNAKE_CASE						(	unittest.TestCase	):







  '''simple docstring'''






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       int       =				['a', 'b', 'c']

       # Defaults to last layer if both are None
       UpperCAmelCase_       ,      UpperCAmelCase_							:       List[str]       =				get_aligned_output_features_output_indices(snake_case_				,		snake_case_				,		snake_case_						)
       self.assertEqual(snake_case_				,		['c']						)
       self.assertEqual(snake_case_				,		[2]						)

       # Out indices set to match out features
       UpperCAmelCase_       ,      UpperCAmelCase_							:       int       =				get_aligned_output_features_output_indices(['a', 'c']				,		snake_case_				,		snake_case_						)
       self.assertEqual(snake_case_				,		['a', 'c']						)
       self.assertEqual(snake_case_				,		[0, 2]						)

       # Out features set to match out indices
       UpperCAmelCase_       ,      UpperCAmelCase_							:       List[Any]       =				get_aligned_output_features_output_indices(snake_case_				,		[0, 2]				,		snake_case_						)
       self.assertEqual(snake_case_				,		['a', 'c']						)
       self.assertEqual(snake_case_				,		[0, 2]						)

       # Out features selected from negative indices
       UpperCAmelCase_       ,      UpperCAmelCase_							:       int       =				get_aligned_output_features_output_indices(snake_case_				,		[-3, -1]				,		snake_case_						)
       self.assertEqual(snake_case_				,		['a', 'c']						)
       self.assertEqual(snake_case_				,		[-3, -1]						)


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       with self.assertRaises(snake_case_						):
            verify_out_features_out_indices(['a', 'b']				,		(0, 1)				,		snake_case_						)

       # Out features must be a list
       with self.assertRaises(snake_case_						):
            verify_out_features_out_indices(('a', 'b')				,		(0, 1)				,		['a', 'b']						)

       # Out features must be a subset of stage names
       with self.assertRaises(snake_case_						):
            verify_out_features_out_indices(['a', 'b']				,		(0, 1)				,		['a']						)

       # Out indices must be a list or tuple
       with self.assertRaises(snake_case_						):
            verify_out_features_out_indices(snake_case_				,		0				,		['a', 'b']						)

       # Out indices must be a subset of stage names
       with self.assertRaises(snake_case_						):
            verify_out_features_out_indices(snake_case_				,		(0, 1)				,		['a']						)

       # Out features and out indices must be the same length
       with self.assertRaises(snake_case_						):
            verify_out_features_out_indices(['a', 'b']				,		(0,)				,		['a', 'b', 'c']						)

       # Out features should match out indices
       with self.assertRaises(snake_case_						):
            verify_out_features_out_indices(['a', 'b']				,		(0, 2)				,		['a', 'b', 'c']						)

       # Out features and out indices should be in order
       with self.assertRaises(snake_case_						):
            verify_out_features_out_indices(['b', 'a']				,		(0, 1)				,		['a', 'b']						)

       # Check passes with valid inputs
       verify_out_features_out_indices(['a', 'b', 'd']				,		(0, 1, -1)				,		['a', 'b', 'c', 'd']						)






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       List[str]       =				BackboneMixin()

       UpperCAmelCase_							:       Any       =				['a', 'b', 'c']
       UpperCAmelCase_							:       str       =				['a', 'c']
       UpperCAmelCase_							:       str       =				[0, 2]

       # Check that the output features and indices are set correctly
       self.assertEqual(backbone.out_features				,		['a', 'c']						)
       self.assertEqual(backbone.out_indices				,		[0, 2]						)

       # Check out features and indices are updated correctly
       UpperCAmelCase_							:       str       =				['a', 'b']
       self.assertEqual(backbone.out_features				,		['a', 'b']						)
       self.assertEqual(backbone.out_indices				,		[0, 1]						)

       UpperCAmelCase_							:       Optional[int]       =				[-3, -1]
       self.assertEqual(backbone.out_features				,		['a', 'c']						)
       self.assertEqual(backbone.out_indices				,		[-3, -1]						)






