Implementation:Isaac sim IsaacGymEnvs DataTree

Knowledge Sources	IsaacGymEnvs
Domains	Data_Management, Motion_Sampling
Last Updated	2026-02-15 11:00 GMT

Overview

DataTree implements a hierarchical motion data sampler that organizes motion capture clips into a tree structure by metadata categories and uses a water floating algorithm for balanced probability-based selection across all branches.

Description

The data_tree class represents a hierarchical tree where motion capture data is organized by metadata attributes such as behavior, direction, type, and style. Each internal node corresponds to a category level, and leaf nodes store the actual motion clip references along with their lengths. The tree is built incrementally via add_node(), which takes a hierarchy of category labels and inserts data at the appropriate leaf position, creating intermediate nodes as needed.

The water_floating_algorithm() method implements a balanced sampling strategy that ensures diverse coverage across all categories. At each level, it selects the child with the minimum accumulated "picked" count (similar to water finding the lowest point), descends recursively to a leaf, and marks the selected data as depleted. This ensures that all categories are sampled proportionally before any category is revisited, preventing the training set from being dominated by categories with more data. Nodes named "mix" are excluded from selection by setting their depleted flag to infinity.

The assign_probability() method provides an alternative sampling approach by distributing a total probability mass equally among children at each level, resulting in uniform per-trajectory probabilities within each category. The parse_dataset() function ties everything together: it shuffles the mocap data list, builds the tree, uses the water floating algorithm to select a training set up to a target size, and saves the resulting train/test splits as TSV files along with a JSON info file containing the verbose tree structure with pick statistics.

Usage

Use this module during the data preparation phase of AMP training to create balanced train/test splits from a collection of motion capture clips organized by behavioral categories. The water floating algorithm is particularly useful when the dataset contains uneven amounts of data per category (e.g., many walking clips but few jumping clips) and you want to ensure all behaviors are represented in the training set. The probability assignment method can be used for weighted sampling during training.

Code Reference

Source Location

• Repository: IsaacGymEnvs
• File: isaacgymenvs/tasks/amp/utils_amp/data_tree.py
• Lines: 36-222

Signature

class data_tree(object):
    def __init__(self, name: str): ...
    def add_node(self, dict_hierarchy: list, mocap_data: list): ...
    def summarize_length(self) -> float: ...
    def to_dict(self, verbose: bool = False) -> dict: ...
    def water_floating_algorithm(self) -> Tuple[str, dict]: ...
    def assign_probability(self, total_prob: float) -> Tuple[list, list]: ...

    @property
    def name(self) -> str: ...
    @property
    def picked(self) -> list: ...
    @property
    def total_length(self) -> float: ...


def parse_dataset(env, args): ...
def save_tsv_files(base_dir: str, name: str, data_dict: list): ...

Import

from isaacgymenvs.tasks.amp.utils_amp.data_tree import data_tree, parse_dataset

I/O Contract

Inputs

Name	Type	Required	Description
name	str	Yes	Name identifier for the tree node (e.g., "root", category label)
dict_hierarchy	list[str]	Yes (add_node)	Ordered list of category labels defining the path from root to leaf
mocap_data	list	Yes (add_node)	Motion data entry: [clip_path, length, ...metadata]
verbose	bool	No	If True, to_dict() includes pick counts and lengths in keys (default: False)
total_prob	float	Yes (assign_probability)	Total probability mass to distribute across the tree
env	object	Yes (parse_dataset)	Environment object with motion_info, motion, and get_all_motion_length()
args	object	Yes (parse_dataset)	Arguments with parse_dataset_train percentage and mocap_list_file path

Outputs

Name	Type	Description
chosen_data	str	Path to the selected motion clip (from water_floating_algorithm)
data_info	dict	Selection metadata with 'name' (hierarchy path), 'length', and 'all_depleted' flag
total_length	float	Sum of all motion clip lengths in the subtree (from summarize_length)
data_dict	dict/list	Nested dictionary or list representation of the tree (from to_dict)
leaves	list	Flat list of all leaf motion clip paths (from assign_probability)
probs	list	Corresponding probability weights for each leaf (from assign_probability)
train TSV	file	Tab-separated training set file saved to disk (from parse_dataset)
test TSV	file	Tab-separated test set file saved to disk (from parse_dataset)

Usage Examples

from isaacgymenvs.tasks.amp.utils_amp.data_tree import data_tree, parse_dataset

# Build a hierarchical motion data tree
tree = data_tree("root")

# Add motion clips with category hierarchy: [behavior, direction]
# mocap_data format: [clip_path, length, behavior, direction]
tree.add_node(
    ["walk", "forward"],
    ["data/walk_fwd_01.npy", 150, "walk", "forward"]
)
tree.add_node(
    ["walk", "forward"],
    ["data/walk_fwd_02.npy", 200, "walk", "forward"]
)
tree.add_node(
    ["run", "forward"],
    ["data/run_fwd_01.npy", 100, "run", "forward"]
)
tree.add_node(
    ["walk", "backward"],
    ["data/walk_bwd_01.npy", 180, "walk", "backward"]
)

# Summarize total lengths
total = tree.summarize_length()
print(f"Total motion length: {total}")

# Use water floating algorithm to select balanced training data
selected_clips = []
for _ in range(3):
    clip, info = tree.water_floating_algorithm()
    selected_clips.append(clip)
    print(f"Selected: {clip}, from: {info['name']}, length: {info['length']}")

# Get verbose tree structure with pick statistics
tree_dict = tree.to_dict(verbose=True)

# Assign uniform sampling probabilities
leaves, probs = tree.assign_probability(1.0)
for leaf, prob in zip(leaves, probs):
    print(f"{leaf}: probability = {prob:.4f}")