code: def default_models_to_snapshot(
      networks: SACNetworks, spec: specs.EnvironmentSpec
  ):
      """Defines default models to be snapshotted."""
      dummy_obs = utils.zeros_like(spec.observations)
      dummy_action = utils.zeros_like(spec.actions)
      dummy_key = jax.random.PRNGKey(0)
  
      def q_network(
          source: core.VariableSource,
      ) -> types.ModelToSnapshot:
          params = source.get_variables(["critic"])[0]
          return types.ModelToSnapshot(
              networks.q_network.apply,
              params,
              {"obs": dummy_obs, "action": dummy_action},
          )
  
      def default_training_actor(
          source: core.VariableSource,
      ) -> types.ModelToSnapshot:
          params = source.get_variables(["policy"])[0]
          return types.ModelToSnapshot(
              apply_policy_and_sample(networks, False),
              params,
              {"key": dummy_key, "obs": dummy_obs},
          )
  
      def default_eval_actor(
          source: core.VariableSource,
      ) -> types.ModelToSnapshot:
          params = source.get_variables(["policy"])[0]
          return types.ModelToSnapshot(
              apply_policy_and_sample(networks, True),
              params,
              {"key": dummy_key, "obs": dummy_obs},
          )
  
      return {
          "q_network": q_network,
          "default_training_actor": default_training_actor,
          "default_eval_actor": default_eval_actor,
      }
  
file_path: acme/acme/agents/jax/sac/networks.py
imports: import dataclasses
  from typing import Optional, Tuple
  from acme import core
  from acme import specs
  from acme.agents.jax import actor_core
  from acme.jax import networks
  from acme.jax import types
  from acme.jax import utils
  import haiku
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: default_models_to_snapshot(networks, spec) and returns: dict, types, types, types