Skip to content

resources

resources

Modules:

Name Description
generate_user_asset_library_index

Script to generate an asset index for a user-provided asset library.
generate_user_grasp_library_index

Script to generate a grasp index for a user-provided grasp library.

generate_user_asset_library_index

Script to generate an asset index for a user-provided asset library. Optionally, the script can also compute some unprovided metadata for the assets.

The asset library can have any directory structure, but asset XMLs should be named .xml. Their corresponding metadata JSONs should be named .json. If these assets have array-valued metadata (e.g. clip features), they should be stored in an npz file named .npz, keys are forward-slash-delimited representing the hierarchy of the metadata.

Functions:

Name Description
build_asset_index
hydrate_asset_metadata
main

build_asset_index 

build_asset_index(source_dir: Path, args) -> dict[str, UserAssetLibraryIndexEntry]
Source code in molmo_spaces/resources/generate_user_asset_library_index.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def build_asset_index(source_dir: Path, args) -> dict[str, UserAssetLibraryIndexEntry]:
    asset_index: dict[str, UserAssetLibraryIndexEntry] = {}

    for object_xml in sorted(source_dir.rglob("*.xml")):
        asset_id = object_xml.stem
        metadata_json = object_xml.with_suffix(".json")
        metadata_npz = object_xml.with_suffix(".npz")
        if not metadata_json.exists():
            if args.suppress_missing_metadata:
                print(
                    f"Warning: Expected metadata JSON for {object_xml} at {metadata_json}, but it was not found. Skipping."
                )
                continue
            raise FileNotFoundError(
                f"Expected metadata JSON for {object_xml} at {metadata_json}, but it was not found."
            )

        if args.compute_metadata:
            computed_metadata_json = metadata_json.with_name(f"{asset_id}_computed.json")
            if computed_metadata_json.exists() and not args.overwrite:
                print(
                    f"Computed metadata JSON for {object_xml} already exists at {computed_metadata_json}, skipping. Pass --overwrite to recompute."
                )
                continue

            with metadata_json.open() as f:
                metadata = json.load(f)

            hydrate_asset_metadata(object_xml, metadata)

            with computed_metadata_json.open("w") as f:
                json.dump(metadata, f, indent=4)
            metadata_path = computed_metadata_json
        else:
            metadata_path = metadata_json

        if asset_id in asset_index:
            prev = asset_index[asset_id].object_path
            raise ValueError(
                f"Duplicate asset id '{asset_id}' found at "
                f"{object_xml.relative_to(source_dir)} and {prev}."
            )

        asset_index[asset_id] = UserAssetLibraryIndexEntry(
            uid=asset_id,
            object_path=object_xml.relative_to(source_dir),
            metadata_path=metadata_path.relative_to(source_dir),
            metadata_npz_path=(
                metadata_npz.relative_to(source_dir) if metadata_npz.exists() else None
            ),
        )

    return dict(sorted(asset_index.items()))

hydrate_asset_metadata 

hydrate_asset_metadata(asset_path: Path, metadata: dict)
Source code in molmo_spaces/resources/generate_user_asset_library_index.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def hydrate_asset_metadata(asset_path: Path, metadata: dict):
    model = MjModel.from_xml_path(str(asset_path))
    data = MjData(model)
    mujoco.mj_fwdPosition(model, data)

    assert model.nbody == 2, "Expected a single body in the asset (other than worldbody)"
    body = model.body(1)

    if "mass" not in metadata:
        metadata["mass"] = body.mass.item()

    if "boundingBox" not in metadata:
        aabb_center, aabb_size = body_aabb(model, data, body.id, visible_only=False)
        # metadata bbox is centered at origin, so grow noncentered AABB to make centered
        aabb = 2 * (aabb_center + aabb_size / 2)
        metadata["boundingBox"] = {
            "x": aabb[0].item(),
            "y": aabb[1].item(),
            "z": aabb[2].item(),
        }

main 

main() -> None
Source code in molmo_spaces/resources/generate_user_asset_library_index.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "source_dir",
        type=Path,
        help="Asset library directory.",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        help="Output path for assets_index.json. Defaults to <source_dir>/assets_index.json.",
    )
    parser.add_argument(
        "--overwrite", action="store_true", help="Overwrite existing output file if it exists."
    )
    parser.add_argument(
        "--compute-metadata",
        action="store_true",
        help="Compute metadata for the assets, write to a new JSON file next to the original.",
    )
    parser.add_argument(
        "--suppress-missing-metadata",
        action="store_true",
        help="Suppress missing metadata errors, skipping instead.",
    )
    args = parser.parse_args()

    source_dir = args.source_dir.resolve()
    output_path: Path = (
        args.output.resolve() if args.output is not None else source_dir / "assets_index.json"
    )
    if output_path.exists() and not args.overwrite:
        raise FileExistsError(f"{output_path} already exists. Use --overwrite to overwrite.")

    asset_index = build_asset_index(source_dir, args)
    output_json = UserAssetLibraryIndex.dump_json(asset_index, indent=4).decode("utf-8")

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(output_json)

    print(f"Wrote {output_path} with {len(asset_index)} assets.")

generate_user_grasp_library_index

Script to generate a grasp index for a user-provided grasp library. Assumes the following conventions: - Nonarticulated grasp files are of the form //grasps.npz - Articulated grasp files are of the form //joint_grasps_.npz

Where
  • is the asset ID of the object the grasp is for
  • is the name of the robot/gripper the grasp is for (i.e. what it was collected with)
  • is the name of the joint the grasp is for, for articulated grasps

These conventions are required for this indexing script, but not by the rest of MolmoSpaces.

Functions:

Name Description
build_grasp_index
get_articulated_grasp_paths
get_grasp_paths
main

build_grasp_index 

build_grasp_index(source_dir: Path) -> UserGraspLibraryIndex
Source code in molmo_spaces/resources/generate_user_grasp_library_index.py 
61
62
63
64
65
66
67
def build_grasp_index(source_dir: Path) -> UserGraspLibraryIndex:
    grasp_paths = get_grasp_paths(source_dir)
    articulated_grasp_paths = get_articulated_grasp_paths(source_dir)
    return UserGraspLibraryIndex(
        grasp_paths=grasp_paths,
        articulated_grasp_paths=articulated_grasp_paths,
    )

get_articulated_grasp_paths 

get_articulated_grasp_paths(source_dir: Path) -> dict[str, dict[str, dict[str, Path]]]
Source code in molmo_spaces/resources/generate_user_grasp_library_index.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def get_articulated_grasp_paths(source_dir: Path) -> dict[str, dict[str, dict[str, Path]]]:
    articulated_grasp_paths: dict[str, dict[str, dict[str, Path]]] = {}

    for grasp_file in sorted(source_dir.rglob("joint_grasps_*.npz")):
        robot = grasp_file.parent.name
        uid = grasp_file.parent.parent.name
        joint_name = grasp_file.stem.removeprefix("joint_grasps_")
        assert joint_name

        rel_path = grasp_file.relative_to(source_dir)
        robot_dict = articulated_grasp_paths.setdefault(robot, {})
        uid_dict = robot_dict.setdefault(uid, {})
        uid_dict[joint_name] = rel_path

    # Keep output stable for easier diffs/review.
    sorted_articulated_grasp_paths = {
        robot: dict(sorted(uid_to_path.items()))
        for robot, uid_to_path in sorted(articulated_grasp_paths.items())
    }
    return sorted_articulated_grasp_paths

get_grasp_paths 

get_grasp_paths(source_dir: Path) -> dict[str, dict[str, Path]]
Source code in molmo_spaces/resources/generate_user_grasp_library_index.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def get_grasp_paths(source_dir: Path) -> dict[str, dict[str, Path]]:
    grasp_paths: dict[str, dict[str, Path]] = {}

    for grasp_file in sorted(source_dir.rglob("grasps.npz")):
        robot = grasp_file.parent.name
        uid = grasp_file.parent.parent.name

        rel_path = grasp_file.relative_to(source_dir)
        grasp_paths.setdefault(robot, {})[uid] = rel_path

    # Keep output stable for easier diffs/review.
    sorted_grasp_paths = {
        robot: dict(sorted(uid_to_path.items()))
        for robot, uid_to_path in sorted(grasp_paths.items())
    }
    return sorted_grasp_paths

main 

main() -> None
Source code in molmo_spaces/resources/generate_user_grasp_library_index.py 
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "source_dir",
        type=Path,
        help="Grasp library directory.",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        help="Output path for grasps_index.json. Defaults to <source_dir>/grasps_index.json.",
    )
    parser.add_argument(
        "--overwrite", action="store_true", help="Overwrite existing output file if it exists."
    )
    args = parser.parse_args()

    source_dir = args.source_dir.resolve()
    output_path: Path = (
        args.output.resolve() if args.output is not None else source_dir / "grasps_index.json"
    )
    if output_path.exists() and not args.overwrite:
        raise FileExistsError(f"{output_path} already exists. Use --overwrite to overwrite.")

    grasp_index = build_grasp_index(source_dir=source_dir)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(grasp_index.model_dump_json(indent=4))

    num_robots = len(grasp_index.grasp_paths)
    num_uids = sum(len(uid_to_path) for uid_to_path in grasp_index.grasp_paths.values())
    print(f"Wrote {output_path} with {num_robots} robots / {num_uids} grasp files.")