| 25 | mu_0ts = jnp.concatenate([mu_0_random[None], mu_0ts], axis=0) |
| 26 | |
| 27 | def dumps(sys, statess) -> str: |
| 28 | |
| 29 | d = _to_dict(sys) |
| 30 | |
| 31 | # Fill in empty link names |
| 32 | link_names = [n or f"link {i}" for i, n in enumerate(sys.link_names)] |
| 33 | link_names += ["world"] |
| 34 | |
| 35 | # Unpack geoms into a dict for the visualizer |
| 36 | link_geoms = {} |
| 37 | for id_ in range(sys.ngeom): |
| 38 | link_idx = sys.geom_bodyid[id_] - 1 |
| 39 | rgba = sys.geom_rgba[id_] |
| 40 | |
| 41 | geom = { |
| 42 | "name": _GEOM_TYPE_NAMES[sys.geom_type[id_]], |
| 43 | "link_idx": link_idx, |
| 44 | "pos": sys.geom_pos[id_], |
| 45 | "rot": sys.geom_quat[id_], |
| 46 | "rgba": rgba, |
| 47 | "size": sys.geom_size[id_], |
| 48 | } |
| 49 | |
| 50 | link_geoms.setdefault(link_names[link_idx], []).append(_to_dict(geom)) |
| 51 | |
| 52 | # repeat link_geoms for each body across all timesteps |
| 53 | all_link_geoms = {} |
| 54 | all_link_names = [] |
| 55 | traj_len = len(statess[0]) |
| 56 | plot_idx = jnp.arange(0, traj_len, plot_interval) |
| 57 | plot_idx = jnp.append(plot_idx, traj_len - 1) |
| 58 | for k in range(traj_len): |
| 59 | for _, (name, geoms) in enumerate(link_geoms.items()): |
| 60 | name = f"{name}_{k}" if k > 0 else name |
| 61 | geoms_new = [] |
| 62 | for geom in geoms: |
| 63 | geom_new = geom.copy() |
| 64 | if "world" in name: |
| 65 | geom_new["link_idx"] = -1 |
| 66 | elif "goal" in name: |
| 67 | geom_new["rgba"] = [0.0, 1.0, 0.0, 1.0] |
| 68 | elif "_ref" in name: |
| 69 | if "torso" in name or "thigh" in name: |
| 70 | geom_new["link_idx"] = geom["link_idx"] + k * (len(link_names) - 1) |
| 71 | a = k / traj_len * 0.8 + 0.2 |
| 72 | geom_new["rgba"] = [(1-a), 1.0, (1-a), 1.0] |
| 73 | else: |
| 74 | geom_new["rgba"] = [1.0, 1.0, 1.0, 0.0] |
| 75 | else: |
| 76 | geom_new["link_idx"] = geom["link_idx"] + k * (len(link_names) - 1) |
| 77 | a = k / traj_len * 0.8 + 0.2 |
| 78 | geom_new["rgba"] = [1, (1 - a), (1 - a), 1.0] |
| 79 | geoms_new.append(geom_new) |
| 80 | all_link_geoms[name] = geoms_new |
| 81 | all_link_names.append(name) |
| 82 | d["geoms"] = all_link_geoms |
| 83 | d["link_names"] = all_link_names |
| 84 | |