class_docstr: An encoder network for soft actor critic.
  
  This is the network described in the SAC-AE paper "Improving sample
  efficiency in model-free reinforcement learning from images"
  (Yarats, Zhang, Kostrikov, Amos, Pineau, & Fergus, 2019), and used in
  experiments in the DrQ paper "Image Augmentation Is All You Need:
  Regularizing Deep Reinforcement Learning from Pixels" (Kostrikov, Yarats,
  & Fergus, 2020).
class_docstr_tok: ["An","encoder","network","for","soft","actor","critic",".","This","is","the","network","described","in","the","SAC-AE","paper","``","Improving","sample","efficiency","in","model-free","reinforcement","learning","from","images","''","(","Yarats",",","Zhang",",","Kostrikov",",","Amos",",","Pineau",",","&","Fergus",",","2019",")",",","and","used","in","experiments","in","the","DrQ","paper","``","Image","Augmentation","Is","All","You","Need",":","Regularizing","Deep","Reinforcement","Learning","from","Pixels","''","(","Kostrikov",",","Yarats",",","&","Fergus",",","2020",")","."]
code: def __call__(self, x):
      # Images are stored in the replay buffer as uint8.
      x = x.astype(jnp.float32) / 255.0
  
      # Flatten the last dimension (normally to deal with stacked rgb frames)
      if len(x.shape) > 3:
          x = x.reshape((*x.shape[:2], -1))
  
      kernel_init = nn.initializers.orthogonal()
      x = nn.Conv(
          features=32,
          kernel_size=(3, 3),
          strides=(2, 2),
          kernel_init=kernel_init,
      )(x)
      x = nn.relu(x)
      x = nn.Conv(
          features=32,
          kernel_size=(3, 3),
          strides=(1, 1),
          kernel_init=kernel_init,
      )(x)
      x = nn.relu(x)
      x = nn.Conv(
          features=32,
          kernel_size=(3, 3),
          strides=(1, 1),
          kernel_init=kernel_init,
      )(x)
      x = nn.relu(x)
      x = nn.Conv(
          features=32,
          kernel_size=(3, 3),
          strides=(1, 1),
          kernel_init=kernel_init,
      )(x)
      x = nn.relu(x)
      x = jnp.reshape(x, -1)  # Flatten
  
      critic_z = nn.Dense(features=50, kernel_init=kernel_init)(x)
      critic_z = nn.LayerNorm()(critic_z)
      critic_z = nn.tanh(critic_z)
  
      # Only the critic should train the convolution layers, so stop the
      # gradients from the actor.
      actor_z = nn.Dense(features=50, kernel_init=kernel_init)(
          jax.lax.stop_gradient(x)
      )
      actor_z = nn.LayerNorm()(actor_z)
      actor_z = nn.tanh(actor_z)
  
      return SACEncoderOutputs(critic_z, actor_z)
  
file_path: dopamine/dopamine/labs/sac_from_pixels/continuous_networks.py
imports: import dataclasses
  from typing import Optional, Tuple
  from dopamine.jax import continuous_networks
  from flax import linen
  import gin
  import jax
  from jax import numpy
label_desc: Use image node_id 1 for calling the SACEncoderNetwork obj's underlying member method code with example usage: obj.__call__(x) and returns: SACEncoderOutputs