add arrows

app.py

@@ -9,15 +9,6 @@ import random
 import math
 import hydra
 import numpy as np
-
-# import importlib
-# import torch
-# import torch.nn as nn
-# print(torch.__version__)
-# print(torch.version.cuda)
-# import torch.backends.cudnn
-# from torch.utils.data import DataLoader
-
 import glob
 import os
 import subprocess
@@ -32,69 +23,22 @@ import spaces
 from spaces import zero
 zero.startup()
 
-# def install_cuda_toolkit():
-#     # CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"
-#     # CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_linux.run"
-#     CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run"
-#     CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
-#     subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
-#     subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
-#     subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
-
-#     os.environ["CUDA_HOME"] = "/usr/local/cuda"
-#     os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
-#     os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
-#         os.environ["CUDA_HOME"],
-#         "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
-#     )
-#     # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range
-#     os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
-
-# install_cuda_toolkit()
-
-# gs_path = Path(__file__).parent / "src/third-party/diff-gaussian-rasterization-w-depth"
-# subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", str(gs_path)])
-# site.main()  # re-processes every *.pth in site-packages
-# importlib.invalidate_caches()
-# diff_gaussian_rasterization = importlib.import_module("diff_gaussian_rasterization")
-
-# subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-deps", "dgl", "-f", "https://data.dgl.ai/wheels/torch-2.4/cu124/repo.html"])
-
-# subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "-y", "torch", "torchvision", "torchaudio"])
-# subprocess.check_call([sys.executable, "-m", "pip", "install", "torch==2.4.0", "torchvision==0.19.0", "torchaudio==2.4.0", "--index-url", "https://download.pytorch.org/whl/cu124"])
-# os.system('conda install conda-forge::ffmpeg')
-
-# site.main()  # re-processes every *.pth in site-packages
-# importlib.invalidate_caches()
-# torch = importlib.import_module("torch")
-
 import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader
 import kornia
-# import warp as wp
 
 from diff_gaussian_rasterization import GaussianRasterizer
 from diff_gaussian_rasterization import GaussianRasterizationSettings as Camera
 
-# print(torch.__version__)
-# print(torch.version.cuda)
-
 import sys
 sys.path.insert(0, str(Path(__file__).parent / "src"))
 sys.path.append(str(Path(__file__).parent / "src" / "experiments"))
 
-# from pgnd.sim import Friction, CacheDiffSimWithFrictionBatch, StaticsBatch, CollidersBatch
-# from pgnd.material import PGNDModel
-# from pgnd.utils import Logger, get_root, mkdir
-
 from real_world.utils.render_utils import interpolate_motions
 from real_world.gs.helpers import setup_camera
 from real_world.gs.convert import save_to_splat, read_splat
 
-# from diff_gaussian_rasterization import GaussianRasterizer
-# from diff_gaussian_rasterization import GaussianRasterizationSettings as Camera
-
 root = Path(__file__).parent / "src" / "experiments"
 
 def make_video(
@@ -172,6 +116,7 @@ class DynamicsVisualizer:
         self.k_wgt = 16  # knn for weights
         self.with_bg = True
         self.render_gripper = True
+        self.render_direction = True
         self.verbose = False
 
         self.dt_base = cfg.sim.dt
@@ -840,6 +785,65 @@ class DynamicsVisualizer:
                 opa = torch.cat([opa, g_opacities], dim=0)
                 scales = torch.cat([scales, g_scales], dim=0)
 
+            if self.render_direction:
+                gripper_direction = self.gripper_direction
+                gripper_direction = gripper_direction / (torch.norm(gripper_direction, dim=-1, keepdim=True) + 1e-10)  # normalize
+
+                R = self.preprocess_metadata['R']
+                # model frame to data frame
+                direction = gripper_direction @ R.T
+
+                n_grippers = 1
+                N = 200
+                length = 0.2
+                kk = 5
+                xyz_test = torch.zeros((n_grippers, N + N // kk + N // kk, 3), device=self.torch_device, dtype=xyz.dtype)
+
+                gripper_now_inv_xyz = self.inverse_preprocess_gripper(self.state['prev_key_pos'][None].clone())
+                gripper_now_inv_rot = torch.eye(3, device=self.torch_device).unsqueeze(0).repeat(n_grippers, 1, 1)
+
+                center_point = torch.tensor([0.0, 0.0, 0.10], device=self.torch_device, dtype=xyz.dtype).reshape(1, 3)  # center point in gripper frame
+                gripper_center_inv_xyz = gripper_now_inv_xyz + \
+                        torch.einsum('ijk,ik->ij', gripper_now_inv_rot, center_point)  # (n_grippers, 3)
+
+                for i in range(N):
+                    offset = i / N * length * direction
+                    xyz_test[:, i] = gripper_center_inv_xyz + offset
+                
+                if direction[0, 2] < 0.9 and direction[0, 2] > -0.9:  # not vertical
+                    direction_up = -direction + torch.tensor([0.0, 0.0, 0.5], device=self.torch_device, dtype=xyz.dtype)
+                    direction_up = direction_up / (torch.norm(direction_up, dim=-1, keepdim=True) + 1e-10)  # normalize
+                    direction_down = -direction + torch.tensor([0.0, 0.0, -0.5], device=self.torch_device, dtype=xyz.dtype)
+                    direction_down = direction_down / (torch.norm(direction_down, dim=-1, keepdim=True) + 1e-10)  # normalize
+                else:
+                    direction_up = -direction + torch.tensor([0.0, 0.5, 0.0], device=self.torch_device, dtype=xyz.dtype)
+                    direction_up = direction_up / (torch.norm(direction_up, dim=-1, keepdim=True) + 1e-10)  # normalize
+                    direction_down = -direction + torch.tensor([0.0, -0.5, 0.0], device=self.torch_device, dtype=xyz.dtype)
+                    direction_down = direction_down / (torch.norm(direction_down, dim=-1, keepdim=True) + 1e-10)  # normalize
+
+                for i in range(N, N + N // kk):
+                    offset = length * direction + (i - N) / N * length * direction_up
+                    xyz_test[:, i] = gripper_center_inv_xyz + offset
+                
+                for i in range(N + N // kk, N + N // kk + N // kk):
+                    offset = length * direction + (i - N - N // kk) / N * length * direction_down
+                    xyz_test[:, i] = gripper_center_inv_xyz + offset
+                
+                color_test = torch.zeros_like(xyz_test, device=self.torch_device, dtype=xyz.dtype)
+                color_test[:, :, 0] = 255 / 255  # red
+                color_test[:, :, 1] = 80 / 255  # green
+                color_test[:, :, 2] = 110 / 255  # blue
+                quat_test = torch.zeros((n_grippers, N + N // kk + N // kk, 4), device=self.torch_device, dtype=xyz.dtype)
+                quat_test[:, :, 0] = 1.0  # identity quaternion
+                opa_test = torch.ones((n_grippers, N + N // kk + N // kk, 1), device=self.torch_device, dtype=xyz.dtype)
+                scales_test = torch.ones((n_grippers, N + N // kk + N // kk, 3), device=self.torch_device, dtype=xyz.dtype) * 0.002
+                
+                xyz = torch.cat([xyz, xyz_test.reshape(-1, 3)], dim=0)
+                rgb = torch.cat([rgb, color_test.reshape(-1, 3)], dim=0)
+                quat = torch.cat([quat, quat_test.reshape(-1, 4)], dim=0)
+                opa = torch.cat([opa, opa_test.reshape(-1, 1)], dim=0)
+                scales = torch.cat([scales, scales_test.reshape(-1, 3)], dim=0)
+
             # normalize
             quat = torch.nn.functional.normalize(quat, dim=-1)
 
@@ -944,6 +948,7 @@ class DynamicsVisualizer:
         num_particles = self.cfg.sim.n_particles
         self.load_scaniverse(in_dir)
         self.init_model(batch_size, num_steps, num_particles, ckpt_path=None)
+        self.render_direction = False
 
         params = self.preprocess_gs(self.params)
         if self.with_bg:
@@ -1079,6 +1084,7 @@ class DynamicsVisualizer:
         batch_size = 1
         num_particles = self.cfg.sim.n_particles
         self.init_model(batch_size, num_steps, num_particles, ckpt_path=None)
+        self.render_direction = True
 
         # im_list = []
         for i in range(num_steps):
@@ -1090,6 +1096,8 @@ class DynamicsVisualizer:
             if self.verbose:
                 print('command:', command.cpu().numpy().tolist())
             
+            self.gripper_direction = command.clone()
+
             assert self.state['sub_pos'] is None
 
             if self.state['sub_pos'] is None:

src/experiments/log/gs/temp/gs_pred.splat

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:111f09eae16a77a8e63af83f4adcb78c21ac301113a18fae10bcaa89432aff5f
-size 7675392
+oid sha256:753a59d8cb6fb82e29d233db6865e6c6a87fabc5b469619e4310d3fbee619f6c
+size 7684352