(n_frames=4, hidden_sizes=(32, 64, 64),
                 output_size=512, mode='train',
                 kernel_initializer=None, padding='VALID')
return tl.Serial(
      _BytesToFloats(),
      _FrameStack(n_frames=n_frames),  # (B, T, H, W, 4C)
      tl.Conv(hidden_sizes[0], (8, 8), (4, 4), padding=padding,
              kernel_initializer=kernel_initializer),
      tl.Relu(),
      tl.Conv(hidden_sizes[1], (4, 4), (2, 2), padding=padding,
              kernel_initializer=kernel_initializer),
      tl.Relu(),
      tl.Conv(hidden_sizes[2], (3, 3), (1, 1), padding=padding,
              kernel_initializer=kernel_initializer),
      tl.Relu(),
      tl.Flatten(n_axes_to_keep=2),  # B, T and rest.
      tl.Dense(output_size),
      tl.Relu(),
  )
An Atari CNN.
An Atari CNN.
An Atari CNN .
def AtariCnnBody(n_frames=4, hidden_sizes=(32, 64, 64),
                 output_size=512, mode='train',
                 kernel_initializer=None, padding='VALID'):
  """An Atari CNN."""
  del mode

  # TODO(jonni): Include link to paper?
  # Input shape: (B, T, H, W, C)
  # Output shape: (B, T, output_size)
  return tl.Serial(
      _BytesToFloats(),
      _FrameStack(n_frames=n_frames),  # (B, T, H, W, 4C)
      tl.Conv(hidden_sizes[0], (8, 8), (4, 4), padding=padding,
              kernel_initializer=kernel_initializer),
      tl.Relu(),
      tl.Conv(hidden_sizes[1], (4, 4), (2, 2), padding=padding,
              kernel_initializer=kernel_initializer),
      tl.Relu(),
      tl.Conv(hidden_sizes[2], (3, 3), (1, 1), padding=padding,
              kernel_initializer=kernel_initializer),
      tl.Relu(),
      tl.Flatten(n_axes_to_keep=2),  # B, T and rest.
      tl.Dense(output_size),
      tl.Relu(),
  )
def AtariCnnBody ( n_frames = 4 , hidden_sizes = ( 32 , 64 , 64 ) , output_size = 512 , mode = 'train' , kernel_initializer = None , padding = 'VALID' ) : del mode # TODO(jonni): Include link to paper? # Input shape: (B, T, H, W, C) # Output shape: (B, T, output_size) return tl . Serial ( _BytesToFloats ( ) , _FrameStack ( n_frames = n_frames ) , # (B, T, H, W, 4C) tl . Conv ( hidden_sizes [ 0 ] , ( 8 , 8 ) , ( 4 , 4 ) , padding = padding , kernel_initializer = kernel_initializer ) , tl . Relu ( ) , tl . Conv ( hidden_sizes [ 1 ] , ( 4 , 4 ) , ( 2 , 2 ) , padding = padding , kernel_initializer = kernel_initializer ) , tl . Relu ( ) , tl . Conv ( hidden_sizes [ 2 ] , ( 3 , 3 ) , ( 1 , 1 ) , padding = padding , kernel_initializer = kernel_initializer ) , tl . Relu ( ) , tl . Flatten ( n_axes_to_keep = 2 ) , # B, T and rest. tl . Dense ( output_size ) , tl . Relu ( ) , )