84d517d540204dcaa20237dc05bdd4a63a958923,reagent/training/world_model/mdnrnn_trainer.py,MDNRNNTrainer,get_loss,#MDNRNNTrainer#Any#Any#,59

Before Change

        assert isinstance(training_batch, rlt.PreprocessedMemoryNetworkInput)
        // mdnrnn"s input should have seq_len as the first dimension

        mdnrnn_input = rlt.PreprocessedStateAction.from_tensors(
            training_batch.state.float_features, training_batch.action
        )
        mdnrnn_output = self.memory_network(mdnrnn_input)
        // mus, sigmas: [seq_len, batch_size, num_gaussian, state_dim]
        mus, sigmas, logpi, rs, nts = (
            mdnrnn_output.mus,

After Change

        // mdnrnn"s input should have seq_len as the first dimension

        mdnrnn_output = self.memory_network(
            training_batch.state, rlt.FeatureData(training_batch.action)
        )
        // mus, sigmas: [seq_len, batch_size, num_gaussian, state_dim]
        mus, sigmas, logpi, rs, nts = (
            mdnrnn_output.mus,

Instances