class_docstr: A seq2seq policy that deals with simple additions.
code: def __init__(
      self,
      hidden_size: int = 256,
      teacher_force: bool = False,
      max_len_query_digit: int = 3,
      logger: logging.Logger = None,
  ):
      if logger is None:
          self._logger = create_logger("Seq2seqPolicy")
      else:
          self._logger = logger
  
      max_input_len = max_len_query_digit + 2 + 2
      max_output_len = max_len_query_digit + 3
      encoder_shape = jnp.ones(
          (1, max_input_len, char_table.vocab_size), dtype=jnp.float32
      )
      decoder_shape = jnp.ones(
          (1, max_output_len, char_table.vocab_size), dtype=jnp.float32
      )
      model = Seq2seq(
          hidden_size=hidden_size, teacher_force=teacher_force
      )
      key = random.PRNGKey(0)
      params = model.init(
          {"params": key, "lstm": key}, encoder_shape, decoder_shape
      )["params"]
      self.num_params, format_params_fn = get_params_format_fn(params)
      self._logger.info(
          "Seq2seqPolicy.num_params = {}".format(self.num_params)
      )
      self._format_params_fn = jax.vmap(format_params_fn)
  
      def forward_fn(p, o):
          x = jax.nn.one_hot(
              char_table.encode(jnp.array([11]))[0:1],
              char_table.vocab_size,
              dtype=jnp.float32,
          )
          x = jnp.tile(x, (o.shape[0], max_output_len, 1))
          logits, predictions = model.apply({"params": p}, o, x)
          return logits
  
      self._forward_fn = jax.vmap(forward_fn)
  
imports: import logging
  import numpy
  from typing import Tuple
  from typing import Any
  from functools import partial
  import jax
  import jax.numpy
  from jax import random
  from flax import linen
  from evojax.policy.base import PolicyNetwork
  from evojax.policy.base import PolicyState
  from evojax.task.base import TaskState
  from evojax.util import create_logger
  from evojax.util import get_params_format_fn
label_desc: Use image node_id 1 to create a new Seq2seqPolicy object from inherited base classes: PolicyNetwork with example: obj = Seq2seqPolicy(hidden_size, teacher_force, max_len_query_digit, logger)