codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch decisiontransformer model import inspect import unittest from transformers import decisiontransformerconfig istorchavailable from transformers testingutils import requiretorch slow torchdevice from generation testutils import generationtestermixin from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin floatstensor idstensor randomattentionmask from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from transformers import decisiontransformermodel from transformers models decisiontransformer modelingdecisiontransformer import decisiontransformerpretrainedmodelarchivelist class decisiontransformermodeltester def init self parent batchsize13 seqlength7 actdim6 statedim17 hiddensize23 maxlength11 istrainingtrue self parent parent self batchsize batchsize self seqlength seqlength self actdim actdim self statedim statedim self hiddensize hiddensize self maxlength maxlength self istraining istraining def prepareconfigandinputsself states floatstensorself batchsize self seqlength self statedim actions floatstensorself batchsize self seqlength self actdim rewards floatstensorself batchsize self seqlength 1 returnstogo floatstensorself batchsize self seqlength 1 timesteps idstensorself batchsize self seqlength vocabsize1000 attentionmask randomattentionmaskself batchsize self seqlength config self getconfig return config states actions rewards returnstogo timesteps attentionmask def getconfigself return decisiontransformerconfig batchsizeself batchsize seqlengthself seqlength actdimself actdim statedimself statedim hiddensizeself hiddensize maxlengthself maxlength def createandcheckmodel self config states actions rewards returnstogo timesteps attentionmask model decisiontransformermodelconfigconfig model totorchdevice model eval result modelstates actions rewards returnstogo timesteps attentionmask self parent assertequalresult statepreds shape states shape self parent assertequalresult actionpreds shape actions shape self parent assertequalresult returnpreds shape returnstogo shape self parent assertequal result lasthiddenstate shape self batchsize self seqlength 3 self hiddensize seq length 3 as there are 3 modelities states returns and actions def prepareconfigandinputsforcommonself configandinputs self prepareconfigandinputs config states actions rewards returnstogo timesteps attentionmask configandinputs inputsdict states states actions actions rewards rewards returnstogo returnstogo timesteps timesteps attentionmask attentionmask return config inputsdict requiretorch class decisiontransformermodeltestmodeltestermixin generationtestermixin pipelinetestermixin unittest testcase allmodelclasses decisiontransformermodel if istorchavailable else allgenerativemodelclasses pipelinemodelmapping featureextraction decisiontransformermodel if istorchavailable else ignoring of a failing test from generationtestermixin as the model does not use inputsids testgeneratewithoutinputids false ignoring of a failing tests from modeltestermixin as the model does not implement these features testpruning false testresizeembeddings false testheadmasking false testattentionoutputs false testhiddenstatesoutput false testinputsembeds false testmodelcommonattributes false testgradientcheckpointing false testtorchscript false def setupself self modeltester decisiontransformermodeltesterself self configtester configtesterself configclassdecisiontransformerconfig hiddensize37 def testconfigself self configtester runcommontests def testmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelconfigandinputs slow def testmodelfrompretrainedself for modelname in decisiontransformerpretrainedmodelarchivelist 1 model decisiontransformermodel frompretrainedmodelname self assertisnotnonemodel def testforwardsignatureself config self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig signature inspect signaturemodel forward signature parameters is an ordereddict so argnames order is deterministic argnames signature parameters keys expectedargnames states actions rewards returnstogo timesteps attentionmask self assertlistequalargnames lenexpectedargnames expectedargnames requiretorch class decisiontransformermodelintegrationtestunittest testcase slow def testautoregressivepredictionself numsteps 2 number of steps of autoregressive prediction we will perform targetreturn 10 defined by the rl environment may be normalized model decisiontransformermodel frompretrainededbeechingdecisiontransformergymhopperexpert model model totorchdevice config model config torch manualseed0 state torch randn1 1 config statedim todevicetorchdevice dtypetorch float32 env reset expectedoutputs torch tensor 0 242793 0 28693074 0 8742613 0 67815274 0 08101085 0 12952147 devicetorchdevice returnstogo torch tensortargetreturn devicetorchdevice dtypetorch float32 reshape1 1 1 states state actions torch zeros1 0 config actdim devicetorchdevice dtypetorch float32 rewards torch zeros1 0 devicetorchdevice dtypetorch float32 timesteps torch tensor0 devicetorchdevice dtypetorch long reshape1 1 for step in rangenumsteps actions torch catactions torch zeros1 1 config actdim devicetorchdevice dim1 rewards torch catrewards torch zeros1 1 devicetorchdevice dim1 attentionmask torch ones1 states shape1 todtypetorch long devicestates device with torch nograd actionpred model statesstates actionsactions rewardsrewards returnstogoreturnstogo timestepstimesteps attentionmaskattentionmask returndictfalse self assertequalactionpred shape actions shape self asserttruetorch allcloseactionpred0 1 expectedoutputsstep atol1e4 state reward env stepaction torch randn1 1 config statedim todevicetorchdevice dtypetorch float32 1 0 false actions1 actionpred0 1 states torch catstates state dim1 predreturn returnstogo0 1 reward returnstogo torch catreturnstogo predreturn reshape1 1 1 dim1 timesteps torch cat timesteps torch ones1 1 devicetorchdevice dtypetorch long step 1 dim1 coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch decisiontransformer model seq length 3 as there are 3 modelities states returns and actions ignoring of a failing test from generationtestermixin as the model does not use inputs_ids ignoring of a failing tests from modeltestermixin as the model does not implement these features signature parameters is an ordereddict so arg_names order is deterministic an integration test that performs autoregressive prediction of state action and return from a sequence of state actions and returns test is performed over two timesteps number of steps of autoregressive prediction we will perform defined by the rl environment may be normalized env reset env step action
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 DecisionTransformerModelTester:
    def __init__(
        self,
        parent,
        batch_size=13,
        seq_length=7,
        act_dim=6,
        state_dim=17,
        hidden_size=23,
        max_length=11,
        is_training=True,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.act_dim = act_dim
        self.state_dim = state_dim
        self.hidden_size = hidden_size
        self.max_length = max_length
        self.is_training = is_training

    def prepare_config_and_inputs(self):
        states = floats_tensor((self.batch_size, self.seq_length, self.state_dim))
        actions = floats_tensor((self.batch_size, self.seq_length, self.act_dim))
        rewards = floats_tensor((self.batch_size, self.seq_length, 1))
        returns_to_go = floats_tensor((self.batch_size, self.seq_length, 1))
        timesteps = ids_tensor((self.batch_size, self.seq_length), vocab_size=1000)
        attention_mask = random_attention_mask((self.batch_size, self.seq_length))

        config = self.get_config()

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

    def get_config(self):
        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 create_and_check_model(
        self,
        config,
        states,
        actions,
        rewards,
        returns_to_go,
        timesteps,
        attention_mask,
    ):
        model = DecisionTransformerModel(config=config)
        model.to(torch_device)
        model.eval()
        result = model(states, actions, rewards, returns_to_go, timesteps, attention_mask)

        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)
        )  

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        (
            config,
            states,
            actions,
            rewards,
            returns_to_go,
            timesteps,
            attention_mask,
        ) = config_and_inputs
        inputs_dict = {
            "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 DecisionTransformerModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (DecisionTransformerModel,) if is_torch_available() else ()
    all_generative_model_classes = ()
    pipeline_model_mapping = {"feature-extraction": DecisionTransformerModel} if is_torch_available() else {}

    
    test_generate_without_input_ids = False

    
    test_pruning = False
    test_resize_embeddings = False
    test_head_masking = False
    test_attention_outputs = False
    test_hidden_states_output = False
    test_inputs_embeds = False
    test_model_common_attributes = False
    test_gradient_checkpointing = False
    test_torchscript = False

    def setUp(self):
        self.model_tester = DecisionTransformerModelTester(self)
        self.config_tester = ConfigTester(self, config_class=DecisionTransformerConfig, hidden_size=37)

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

    def test_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*config_and_inputs)

    @slow
    def test_model_from_pretrained(self):
        for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            model = DecisionTransformerModel.from_pretrained(model_name)
            self.assertIsNotNone(model)

    def test_forward_signature(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            signature = inspect.signature(model.forward)
            
            arg_names = [*signature.parameters.keys()]

            expected_arg_names = [
                "states",
                "actions",
                "rewards",
                "returns_to_go",
                "timesteps",
                "attention_mask",
            ]

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


@require_torch
class DecisionTransformerModelIntegrationTest(unittest.TestCase):
    @slow
    def test_autoregressive_prediction(self):
        

        NUM_STEPS = 2  
        TARGET_RETURN = 10  
        model = DecisionTransformerModel.from_pretrained("edbeeching/decision-transformer-gym-hopper-expert")
        model = model.to(torch_device)
        config = model.config
        torch.manual_seed(0)
        state = torch.randn(1, 1, config.state_dim).to(device=torch_device, dtype=torch.float32)  

        expected_outputs = torch.tensor(
            [[0.242793, -0.28693074, 0.8742613], [0.67815274, -0.08101085, -0.12952147]], device=torch_device
        )

        returns_to_go = torch.tensor(TARGET_RETURN, device=torch_device, dtype=torch.float32).reshape(1, 1, 1)
        states = state
        actions = torch.zeros(1, 0, config.act_dim, device=torch_device, dtype=torch.float32)
        rewards = torch.zeros(1, 0, device=torch_device, dtype=torch.float32)
        timesteps = torch.tensor(0, device=torch_device, dtype=torch.long).reshape(1, 1)

        for step in range(NUM_STEPS):
            actions = torch.cat([actions, torch.zeros(1, 1, config.act_dim, device=torch_device)], dim=1)
            rewards = torch.cat([rewards, torch.zeros(1, 1, device=torch_device)], dim=1)

            attention_mask = torch.ones(1, states.shape[1]).to(dtype=torch.long, device=states.device)

            with torch.no_grad():
                _, action_pred, _ = model(
                    states=states,
                    actions=actions,
                    rewards=rewards,
                    returns_to_go=returns_to_go,
                    timesteps=timesteps,
                    attention_mask=attention_mask,
                    return_dict=False,
                )

            self.assertEqual(action_pred.shape, actions.shape)
            self.assertTrue(torch.allclose(action_pred[0, -1], expected_outputs[step], atol=1e-4))
            state, reward, _, _ = (  
                torch.randn(1, 1, config.state_dim).to(device=torch_device, dtype=torch.float32),
                1.0,
                False,
                {},
            )

            actions[-1] = action_pred[0, -1]
            states = torch.cat([states, state], dim=1)
            pred_return = returns_to_go[0, -1] - reward
            returns_to_go = torch.cat([returns_to_go, pred_return.reshape(1, 1, 1)], dim=1)
            timesteps = torch.cat(
                [timesteps, torch.ones((1, 1), device=torch_device, dtype=torch.long) * (step + 1)], dim=1
            )