release repository

tests/interpolation_test.py

@@ -0,0 +1,67 @@
+#
+# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+#
+
+# Third Party
+import torch
+
+# CuRobo
+from curobo.types.base import TensorDeviceType
+from curobo.types.robot import JointState
+from curobo.util.trajectory import InterpolateType, get_batch_interpolated_trajectory
+
+
+def test_linear_interpolation():
+    b, h, dof = 1, 24, 1
+    raw_dt = 0.4
+    int_dt = 0.01
+    tensor_args = TensorDeviceType()
+    # initialize raw trajectory:
+    in_traj = JointState.zeros((b, h, dof), tensor_args)
+    in_traj.position = torch.zeros((b, h, dof), device=tensor_args.device)
+    in_traj.position[:, 1, :] = 0.1
+
+    in_traj.position[:, -2, :] = -0.01
+    in_traj.position[:, 10, :] = -0.01
+
+    in_traj.position[:, -1, :] = 0.01
+    in_traj.velocity = in_traj.position - torch.roll(in_traj.position, -1, dims=1)
+    in_traj.velocity[:, 0, :] = 0.0
+    in_traj.velocity[:, -1, :] = 0.0
+
+    max_vel = torch.ones((1, 1, dof), device=tensor_args.device, dtype=tensor_args.dtype)
+    max_acc = torch.ones((1, 1, dof), device=tensor_args.device, dtype=tensor_args.dtype) * 25
+    max_jerk = torch.ones((1, 1, dof), device=tensor_args.device, dtype=tensor_args.dtype) * 500
+
+    # create max_velocity buffer:
+    out_traj_gpu, _, _ = get_batch_interpolated_trajectory(
+        in_traj, int_dt, max_vel, max_acc=max_acc, max_jerk=max_jerk, raw_dt=raw_dt
+    )
+    #
+    out_traj_gpu = out_traj_gpu.clone()
+
+    out_traj_cpu, _, _ = get_batch_interpolated_trajectory(
+        in_traj,
+        int_dt,
+        max_vel,
+        raw_dt=raw_dt,
+        kind=InterpolateType.LINEAR,
+        max_acc=max_acc,
+        max_jerk=max_jerk,
+    )
+    assert (
+        torch.max(
+            torch.abs(out_traj_gpu.position[:, -5:, :] - out_traj_cpu.position[:, -5:, :])
+        ).item()
+        < 0.05
+    )
+
+
+# test_linear_interpolation()