text: import itertools
  import numpy as np
  from copy import deepcopy
  
  class Node():
      def __init__(self, label, qpos_ids, qvel_ids, act_ids, body_fn=None, bodies=None, extra_obs=None, tendons=None):
          self.label = label
          self.qpos_ids = qpos_ids
          self.qvel_ids = qvel_ids
          self.act_ids = act_ids
          self.bodies = bodies
          self.extra_obs = {} if extra_obs is None else extra_obs
          self.body_fn = body_fn
          self.tendons = tendons
          pass
  
      def __str__(self):
          return self.label
  
      def __repr__(self):
          return self.label
  
  
  class HyperEdge():
      def __init__(self, *edges):
          self.edges = set(edges)
  
      def __contains__(self, item):
          return item in self.edges
  
      def __str__(self):
          return "HyperEdge({})".format(self.edges)
  
      def __repr__(self):
          return "HyperEdge({})".format(self.edges)
  
  
  def get_joints_at_kdist(agent_id, agent_partitions, hyperedges, k=0, kagents=False,):
      """ Identify all joints at distance <= k from agent agent_id
  
      :param agent_id: id of agent to be considered
      :param agent_partitions: list of joint tuples in order of agentids
      :param edges: list of tuples (joint1, joint2)
      :param k: kth degree
      :param kagents: True (observe all joints of an agent if a single one is) or False (individual joint granularity)
      :return:
          dict with k as key, and list of joints at that distance
      """
      assert not kagents, "kagents not implemented!"
  
      agent_joints = agent_partitions[agent_id]
  
      def _adjacent(lst, kagents=False):
          # return all sets adjacent to any element in lst
          ret = set([])
          for l in lst:
              ret = ret.union(set(itertools.chain(*[e.edges.difference({l}) for e in hyperedges if l in e])))
          return ret
  
      seen = set([])
      new = set([])
      k_dict = {}
      for _k in range(k+1):
          if not _k:
              new = set(agent_joints)
          else:
              print(hyperedges)
              new = _adjacent(new) - seen
          seen = seen.union(new)
          k_dict[_k] = sorted(list(new), key=lambda x:x.label)
      return k_dict
  
  
  def build_obs(env, k_dict, k_categories, global_dict, global_categories, vec_len=None):
      """Given a k_dict from get_joints_at_kdist, extract observation vector.
  
      :param k_dict: k_dict
      :param qpos: qpos numpy array
      :param qvel: qvel numpy array
      :param vec_len: if None no padding, else zero-pad to vec_len
      :return:
      observation vector
      """
  
      # TODO: This needs to be fixed, it was designed for half-cheetah only!
      #if add_global_pos:
      #    obs_qpos_lst.append(global_qpos)
      #    obs_qvel_lst.append(global_qvel)
  
  
      body_set_dict = {}
      obs_lst = []
      # Add parts attributes
      for k in sorted(list(k_dict.keys())):
          cats = k_categories[k]
          for _t in k_dict[k]:
              for c in cats:
                  if c in _t.extra_obs:
                      items = _t.extra_obs[c](env).tolist()
                      obs_lst.extend(items if isinstance(items, list) else [items])
                  else:
                      if c in ["qvel","qpos"]: # this is a "joint position/velocity" item
                          items = getattr(env.sim.data, c)[getattr(_t, "{}_ids".format(c))]
                          obs_lst.extend(items if isinstance(items, list) else [items])
                      elif c in ["qfrc_actuator"]: # this is a "vel position" item
                          items = getattr(env.sim.data, c)[getattr(_t, "{}_ids".format("qvel"))]
                          obs_lst.extend(items if isinstance(items, list) else [items])
                      elif c in ["cvel", "cinert", "cfrc_ext"]:  # this is a "body position" item
                          if _t.bodies is not None:
                              for b in _t.bodies:
                                  if c not in body_set_dict:
                                      body_set_dict[c] = set()
                                  if b not in body_set_dict[c]:
                                      items = getattr(env.sim.data, c)[b].tolist()
                                      items = getattr(_t, "body_fn", lambda _id,x:x)(b, items)
                                      obs_lst.extend(items if isinstance(items, list) else [items])
                                      body_set_dict[c].add(b)
  
      # Add global attributes
      body_set_dict = {}
      for c in global_categories:
          if c in ["qvel", "qpos"]:  # this is a "joint position" item
              for j in global_dict.get("joints", []):
                  items = getattr(env.sim.data, c)[getattr(j, "{}_ids".format(c))]
                  obs_lst.extend(items if isinstance(items, list) else [items])
          else:
              for b in global_dict.get("bodies", []):
                  if c not in body_set_dict:
                      body_set_dict[c] = set()
                  if b not in body_set_dict[c]:
                      obs_lst.extend(getattr(env.sim.data, c)[b].tolist())
                      body_set_dict[c].add(b)
  
      if vec_len is not None:
          pad = np.array((vec_len - len(obs_lst))*[0])
          if len(pad):
              return np.concatenate([np.array(obs_lst), pad])
      return np.array(obs_lst)
  
  
  def build_actions(agent_partitions, k_dict):
      # Composes agent actions output from networks
      # into coherent joint action vector to be sent to the env.
      pass
  
  def get_parts_and_edges(label, partitioning):
      if label in ["half_cheetah", "HalfCheetah-v2"]:
  
          # define Mujoco graph
          bthigh = Node("bthigh", -6, -6, 0)
          bshin = Node("bshin", -5, -5, 1)
          bfoot = Node("bfoot", -4, -4, 2)
          fthigh = Node("fthigh", -3, -3, 3)
          fshin = Node("fshin", -2, -2, 4)
          ffoot = Node("ffoot", -1, -1, 5)
  
          edges = [HyperEdge(bfoot, bshin),
                   HyperEdge(bshin, bthigh),
                   HyperEdge(bthigh, fthigh),
                   HyperEdge(fthigh, fshin),
                   HyperEdge(fshin, ffoot)]
  
          root_x = Node("root_x", 0, 0, -1,
                        extra_obs={"qpos": lambda env: np.array([])})
          root_z = Node("root_z", 1, 1, -1)
          root_y = Node("root_y", 2, 2, -1)
          globals = {"joints":[root_x, root_y, root_z]}
  
          if partitioning == "2x3":
              parts = [(bfoot, bshin, bthigh),
                       (ffoot, fshin, fthigh)]
          elif partitioning == "6x1":
              parts = [(bfoot,), (bshin,), (bthigh,), (ffoot,), (fshin,), (fthigh,)]
          else:
              raise Exception("UNKNOWN partitioning config: {}".format(partitioning))
  
          return parts, edges, globals
  
      elif label in ["Ant-v2"]:
  
          # define Mujoco graph
          torso = 1
          front_left_leg = 2
          aux_1 = 3
          ankle_1 = 4
          front_right_leg = 5
          aux_2 = 6
          ankle_2 = 7
          back_leg = 8
          aux_3 = 9
          ankle_3 = 10
          right_back_leg = 11
          aux_4 = 12
          ankle_4 = 13
  
          hip1 = Node("hip1", -8, -8, 2, bodies=[torso, front_left_leg], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist()) #
          ankle1 = Node("ankle1", -7, -7, 3, bodies=[front_left_leg, aux_1, ankle_1], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())#,
          hip2 = Node("hip2", -6, -6, 4, bodies=[torso, front_right_leg], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())#,
          ankle2 = Node("ankle2", -5, -5, 5, bodies=[front_right_leg, aux_2, ankle_2], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())#,
          hip3 = Node("hip3", -4, -4, 6, bodies=[torso, back_leg], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())#,
          ankle3 = Node("ankle3", -3, -3, 7, bodies=[back_leg, aux_3, ankle_3], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())#,
          hip4 = Node("hip4", -2, -2, 0, bodies=[torso, right_back_leg], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())#,
          ankle4 = Node("ankle4", -1, -1, 1, bodies=[right_back_leg, aux_4, ankle_4], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())#,
  
          edges = [HyperEdge(ankle4, hip4),
                   HyperEdge(ankle1, hip1),
                   HyperEdge(ankle2, hip2),
                   HyperEdge(ankle3, hip3),
                   HyperEdge(hip4, hip1, hip2, hip3),
                   ]
  
          free_joint = Node("free", 0, 0, -1, extra_obs={"qpos": lambda env: env.sim.data.qpos[:7],
                                                         "qvel": lambda env: env.sim.data.qvel[:6],
                                                         "cfrc_ext": lambda env: np.clip(env.sim.data.cfrc_ext[0:1], -1, 1)})
          globals = {"joints": [free_joint]}
  
          if partitioning == "2x4": # neighbouring legs together
              parts = [(hip1, ankle1, hip2, ankle2),
                       (hip3, ankle3, hip4, ankle4)]
          elif partitioning == "2x4d": # diagonal legs together
              parts = [(hip1, ankle1, hip3, ankle3),
                       (hip2, ankle2, hip4, ankle4)]
          elif partitioning == "4x2":
              parts = [(hip1, ankle1),
                       (hip2, ankle2),
                       (hip3, ankle3),
                       (hip4, ankle4)]
          else:
              raise Exception("UNKNOWN partitioning config: {}".format(partitioning))
  
          return parts, edges, globals
  
      elif label in ["Hopper-v2"]:
  
          # define Mujoco-Graph
          thigh_joint = Node("thigh_joint", -3, -3, 0,
                             extra_obs={"qvel": lambda env: np.clip(np.array([env.sim.data.qvel[-3]]), -10, 10)})
          leg_joint = Node("leg_joint", -2, -2, 1,
                           extra_obs={"qvel": lambda env: np.clip(np.array([env.sim.data.qvel[-2]]), -10, 10)})
          foot_joint = Node("foot_joint", -1, -1, 2,
                            extra_obs={"qvel": lambda env: np.clip(np.array([env.sim.data.qvel[-1]]), -10, 10)})
  
          edges = [HyperEdge(foot_joint, leg_joint),
                   HyperEdge(leg_joint, thigh_joint)]
  
          root_x = Node("root_x", 0, 0, -1, extra_obs={"qpos": lambda env: np.array([]),
                                                       "qvel": lambda env: np.clip(np.array([env.sim.data.qvel[1]]), -10, 10)})
          root_z = Node("root_z", 1, 1, -1, extra_obs={"qvel": lambda env: np.clip(np.array([env.sim.data.qvel[1]]), -10, 10)})
          root_y = Node("root_y", 2, 2, -1, extra_obs={"qvel": lambda env: np.clip(np.array([env.sim.data.qvel[2]]), -10, 10)})
          globals = {"joints":[root_x, root_y, root_z]}
  
          if partitioning == "3x1":
              parts = [(thigh_joint,),
                       (leg_joint,),
                       (foot_joint,)]
  
          else:
              raise Exception("UNKNOWN partitioning config: {}".format(partitioning))
  
          return parts, edges, globals
  
      elif label in ["Humanoid-v2", "HumanoidStandup-v2"]:
  
          # define Mujoco-Graph
          abdomen_y = Node("abdomen_y", -16, -16, 0) # act ordering bug in env -- double check!
          abdomen_z = Node("abdomen_z", -17, -17, 1)
          abdomen_x = Node("abdomen_x", -15, -15, 2)
          right_hip_x = Node("right_hip_x", -14, -14, 3)
          right_hip_z = Node("right_hip_z", -13, -13, 4)
          right_hip_y = Node("right_hip_y", -12, -12, 5)
          right_knee = Node("right_knee", -11, -11, 6)
          left_hip_x = Node("left_hip_x", -10, -10, 7)
          left_hip_z = Node("left_hip_z", -9, -9, 8)
          left_hip_y = Node("left_hip_y", -8, -8, 9)
          left_knee = Node("left_knee", -7, -7, 10)
          right_shoulder1 = Node("right_shoulder1", -6, -6, 11)
          right_shoulder2 = Node("right_shoulder2", -5, -5, 12)
          right_elbow = Node("right_elbow", -4, -4, 13)
          left_shoulder1 = Node("left_shoulder1", -3, -3, 14)
          left_shoulder2 = Node("left_shoulder2", -2, -2, 15)
          left_elbow = Node("left_elbow", -1, -1, 16)
  
          edges = [HyperEdge(abdomen_x, abdomen_y, abdomen_z),
                   HyperEdge(right_hip_x, right_hip_y, right_hip_z),
                   HyperEdge(left_hip_x, left_hip_y, left_hip_z),
                   HyperEdge(left_elbow, left_shoulder1, left_shoulder2),
                   HyperEdge(right_elbow, right_shoulder1, right_shoulder2),
                   HyperEdge(left_knee, left_hip_x, left_hip_y, left_hip_z),
                   HyperEdge(right_knee, right_hip_x, right_hip_y, right_hip_z),
                   HyperEdge(left_shoulder1, left_shoulder2, abdomen_x, abdomen_y, abdomen_z),
                   HyperEdge(right_shoulder1, right_shoulder2, abdomen_x, abdomen_y, abdomen_z),
                   HyperEdge(abdomen_x, abdomen_y, abdomen_z, left_hip_x, left_hip_y, left_hip_z),
                   HyperEdge(abdomen_x, abdomen_y, abdomen_z, right_hip_x, right_hip_y, right_hip_z),
                   ]
  
          globals = {}
  
          if partitioning == "9|8": # 17 in total, so one action is a dummy (to be handled by pymarl)
              # isolate upper and lower body
              parts = [(left_shoulder1, left_shoulder2, abdomen_x, abdomen_y, abdomen_z,
                        right_shoulder1, right_shoulder2,
                        right_elbow, left_elbow),
                       (left_hip_x, left_hip_y, left_hip_z,
                        right_hip_x, right_hip_y, right_hip_z,
                        right_knee, left_knee)]
              # TODO: There could be tons of decompositions here
  
          else:
              raise Exception("UNKNOWN partitioning config: {}".format(partitioning))
  
          return parts, edges, globals
  
      elif label in ["Reacher-v2"]:
  
          # define Mujoco-Graph
          body0 = 1
          body1 = 2
          fingertip = 3
          joint0 = Node("joint0", -4, -4, 0,
                        bodies=[body0, body1],
                        extra_obs={"qpos":(lambda env:np.array([np.sin(env.sim.data.qpos[-4]),
                                                                np.cos(env.sim.data.qpos[-4])]))})
          joint1 = Node("joint1", -3, -3, 1,
                        bodies=[body1, fingertip],
                        extra_obs={"fingertip_dist":(lambda env:env.get_body_com("fingertip") - env.get_body_com("target")),
                                   "qpos":(lambda env:np.array([np.sin(env.sim.data.qpos[-3]),
                                                                np.cos(env.sim.data.qpos[-3])]))})
          edges = [HyperEdge(joint0, joint1)]
  
          worldbody = 0
          target = 4
          target_x = Node("target_x", -2, -2, -1, extra_obs={"qvel":(lambda env:np.array([]))})
          target_y = Node("target_y", -1, -1, -1, extra_obs={"qvel":(lambda env:np.array([]))})
          globals = {"bodies":[worldbody, target],
                     "joints":[target_x, target_y]}
  
          if partitioning == "2x1":
              # isolate upper and lower arms
              parts = [(joint0,), (joint1,)]
              # TODO: There could be tons of decompositions here
  
          else:
              raise Exception("UNKNOWN partitioning config: {}".format(partitioning))
  
          return parts, edges, globals
  
      elif label in ["Swimmer-v2"]:
  
          # define Mujoco-Graph
          joint0 = Node("rot2", -2, -2, 0) # TODO: double-check ids
          joint1 = Node("rot3", -1, -1, 1)
  
          edges = [HyperEdge(joint0, joint1)]
          globals = {}
  
          if partitioning == "2x1":
              # isolate upper and lower body
              parts = [(joint0,), (joint1,)]
              # TODO: There could be tons of decompositions here
  
          else:
              raise Exception("UNKNOWN partitioning config: {}".format(partitioning))
  
          return parts, edges, globals
  
      elif label in ["Walker2d-v2"]:
  
          # define Mujoco-Graph
          thigh_joint = Node("thigh_joint", -6, -6, 0)
          leg_joint = Node("leg_joint", -5, -5, 1)
          foot_joint = Node("foot_joint", -4, -4, 2)
          thigh_left_joint = Node("thigh_left_joint", -3, -3, 3)
          leg_left_joint = Node("leg_left_joint", -2, -2, 4)
          foot_left_joint = Node("foot_left_joint", -1, -1, 5)
  
          edges = [HyperEdge(foot_joint, leg_joint),
                   HyperEdge(leg_joint, thigh_joint),
                   HyperEdge(foot_left_joint, leg_left_joint),
                   HyperEdge(leg_left_joint, thigh_left_joint),
                   HyperEdge(thigh_joint, thigh_left_joint)
                   ]
          globals = {}
  
          if partitioning == "2x3":
              # isolate upper and lower body
              parts = [(foot_joint, leg_joint, thigh_joint),
                       (foot_left_joint, leg_left_joint, thigh_left_joint,)]
              # TODO: There could be tons of decompositions here
  
          else:
              raise Exception("UNKNOWN partitioning config: {}".format(partitioning))
  
          return parts, edges, globals
  
      elif label in ["coupled_half_cheetah"]:
  
          # define Mujoco graph
          tendon = 0
  
          bthigh = Node("bthigh", -6, -6, 0,
                       tendons=[tendon],
                       extra_obs = {"ten_J": lambda env: env.sim.data.ten_J[tendon],
                                    "ten_length": lambda env: env.sim.data.ten_length,
                                    "ten_velocity": lambda env: env.sim.data.ten_velocity})
          bshin = Node("bshin", -5, -5, 1)
          bfoot = Node("bfoot", -4, -4, 2)
          fthigh = Node("fthigh", -3, -3, 3)
          fshin = Node("fshin", -2, -2, 4)
          ffoot = Node("ffoot", -1, -1, 5)
  
          bthigh2 = Node("bthigh2", -6, -6, 0,
                        tendons=[tendon],
                        extra_obs={"ten_J": lambda env: env.sim.data.ten_J[tendon],
                                   "ten_length": lambda env: env.sim.data.ten_length,
                                   "ten_velocity": lambda env: env.sim.data.ten_velocity})
          bshin2 = Node("bshin2", -5, -5, 1)
          bfoot2 = Node("bfoot2", -4, -4, 2)
          fthigh2 = Node("fthigh2", -3, -3, 3)
          fshin2 = Node("fshin2", -2, -2, 4)
          ffoot2 = Node("ffoot2", -1, -1, 5)
  
  
          edges = [HyperEdge(bfoot, bshin),
                   HyperEdge(bshin, bthigh),
                   HyperEdge(bthigh, fthigh),
                   HyperEdge(fthigh, fshin),
                   HyperEdge(fshin, ffoot),
                   HyperEdge(bfoot2, bshin2),
                   HyperEdge(bshin2, bthigh2),
                   HyperEdge(bthigh2, fthigh2),
                   HyperEdge(fthigh2, fshin2),
                   HyperEdge(fshin2, ffoot2)
                   ]
          globals = {}
  
          root_x = Node("root_x", 0, 0, -1,
                        extra_obs={"qpos": lambda env: np.array([])})
          root_z = Node("root_z", 1, 1, -1)
          root_y = Node("root_y", 2, 2, -1)
          globals = {"joints":[root_x, root_y, root_z]}
  
          if partitioning == "1p1":
              parts = [(bfoot, bshin, bthigh, ffoot, fshin, fthigh),
                       (bfoot2, bshin2, bthigh2, ffoot2, fshin2, fthigh2)
                       ]
          else:
              raise Exception("UNKNOWN partitioning config: {}".format(partitioning))
  
          return parts, edges, globals
  
      elif label in ["manyagent_swimmer"]:
  
          # Generate asset file
          try:
              n_agents = int(partitioning.split("x")[0])
              n_segs_per_agents = int(partitioning.split("x")[1])
              n_segs = n_agents * n_segs_per_agents
          except Exception as e:
              raise Exception("UNKNOWN partitioning config: {}".format(partitioning))
  
          # Note: Default Swimmer corresponds to n_segs = 3
  
          # define Mujoco-Graph
          joints = [Node("rot{:d}".format(i), -n_segs + i, -n_segs + i, i) for i in range(0, n_segs)]
          edges = [HyperEdge(joints[i], joints[i+1]) for i in range(n_segs-1)]
          globals = {}
  
          parts = [tuple(joints[i * n_segs_per_agents:(i + 1) * n_segs_per_agents]) for i in range(n_agents)]
          return parts, edges, globals
  
      elif label in ["manyagent_ant"]: # TODO: FIX!
  
          # Generate asset file
          try:
              n_agents = int(partitioning.split("x")[0])
              n_segs_per_agents = int(partitioning.split("x")[1])
              n_segs = n_agents * n_segs_per_agents
          except Exception as e:
              raise Exception("UNKNOWN partitioning config: {}".format(partitioning))
  
  
          # # define Mujoco graph
          # torso = 1
          # front_left_leg = 2
          # aux_1 = 3
          # ankle_1 = 4
          # right_back_leg = 11
          # aux_4 = 12
          # ankle_4 = 13
          #
          # off = -4*(n_segs-1)
          # hip1 = Node("hip1", -4-off, -4-off, 2, bodies=[torso, front_left_leg], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist()) #
          # ankle1 = Node("ankle1", -3-off, -3-off, 3, bodies=[front_left_leg, aux_1, ankle_1], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())#,
          # hip4 = Node("hip4", -2-off, -2-off, 0, bodies=[torso, right_back_leg], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())#,
          # ankle4 = Node("ankle4", -1-off, -1-off, 1, bodies=[right_back_leg, aux_4, ankle_4], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())#,
          #
          # edges = [HyperEdge(ankle4, hip4),
          #          HyperEdge(ankle1, hip1),
          #          HyperEdge(hip4, hip1),
          #          ]
  
          edges = []
          joints = []
          for si in range(n_segs):
  
              torso = 1 + si*7
              front_right_leg = 2 + si*7
              aux1 = 3 + si*7
              ankle1 = 4 + si*7
              back_leg = 5 + si*7
              aux2 = 6 + si*7
              ankle2 = 7 + si*7
  
              off = -4 * (n_segs - 1 - si)
              hip1n = Node("hip1_{:d}".format(si), -4-off, -4-off, 2+4*si, bodies=[torso, front_right_leg], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())
              ankle1n = Node("ankle1_{:d}".format(si), -3-off, -3-off, 3+4*si, bodies=[front_right_leg, aux1, ankle1], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())
              hip2n = Node("hip2_{:d}".format(si), -2-off, -2-off, 0+4*si, bodies=[torso, back_leg], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())
              ankle2n = Node("ankle2_{:d}".format(si), -1-off, -1-off, 1+4*si, bodies=[back_leg, aux2, ankle2], body_fn=lambda _id, x:np.clip(x, -1, 1).tolist())
  
              edges += [HyperEdge(ankle1n, hip1n),
                        HyperEdge(ankle2n, hip2n),
                        HyperEdge(hip1n, hip2n)]
              if si:
                  edges += [HyperEdge(hip1m, hip2m, hip1n, hip2n)]
  
              hip1m = deepcopy(hip1n)
              hip2m = deepcopy(hip2n)
              joints.append([hip1n,
                             ankle1n,
                             hip2n,
                             ankle2n])
  
          free_joint = Node("free", 0, 0, -1, extra_obs={"qpos": lambda env: env.sim.data.qpos[:7],
                                                         "qvel": lambda env: env.sim.data.qvel[:6],
                                                         "cfrc_ext": lambda env: np.clip(env.sim.data.cfrc_ext[0:1], -1, 1)})
          globals = {"joints": [free_joint]}
  
          parts =  [[x for sublist in joints[i * n_segs_per_agents:(i + 1) * n_segs_per_agents] for x in sublist] for i in range(n_agents)]
  
          return parts, edges, globals