class_docstr: Behavior descriptor extractor.
docstr: Augment the original task state definition with more entries.
  
  Args:
      task_state_def - This should be a flax.struct.dataclass instance.
  Returns:
      A flax.struct.dataclass type of class definition with extra entries.
docstr_tok: ["Augment","the","original","task","state","definition","with","more","entries",".","Args",":","task_state_def","-","This","should","be","a","flax.struct.dataclass","instance",".","Returns",":","A","flax.struct.dataclass","type","of","class","definition","with","extra","entries","."]
code: def get_extended_task_state_def(self, task_state_def: T) -> T:
      """Augment the original task state definition with more entries.
  
      Args:
          task_state_def - This should be a flax.struct.dataclass instance.
      Returns:
          A flax.struct.dataclass type of class definition with extra entries.
      """
      if self.bd_spec is None:
          return task_state_def
      else:
          fields = []
          # Keep what is in the original task_state.
          data_fields = task_state_def.__dict__["__annotations__"]
          for name, field in data_fields.items():
              fields.append((name, field))
          # Add bd fields.
          fields.extend([(x[0], jnp.int32) for x in self.bd_spec])
          # Add extra bd state fields to help the calculation.
          fields.extend(self.bd_state_spec)
          return dataclass(
              dataclasses.make_dataclass(
                  type(task_state_def).__name__,
                  fields=fields,
                  bases=task_state_def.__bases__,
                  init=False,
              )
          )
  
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 2 for calling the BDExtractor obj's underlying member method code with example usage: obj.get_extended_task_state_def(task_state_def) and returns: task_state_def, dataclass