class_docstr: Base class for Task Embedding policies in TensorFlow.
  
  This policy needs a task id in addition to observation to sample an action.
class_docstr_tok: ["Base","class","for","Task","Embedding","policies","in","TensorFlow",".","This","policy","needs","a","task","id","in","addition","to","observation","to","sample","an","action","."]
docstr: Splits up observation into one-hot task and environment observation.
  
  Args:
      collated (np.ndarray): Environment observation concatenated with
          task one-hot, with shape :math:`(O+N, )`. O is the dimension of
          observation, N is the number of tasks.
  
  Returns:
      np.ndarray: Vanilla environment observation,
          with shape :math:`(O, )`. O is the dimension of observation.
      np.ndarray: Task one-hot, with shape :math:`(N, )`. N is the number
          of tasks.
docstr_tok: ["Splits","up","observation","into","one-hot","task","and","environment","observation",".","Args",":","collated","(","np.ndarray",")",":","Environment","observation","concatenated","with","task","one-hot",",","with","shape",":","math",":","`","(","O+N",",",")","`",".","O","is","the","dimension","of","observation",",","N","is","the","number","of","tasks",".","Returns",":","np.ndarray",":","Vanilla","environment","observation",",","with","shape",":","math",":","`","(","O",",",")","`",".","O","is","the","dimension","of","observation",".","np.ndarray",":","Task","one-hot",",","with","shape",":","math",":","`","(","N",",",")","`",".","N","is","the","number","of","tasks","."]
code: def split_augmented_observation(self, collated):
      """Splits up observation into one-hot task and environment observation.
  
      Args:
          collated (np.ndarray): Environment observation concatenated with
              task one-hot, with shape :math:`(O+N, )`. O is the dimension of
              observation, N is the number of tasks.
  
      Returns:
          np.ndarray: Vanilla environment observation,
              with shape :math:`(O, )`. O is the dimension of observation.
          np.ndarray: Task one-hot, with shape :math:`(N, )`. N is the number
              of tasks.
  
      """
      task_dim = self.task_space.flat_dim
      return collated[:-task_dim], collated[-task_dim:]
  
file_path: garage/src/garage/tf/policies/task_embedding_policy.py
imports: import abc
  from garage.tf.policies.policy import Policy
label_desc: Use image node_id 13 for calling the TaskEmbeddingPolicy obj's underlying member method code with example usage: obj.split_augmented_observation(collated) and returns: collated,  collated