class_docstr: Behavior descriptor extractor.
docstr: Return an extended task_state that includes bd_state fields.
  
  Args:
      task_state - Original task state.
  Returns:
      An instance of the extended task state.
docstr_tok: ["Return","an","extended","task_state","that","includes","bd_state","fields",".","Args",":","task_state","-","Original","task","state",".","Returns",":","An","instance","of","the","extended","task","state","."]
code: def init_extended_state(self, task_state: TaskState) -> T:
      """Return an extended task_state that includes bd_state fields.
  
      Args:
          task_state - Original task state.
      Returns:
          An instance of the extended task state.
      """
      bd_fields = {
          x[0]: jnp.zeros((), dtype=jnp.int32) for x in self.bd_spec
      }
      bd_state = self.init_state(task_state)
      return self.extended_task_state(
          **bd_state,  # These keep track of the stats we need.
          **bd_fields,  # These will be updated in self.summarize.
          **task_state.__dict__,
      )
  
imports: from abc import ABC
  from abc import abstractmethod
  from typing import Dict
  from typing import List
  from typing import Tuple
  from typing import TypeVar
  import dataclasses
  import jax.numpy
  from flax.struct import dataclass
label_desc: Use image node_id 3 for calling the BDExtractor obj's underlying member method code with example usage: obj.init_extended_state(task_state) and returns: self