class_docstr: This uses MLPs to model (s, a) -> (r, d, s').
class_docstr_tok: ["This","uses","MLPs","to","model","(","s",",","a",")","-",">","(","r",",","d",",","s","'",")","."]
code: def __call__(
      self, state: tf.Tensor, action: tf.Tensor
  ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
      embedded_state = snt.Flatten()(state)
      embedded_action = tf.one_hot(action, depth=self._num_actions)
  
      embedding = tf.concat([embedded_state, embedded_action], axis=-1)
  
      # Predict the next state, reward, and termination.
      next_state = self._state_network(embedding)
      reward = self._reward_network(embedding)
      discount_logits = self._discount_network(embedding)
  
      return next_state, reward, discount_logits
  
imports: from typing import Optional, Tuple
  from acme import specs
  from acme.agents.tf.mcts import types
  from acme.agents.tf.mcts.models import base
  from acme.tf import utils
  from bsuite.baselines.utils import replay
  import dm_env
  import numpy
  from scipy import special
  import sonnet
  import tensorflow
label_desc: Use image node_id 2 for calling the MLPTransitionModel obj's underlying member method code with example usage: obj.__call__(state, action) and returns: next_state,  reward,  discount_logits