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 __init__(
      self,
      environment_spec: specs.EnvironmentSpec,
      hidden_sizes: Tuple[int, ...],
  ):
      super(MLPTransitionModel, self).__init__(
          name="mlp_transition_model"
      )
  
      # Get num actions/observation shape.
      self._num_actions = environment_spec.actions.num_values
      self._input_shape = environment_spec.observations.shape
      self._flat_shape = int(np.prod(self._input_shape))
  
      # Prediction networks.
      self._state_network = snt.Sequential(
          [
              snt.nets.MLP(hidden_sizes + (self._flat_shape,)),
              snt.Reshape(self._input_shape),
          ]
      )
      self._reward_network = snt.Sequential(
          [
              snt.nets.MLP(hidden_sizes + (1,)),
              lambda r: tf.squeeze(r, axis=-1),
          ]
      )
      self._discount_network = snt.Sequential(
          [
              snt.nets.MLP(hidden_sizes + (1,)),
              lambda d: tf.squeeze(d, axis=-1),
          ]
      )
  
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 1 to create a new MLPTransitionModel object from inherited base classes: snt with example: obj = MLPTransitionModel(environment_spec, hidden_sizes)