trajectory.py

#!/usr/bin/env python

import argparse
import sys
from typing import Any

import numpy as np
import yaml
from idyntree import bindings as iDynTree

from excitation.postureOptimizer import PostureOptimizer
from excitation.trajectoryGenerator import FixedPositionTrajectory, PulsedTrajectory, simulateTrajectory
from excitation.trajectoryOptimizer import TrajectoryOptimizer
from identification.model import Model
from identifier import Identification

parser = argparse.ArgumentParser(description="Generate excitation trajectories, save to <filename>.")
parser.add_argument(
    "--filename",
    type=str,
    help="the filename to save the trajectory to, otherwise <model>.trajectory.npz",
)
parser.add_argument("--config", required=True, type=str, help="use options from given config file")
parser.add_argument("--model", required=True, type=str, help="the file to load the robot model from")
parser.add_argument(
    "--model_real",
    required=False,
    type=str,
    help='the file to load the "real" robot model from',
)
parser.add_argument("--world", required=False, type=str, help="the file to load world links from")
args = parser.parse_args()

with open(args.config) as stream:
    try:
        config = yaml.load(stream, Loader=yaml.SafeLoader)
    except yaml.YAMLError as exc:
        print(exc)

config["urdf"] = args.model
config["urdf_real"] = args.model_real
if config["useStaticTrajectories"] and not config["urdf_real"]:
    print("When optimizing static postures, need model_real argument!")
    sys.exit()
config["jointNames"] = iDynTree.StringVector([])
if not iDynTree.dofsListFromURDF(config["urdf"], config["jointNames"]):
    sys.exit()
config["num_dofs"] = len(config["jointNames"])
config["skipSamples"] = 0


def main():
    # save either optimized or random trajectory parameters to filename
    if args.filename:
        traj_file = args.filename
    else:
        traj_file = config["urdf"] + ".trajectory.npz"

    if config["optimizeTrajectory"]:
        # find trajectory params by optimization
        old_sim = config["simulateTorques"]
        config["simulateTorques"] = True
        model = Model(config, config["urdf"])
        if config["useStaticTrajectories"]:
            old_gravity = config["identifyGravityParamsOnly"]
            idf = Identification(
                config,
                config["urdf"],
                config["urdf_real"],
                measurements_files=None,
                regressor_file=None,
                validation_file=None,
            )
            trajectoryOptimizer: PostureOptimizer | TrajectoryOptimizer = PostureOptimizer(
                config, idf, model, simulation_func=simulateTrajectory, world=args.world
            )
            config["identifyGravityParamsOnly"] = old_gravity
        else:
            idf = Identification(
                config,
                config["urdf"],
                urdf_file_real=None,
                measurements_files=None,
                regressor_file=None,
                validation_file=None,
            )
            trajectoryOptimizer = TrajectoryOptimizer(
                config, idf, model, simulation_func=simulateTrajectory, world=args.world
            )

        trajectory = trajectoryOptimizer.optimizeTrajectory()
        config["simulateTorques"] = old_sim
    else:
        # use some random params
        print("no optimized trajectory found, generating random one")
        trajectory = PulsedTrajectory(config["num_dofs"], use_deg=config["useDeg"]).initWithRandomParams()
        print(f"a {[t_a.tolist() for t_a in trajectory.a]}")
        print(f"b {[t_b.tolist() for t_b in trajectory.b]}")
        print(f"q {trajectory.q.tolist()}")
        print(f"nf {trajectory.nf.tolist()}")
        print(f"wf {trajectory.w_f_global}")

    print(f"Saving found trajectory to {traj_file}")

    if config["useStaticTrajectories"]:
        # always saved with rad angles
        if not isinstance(trajectory, FixedPositionTrajectory) or trajectory.angles is None:
            raise RuntimeError("Expected initialized FixedPositionTrajectory for static trajectories")
        np.savez(traj_file, static=True, angles=np.array(trajectory.angles, dtype=object))
    else:
        # TODO: remove degrees option
        if not isinstance(trajectory, PulsedTrajectory):
            raise RuntimeError("Expected PulsedTrajectory for non-static trajectories")
        # a and b may be ragged (per-joint nf), so save as object arrays
        a_arr = np.array(trajectory.a, dtype=object)
        b_arr = np.array(trajectory.b, dtype=object)
        save_dict: dict[str, Any] = {
            "use_deg": trajectory.use_deg,
            "static": False,
            "a": a_arr,
            "b": b_arr,
            "q": trajectory.q,
            "nf": trajectory.nf,
            "wf": trajectory.w_f_global,
        }
        # save joint limits if using bounded trajectory mode
        if trajectory.joint_limits is not None:
            save_dict["joint_limits"] = np.array(trajectory.joint_limits)
        np.savez(traj_file, **save_dict)


if __name__ == "__main__":
    main()