traj_opt/display_traj.py at main · Caltech-AMBER/traj_opt · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
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
97
98
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
#!/usr/bin/env python3
"""
Script to display trajectory from YAML files.
Supports both flat ground and stair configurations.
"""

import os
import numpy as np
import argparse
import yaml

from src.g1_robot import G1Robot
from src.trajectory_manager import TrajectoryManager
from src.hybrid_systems.config_parser import ConfigParser


def main():
    """Display trajectory from YAML file."""

    # Parse command line arguments
    parser = argparse.ArgumentParser(
        description='Display G1 Robot Trajectory from YAML file'
    )
    parser.add_argument(
        '--solution-path', '-t',
        required=True,
        help='Path to trajectory YAML file to display'
    )

    parser.add_argument(
        '--config-file', '-c',
        help='Path to configuration file (required for loading trajectory)'
    )
    parser.add_argument(
        '--step-height', '-sh',
        type=float,
        help='Step height in meters (for stair visualization)'
    )
    parser.add_argument(
        '--save-video', '-s',
        action='store_true',
        help='Save animation to video file'
    )
    parser.add_argument(
        '--video-filename',
        default="trajectory_animation.mp4",
        help='Output video filename (default: trajectory_animation.mp4)'
    )
    args = parser.parse_args()

    # --------------------------------- #
    # ---------- Setup robot ---------- #
    # --------------------------------- #
    repo_root = os.environ.get("AMBER_TRAJ_OPT")
    urdf_path = repo_root + "/rsc/g1/g1_21j.urdf"
    mesh_dir = repo_root + "/rsc/g1"

    # Use G1Robot instead of base Robot class
    g1 = G1Robot(urdf_path=urdf_path, mesh_dir=mesh_dir)

    # ------------------------------------------ #
    # ---------- Load Configuration ---------- #
    # ------------------------------------------ #
    if not args.config_file:
        print("Error: --config-file is required for loading trajectory")
        return

    config_parser = ConfigParser(g1)
    config = config_parser.load_config(args.config_file)
    print(f"Loaded configuration from {args.config_file}")

    # ------------------------------------------ #
    # ---------- Stair Configuration ---------- #
    # ------------------------------------------ #
    step_height = args.step_height
    if step_height is not None:
        print(f"Using step height: {step_height:.3f}m")

    # ------------------------------------------ #
    # ---------- Load Trajectory ---------- #
    # ------------------------------------------ #
    traj_manager = TrajectoryManager(g1, config)


    trajectory_data = traj_manager.load_trajectory(args.solution_path)
    print(f"Successfully loaded trajectory from {args.solution_path}")


    # ---------------------------------------- #
    # ---------- Display Trajectory ---------- #
    # ---------------------------------------- #
    for domain_name, domain_data in trajectory_data.items():
        N = domain_data['x'].shape[1]

        # Extract time from auxiliary variables (first element is time T)
        S = domain_data['S']

        if isinstance(S, (list, np.ndarray)) and len(S) > 0:
            time_value = float(S[0])
        else:
            time_value = float(S)
        dt = time_value / N

        dt_vec = np.ones(N) *dt
        print(f"Domain {domain_name}: {N} nodes, dt={dt:.6f}s, total_time={time_value:.6f}s")

        # Initialize visualizer
        g1.display_configuration(domain_data['x'][:g1.nq, 0])

        # Add stair boxes if step height is provided
        if step_height is not None:
            print(f"Adding stair box visualization with height {step_height:.3f}m...")
            g1.add_three_stair_boxes(stair_height=step_height, stair_config=config['stair_config'])

        # Display the trajectory
        print("Displaying trajectory...")
        g1.display_trajectory(domain_data['x'], dt_vec)

        # Optionally save the animation to a video file
        if args.save_video:
            print(f"Saving animation to {args.video_filename}...")
            g1.save_trajectory_animation(
                domain_data['x'], dt, args.video_filename
            )
            print(f"Animation saved to {args.video_filename}")


if __name__ == "__main__":
    main()