class_docstr: Interface for all policy networks in EvoJAX.
class_docstr_tok: ["Interface","for","all","policy","networks","in","EvoJAX","."]
docstr: Reset the policy.
  
  Args:
      TaskState - Initial observations.
  Returns:
      PolicyState. Policy internal states.
docstr_tok: ["Reset","the","policy",".","Args",":","TaskState","-","Initial","observations",".","Returns",":","PolicyState",".","Policy","internal","states","."]
code: def reset(self, states: TaskState) -> PolicyState:
      """Reset the policy.
  
      Args:
          TaskState - Initial observations.
      Returns:
          PolicyState. Policy internal states.
      """
      keys = jax.random.split(
          jax.random.PRNGKey(0), states.obs.shape[0]
      )
      return PolicyState(keys=keys)
  
file_path: evojax/evojax/policy/base.py
imports: from abc import ABC
  from abc import abstractmethod
  from typing import Tuple
  import jax.numpy
  import jax.random
  from flax.struct import dataclass
  from evojax.task.base import TaskState
label_desc: Use image node_id 1 for calling the PolicyNetwork obj's underlying member method code with example usage: obj.reset(states) and returns: PolicyState