Tool-use Environment Benchmark built with Drake Simulation
import os
import cv2
import sys


for file in sorted(os.listdir("textures")):
    new_path = os.path.join("new_textures", file)
    old_path = os.path.join("textures", file)
    img = cv2.imread(old_path)
    print(new_path.replace(".jpg", ".png"))
    cv2.imwrite(new_path.replace(".jpg", ".png"), img)

assets/textures/change_image_format.py
liruiw/Fleet-Tools
# --------------------------------------------------------
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------

import os

import numpy as np
from core.utils.utils import *
import IPython
import time
from core.agent.replay_memory import ReplayMemory
import os.path as osp
import ray
import gym


# from core.utils.torch_utils import *


class Agent(object):
    """
    A general agent class
    """

    def __init__(self, config, logdir, agent_id=-1, args=None):
        self.env_name = config.env_name
        self.logdir = logdir
        self.device = "cuda"
        self.agent_id = agent_id
        self.image_input = config.task.env.use_image
        self.keypoint_input = config.task.env.use_keypoint

        self.log_freq = config.log_freq
        self.batch_size = config.train.learn.minibatches
        self.latent_size = config.train.policy.latent_size

        self.observation_space = gym.spaces.Box(low=-1, high=1, shape=config.train.numObservations)
        self.action_space = PandaTaskSpace6D(*config.task.env.action_scale)  # action_space(config.task.env.numActions)
        self.config = config
        self.train_iter = 0
        self.train_time = 0
        self.memory = ReplayMemory(config.train.replaybuffer_size, config.train.seed, config)

        if not config.run_expert:  # no network is needed
            self.init_policies()

    def get_self(self):
        return self

    def get_exp_stats(self):
        """get experiment details"""
        return self.memory.episode_position

    def select_action(self, state):
        """policy generates action. not used."""
        extra = {}
        self.set_mode(True)
        state = self.select_state_test(state)
        actions = self.policy.sample(state)[0][0].detach().cpu().numpy()
        return actions, extra

    def train(self, step):
        """training loop for the policy. not used."""
        start_time = time.time()
        self.set_mode(test=False)

        for _ in range(self.config.train.train_frequency):
            self.train_iter += 1
            iter_start_time = time.time()
            if not hasattr(self, "memory_dataloader_iter"):
                data = self.memory.sample(batch_size=self.batch_size)
            else:
                data = next(self.memory_dataloader_iter)
            self.data_time = time.time() - iter_start_time

        self.train_time += time.time() - start_time
        self.get_lr()
        info = {k: float(getattr(self, k)) for k in self.info if hasattr(self, k)}
        return info

    def train_step(self, memory, batch_size):
        pass

    def step(self):
        """step the optimizer and scheduler"""
        self.log_stat()

    def log_stat(self):
        """log grad and param statistics for tensorboard"""
        self.train_batch_size = len(self.state_batch)

    def set_memory_position(self, start_episode_position):
        self.memory.episode_position = start_episode_position

    def get_memory_position(self):
        return self.memory.size

    # functions for ray remote calls
    def append_traj_data(self, traj_infos, scene_info, save_demo=False):
        """add a full trajectory to the replay buffer and then save it to disk"""
        episode_idx = self.memory.episode_position + 1
        success = True
        status = False

        try:
            start_time = time.time()
            for t in traj_infos:
                if t[0] is not None:
                    obs, reward, done, info = t[0]
                    info["reward"] = reward
                    info["obs"] = obs
                    info["done"] = done
                success = self.memory.push(**info)
                if not success:
                    print("push failed")
                    status = False
                    break

            if reward > 0 and save_demo and success:
                # save an episode in the buffer to local disk
                print(f"saving traj data: {episode_idx}")
                demo_path = "assets/demonstrations/{}".format(
                    self.config.task.tool_class_name + "for" + self.config.task.task_name + self.config.save_demo_suffix
                )
                start_time = time.time()
                status = self.memory.save_episode(demo_path, scene_info, episode_idx)
                curr_pos = self.memory.position
                self.memory.episode_end_map[curr_pos - len(traj_infos) : curr_pos] = curr_pos - 1
                print(f"save episode time: {time.time() - start_time:.3f} to {demo_path}")

        except Exception as e:
            print(e)
            print("saving memory failed:", episode_idx)
        return status


@ray.remote(num_cpus=1, num_gpus=0.1)
class AgentWrapper(Agent):
    pass

core/agent/agent.py
liruiw/Fleet-Tools
"""
Built on on SAC implementation from
https://github.com/pranz24/pytorch-soft-actor-critic
"""
import random
import numpy as np
from operator import itemgetter
import IPython
import ray
from core.utils.utils import *
import os
import cv2
import json
from env.env_util import color_print


class ReplayMemory:
    """
    Replay buffer class
    """

    def __init__(self, capacity, seed, config):
        random.seed(seed)
        self.capacity = capacity

        self.task_config = config.task
        self.train_config = config.train

        # self.buffer = []
        for idx, key in enumerate(key_list):
            if key == "obs" and not self.task_config.env.use_image:
                self.obs = np.zeros(self.capacity)
            else:
                setattr(
                    self,
                    key,
                    np.zeros((self.capacity, *key_shape_list[idx]), dtype=np.float32),
                )
        self.position = 0
        self.size = 0
        self.episode_position = 0
        self.buffer_idx = 200
        self.pos_idx = np.zeros(self.capacity)
        self.seed = seed
        self.episode_positions = np.zeros(self.capacity)
        self.episode_rewards = np.zeros(self.capacity)
        self.episode_end_map = np.zeros(self.capacity)

    def __len__(self):
        return self.upper_idx()

    def reset(self):
        """reset the replay buffer"""
        for idx, key in enumerate(key_list):
            setattr(
                self,
                key,
                np.zeros((self.capacity, *key_shape_list[idx]), dtype=np.float32),
            )
        self.position = 0
        self.size = 0
        self.episode_position = 0
        self.pos_idx = np.zeros(self.capacity)
        self.episode_positions = np.zeros(self.capacity)
        self.episode_rewards = np.zeros(self.capacity)

    def push(self, **kwargs):
        """push a single data item into the replay buffer"""
        for idx, key in enumerate(key_list):
            if key == "obs":
                if type(kwargs[key]) == list and len(kwargs[key]) > 0:
                    kwargs[key] = np.concatenate(kwargs[key], axis=-1)
                if not self.task_config.env.use_image:  # dummy
                    kwargs[key] = 0.0
            try:
                getattr(self, key)[self.position] = kwargs[key]
            except Exception as e:  # would be overwrite
                print(e)
                print(
                    "push to memory failure:",
                    self.position,
                    key,
                    getattr(self, key).shape,
                    kwargs[key].shape,
                )
                return False

        timestep = self.timestep[self.position]
        reward = self.reward[self.position]
        self.pos_idx[self.position] = float(reward > 0.5)

        if timestep == 0:
            # print("increments index.")
            self.episode_position += 1

        self.episode_positions[self.position] = self.episode_position
        self.position = (self.position + 1) % self.capacity
        self.size = min(self.size + 1, self.capacity)
        return True

    def save_scene_descriptions(self, path, info_dict):
        """save object and tool related poses into a scene descriptions"""
        arg_json = os.path.join(path, "scene_descriptions.json")
        info_dict["tool_pose"] = np.array(info_dict["tool_pose"]).tolist()
        info_dict["object_pose"] = np.array(info_dict["object_pose"]).tolist()
        info_dict["robot_joints"] = np.array(info_dict["robot_joints"]).tolist()

        with open(arg_json, "w") as f:
            json.dump(info_dict, f, indent=4, sort_keys=True)

    def load_offline(self, saved_path):
        """load offline data. Can add train/test split here."""
        for idx, key in enumerate(key_list):
            delattr(self, key)  # reset attributes

        episode_dirs = sorted(os.listdir(saved_path), key=lambda s: int(s.split("_")[1]))
        print("number of saved trajectories:", len(episode_dirs))

        traj_datas = {}
        for path in episode_dirs:
            episode_path = os.path.join(saved_path, path)
            if "traj_data.npz" in os.listdir(episode_path):
                traj_data = np.load(os.path.join(episode_path, "traj_data.npz"))
                try:
                    traj_datas = {
                        k: np.concatenate((traj_datas[k], v), axis=0) if k in traj_datas else v
                        for k, v in traj_data.items()
                    }
                except Exception as e:  # would be overwrite
                    print(e)
                    print("load trajectory folder: {} failed".format(episode_path))

        for k in key_list:
            if k in traj_datas:
                setattr(self, k, traj_datas[k])

        # compute episode index
        self.position = len(self.state) - 1
        self.size = self.position + 1
        self.pos_idx = np.zeros(len(self.state))
        self.episode_positions = traj_datas["episode_positions"]
        self.episode_position = traj_datas["episode_positions"][-1]
        self.episode_rewards = np.zeros(len(self.state))
        self.episode_dirs = episode_dirs

        print(
            "total number of data:",
            self.position,
            key_list,
            traj_datas.keys(),
            [len(traj_datas[k]) for k in key_list if k in traj_datas],
        )

    def save_episode(self, saved_path, scene_info, episode_idx):
        """saving an episode to the disk"""
        prev_episode_idx = np.where(self.episode_positions[: self.position] == episode_idx)[0]

        # skip the first episode due to boundary issues
        if len(prev_episode_idx) == 0 or prev_episode_idx[0] == 0 or prev_episode_idx[-1] == self.capacity - 1:
            # check the last episode step and first episode step
            print("skip empty traj")
            return False

        saved_path = os.path.join(saved_path, "episode_{}".format(episode_idx))
        mkdir_if_missing(saved_path)

        color_print(
            (f"saving episode {prev_episode_idx} | path: {saved_path}"),
            "light_slate_blue",
        )
        self.save_scene_descriptions(saved_path, scene_info)

        # separate out image and the rest
        try:
            images = self.obs[prev_episode_idx]
            data = {k: getattr(self, k)[prev_episode_idx] for k in key_list if k != "obs"}
            data["episode_positions"] = self.episode_positions
            data["obs"] = [f"{saved_path}/{i}" for i in range(len(prev_episode_idx))]
            np.savez(os.path.join(saved_path, "traj_data.npz"), **data)

            if self.task_config.env.use_image:
                for i, img in enumerate(images):
                    overhead_img, wrist_img = img[..., :5], img[..., 5:]
                    cv2.imwrite(
                        f"{saved_path}/{i}_wrist_color.png",
                        (wrist_img[:, :, [2, 1, 0]] * 255).astype(np.uint8),
                    )
                    cv2.imwrite(
                        f"{saved_path}/{i}_wrist_depth.png",
                        (wrist_img[..., [3]] * 1000).astype(np.uint16),
                    )
                    cv2.imwrite(f"{saved_path}/{i}_wrist_mask.png", wrist_img[..., [4]])
                    cv2.imwrite(
                        f"{saved_path}/{i}_overhead_color.png",
                        (overhead_img[:, :, [2, 1, 0]] * 255).astype(np.uint8),
                    )
                    cv2.imwrite(
                        f"{saved_path}/{i}_overhead_depth.png",
                        (overhead_img[..., [3]] * 1000).astype(np.uint16),
                    )
                    cv2.imwrite(f"{saved_path}/{i}_overhead_mask.png", overhead_img[..., [4]])
            return True

        except Exception as e:  # would be overwrite
            print(e)
            print("writing image failed")
            return False

    def save(self, save_dir="."):
        """Saves the current buffer to memory."""
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)

        s = time.time()
        save_dict = {}
        save_attrs = []
        for name in save_attrs:
            save_dict[name] = getattr(self, name)

        np.savez(os.path.join(save_dir, self.save_data_name), **save_dict)

    def __len__(self):
        return len(self.state)


@ray.remote(num_cpus=1)
class ReplayMemoryWrapper(ReplayMemory):
    pass

core/agent/replay_memory.py
liruiw/Fleet-Tools
import datetime
import os
import os.path as osp
import numpy as np
import cv2

from typing import Dict
import time
from env import *
from dotmap import DotMap
import ray
from tqdm import trange
from core.agent.agent import Agent, AgentWrapper
from collections import deque

from core.utils import *
from visdom import Visdom
from core.expert.expert import AnalyticExpert
from core.expert.human_expert import HumanExpert
import copy
import gc


def parse_task(cfg):
    register_env(cfg.task.name)
    env = make_env(cfg.task.name, task_config=cfg.task)
    return env


class ActorWrapper:
    """
    Wrapper class for agent training and logging
    Fleet Agent is the model for the fleet
    """

    def __init__(self, config, fleet_agent_id, robot_id):
        self.config = copy.deepcopy(config)
        self.robot_id = robot_id

        self.experiment_setup()
        if config.train.local_agent:
            print("build local agent")
            self.agent = eval(config.agent)(self.config, None, robot_id, self.config)

            if len(self.config.pretrained) > 0:
                # loading pretrained
                self.agent.load(self.config.pretrained, "fleet", self.config.load_pretrained_mode)
            self.agent.memory.episode_position = self.config.start_episode_position

        self.config.task.random_seed = self.config.start_episode_position + robot_id

        # setting the memory episode
        self.local_train_steps = 0
        self.fleet_agent_id = fleet_agent_id

        self.active = True
        self.parallel = self.config.parallel
        _, _, sim_fail, _ = self.env.reset()
        while sim_fail:
            _, _, sim_fail, _ = self.env.reset()

        self.local_episode = 0
        self.success_episode = 0
        self.plan_success_episode = 0
        self.avg_rewards = deque([], maxlen=1000)

        if self.config.visdom and len(self.config.task.env.numObservations) > 1:
            self.vis = Visdom(port=8097)
            _, w, h = self.config.task.env.numObservations
            self.win_id_wrist = self.vis.image(np.zeros([3, w, h]))
            self.win_id_overhead = self.vis.image(np.zeros([3, w, h]))
        self.curr_traj_for_replaybuffer = []

    def experiment_setup(self):
        self.env = parse_task(self.config)

    def reset_sim(self):
        """rebuild the manipulation station"""
        if hasattr(self, "env"):
            self.env.cleanup_meshcat()

        # c++ backend is leaking
        del self.env
        gc.collect()
        self.experiment_setup()

    def add_trajectory_transition(self, scene_info):
        """add the full trajectory into the memory"""
        save_demo = True
        if self.parallel:
            # should be blocking
            success = ray.get(
                self.fleet_agent_id.append_traj_data.remote(
                    self.curr_traj_for_replaybuffer,
                    scene_info,
                    save_demo and (self.local_episode < self.config.demo_number),
                )
            )
        else:
            success = self.fleet_agent_id.append_traj_data(
                self.curr_traj_for_replaybuffer,
                scene_info,
                save_demo and (self.local_episode < self.config.demo_number),
            )

        self.curr_traj_for_replaybuffer = []
        return success

    def reset(self):
        """reset the environment"""
        return self.env.reset()

    def visdom_visualize(self, ret):
        """publish rendered images to visdom"""
        if not len(self.config.task.env.numObservations) > 1:
            return
        head_cam_img = ret[0][1]
        wrist_cam_img = ret[0][0]  #
        self.vis.image(wrist_cam_img.transpose([2, 0, 1])[:3], win=self.win_id_wrist)
        self.vis.image(head_cam_img.transpose([2, 0, 1])[:3], win=self.win_id_overhead)

    def get_expert(self):
        """get the current expert type"""
        return eval(self.config.expert)(self.env, self.config, self.config, self.robot_id)

    def get_flags(self, time):
        """get different booleans for the current step"""
        apply_dart = (
            (time > 1)
            and (time < self.config.task.task_completion_time - 3)
            and (np.random.uniform() < self.config.task.inject_noise_to_expert)
        )
        return apply_dart

    def reset_task_episode(self, info):
        self.curr_traj_for_replaybuffer = []
        return {
            "tool_class_name": info["tool_class_name"],
            "tool_name": info["tool_name"],
            "object_name": info["tool_object_name"],
            "tool_pose": [info["tool_pose"].tolist()],
            "object_pose": [info["object_pose"].tolist()],
            "robot_joints": [info["joint_pos"].tolist()],
        }

    def rollout(self, explore=False, dagger=False, test=False):
        """policy rollout and save data"""
        self.reset_sim()  # reset the entire simulation

        episode_info = {}
        start_time = time.time()

        # more checks on this
        max_loop_iter = 10
        for idx in range(max_loop_iter):
            ret = self.env.reset()
            if not ret[-2]:
                break

        if ret[-2]:
            return episode_info, False

        if self.config.record_video:
            video_writer = cv2.VideoWriter(
                f"assets/demonstrations/{self.config.task.tool_class_name}for{self.config.task.task_name}/episode_{self.success_episode}/video_{self.local_episode}_{self.config.task.tool_fix_idx}.mp4",
                cv2.VideoWriter_fourcc("m", "p", "4", "v"),
                5,
                (640, 480),
            )

        self.expert = self.get_expert()
        onpolicy_agent = self.agent if not self.config.run_expert else self.expert
        context = self.env.export_context()
        info = ret[-1]
        scene_info = self.reset_task_episode(info)

        self.config.task.scene_info = scene_info

        rews = 0
        planned_success = 0
        self.local_episode += 1

        for step in range(self.config.task.env.episodeLength):
            if self.config.task.dart_demonstration and self.get_flags(self.env.time):
                t = np.random.uniform(-0.02, 0.02, size=(3,))
                r = np.random.uniform(-0.2, 0.2, size=(3,))
                perturb_action = np.hstack([t, r])
                perturb_action = self.env.process_action(perturb_action)
                self.env.step(perturb_action)

                # self.env.random_perturb()
                self.expert.reset_expert()
                scene_info = self.reset_task_episode(info)

            action, extra = onpolicy_agent.select_action(ret)
            planned_success = self.config.run_expert and (self.config.teleop or not self.expert.check_plan_empty())
            if self.config.run_expert and not planned_success:  # only run planned succeeded episode
                return episode_info

            if "reset" in extra and extra["reset"]:
                print("reset env")
                return episode_info, False

            action = self.env.process_action(action)
            ret[-1]["action"] = action  # to save
            self.curr_traj_for_replaybuffer.append(copy.deepcopy([ret]))
            ret = self.env.step(action)

            # append trajectory to the replay buffer
            obs, rew, done, info = ret
            if "collision" in info or "sim_failure" in info:
                print(f"env collision: {'collision' in info} or sim terminate request: {'sim_failure' in info}")
                return episode_info, False

            if self.config.record_video:
                vis_img = self.env.img_obs[0]
                vis_img = (vis_img[..., :3] * 255).astype(np.uint8)
                video_writer.write(vis_img)

            rews += rew
            if "tool_pose" in info:  # avoid the exception case
                scene_info["tool_pose"].append(info["tool_pose"])
                scene_info["object_pose"].append(info["object_pose"])
                scene_info["robot_joints"].append(info["joint_pos"])

                if self.config.visdom:
                    self.visdom_visualize(ret)

            if done:
                break

        total_time = time.time() - start_time
        self.plan_success_episode += float(planned_success)
        self.curr_traj_for_replaybuffer.append(copy.deepcopy([ret]))
        start_time = time.time()
        save_success = self.add_trajectory_transition(scene_info)
        print(f"add / save trajectory to buffer time: {time.time() - start_time}")

        self.log_info(episode_info, scene_info, total_time, step, rews)
        return episode_info, save_success

    def log_info(self, episode_info, scene_info, total_time, step, rews):
        episode_info["episode_length"] = step
        episode_info["local_episode_reward"] = rews

        episode_info["fps"] = total_time / episode_info["episode_length"] if episode_info["episode_length"] > 0 else 0
        if self.parallel:
            episode_info["buffer_size"] = ray.get(self.fleet_agent_id.get_memory_position.remote())
        else:
            episode_info["buffer_size"] = self.fleet_agent_id.get_memory_position()

        # episode_info['global_buffer_size'] = fleet_agent.memory.size
        episode_info["local_episode"] = self.local_episode
        episode_info["object_name"] = self.local_episode
        episode_info["tool_class_name"] = self.local_episode
        episode_info["tool_name"] = self.local_episode

        # measure expert's success rates
        self.avg_rewards.append(float(rews > 0))
        self.success_episode += float(rews > 0)
        print(
            f" episode time: {total_time:.3f}"
            f" traj length: {step}"
            f" success rates: {self.success_episode}/{int(self.plan_success_episode)}"
            f" {self.success_episode/int(self.plan_success_episode):.3f}"
            f" plan success rates: {int(self.plan_success_episode)}/{self.local_episode}"
            f" {(self.plan_success_episode)/self.local_episode:.3f}"
        )

    def get_demos(self):
        return self.env.get_offline_data(self.exp_cfg.agent_cfg.num_task_transitions)

    def set_stat(self, stat):
        (self.local_episode, self.success_episode, self.plan_success_episode) = stat

    def return_stat(self):
        return (
            int(self.local_episode),
            int(self.success_episode),
            int(self.plan_success_episode),
        )


class RobotFleet:
    def __init__(self, exp_cfg):
        """Fleet of robots for running experiments"""
        self.config = exp_cfg
        self.exp_cfg = DotMap(exp_cfg)

        # Logging setup
        if len(self.config.save_dir) == 0:
            self.logdir = os.path.join(
                self.config.logdir,
                "{}_{}_{}".format(
                    datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
                    self.config.env_name + self.config.config_suffix,
                    self.config.logdir_suffix,
                ),
            )
        else:
            self.logdir = exp_cfg.save_dir

        print("LOGDIR:", self.logdir)

        self.log_freq = self.config.log_freq
        self.parallel = self.config.parallel
        print(f"number of robots in the fleet: {self.config.num_envs}")

        # Experiment setup
        if self.parallel:
            self.fleet_agent = eval(self.exp_cfg.agent + "Wrapper").remote(
                self.exp_cfg, self.logdir, -1, self.config
            )  #
            self.robots = [
                eval(self.exp_cfg.actor_wrapper).remote(self.exp_cfg, self.fleet_agent, robot_idx)
                for robot_idx in range(self.config.num_envs)
            ]
            self.fleet_agent.set_memory_position.remote(self.config.start_episode_position)

        else:
            # if not self.exp_cfg.cpu_only:
            self.fleet_agent = eval(self.exp_cfg.agent)(self.exp_cfg, self.logdir, -1, self.config)
            self.fleet_agent.set_memory_position(self.config.start_episode_position)

            self.robots = ActorWrapper(self.exp_cfg, self.fleet_agent, 0)
        self.fleet_steps = 0
        self.episode_steps = [0 for _ in range(self.config.num_envs)]
        self.num_actions = self.exp_cfg.task.env.numActions
        self.start_time = time.time()


@ray.remote(num_cpus=1, num_gpus=0.03)
class ActorWrapperGPU05(ActorWrapper):
    """
    A ray wrapper over an agent
    """

    pass

core/experiment.py
liruiw/Fleet-Tools
from .base_expert import TaskExpert
from .expert import AnalyticExpert
from .human_expert import HumanExpert


core/expert/__init__.py
liruiw/Fleet-Tools
from core.utils import *
import json
import core.expert.kpam.SE3_utils as SE3_utils
from env.env_util import *

# The specification of optimization problem
import core.expert.kpam.term_spec as term_spec
import core.expert.kpam.mp_terms as mp_terms
from core.expert.kpam.optimization_problem import OptimizationProblemkPAM, solve_kpam
from core.expert.kpam.optimization_spec import OptimizationProblemSpecification

from colored import fg
import IPython


def dense_sample_traj_times(sample_times, task_completion_time, start_idx=1):
    """densify the waypoints for IK trajectory"""
    ik_times = list(sample_times)
    for i in range(start_idx, len(sample_times)):
        ratio = 1 if sample_times[i] < task_completion_time else 2
        N = int((sample_times[i] - sample_times[i - 1]) * ratio)  #
        for j in range(N):
            ik_times.insert(i, get_interp_time(sample_times[i - 1], sample_times[i], j / N))
    ik_times = sorted(np.unique(ik_times))
    return ik_times


def compose_circular_key_frames(
    center_pose,
    start_time=0,
    r=0.02,
    discrete_frame_num=5,
    total_degrees=np.pi * 2,
    total_time=3,
):
    """generating the circular motion for whisk for instance"""
    poses = []
    times = [start_time + (i + 1) / discrete_frame_num * total_time for i in range(discrete_frame_num)]
    for i in range(discrete_frame_num):
        pose = center_pose.copy()
        pose[:2, 3] += (rotZ(i / discrete_frame_num * total_degrees)[:3, :3] @ np.array([r, 0, 0]))[:2]
        poses.append(pose)
    return times, poses


def compose_rotate_y_key_frames(
    center_pose,
    start_time=0,
    r=0.02,
    discrete_frame_num=5,
    total_degrees=np.pi * 2,
    total_time=3,
):
    """generating the circular motion for whisk for instance"""
    poses = []
    times = [start_time + (i + 1) / discrete_frame_num * total_time for i in range(discrete_frame_num)]
    for i in range(discrete_frame_num):
        pose = center_pose.copy()
        pose[:2, 3] += (rotZ(i / discrete_frame_num * total_degrees)[:3, :3] @ np.array([r, 0, 0]))[:2]
        poses.append(pose)
    return times, poses


def compose_rotating_key_frames(
    ef_pose,
    center_point,
    start_time=0,
    discrete_frame_num=5,
    total_degrees=np.pi * 2,
    total_time=3,
):
    """generating the circular motion with facing center for wrench for instance"""
    poses = []
    times = [start_time + (i + 1) / discrete_frame_num * total_time for i in range(discrete_frame_num)]
    rel_pose = ef_pose.copy()
    rel_pose[:3, 3] -= center_point

    for i in range(discrete_frame_num):  # clockwise
        pose = rotZ(-(i + 1) / discrete_frame_num * total_degrees) @ rel_pose.copy()
        pose[:3, 3] += center_point
        poses.append(pose)
    return times, poses


class TaskExpert:
    # keypoint-based method for tool use
    def __init__(self, envs, args, cfg, robot_id=-1):
        self.actuation_time = 3.0

    def make_vector_vis(self, path, vector, position, color):
        from env.env_util import meshcat

        length_scale = 3
        radius = 0.001
        opacity = 1.0
        vector_norm = np.linalg.norm(vector)
        vector_dir = vector / vector_norm
        transform = RigidTransform()
        test_vec = np.array([0, 0, 1])
        rot_vec_1 = np.cross(test_vec, vector_dir)
        rot_vec_2 = np.cross(vector_dir, rot_vec_1)
        rot_matrix = np.stack([rot_vec_2, rot_vec_1, vector_dir], axis=1)
        transform.set_rotation(RotationMatrix(rot_matrix))
        transform.set_translation(position)

        AddMeshcatVector(
            meshcat,
            path + "_feedforward",
            Rgba(*color),
            length=length_scale * vector_norm,
            radius=radius,
            opacity=opacity,
            X_PT=transform,
        )

    def generate_spatula_poses(self, reach_ratio=0.3, reach_dist=0.2):
        """hard-coded post-activation trajectory:
        (1) reach to the pre-actuation
        (2) actuate it and lift it forward to flip the object
        """
        if self.env.task_config.expert_augment_goals:
            self.augment_goal_pose_and_select()

        # 1 cm lower
        self.task_goal_tool_pose[2, 3] -= 0.002
        tool_forward_dir = (self.task_goal_tool_pose[:3, :3] @ np.array([0, 0, 1]))[:2]
        tool_forward_dir = safe_div(tool_forward_dir, np.linalg.norm(tool_forward_dir))
        # self.task_goal_tool_pose[:2, 3] += tool_forward_dir * 0.05 # avoid reaching the middle point to maintain stick contact
        self.task_goal_tool_pose[:2, 3] += tool_forward_dir * 0.01
        # avoid reaching the middle point to maintain stick contact

        # define task poses
        self.task_goal_hand_pose = self.task_goal_tool_pose @ se3_inverse(self.env.tool_rel_pose)
        hand_forward_dir = tool_forward_dir  # (self.tool_forward_dir[:3, :3] @ np.array([0, 0, 1]))[:2]
        hand_forward_dir = safe_div(hand_forward_dir, np.linalg.norm(hand_forward_dir))
        trans_standoff1 = hand_forward_dir * -reach_dist  # / 2

        self.standoff_pose = self.task_goal_hand_pose.copy()
        self.standoff_pose[:2, 3] += trans_standoff1
        self.standoff_pose3 = self.standoff_pose.copy()
        self.standoff_pose3[2, 3] += 0.15

        # also add a bit of rotation
        self.lift_pose = self.task_goal_hand_pose.copy()
        self.lift_pose[:3, :3] = self.lift_pose[:3, :3] @ rotY(np.pi / 6)[:3, :3]
        self.lift_pose[2, 3] += 0.1  # 0.15
        # self.lift_pose[:2, 3] += hand_forward_dir * 0.12  #  * 3 # no need to flip
        self.lift_pose[:2, 3] -= hand_forward_dir * 0.02  # 13 # lift

        curr_time = self.env.time
        finish_time = self.cfg.task.task_completion_time
        standoff1_time = get_interp_time(curr_time, finish_time, 1 - reach_ratio)
        standoff3_time = get_interp_time(curr_time, finish_time, 1 - 2 * reach_ratio)
        self.actuate_times = finish_time + self.actuation_time
        self.sample_times = [
            curr_time,
            standoff1_time,
            finish_time,
            self.actuate_times,
        ]

        # standoff3_time, self.standoff_pose3,
        self.traj_keyframes = [
            self.env.ee_pose.reshape(4, 4),
            self.standoff_pose,
            self.task_goal_hand_pose,
            self.lift_pose,
        ]

    def generate_scoop_poses(self, reach_ratio=0.3, reach_dist=0.2):
        """hard-coded post-activation trajectory:
        (1) reach the center of the container
        (2) rotate the spoon while moving forward to scoop the ball
        """
        if self.env.task_config.expert_augment_goals:
            self.augment_goal_pose_and_select()

        tool_forward_dir = (self.task_goal_tool_pose[:3, :3] @ np.array([1, 0, 0]))[:2]
        tool_forward_dir = tool_forward_dir[:2] / np.linalg.norm(tool_forward_dir[:2])
        self.task_goal_tool_pose[:2, 3] += tool_forward_dir * 0.015
        self.task_goal_tool_pose[2, 3] += 0.01

        # define task poses
        self.task_goal_hand_pose = self.task_goal_tool_pose @ se3_inverse(self.env.tool_rel_pose)
        self.standoff_pose = self.task_goal_hand_pose.copy()  # in tool coordinate
        self.standoff_pose[2, 3] += reach_dist

        # scoop motion requires some efforts
        hand_forward_dir = (self.task_goal_hand_pose[:3, :3] @ np.array([0, -1, 0]))[:2]
        hand_forward_dir = hand_forward_dir[:2] / np.linalg.norm(hand_forward_dir[:2])
        self.lift_pose = self.task_goal_hand_pose.copy()  # in tool coordinate
        self.lift_pose[:3, :3] = self.lift_pose[:3, :3] @ rotX(np.pi / 3)[:3, :3]
        self.lift_pose[2, 3] += 0.07  # in the hand coordinate, so it rotates the tool a lot
        self.lift_pose[:2, 3] -= hand_forward_dir * 0.14

        # how to measure balls
        curr_time = self.env.time
        finish_time = self.cfg.task.task_completion_time
        standoff_time = get_interp_time(curr_time, finish_time, 1 - reach_ratio)
        self.actuate_times, self.actuate_poses = [finish_time + self.actuation_time], [self.lift_pose]
        self.sample_times = [curr_time, standoff_time, finish_time] + self.actuate_times
        self.traj_keyframes = [
            self.env.ee_pose.reshape(4, 4),
            self.standoff_pose,
            self.task_goal_hand_pose,
        ] + self.actuate_poses

    def generate_hammer_poses(self, reach_ratio=0.3, reach_dist=0.1):
        """hard-coded post-activation trajectory:
        (1) reach above the screw
        (2) hammer it by moving downward
        """
        if self.env.task_config.expert_augment_goals:
            self.augment_goal_pose_and_select()

        # define task poses
        self.task_goal_hand_pose = self.task_goal_tool_pose @ se3_inverse(self.env.tool_rel_pose)
        self.standoff_pose = self.task_goal_hand_pose.copy()  # in tool coordinate
        self.standoff_pose[2, 3] += reach_dist / 2

        curr_time = self.env.time
        finish_time = self.cfg.task.task_completion_time
        standoff_time = get_interp_time(curr_time, finish_time, 1 - reach_ratio)

        # maybe add a rotation in the tool space for both standoff and the task goal pose
        self.hammer_pose = self.task_goal_hand_pose.copy()
        self.hammer_pose[2, 3] -= 0.05

        self.task_goal_hand_pose_rotated = self.task_goal_hand_pose.copy()
        self.lift_pose = self.task_goal_hand_pose.copy()
        self.lift_pose[2, 3] += reach_dist

        self.actuate_times = [
            finish_time + self.actuation_time * 0.6,
            finish_time + self.actuation_time,
        ]
        self.actuate_poses = [self.hammer_pose, self.lift_pose]
        self.sample_times = [curr_time, standoff_time, finish_time] + self.actuate_times
        self.traj_keyframes = [
            self.env.ee_pose.reshape(4, 4),
            self.standoff_pose,
            self.task_goal_hand_pose,
        ] + self.actuate_poses

    def generate_whisk_poses(self, reach_ratio=0.3, reach_dist=0.2):
        """hard-coded post-activation trajectory:
        (1) reach the container with whisk
        (2) whisk around the container
        """
        if self.env.task_config.expert_augment_goals:
            self.augment_goal_pose_and_select()

        # define task poses
        self.task_goal_hand_pose = self.task_goal_tool_pose @ se3_inverse(self.env.tool_rel_pose)
        self.standoff_pose = self.task_goal_hand_pose.copy()  # in tool coordinate
        self.standoff_pose[2, 3] += reach_dist

        curr_time = self.env.time
        finish_time = self.cfg.task.task_completion_time
        standoff_time = get_interp_time(curr_time, finish_time, 1 - reach_ratio)

        self.actuate_times, self.actuate_poses = compose_circular_key_frames(self.task_goal_hand_pose, finish_time)
        self.sample_times = [curr_time, standoff_time, finish_time] + self.actuate_times
        self.traj_keyframes = [
            self.env.ee_pose.reshape(4, 4),
            self.standoff_pose,
            self.task_goal_hand_pose,
        ] + self.actuate_poses

    def generate_knife_poses(self, reach_ratio=0.3, reach_dist=0.2):
        """hard-coded post-activation trajectory:
        (1) reach the middle of the two objects
        (2) push toward one of them with the knife side to separate
        """

        self.task_goal_tool_pose[2, 3] -= 0.005
        self.task_goal_hand_pose = self.task_goal_tool_pose @ se3_inverse(self.env.tool_rel_pose)
        self.standoff_pose = self.task_goal_hand_pose.copy()  # in tool coordinate
        self.standoff_pose[2, 3] += reach_dist / 2

        curr_time = self.env.time
        finish_time = self.cfg.task.task_completion_time
        standoff_time = get_interp_time(curr_time, finish_time, 1 - reach_ratio)

        # need two points on both objects. push from the current tool tip point towards one of the object
        tool_forward_dir1 = (self.env.obj_pose1[:3, 3] - self.curr_solution_tool_keypoint_head)[:2]
        tool_forward_dir1 = tool_forward_dir1[:2] / np.linalg.norm(tool_forward_dir1[:2])

        tool_forward_dir2 = (self.env.obj_pose2[:3, 3] - self.curr_solution_tool_keypoint_head)[:2]
        tool_forward_dir2 = tool_forward_dir2[:2] / np.linalg.norm(tool_forward_dir2[:2])

        # choose based on which one's x direction is more negative
        tool_forward_dir = tool_forward_dir2 if tool_forward_dir2[0] < tool_forward_dir1[0] else tool_forward_dir1
        self.lift_pose = self.task_goal_hand_pose.copy()
        self.lift_pose[:2, 3] -= tool_forward_dir * 0.2

        self.actuate_times = [finish_time + self.actuation_time]
        self.actuate_poses = [self.lift_pose]
        self.sample_times = [curr_time, standoff_time, finish_time] + self.actuate_times
        self.traj_keyframes = [
            self.env.ee_pose.reshape(4, 4),
            self.standoff_pose,
            self.task_goal_hand_pose,
        ] + self.actuate_poses

    def generate_mug_poses(self, reach_ratio=0.3, reach_dist=0.2):
        """hard-coded post-activation trajectory:
        (1) reach right to the screw
        (2) rotate the wrench around the screw
        """
        # has to be the tool keypoint axis
        tool_forward_dir = (self.curr_solution_tool_keypoint_head - self.curr_solution_tool_keypoint_tail)[:2]

        # tool_forward_dir = (self.task_goal_tool_pose[:3, :3] @ np.array([0, 0, 1]))[:2]
        tool_forward_dir = safe_div(tool_forward_dir, np.linalg.norm(tool_forward_dir))

        self.task_goal_tool_pose[:2, 3] += tool_forward_dir * 0.04
        self.task_goal_hand_pose = self.task_goal_tool_pose @ se3_inverse(self.env.tool_rel_pose)
        self.task_goal_hand_pose[2, 3] += 0.04
        self.standoff_pose = self.task_goal_hand_pose.copy()  # in tool coordinate
        self.standoff_pose[2, 3] += 0.03  # a bit higher

        curr_time = self.env.time
        finish_time = self.cfg.task.task_completion_time
        standoff_time = get_interp_time(curr_time, finish_time, 1 - reach_ratio)  # self.task_goal_tool_pose[:3, 3]
        # then rotate around Z axis local

        self.actuate_times = [self.actuation_time + finish_time]
        rotated_pose = self.task_goal_hand_pose.copy()
        rotated_pose[:3, :3] = rotated_pose[:3, :3] @ rotY(-np.pi / 6)[:3, :3]
        self.actuate_poses = [rotated_pose]  #

        self.sample_times = [curr_time, standoff_time, finish_time] + self.actuate_times
        self.traj_keyframes = [
            self.env.ee_pose.reshape(4, 4),
            self.standoff_pose,
            self.task_goal_hand_pose,
        ] + self.actuate_poses

    def generate_wrench_poses(self, reach_ratio=0.3, reach_dist=0.2):
        """hard-coded post-activation trajectory:
        (1) reach right to the screw
        (2) rotate the wrench around the screw
        """

        # has to be the tool keypoint axis
        tool_forward_dir = (self.curr_solution_tool_keypoint_head - self.curr_solution_tool_keypoint_tail)[:2]
        tool_forward_dir = safe_div(tool_forward_dir, np.linalg.norm(tool_forward_dir))
        self.task_goal_hand_pose = self.task_goal_tool_pose @ se3_inverse(self.env.tool_rel_pose)
        self.standoff_pose = self.task_goal_hand_pose.copy()  # in tool coordinate
        hand_forward_dir = (self.task_goal_hand_pose[:3, :3] @ np.array([0, 0, 1]))[:2]
        hand_forward_dir = safe_div(hand_forward_dir, np.linalg.norm(hand_forward_dir))
        self.standoff_pose[:2, 3] -= tool_forward_dir * 0.08  #

        curr_time = self.env.time
        finish_time = self.cfg.task.task_completion_time
        standoff_time = get_interp_time(curr_time, finish_time, 1 - reach_ratio)  # self.task_goal_tool_pose[:3, 3]
        self.actuate_times, self.actuate_poses = compose_rotating_key_frames(
            self.task_goal_hand_pose,
            self.curr_solution_tool_keypoint_head,
            total_degrees=np.pi / 2,
            total_time=self.actuation_time * 2,
            start_time=finish_time,
        )
        self.lift_pose = self.actuate_poses[-1]
        self.sample_times = [curr_time, standoff_time, finish_time] + self.actuate_times
        self.traj_keyframes = [
            self.env.ee_pose.reshape(4, 4),
            self.standoff_pose,
            self.task_goal_hand_pose,
        ] + self.actuate_poses

core/expert/base_expert.py
liruiw/Fleet-Tools
from core.utils import *
import json
import core.expert.kpam.SE3_utils as SE3_utils
from env.env_util import *

# The specification of optimization problem
import core.expert.kpam.term_spec as term_spec
import core.expert.kpam.mp_terms as mp_terms
from core.expert.kpam.optimization_problem import OptimizationProblemkPAM, solve_kpam
from core.expert.kpam.optimization_spec import OptimizationProblemSpecification
from core.expert.kpam.mp_builder import OptimizationBuilderkPAM

from colored import fg
import IPython
from core.expert.base_expert import *


class AnalyticExpert(TaskExpert):
    # keypoint-based method for tool use
    def __init__(self, envs, args, cfg, robot_id=-1):
        super(AnalyticExpert, self).__init__(envs, args, cfg, robot_id=-1)
        self.env = envs
        self.args = args
        self.cfg = cfg
        self.robot_id = robot_id
        self.tool_name = self.env.task_config.tool_class_name
        self.task_name = self.env.task_config.task_name

        self.plan_time = 0
        self.expert_differentialIK = cfg.task.expert_differentialIK
        self.goal_joint = None
        self.reset_expert()
        self.setup()

    def setup(self):
        # load keypoints
        self.solved_ik_times = []
        self.joint_traj_waypoints = []

    def reset_expert(self):
        """reinitialize expert state"""
        self.joint_space_traj = None
        self.task_space_traj = None
        self.contact_task_space_traj = None
        self.plan_succeeded = False

        self.plan_time = 0.0
        self.solved_expert_traj = []

    def check_plan_empty(self):
        """check if already have a plan"""
        return self.joint_space_traj is None

    def check_replan(self):
        """check if needs to replan"""
        return self.env.time < self.standoff_time and (
            self.env.time - self.plan_time > self.env.task_config.replan_freq
        )

    def use_contact_expert(self):
        """compute the actuation trajectory by planning through contact"""
        if self.contact_task_space_traj is None:  #  and self.env.time > 1 self.cfg.task.task_completion_time:
            # self.contact_expert = ContactExpert(self.env, self.args, self.cfg)
            self.build_tool_object_env()
            self.contact_task_space_traj = self.trajopt_through_contact()

    def select_action(self, state, env_idx=0):
        """the main calling function"""
        """select action for the current state, base on the env time and the solved trajectory"""
        s = time.time()
        if self.check_plan_empty() or self.check_replan():
            self.standoff_time = get_interp_time(self.env.time, self.cfg.task.task_completion_time, 0.7)
            if self.env.task_config.joint_space_kpam:
                self.solve_kpam_joint()
            else:
                self.solve_kpam_pose()

            print("env time: {:.3f} plan generation time: {:.3f}".format(self.env.time, time.time() - s))
            self.solve_postactuation_traj()
            self.solve_joint_traj()

        if self.env.task_config.use_meshcat and hasattr(self, "curr_solution_tool_keypoint_head"):
            self.env.visualize_task_keypoints(
                [
                    self.curr_solution_tool_keypoint_head,
                    self.curr_solution_tool_keypoint_tail,
                ],
                "solution_points",
            )
        extra = {"reset": self.check_plan_empty() and not self.expert_differentialIK}
        return self.get_action(), extra

    def calculate_distance_from_expert_plan(self):
        """compute the distance from the current environment state from what the expert plan is"""
        target_pose = self.task_space_traj.value(self.env.time)
        curr_state = self.env.ee_pose.reshape(4, 4)
        # compute point distance from the current and the target states

    def create_opt_problem(self, optimization_spec):
        """create a keypoint optimization problem from the current keypoint state"""
        if len(self.env.task_config.overwrite_kpam_config_file) > 0:
            optimization_spec.load_from_yaml(
                "core/expert/kpam/config/{}.yaml".format(self.env.task_config.overwrite_kpam_config_file)
            )
        else:
            optimization_spec.load_from_yaml("core/expert/kpam/config/{}.yaml".format(self.task_name))
        print(f"load tool keypoint file from {self.task_name}.yaml")

        # match the table
        constraint_update_keypoint_target = {"tool_head": self.env.curr_tool_obj_keypoints}

        # minimize movement
        cost_update_keypoint_target = {
            "tool_head": self.env.curr_tool_keypoint_head,
            "tool_tail": self.env.curr_tool_keypoint_tail,
            "tool_side": self.env.curr_tool_keypoint_side,
        }

        for term in optimization_spec._cost_list:
            if hasattr(term, "keypoint_name") and term.keypoint_name in cost_update_keypoint_target.keys():
                term.target_position = cost_update_keypoint_target[term.keypoint_name]

        for term in optimization_spec._constraint_list:
            if hasattr(term, "keypoint_name") and term.keypoint_name in constraint_update_keypoint_target.keys():
                term.target_position = constraint_update_keypoint_target[term.keypoint_name]

        return optimization_spec

    # tool use related
    def solve_kpam_joint(self, generate_traj=True):
        """solve the formulated kpam problem and get goal joint"""
        s = time.time()
        self.env.get_obs(render=False)

        # solve for the goal end effector pose to match keypoints
        keypoint_loc = np.stack(
            (
                self.env.curr_tool_keypoint_head,
                self.env.curr_tool_keypoint_tail,
                self.env.curr_tool_keypoint_side,
            ),
            axis=0,
        )

        keypoint_loc_in_hand = np.stack(
            (
                self.env.tool_keypoint_head_in_hand,
                self.env.tool_keypoint_tail_in_hand,
                self.env.tool_keypoint_side_in_hand,
            ),
            axis=0,
        )

        optimization_spec = OptimizationProblemSpecification()
        optimization_spec = self.create_opt_problem(optimization_spec)

        constraint_dicts = [c.to_dict() for c in optimization_spec._constraint_list]
        cost_dicts = [c.to_dict() for c in optimization_spec._cost_list]

        # need to parse the kpam config file and create a kpam problem
        indexes = np.random.randint(len(anchor_seeds), size=(8,))
        random_seeds = [self.env.joint_positions.copy()] + [anchor_seeds[idx] for idx in indexes]
        solutions = []

        for seed in random_seeds:
            res = solve_ik_kpam(
                [constraint_dicts, cost_dicts],
                self.env.controller_plant,
                self.env.controller_plant.GetFrameByName("panda_hand"),
                keypoint_loc_in_hand,
                self.env.curr_tool_obj_keypoints,
                RigidTransform(self.env.ee_pose.reshape(4, 4)),
                seed.reshape(-1, 1),
                self.env.joint_positions.copy(),
                rot_tol=0.01,
                timeout=True,
                consider_collision=False,
            )

            if res is not None:
                solutions.append(res.get_x_val()[:9])

        diff_ik_context = self.env.differential_ik.GetMyMutableContextFromRoot(self.env.context)
        if len(solutions) == 0:
            print("empty solution in kpam")
            self.env.need_termination = True
            self.goal_joint = self.env.joint_positions[:9].copy()
        else:
            solutions = np.array(solutions)
            joint_positions = self.env.joint_positions[:9]
            dist_to_init_joints = np.linalg.norm(solutions - joint_positions.copy(), axis=-1)
            res = solutions[np.argmin(dist_to_init_joints)]
            self.goal_joint = res

            # not used anyways
            self.env.differential_ik.SetPositions(diff_ik_context, res)

        self.task_goal_hand_pose = self.env.differential_ik.ForwardKinematics(diff_ik_context)
        self.task_goal_hand_pose = np.array(self.task_goal_hand_pose.GetAsMatrix4())
        self.task_goal_tool_pose = self.task_goal_hand_pose @ self.env.tool_rel_pose

        # Transform the keypoint
        self.curr_solution_tool_keypoint_head = SE3_utils.transform_point(
            self.task_goal_hand_pose, keypoint_loc_in_hand[0, :]
        )
        self.curr_solution_tool_keypoint_tail = SE3_utils.transform_point(
            self.task_goal_hand_pose, keypoint_loc_in_hand[1, :]
        )
        self.curr_solution_tool_keypoint_side = SE3_utils.transform_point(
            self.task_goal_hand_pose, keypoint_loc_in_hand[2, :]
        )

        self.plan_time = self.env.time
        self.goal_keypoint = np.stack(
            (
                self.curr_solution_tool_keypoint_head,
                self.curr_solution_tool_keypoint_tail,
                self.curr_solution_tool_keypoint_side,
            ),
            axis=0,
        )

        self.env.info["goal_keypoint"] = self.goal_keypoint
        self.env.info["goal_pose"] = self.task_goal_hand_pose

    # tool use related
    def solve_kpam_pose(self, generate_traj=True):
        """solve the formulated kpam problem and get goal pose
        and then solve for the trajectory"""
        self.env.get_obs(render=False)

        # solve for the goal end effector pose to match keypoints
        keypoint_loc = np.stack(
            (
                self.env.curr_tool_keypoint_head,
                self.env.curr_tool_keypoint_tail,
                self.env.curr_tool_keypoint_side,
            ),
            axis=0,
        )

        # create a kpam problem
        optimization_spec = OptimizationProblemSpecification()
        optimization_spec = self.create_opt_problem(optimization_spec)
        builder = OptimizationBuilderkPAM(optimization_spec)
        problem = builder.build_optimization(keypoint_loc)
        solve_kpam(problem)
        self.task_goal_tool_pose = problem.T_action.dot(self.env.tool_pose)

        # Transform the keypoint
        self.curr_solution_tool_keypoint_head = SE3_utils.transform_point(problem.T_action, keypoint_loc[0, :])
        self.curr_solution_tool_keypoint_tail = SE3_utils.transform_point(problem.T_action, keypoint_loc[1, :])
        self.curr_solution_tool_keypoint_side = SE3_utils.transform_point(problem.T_action, keypoint_loc[2, :])

        self.plan_time = self.env.time
        self.goal_keypoint = np.stack(
            (
                self.curr_solution_tool_keypoint_head,
                self.curr_solution_tool_keypoint_tail,
                self.curr_solution_tool_keypoint_side,
            ),
            axis=0,
        )

        self.env.info["goal_keypoint"] = self.goal_keypoint

    def solve_goal_joint(self):
        # solve a joint to reach task_goal_hand_pose

        seeds = [self.env.joint_positions.copy()]
        indexes = np.random.randint(len(anchor_seeds), size=(8,))
        seeds = seeds + [anchor_seeds[idx] for idx in indexes]
        solved_iks = []

        for seed_ in seeds:
            rot_tol = 0.005
            res = self.env.solve_ik(self.task_goal_hand_pose, seed_, rot_tol)

            if res is not None:
                res = res.get_x_val()[:9]
                solved_iks.append(res)

        if len(solved_iks) == 0:
            print("expert solve goal IK failed")

        else:
            # solve multiple ones and pick the closest one
            diff = np.array(solved_iks) - self.env.joint_positions[None]
            dist = np.linalg.norm(diff, axis=-1)
            self.goal_joint = solved_iks[np.argmin(dist)]

    def set_plan(self, joint_plan):
        self.joint_traj_waypoints = joint_plan
        self.joint_space_traj = PiecewisePolynomial.CubicShapePreserving(
            self.dense_traj_times, np.array(self.joint_traj_waypoints).T
        )
        self.get_task_traj_from_joint_traj()

    def solve_traj_endpoints(self):
        """
        solve for the IKs for each individual waypoint as an initial guess, and then
        solve for the whole trajectory with smoothness cost
        """
        keyposes = self.traj_keyframes
        keytimes = self.sample_times

        self.joint_traj_waypoints = [self.env.joint_positions.copy()]
        self.joint_space_traj = PiecewisePolynomial.FirstOrderHold(
            [self.env.time, self.cfg.task.task_completion_time],
            np.array([self.env.joint_positions.copy(), self.goal_joint]).T,
        )

        self.dense_traj_times = dense_sample_traj_times(self.sample_times, self.cfg.task.task_completion_time)

        # traj_pose = self.env.set_diffik_state(self.env.context, set_joints=self.env.joint_positions)
        print("solve traj endpoint")

        # interpolated joint
        res = solve_ik_traj_with_standoff(
            self.env.controller_plant,
            [self.env.ee_pose.reshape(4, 4), self.task_goal_hand_pose],
            np.array([self.env.joint_positions.copy(), self.goal_joint]).T,
            q_traj=self.joint_space_traj,
            waypoint_times=self.dense_traj_times,
            keyposes=keyposes,
            keytimes=keytimes,
        )

        # solve the standoff and the remaining pose use the goal as seed.
        # stitch the trajectory

        if res is not None:
            #  use the joint trajectory to build task trajectory
            self.joint_traj_waypoints = res.get_x_val().reshape(-1, 9)
            self.joint_traj_waypoints[:, -2:] = 0.04
            self.joint_traj_waypoints = list(self.joint_traj_waypoints)
            self.joint_space_traj = PiecewisePolynomial.CubicShapePreserving(
                self.dense_traj_times, np.array(self.joint_traj_waypoints).T
            )
            self.get_task_traj_from_joint_traj()
        else:
            print("endpoint trajectory not solved! environment termination")
            self.env.need_termination = True
            self.get_task_traj_from_joint_traj()

    def get_task_traj_from_joint_traj(self):
        """forward kinematics the joint trajectory to get the task trajectory"""
        self.pose_traj = []
        ik_times = dense_sample_traj_times(self.sample_times, self.cfg.task.task_completion_time)
        self.dense_ik_times = ik_times
        for traj_time in ik_times:
            diff_ik_context = self.env.differential_ik.GetMyMutableContextFromRoot(self.env.context)
            set_joints = self.joint_space_traj.value(traj_time)
            self.env.differential_ik.SetPositions(diff_ik_context, set_joints)
            pose = self.env.differential_ik.ForwardKinematics(diff_ik_context)
            self.pose_traj.append(pose.GetAsMatrix4())

        self.task_space_traj = PiecewisePose.MakeLinear(ik_times, [RigidTransform(p) for p in self.pose_traj])
        self.env.set_diffik_state(self.env.context)

    def solve_postactuation_traj(self, reach_ratio=0.3, reach_dist=0.12):
        """
        generate the full task trajectory with a FirstOrderHold
        """
        getattr(self, f"generate_{self.task_name}_poses")(reach_ratio, reach_dist)

        if self.env.task_config.contact_planner:
            print("use trajopt to improve the actuation trajectory")
            self.use_contact_expert()

        if self.env.task_config.env.vis_traj_keyframe and self.env.task_config.use_meshcat:
            self.vis_keyframes()

        self.env.info["goal_pose"] = self.task_goal_hand_pose

    def solve_traj_pose_interp(self, ik_times):
        """
        solve for the IKs for each individual waypoint as an initial guess, and then
        solve for the whole trajectory with smoothness cost
        """
        self.task_space_traj = PiecewisePose.MakeLinear(
            self.sample_times, [RigidTransform(p) for p in self.traj_keyframes]
        )
        seed = self.env.joint_positions.copy()
        solved_poses = [self.env.ee_pose.reshape(4, 4)]
        self.joint_traj_waypoints = [self.env.joint_positions.copy()]

        for idx, t in enumerate(ik_times[1:]):
            target_pose = self.task_space_traj.value(t)

            # try different seeds for solving IK
            seeds = [seed, self.goal_joint]
            per_joint_solved_iks = []
            indexes = np.random.randint(len(anchor_seeds), size=(2,))
            seeds = seeds + [anchor_seeds[idx] for idx in indexes]

            for seed_ in seeds:
                rot_tol = 0.1 if t < self.standoff_time else 0.005
                res = self.env.solve_ik(target_pose, seed_.reshape(-1, 1), rot_tol)

                if res is not None:
                    res = res.get_x_val()[:9]
                    per_joint_solved_iks.append(res)

            if len(per_joint_solved_iks) == 0:
                color_print(("expert solve IK failed {:.3f}".format(t)), "sea_green_1b")
                if len(self.joint_traj_waypoints) < idx:
                    color_print("early terminate", "sea_green_1b")
                    break
            else:
                # solve multiple ones and pick the closest one
                diff = np.array(per_joint_solved_iks) - seed[None]
                dist = np.linalg.norm(diff, axis=-1)

                if np.min(dist) > 1 and not self.env.task_config.skip_pose_interpolate:
                    print(f"solve ik min dist: {t:.3f} {np.min(dist):.3f}")
                    continue

                seed = per_joint_solved_iks[np.argmin(dist)]
                res = seed  # use seed for the next joint
                self.solved_ik_times.append(t)
                self.joint_traj_waypoints.append(res[:9])
                solved_poses.append(target_pose)

        solved_poses = np.array(solved_poses)
        self.joint_traj_waypoints = np.array(self.joint_traj_waypoints)

        smoothness_curr = smoothness(self.joint_traj_waypoints.T)
        print("smoothness before refine {:.3f}".format(smoothness_curr))

        res = refine_ik_traj(
            self.env.controller_plant,
            solved_poses,
            np.array(self.joint_traj_waypoints).T,
        )

        if res is not None:
            refined_waypoints = res.get_x_val().reshape(-1, 9)
            smoothness_new = smoothness(refined_waypoints.T)
            print("smoothness after refine {:.3f}".format(smoothness_new))
            if smoothness_new <= smoothness_curr:
                self.joint_traj_waypoints = refined_waypoints

    def solve_joint_traj(self):
        """
        solve single ik of each pose in the pose trajectory
        and then solve a small smoothness trajopt
        """
        if self.env.task_config.pose_interp_traj:
            self.joint_traj_waypoints = [self.env.joint_positions.copy()]
            self.solved_ik_times = [self.env.time]
            print("pose interpolation keyframe sampled times:", self.sample_times)

            ik_times = dense_sample_traj_times(self.sample_times, self.standoff_time)
            exported_context = self.env.export_context()
            """ interpolate pose trajectory and then sole and refine """
            self.solve_traj_pose_interp(ik_times)
            self.env.load_context(exported_context)
            if self.check_plan_empty():  #  self.sample_times[-1] not in self.solved_ik_times
                if len(self.solved_ik_times) < len(ik_times) - 3:
                    color_print((("expert failed to solve for a trajectory. reset env")), "sea_green_1b")
                    self.reset_expert()
                    self.env.need_termination = True

            # each column is a sample straightline trajectory example FirstOrderHold
            self.joint_space_traj = PiecewisePolynomial.CubicShapePreserving(
                self.solved_ik_times, np.array(self.joint_traj_waypoints).T
            )
        else:
            # directly solve the joint end point problem
            self.solve_traj_endpoints()

    def feedback_loop(self, action):
        """
        do a simple feedback loop for each task
        return 6dof delta pose
        """
        return getattr(self, f"reactive_{self.tool_name}_motion")(action)

    def vis_keyframes(self):
        """visualize keyframes specified to execute the task trajectory"""
        from env.env_util import meshcat, visualizer

        for idx, p in enumerate(self.traj_keyframes):
            AddMeshcatTriad(
                meshcat,
                f"keyframe_{idx}",
                X_PT=RigidTransform(p),
                length=0.03,
                radius=0.001,
            )

        vis_tool_pose(meshcat, self.task_goal_tool_pose, "tool_goal_pose")

    def get_joint_action(self):
        """get the joint space action"""
        if self.check_plan_empty():
            print("no joint trajectory")
            return self.env.reset()

        # look ahead
        return self.joint_space_traj.value(self.env.time + self.env.env_dt).reshape(-1)

    def get_pose_action(self, curr_pose, traj_eff_pose):
        """get the task space action"""
        # feedforward position action
        action = pack_pose(se3_inverse(curr_pose).dot(traj_eff_pose))
        return action

    def get_standoff_instant(self):
        # return True if the robot is before the standoff
        return (self.env.time - self.standoff_time) >= 0 and (self.env.time - self.standoff_time) <= (
            self.env.env_dt - 1e-5
        )

    def get_action(self):
        """get the task-space action from the kpam expert joint trajectory"""
        if self.task_space_traj is None:
            print("no task trajectory")
            return self.env.reset()

        curr_pose = self.env.ee_pose.reshape(4, 4)
        traj_eff_pose = self.task_space_traj.value(self.env.time + self.env.env_dt)
        pose_action = self.get_pose_action(curr_pose, traj_eff_pose)
        joint_action = self.get_joint_action()

        if self.env.task_config.reactive_planner and self.env.at_contact:
            new_traj_eff_pose = curr_pose @ unpack_pose(pose_action, rot_type="euler")
            new_ik = self.env.solve_ik(new_traj_eff_pose, joint_action)
            if new_ik is not None:
                joint_action = new_ik.get_x_val()[:9]

        if self.expert_differentialIK or (
            self.env.task_config.contact_shooting_planner and self.env.time > self.standoff_time
        ):  # use workspace motion as supervison

            return pose_action
        return np.concatenate((joint_action, pose_action))

core/expert/expert.py
liruiw/Fleet-Tools
from core.utils import *
import json
from env.env_util import *

from colored import fg
import IPython
from core.expert.expert import *
import cv2
from env.env_util import meshcat, visualizer

###############################################################################################
# Human expert control
class HumanExpert(AnalyticExpert):
    def __init__(self, envs, args, cfg, robot_id=-1):
        super(HumanExpert, self).__init__(envs, args, cfg, robot_id=-1)
        self.env = envs
        self.args = args
        self.robot_id = robot_id
        self.cfg = cfg
        self.desired_pose = self.env.ee_pose.reshape(4, 4)
        self.desired_joint = self.env.joint_positions
        self.task_transform_pose = np.eye(4)
        self.teleop_cam_pose = None

    def select_controller(self, name="pose"):
        if name == "pose":
            self.env.task_config.pose_controller = True
            self.env.controller_switch_port.FixValue(self.env.context, InputPortIndex(1))
        else:
            self.env.task_config.pose_controller = False
            self.env.controller_switch_port.FixValue(self.env.context, InputPortIndex(2))

    def select_action(self, curr_pose, env_idx=0):
        """control the relative delta action for the end effector. simulation environment only."""
        extra = {}
        if self.args.teleop_type == "keyboard":
            action, finger_state = get_keyboard_input(self.env.action_space.sample(), self.robot_id)
            action_pose = unpack_action(action)
            reset = False

        elif self.args.teleop_type == "vr":
            (
                dX_UinitU_W,  # current delta
                finger_state,  # Right Grip
                reset,  # B
                teleop_hold,  # Right Trig
                X_WGdesired,  # A
                exit_loop,
            ) = get_vr_input()

            action_pose = se3_inverse(curr_pose) @ X_WGdesired  # delta pose as action labels
            action = pack_action(action_pose)

            # visualize the target_X_TB using triad
            meshcat.SetTransform("hand_desired_coord", RigidTransform(X_WGdesired))
            meshcat.SetTransform("hand_actual_coord", RigidTransform(curr_pose))

        return action, extra

core/expert/human_expert.py
liruiw/Fleet-Tools
import numpy as np
import functools
from core.expert.kpam.transformations import affine_matrix_from_points
import core.expert.kpam.SE3_utils as SE3_utils

# The specification of optimization problem
import core.expert.kpam.term_spec as term_spec
import core.expert.kpam.mp_terms as mp_terms
from core.expert.kpam.optimization_problem import OptimizationProblemkPAM, solve_kpam
import IPython


class OptimizationBuilderkPAM(object):
    def __init__(self, specification):  # type: (OptimizationProblemSpecification) -> None
        self._specification = specification

    def _check_and_get_keypoint_idx(self, keypoint_name, input_idx=-1):
        # type: (str, int) -> int
        """
        Given the keypoint name, return the index for that keypoint.
        If input_idx is valid, also verify two index match each other
        :param keypoint_name:
        :param input_idx:
        :return:
        """
        assert keypoint_name in self._specification.keypoint_name2idx
        idx = self._specification.keypoint_name2idx[keypoint_name]

        if input_idx >= 0:
            assert idx == input_idx
        return idx

    def build_optimization(self, keypoint_observation):  # type: (np.ndarray) -> OptimizationProblemkPAM
        """
        Given the observed keypoint, build a mathematical program whose
        solution is the transform that map keypoint_observation to target
        :param keypoint_observation: (N, 3) keypoint location
        :return:
        """
        # Basic check
        assert keypoint_observation.shape[1] == 3
        assert keypoint_observation.shape[0] == len(self._specification.keypoint_name2idx)

        # Empty initialization of optimization problem
        problem = OptimizationProblemkPAM()
        problem.T_init = np.eye(4)
        problem.mp = None
        problem.xyzrpy = None
        problem.has_solution = False

        # Build the initial transformation
        problem.T_init = self._build_initial_transformation(keypoint_observation)

        # Build the mathematical_program with decision variable
        problem.build_empty_mp()

        # Process the cost terms
        for cost in self._specification.cost_list:
            self._process_cost_term(problem, keypoint_observation, cost)

        # Process the constraint terms
        for constraint in self._specification.constraint_list:
            self._process_constraint_term(problem, keypoint_observation, constraint)

        return problem

    def _build_initial_transformation(self, keypoint_observation):  # type: (np.ndarray) -> np.ndarray
        # The container for all keypoint that has nominal target
        from_keypoint = []
        to_keypoint = []

        # Get all the keypoint that has nominal target
        keypoint_name2idx = self._specification.keypoint_name2idx
        for keypoint_name in keypoint_name2idx:
            exist_target, location = self._specification.get_nominal_target_position(keypoint_name)
            if exist_target:
                to_keypoint.append(location)
                idx = keypoint_name2idx[keypoint_name]
                from_keypoint.append(
                    [
                        keypoint_observation[idx, 0],
                        keypoint_observation[idx, 1],
                        keypoint_observation[idx, 2],
                    ]
                )

        # Depends on whether we have enough keypoints
        assert len(from_keypoint) == len(to_keypoint)
        if len(from_keypoint) < 3:
            return np.eye(4)

        # There exists enough keypoint, let's compute using karbas algorithm
        # The np array are (3, N) format of keypoints
        from_keypoint_np = np.zeros(shape=(3, len(from_keypoint)))
        to_keypoint_np = np.zeros_like(from_keypoint_np)
        for i in range(len(from_keypoint)):
            for j in range(3):
                from_keypoint_np[j, i] = from_keypoint[i][j]
                to_keypoint_np[j, i] = to_keypoint[i][j]
        T = affine_matrix_from_points(from_keypoint_np, to_keypoint_np, shear=False, scale=False, usesvd=True)
        return T

    def _process_cost_term(self, problem, keypoint_observation, cost):
        # type: (OptimizationProblemkPAM, np.ndarray, term_spec.OptimizationTermSpec) -> None
        """
        Global dispatcher for the cost term
        :param problem:
        :param cost:
        :return:
        """
        if isinstance(cost, term_spec.Point2PointCostL2Spec):
            self._process_keypoint_l2_cost_term(problem, keypoint_observation, cost)
        elif isinstance(cost, term_spec.Point2PlaneCostSpec):
            self._process_point2plane_cost_term(problem, keypoint_observation, cost)
        else:
            raise RuntimeError("The cost term type is unknown/not implemented")

    def _process_constraint_term(self, problem, keypoint_observation, constraint):
        # type: (OptimizationProblemkPAM, np.ndarray, term_spec.OptimizationTermSpec) -> None
        """
        Global dispatcher for constraint term
        :param problem:
        :param keypoint_observation:
        :param constraint:
        :return:
        """
        if isinstance(constraint, term_spec.Point2PointConstraintSpec):
            self._process_point2point_constraint_term(problem, keypoint_observation, constraint)
        elif isinstance(constraint, term_spec.AxisAlignmentConstraintSpec):
            self._process_axis_alignment_constraint_term(problem, keypoint_observation, constraint)
        elif isinstance(constraint, term_spec.KeypointAxisAlignmentConstraintSpec):
            self._process_keypoint_axisalign_constraint_term(problem, keypoint_observation, constraint)
        elif isinstance(constraint, term_spec.KeypointAxisOrthogonalConstraintSpec):
            self._process_keypoint_axisorthogonal_constraint_term(problem, keypoint_observation, constraint)
        else:
            raise RuntimeError("The constraint type is unknown/not implemented")

    def _process_keypoint_l2_cost_term(self, problem, keypoint_observation, cost):
        # type: (OptimizationProblemkPAM, np.ndarray, term_spec.Point2PointCostL2Spec) -> None
        idx = self._check_and_get_keypoint_idx(cost.keypoint_name, cost.keypoint_idx)

        # The current location
        from_keypoints = keypoint_observation[idx, :]
        from_keypoints = np.reshape(from_keypoints, (1, 3))

        # The target location
        to_keypoints = np.zeros(shape=(1, 3))
        to_keypoints[0, 0] = cost.target_position[0]
        to_keypoints[0, 1] = cost.target_position[1]
        to_keypoints[0, 2] = cost.target_position[2]

        # The cost function
        # It is wired that the cost term cannot use symbolic expression
        cost_func = functools.partial(
            mp_terms.keypoint_l2_cost,
            problem,
            from_keypoints,
            to_keypoints,
            cost.penalty_weight,
        )
        problem.mp.AddCost(cost_func, problem.xyzrpy)

    def _process_point2plane_cost_term(self, problem, keypoint_observation, cost):
        # type: (OptimizationProblemkPAM, np.ndarray, term_spec.Point2PlaneCostSpec) -> None
        idx = self._check_and_get_keypoint_idx(cost.keypoint_name, cost.keypoint_idx)

        # The current location
        from_keypoints = keypoint_observation[idx, :]
        from_keypoints = np.reshape(from_keypoints, (1, 3))

        # The target location
        to_keypoints = np.zeros(shape=(1, 3))
        to_keypoints[0, 0] = cost.target_position[0]
        to_keypoints[0, 1] = cost.target_position[1]
        to_keypoints[0, 2] = cost.target_position[2]

        # The cost function
        # It is wired that the cost term cannot use symbolic expression
        cost_func = functools.partial(
            mp_terms.point2plane_cost,
            problem,
            from_keypoints,
            to_keypoints,
            cost.penalty_weight,
            cost.plane_normal,
        )
        problem.mp.AddCost(cost_func, problem.xyzrpy)

    def _process_point2point_constraint_term(self, problem, keypoint_observation, constraint):
        # type: (OptimizationProblemkPAM, np.ndarray, term_spec.Point2PointConstraintSpec) -> None
        idx = self._check_and_get_keypoint_idx(constraint.keypoint_name, constraint.keypoint_idx)

        # The current location
        from_keypoint = keypoint_observation[idx, :]
        transformed_point_symbolic_val = mp_terms.transform_point_symbolic(problem, from_keypoint, problem.xyzrpy)

        # The target location
        to_keypoints = np.asarray(constraint.target_position).copy()
        to_keypoints_lb = to_keypoints - constraint.tolerance
        to_keypoints_ub = to_keypoints + constraint.tolerance
        for j in range(3):
            problem.mp.AddConstraint(transformed_point_symbolic_val[j] >= to_keypoints_lb[j])
            problem.mp.AddConstraint(transformed_point_symbolic_val[j] <= to_keypoints_ub[j])

    @staticmethod
    def _process_axis_alignment_constraint_term(problem, keypoint_observation, constraint):
        # type: (OptimizationProblemkPAM, np.ndarray, term_spec.AxisAlignmentConstraintSpec) -> None
        assert constraint.tolerance < 1
        assert constraint.tolerance > 0

        # Check the axis
        from_axis = np.asarray(constraint.from_axis, dtype=np.float32).copy()  # type: np.ndarray
        to_axis = np.asarray(constraint.target_axis, dtype=np.float32).copy()  # type: np.ndarray
        norm_from = np.linalg.norm(from_axis)
        norm_to = np.linalg.norm(to_axis)
        if norm_from < 1e-4 or norm_to < 1e-4:
            print("Warning: the axis is zero, the constraint is not added to optimization.")
            return

        # Do normalization
        from_axis *= 1.0 / norm_from
        to_axis *= 1.0 / norm_to

        # Build the term
        vector_dot_symbolic_val = mp_terms.vector_dot_symbolic(problem, from_axis, to_axis, problem.xyzrpy)

        # Add to mp
        problem.mp.AddConstraint(
            vector_dot_symbolic_val >= constraint.target_inner_product - constraint.tolerance
        )  #  - constraint.target_inner_product
        problem.mp.AddConstraint(
            vector_dot_symbolic_val <= constraint.target_inner_product + constraint.tolerance
        )  #  - constraint.target_inner_product

    def _process_keypoint_axisalign_constraint_term(self, problem, keypoint_observation, constraint):
        # type: (OptimizationProblemkPAM, np.ndarray, term_spec.KeypointAxisAlignmentConstraintSpec) -> None
        from_idx = self._check_and_get_keypoint_idx(
            constraint.axis_from_keypoint_name, constraint.axis_from_keypoint_idx
        )
        to_idx = self._check_and_get_keypoint_idx(constraint.axis_to_keypoint_name, constraint.axis_to_keypoint_idx)

        # Compute the axis
        from_point = keypoint_observation[from_idx, :]
        to_point = keypoint_observation[to_idx, :]
        vector = to_point - from_point
        norm_vec = np.linalg.norm(vector)
        if norm_vec < 1e-4:
            print("Warning: the axis is zero, the constraint is not added to optimization.")
            return
        vector *= 1.0 / norm_vec

        # Build axis alignment constraint
        axis_constraint = term_spec.AxisAlignmentConstraintSpec()
        axis_constraint.from_axis = [vector[0], vector[1], vector[2]]
        axis_constraint.target_axis = constraint.target_axis
        axis_constraint.tolerance = constraint.tolerance
        axis_constraint.target_inner_product = constraint.target_inner_product

        # OK
        self._process_axis_alignment_constraint_term(problem, keypoint_observation, axis_constraint)

    @staticmethod
    def _process_axis_orthogonal_constraint_term(problem, keypoint_observation, constraint):
        # make the axis orthogonal as the constraint
        # type: (OptimizationProblemkPAM, np.ndarray, term_spec.AxisAlignmentConstraintSpec) -> None
        assert constraint.tolerance < 1
        assert constraint.tolerance > 0

        # Check the axis
        from_axis = np.asarray(constraint.from_axis, dtype=np.float32).copy()  # type: np.ndarray
        to_axis = np.asarray(constraint.target_axis, dtype=np.float32).copy()  # type: np.ndarray
        norm_from = np.linalg.norm(from_axis)
        norm_to = np.linalg.norm(to_axis)
        if norm_from < 1e-4 or norm_to < 1e-4:
            print("Warning: the axis is zero, the constraint is not added to optimization.")
            return

        # Do normalization
        from_axis *= 1.0 / norm_from
        to_axis *= 1.0 / norm_to

        # Build the term
        vector_dot_symbolic_val = mp_terms.vector_dot_symbolic(problem, from_axis, to_axis, problem.xyzrpy)

        # Add to mp
        problem.mp.AddConstraint(
            vector_dot_symbolic_val - constraint.target_inner_product <= constraint.tolerance
        )  # 1.0 -
        problem.mp.AddConstraint(
            vector_dot_symbolic_val - constraint.target_inner_product >= -constraint.tolerance
        )  # 1.0 -

    def _process_keypoint_axisorthogonal_constraint_term(self, problem, keypoint_observation, constraint):
        # type: (OptimizationProblemkPAM, np.ndarray, term_spec.KeypointAxisAlignmentConstraintSpec) -> None
        from_idx = self._check_and_get_keypoint_idx(
            constraint.axis_from_keypoint_name, constraint.axis_from_keypoint_idx
        )
        to_idx = self._check_and_get_keypoint_idx(constraint.axis_to_keypoint_name, constraint.axis_to_keypoint_idx)

        # Compute the axis
        from_point = keypoint_observation[from_idx, :]
        to_point = keypoint_observation[to_idx, :]
        vector = to_point - from_point
        norm_vec = np.linalg.norm(vector)
        if norm_vec < 1e-4:
            print("Warning: the axis is zero, the constraint is not added to optimization.")
            return
        vector *= 1.0 / norm_vec

        # Build axis alignment constraint
        axis_constraint = term_spec.AxisAlignmentConstraintSpec()
        axis_constraint.from_axis = [vector[0], vector[1], vector[2]]
        axis_constraint.target_axis = constraint.target_axis
        axis_constraint.tolerance = constraint.tolerance
        axis_constraint.target_inner_product = constraint.target_inner_product

        # OK
        self._process_axis_orthogonal_constraint_term(problem, keypoint_observation, axis_constraint)


def test_builder():
    from core.kpam.optimization_spec import build_mug2rack, build_mug2shelf

    optimization_spec = OptimizationProblemSpecification()
    optimization_spec.load_from_yaml("/home/wei/catkin_ws/src/kplan_ros/kplan/kpam_opt/config/mug2rack.yaml")
    # optimization_spec = build_mug2shelf()
    builder = OptimizationBuilderkPAM(optimization_spec)

    # Some keypoint location
    keypoint_loc = np.zeros(shape=(3, 3))
    keypoint_loc[0, :] = np.asarray([0.6523959765381968, -0.0028259822512128746, -0.0077028763979177794])
    keypoint_loc[1, :] = np.asarray([0.5961969831659355, -0.0020322277135196576, 0.04569443896150649])
    keypoint_loc[2, :] = np.asarray([0.6477449683371657, -0.005210614473128802, 0.0852574313164387])
    problem = builder.build_optimization(keypoint_loc)

    # Try to solve it
    solve_kpam(problem)
    print(problem.T_action)

    # Transform the keypoint
    print(
        "The transformed bottom center is: ",
        SE3_utils.transform_point(problem.T_action, keypoint_loc[0, :]),
    )
    print(
        "The transformed handle center is: ",
        SE3_utils.transform_point(problem.T_action, keypoint_loc[1, :]),
    )
    print(
        "The transformed top center is: ",
        SE3_utils.transform_point(problem.T_action, keypoint_loc[2, :]),
    )


if __name__ == "__main__":
    test_builder()

core/expert/kpam/mp_builder.py
liruiw/Fleet-Tools
import numpy as np
import core.expert.kpam.SE3_utils as SE3_utils
from core.expert.kpam.optimization_problem import OptimizationProblemkPAM


def keypoint_l2_cost(problem, from_keypoints, goal_keypoints, weight, xyzrpy):
    # type: (OptimizationProblemkPAM, np.ndarray, np.ndarray, float, np.ndarray) -> np.ndarray
    """
    :param from_keypoints: (N, 3) np.ndarray with float
    :param goal_keypoints: (N, 3) np.ndarray with float
    :param weight: (N, 3) np.ndarray with float
    :param xyzrpy: np.ndarray with potentially symbolic variable
    :return: The cost value
    """
    cost = 0.0
    T_eps = SE3_utils.xyzrpy_to_matrix(xyzrpy)
    T = np.dot(problem.T_init, T_eps)

    # The iteration on keypoints
    num_keypoints = from_keypoints.shape[0]
    for i in range(num_keypoints):
        q = SE3_utils.transform_point(T, from_keypoints[i, :])
        q_goal = goal_keypoints[i, :]
        delta = q - q_goal
        cost += np.dot(delta, delta)

    return weight * cost


def point2plane_cost(problem, from_keypoints, goal_keypoints, weight, plane_normal, xyzrpy):
    cost = 0.0
    T_eps = SE3_utils.xyzrpy_to_matrix(xyzrpy)
    T = np.dot(problem.T_init, T_eps)

    # The iteration on keypoints
    num_keypoints = from_keypoints.shape[0]
    for i in range(num_keypoints):
        q = SE3_utils.transform_point(T, from_keypoints[i, :])
        q_goal = goal_keypoints[i, :]
        delta = q - q_goal
        delta_normal = np.dot(delta, plane_normal)
        cost += np.dot(delta_normal, delta_normal)

    return weight * cost


def vector_dot_symbolic(problem, from_axis, to_axis, xyzrpy):
    # type: (OptimizationProblemkPAM, np.ndarray, np.ndarray, np.ndarray) -> np.ndarray
    """
    dot(T.dot(from_axis), to_axis)
    :param problem:
    :param from_axis:
    :param to_axis:
    :param xyzrpy:
    :return:
    """
    T_eps = SE3_utils.xyzrpy_to_matrix_symbolic(xyzrpy)
    T = np.dot(problem.T_init, T_eps)
    R = T[0:3, 0:3]
    transformed_from = np.dot(R, from_axis)
    return np.dot(transformed_from, to_axis)


def transform_point_symbolic(problem, from_keypoint, xyzrpy):
    # type: (OptimizationProblemkPAM, np.ndarray, np.ndarray) -> np.ndarray
    """
    Transform the point using symbolic expression
    :param problem:
    :param from_keypoint: np.ndarray with float
    :param xyzrpy: np.ndarray with symbolic variable
    :return:
    """
    T_eps = SE3_utils.xyzrpy_to_matrix_symbolic(xyzrpy)
    T = np.dot(problem.T_init, T_eps)
    return SE3_utils.transform_point(T, from_keypoint)

core/expert/kpam/mp_terms.py
liruiw/Fleet-Tools
import attr
import numpy as np
import core.expert.kpam.SE3_utils as SE3_utils
from pydrake.all import MathematicalProgram, Solve
import IPython


class OptimizationProblemkPAM(object):
    # The property about initial transformation
    # There is always an initial guess, either identity or from keypoint matching
    T_init = np.ndarray(shape=(4, 4))  # type: np.ndarray

    # The final mathematical program
    mp = None  # type: MathematicalProgram
    xyzrpy = None  # type: np.ndarray

    # The solution to the program
    has_solution = False
    xyzrpy_sol = np.ndarray(shape=(6,))
    # T_action.dot(observed_keypoint) = target_keypoint
    T_action = np.ndarray(shape=(4, 4))

    # Build empty problem
    def build_empty_mp(self):
        # Construct the problem
        mp = MathematicalProgram()
        xyz = mp.NewContinuousVariables(3, "xyz")
        rpy = mp.NewContinuousVariables(3, "rpy")
        xyz_rpy = np.concatenate((xyz, rpy))
        mp.SetInitialGuessForAllVariables(np.zeros(6))

        # Store the result to problem
        self.mp = mp
        self.xyzrpy = xyz_rpy


class OptimizationProblemkPAMJoint(object):
    # The property about initial transformation
    # There is always an initial guess, either identity or from keypoint matching
    T_init = np.ndarray(shape=(4, 4))  # type: np.ndarray

    # The final mathematical program
    mp = None  # type: MathematicalProgram
    xyzrpy = None  # type: np.ndarray

    # The solution to the program
    has_solution = False
    xyzrpy_sol = np.ndarray(shape=(6,))
    # T_action.dot(observed_keypoint) = target_keypoint
    T_action = np.ndarray(shape=(4, 4))

    # Build empty problem
    def build_empty_mp(self):
        # Construct the problem
        mp = MathematicalProgram()
        xyz = mp.NewContinuousVariables(3, "xyz")
        rpy = mp.NewContinuousVariables(3, "rpy")
        xyz_rpy = np.concatenate((xyz, rpy))
        mp.SetInitialGuessForAllVariables(np.zeros(6))

        # Store the result to problem
        self.mp = mp
        self.xyzrpy = xyz_rpy



def solve_kpam(problem):  # type: (OptimizationProblemkPAM) -> bool
    result = Solve(problem.mp)  # mp.
    if not result.is_success():
        problem.has_solution = False
        return False

    # Save the result to problem
    problem.xyzrpy_sol = result.get_x_val()  # problem.mp.GetSolution(problem.xyzrpy)
    T_eps = SE3_utils.xyzrpy_to_matrix(xyzrpy=problem.xyzrpy_sol)
    problem.T_action = np.dot(problem.T_init, T_eps)
    problem.has_solution = True

    # OK
    return True

core/expert/kpam/optimization_problem.py
liruiw/Fleet-Tools
from core.expert.kpam.term_spec import (
    OptimizationTermSpec,
    Point2PointConstraintSpec,
    Point2PointCostL2Spec,
)
from typing import List, Dict
import copy
import yaml
import os
import IPython
import core.expert.kpam.term_spec as term_spec


class OptimizationProblemSpecification(object):
    """
    The class serves as the interface between the
    config file and solver
    """

    def __init__(
        self,
        task_name="",  # type: str
        category_name="",  # type: str
        keypoint_name_list=None,  # type: List[str]
    ):
        """
        The optimization spec can be constructed in python code
        or load from yaml config path. In latter case, these
        parameters can left default and use load_from_yaml method
        :param task_name:
        :param category_name:
        :param keypoint_name_list:
        """
        self._task_name = task_name  # type: str
        self._category_name = category_name  # type: str

        # The default construction of list
        if keypoint_name_list is not None:
            self._keypoint_name_list = keypoint_name_list  # type: List[str]
        else:
            self._keypoint_name_list = []

        # By default, nothing here
        self._cost_list = []  # type: List[OptimizationTermSpec]
        self._constraint_list = []  # type: List[OptimizationTermSpec]

        # Build from keypoint name list
        self._keypoint_name2idx = {}  # type: Dict[str, int]
        self.setup_keypoint_mapping()

        # The container for explicit nominal position
        # The nominal position can either explicit declared
        # Or implicit added using Point2PointCost/Constraint
        self._keypoint_nominal_target_position = {}  # type: Dict[str, List[float]]

    def setup_keypoint_mapping(self):
        self._keypoint_name2idx.clear()
        for i in range(len(self._keypoint_name_list)):
            name = self._keypoint_name_list[i]
            self._keypoint_name2idx[name] = i

    # The access interface
    @property
    def task_name(self):
        return self._task_name

    @property
    def category_name(self):
        return self._category_name

    @property
    def keypoint_name2idx(self):
        return self._keypoint_name2idx

    @property
    def cost_list(self):
        return self._cost_list

    @property
    def constraint_list(self):
        return self._constraint_list

    # The method to manipulate nominal target
    def add_nominal_target_position(self, keypoint_name, nominal_target):  # type: (str, List[float]) -> bool
        # Check the existence of keypoint
        if keypoint_name not in self._keypoint_name2idx:
            return False

        # Check the shape of target
        if len(nominal_target) != 3:
            return False

        # OK
        self._keypoint_nominal_target_position[keypoint_name] = nominal_target
        return True

    def get_nominal_target_position(self, keypoint_name):  # type: (str) -> (bool, List[float])
        # If explicitly defined
        if keypoint_name in self._keypoint_nominal_target_position:
            return True, self._keypoint_nominal_target_position[keypoint_name]

        # Else search for constraints
        for cost_term in self.cost_list:
            if isinstance(cost_term, Point2PointCostL2Spec):
                if cost_term.keypoint_name == keypoint_name:
                    return True, cost_term.target_position
        for constraint_term in self.constraint_list:
            if isinstance(constraint_term, Point2PointConstraintSpec):
                if constraint_term.keypoint_name == keypoint_name:
                    return True, constraint_term.target_position

        # Not available
        return False, []

    # The method to modify the specification from python
    def add_cost(self, cost_term):  # type: (OptimizationTermSpec) -> bool
        if not cost_term.is_cost():
            return False
        copied = copy.deepcopy(cost_term)
        self._cost_list.append(copied)
        return True

    def add_constraint(self, constraint_term):  # type: (OptimizationTermSpec) -> bool
        if constraint_term.is_cost():
            return False
        copied = copy.deepcopy(constraint_term)
        self._constraint_list.append(copied)
        return True

    def add_optimization_term(self, optimization_term):  # type: (OptimizationTermSpec) -> bool
        if optimization_term.is_cost():
            return self.add_cost(optimization_term)
        else:
            return self.add_constraint(optimization_term)

    # The interface from/to yaml
    def write_to_yaml(self, yaml_save_path):  # type: (str) -> None
        data_map = dict()
        data_map["task_name"] = self._task_name
        data_map["category_name"] = self._category_name
        data_map["keypoint_name_list"] = self._keypoint_name_list

        # For cost terms
        cost_map_list = []
        for cost in self._cost_list:
            cost_i_map = cost.to_dict()
            cost_i_map["type"] = cost.type_name()
            cost_map_list.append(cost_i_map)
        data_map["cost_list"] = cost_map_list

        # For constraint terms
        constraint_map_list = []
        for constraint in self._constraint_list:
            constraint_i_map = constraint.to_dict()
            constraint_i_map["type"] = constraint.type_name()
            constraint_map_list.append(constraint_i_map)
        data_map["constraint_list"] = constraint_map_list

        # For nominal location
        data_map["keypoint_nominal_target_position"] = self._keypoint_nominal_target_position

        # Save to yaml
        with open(yaml_save_path, mode="w") as save_file:
            yaml.dump(data_map, save_file)
            save_file.close()

    def load_from_yaml(self, yaml_load_path):  # type: (str) -> bool
        if not os.path.exists(yaml_load_path):
            return False

        # Read the file
        load_file = open(yaml_load_path, mode="r")
        data_map = yaml.load(load_file, Loader=yaml.SafeLoader)
        load_file.close()

        # Basic meta
        self._task_name = data_map["task_name"]
        self._category_name = data_map["category_name"]
        self._keypoint_name_list = data_map["keypoint_name_list"]
        self._keypoint_name2idx.clear()
        self.setup_keypoint_mapping()

        # The nominal location
        self._keypoint_nominal_target_position = data_map["keypoint_nominal_target_position"]

        # For cost terms

        cost_map_list = data_map["cost_list"]
        self._cost_list = []
        for cost in cost_map_list:

            cost_type = cost["type"]  # type: str
            if cost_type == term_spec.Point2PointCostL2Spec.type_name():
                cost_spec = term_spec.Point2PointCostL2Spec()
                cost_spec.from_dict(cost)
                self._cost_list.append(cost_spec)
            elif cost_type == term_spec.Point2PlaneCostSpec.type_name():
                cost_spec = term_spec.Point2PlaneCostSpec()
                cost_spec.from_dict(cost)
                self._cost_list.append(cost_spec)
            else:
                raise RuntimeError("Unknown cost type %s" % cost_type)

        # For constraint terms
        constraint_map_list = data_map["constraint_list"]
        self._constraint_list = []
        for constraint in constraint_map_list:
            constraint_type = constraint["type"]
            if constraint_type == term_spec.Point2PointConstraintSpec.type_name():
                constraint_spec = term_spec.Point2PointConstraintSpec()
                constraint_spec.from_dict(constraint)
                self._constraint_list.append(constraint_spec)
            elif constraint_type == term_spec.AxisAlignmentConstraintSpec.type_name():
                constraint_spec = term_spec.AxisAlignmentConstraintSpec()
                constraint_spec.from_dict(constraint)
                self._constraint_list.append(constraint_spec)
            elif constraint_type == term_spec.KeypointAxisAlignmentConstraintSpec.type_name():
                constraint_spec = term_spec.KeypointAxisAlignmentConstraintSpec()
                constraint_spec.from_dict(constraint)
                self._constraint_list.append(constraint_spec)
            elif constraint_type == term_spec.KeypointAxisOrthogonalConstraintSpec.type_name():
                constraint_spec = term_spec.KeypointAxisOrthogonalConstraintSpec()
                constraint_spec.from_dict(constraint)
                self._constraint_list.append(constraint_spec)
            else:
                raise RuntimeError("Unknown constraint type %s" % constraint_type)

        # OK
        return True


if __name__ == "__main__":
    opt = build_mug2shelf()
    print(opt.get_nominal_target_position("top_center"))

core/expert/kpam/optimization_spec.py
liruiw/Fleet-Tools
import attr
from typing import List


# A very thin base class for all
# the specification of the terms
class OptimizationTermSpec(object):
    # The type of this method
    @staticmethod
    def type_name():  # type: () -> str
        raise NotImplementedError

    @staticmethod
    def is_cost():  # type: () -> bool
        raise NotImplementedError

    # The method for serialization
    def to_dict(self):
        raise NotImplementedError

    def from_dict(self, data_map):
        raise NotImplementedError


# For the specification of constraint types:
@attr.s
class Point2PointConstraintSpec(OptimizationTermSpec):
    # 1 - tolerance <= dot(from, to) <= 1
    keypoint_name = ""  # type: str
    keypoint_idx = -1  # type: int
    target_position = [0, 0, 0]  # type: List[float]
    tolerance = 0.001  # type: float

    @staticmethod
    def type_name():
        return "point2point_constraint"

    @staticmethod
    def is_cost():  # type: () -> bool
        return False

    def to_dict(self):
        constraint_dict = dict()
        constraint_dict["keypoint_name"] = self.keypoint_name
        constraint_dict["keypoint_idx"] = self.keypoint_idx
        constraint_dict["target_position"] = self.target_position
        constraint_dict["tolerance"] = self.tolerance
        return constraint_dict

    def from_dict(self, data_map):
        self.keypoint_name = data_map["keypoint_name"]
        self.keypoint_idx = data_map["keypoint_idx"]
        self.target_position = data_map["target_position"]
        self.tolerance = data_map["tolerance"]


@attr.s
class AxisAlignmentConstraintSpec(OptimizationTermSpec):
    from_axis = [0, 0, 1]  # type: List[float]
    target_axis = [0, 0, 1]  # type: List[float]
    tolerance = 0.01  # type: float
    target_inner_product = 1

    @staticmethod
    def type_name():
        return "axis_alignment"

    @staticmethod
    def is_cost():  # type: () -> bool
        return False

    def to_dict(self):
        constraint_dict = dict()
        constraint_dict["from_axis"] = self.from_axis
        constraint_dict["target_axis"] = self.target_axis
        constraint_dict["tolerance"] = self.tolerance
        # constraint_dict['target_inner_product'] = self.target_inner_product
        return constraint_dict

    def from_dict(self, data_map):
        self.from_axis = data_map["from_axis"]
        self.target_axis = data_map["target_axis"]
        self.tolerance = data_map["tolerance"]
        # self.target_inner_product = data_map['target_inner_product']


@attr.s
class KeypointAxisAlignmentConstraintSpec(OptimizationTermSpec):
    axis_from_keypoint_idx = -1  # type: int
    axis_from_keypoint_name = ""  # type: str
    axis_to_keypoint_idx = -1  # type: int
    axis_to_keypoint_name = ""  # type: str
    target_axis = [0, 0, 1]  # type: List[float]
    tolerance = 0.01  # type: float
    target_inner_product = 1

    @staticmethod
    def type_name():
        return "keypoint_axis_alignment"

    @staticmethod
    def is_cost():  # type: () -> bool
        return False

    def to_dict(self):
        constraint_dict = dict()
        constraint_dict["axis_from_keypoint_name"] = self.axis_from_keypoint_name
        constraint_dict["axis_from_keypoint_idx"] = self.axis_from_keypoint_idx
        constraint_dict["axis_to_keypoint_name"] = self.axis_to_keypoint_name
        constraint_dict["axis_to_keypoint_idx"] = self.axis_to_keypoint_idx
        constraint_dict["target_axis"] = self.target_axis
        constraint_dict["tolerance"] = self.tolerance
        constraint_dict["target_inner_product"] = self.target_inner_product

        return constraint_dict

    def from_dict(self, constraint_dict):
        self.axis_from_keypoint_name = constraint_dict["axis_from_keypoint_name"]
        self.axis_from_keypoint_idx = constraint_dict["axis_from_keypoint_idx"]
        self.axis_to_keypoint_name = constraint_dict["axis_to_keypoint_name"]
        self.axis_to_keypoint_idx = constraint_dict["axis_to_keypoint_idx"]
        self.target_axis = constraint_dict["target_axis"]
        self.tolerance = constraint_dict["tolerance"]
        self.target_inner_product = constraint_dict["target_inner_product"]


@attr.s
class KeypointAxisOrthogonalConstraintSpec(OptimizationTermSpec):
    axis_from_keypoint_idx = -1  # type: int
    axis_from_keypoint_name = ""  # type: str
    axis_to_keypoint_idx = -1  # type: int
    axis_to_keypoint_name = ""  # type: str
    target_axis = [0, 0, 1]  # type: List[float]
    tolerance = 0.01  # type: float
    target_inner_product = 0

    @staticmethod
    def type_name():
        return "keypoint_axis_orthogonal"

    @staticmethod
    def is_cost():  # type: () -> bool
        return False

    def to_dict(self):
        constraint_dict = dict()
        constraint_dict["axis_from_keypoint_name"] = self.axis_from_keypoint_name
        constraint_dict["axis_from_keypoint_idx"] = self.axis_from_keypoint_idx
        constraint_dict["axis_to_keypoint_name"] = self.axis_to_keypoint_name
        constraint_dict["axis_to_keypoint_idx"] = self.axis_to_keypoint_idx
        constraint_dict["target_axis"] = self.target_axis
        constraint_dict["tolerance"] = self.tolerance
        constraint_dict["target_inner_product"] = self.target_inner_product
        return constraint_dict

    def from_dict(self, constraint_dict):
        self.axis_from_keypoint_name = constraint_dict["axis_from_keypoint_name"]
        self.axis_from_keypoint_idx = constraint_dict["axis_from_keypoint_idx"]
        self.axis_to_keypoint_name = constraint_dict["axis_to_keypoint_name"]
        self.axis_to_keypoint_idx = constraint_dict["axis_to_keypoint_idx"]
        self.target_axis = constraint_dict["target_axis"]
        self.tolerance = constraint_dict["tolerance"]
        self.target_inner_product = constraint_dict["target_inner_product"]


# For the specification of cost terms
@attr.s
class Point2PointCostL2Spec(OptimizationTermSpec):
    keypoint_name = ""  # type: str
    keypoint_idx = -1  # type: int
    target_position = [0, 0, 0]  # type: List[float]
    penalty_weight = 1.0  # type: float

    @staticmethod
    def type_name():
        return "point2point_cost"

    @staticmethod
    def is_cost():  # type: () -> bool
        return True

    def to_dict(self):
        constraint_dict = dict()
        constraint_dict["keypoint_name"] = self.keypoint_name
        constraint_dict["keypoint_idx"] = self.keypoint_idx
        constraint_dict["target_position"] = self.target_position
        constraint_dict["penalty_weight"] = self.penalty_weight
        return constraint_dict

    def from_dict(self, data_map):
        self.keypoint_name = data_map["keypoint_name"]
        self.keypoint_idx = data_map["keypoint_idx"]
        self.target_position = data_map["target_position"]
        self.penalty_weight = data_map["penalty_weight"]


# For the specification of point2plane cost
@attr.s
class Point2PlaneCostSpec(OptimizationTermSpec):
    keypoint_name = ""  # type: str
    keypoint_idx = -1  # type: int
    target_position = [0.0, 0.0, 0.0]  # type: List[float]
    plane_normal = [0.0, 0.0, 1.0]  # type: List[float]
    penalty_weight = 1.0  # type: float

    @staticmethod
    def type_name():  # type: () -> str
        return "point2plane_cost"

    @staticmethod
    def is_cost():  # type: () -> bool
        return True

    def to_dict(self):
        constraint_dict = dict()
        constraint_dict["keypoint_name"] = self.keypoint_name
        constraint_dict["keypoint_idx"] = self.keypoint_idx
        constraint_dict["target_position"] = self.target_position
        constraint_dict["plane_normal"] = self.plane_normal
        constraint_dict["penalty_weight"] = self.penalty_weight
        return constraint_dict

    def from_dict(self, data_map):
        self.keypoint_name = data_map["keypoint_name"]
        self.keypoint_idx = data_map["keypoint_idx"]
        self.target_position = data_map["target_position"]
        self.plane_normal = data_map["plane_normal"]
        self.penalty_weight = data_map["penalty_weight"]

core/expert/kpam/term_spec.py
liruiw/Fleet-Tools
import IPython
import warnings

warnings.simplefilter("ignore")

from experiments.args import get_parser
import os.path as osp
from core.experiment import *
from core.expert import *
from core.utils.utils import *

from core.agent.agent import Agent, AgentWrapper
from pydrake.all import RandomGenerator
from collections import deque

import IPython
import numpy as np
import GPUtil
import psutil
import hydra
import os
from omegaconf import OmegaConf

os.environ["WANDB_SILENT"] = "true"
MEMORY_THRE = 75


def choose_setup(args):
    actor_wrapper = "ActorWrapperGPU05"
    max_memory = psutil.virtual_memory().available * 0.5
    print("worker: {} memory: {}".format(args.num_envs, int(max_memory)))
    return actor_wrapper, max_memory


def monitor_usage():
    """monitor whether gpu or cpu memory has reached a state that needs reset"""
    gpu_usage, memory_usage = get_usage()
    GPUs = GPUtil.getGPUs()
    gpu_limit = np.max([GPU.memoryTotal for GPU in GPUs])
    gpu_max = float(gpu_usage) / gpu_limit > 0.98
    memory_max = memory_usage >= MEMORY_THRE
    print(f"==================== Memory: {memory_usage} GPU: {gpu_usage} =====================")

    return memory_max or gpu_max


def reinit(robot):
    """reiniting ray for memory leak"""
    print("============================ Reinit ==========================")
    robot.config.start_episode_position = (
        ray.get(robot.fleet_agent.get_exp_stats.remote())
        if robot.config.parallel
        else robot.fleet_agent.get_exp_stats()
    )
    os.system("nvidia-smi")
    print(f"===================== Ray Reinit Episode: {robot.config.start_episode_position} =================")
    if robot.config.parallel:
        curr_stat = ray.get([r.return_stat.remote() for r in robot.robots])
    else:
        curr_stat = robot.robots.return_stat()

    time.sleep(4)
    ray.shutdown()
    time.sleep(4)
    from env.env_util import g

    # easy way to generate stochasty
    g = RandomGenerator(np.random.randint(1e10) + robot.config.start_episode_position)
    np.random.seed(g())
    robot = RobotFleet(robot.config)  # reinitialize
    print(None, "==============================================================")
    if robot.config.parallel:
        ray.get([r.set_stat.remote(curr_stat[idx]) for idx, r in enumerate(robot.robots)])
    else:
        robot.robots.set_stat(curr_stat)
    return robot


def run_episode(robot):
    """run the training experiment episode by episode"""
    fleet_steps = 0
    saved_episodes = 0

    if robot.config.eval:
        robot.fleet_agent.load(robot.config.eval)

    if robot.config.record_video and fleet_steps <= 3:
        from env.env_util import visualizer, meshcat

        if visualizer is not None:
            animation = visualizer.StartRecording(set_transforms_while_recording=False)

    total_rewards, total_lengths, scene_descriptions = [], [], []

    while fleet_steps < robot.config.num_episode:
        # load weights
        if robot.config.parallel:
            res = ray.get([r.rollout.remote(1) for r in robot.robots])
            info = [r[0] for r in res]
            success = [r[1] for r in res]
            info = {k: np.mean([info_i[k] for info_i in info if k in info_i]) for k in info[0].keys()}
        else:
            info, success = robot.robots.rollout(1)

        if robot.config.training:
            train_info = (
                ray.get(robot.fleet_agent.train.remote(fleet_steps))
                if robot.config.parallel
                else robot.fleet_agent.train(fleet_steps)
            )
            info = merge_two_dicts(info, train_info)

        info["total_time"] = time.time() - robot.start_time
        info["log_dir"] = robot.logdir
        info["global_episode"] = fleet_steps
        saved_episodes += np.sum(success)
        success = 0

        if robot.config.record_video and fleet_steps <= 5 and visualizer is not None:
            # and fleet_steps % 10 == 0 might only work with one runner
            visualizer.StopRecording()
            visualizer.PublishRecording()
            html = meshcat.StaticHtml()

            task_dir = f"assets/demonstrations/{robot.config.task.tool_class_name}for{robot.config.task.task_name}/episode_{fleet_steps}"
            mkdir_if_missing(task_dir)

            # it should be the last one
            with open(f"{task_dir}/recording.html", "w") as f:
                print("write to video path", task_dir)
                f.write(html)
            visualizer.DeleteRecording()  # delete
            animation = visualizer.StartRecording(set_transforms_while_recording=False)

            # restart recording after certain number of episodes
        if "episode_length" in info:
            total_lengths.append(info["episode_length"])
            total_rewards.append(info["local_episode_reward"])
            scene_descriptions.append((info["tool_class_name"], info["tool_name"], info["object_name"]))

        # reiniting
        if fleet_steps % robot.config.reinit_num == 0 or monitor_usage():
            robot = reinit(robot)

        fleet_steps += 1
        if saved_episodes > robot.config.max_episodes:
            break

    print("total number of saved trajectories:", robot.config.max_episodes)
    print(f"log rollout results: {len(total_rewards)} {np.sum(total_rewards)} {np.mean(total_lengths)}")


@hydra.main(config_path="../experiments/config", config_name="config", version_base="1.2")
def main(cfg):

    print(OmegaConf.to_yaml(cfg, resolve=True))
    np.random.seed(cfg.seed)

    cfg.parallel = cfg.num_envs > 1
    actor_wrapper, max_memory = choose_setup(cfg)
    cfg.actor_wrapper = actor_wrapper
    cfg.max_memory = max_memory
    cfg.task.random_seed = cfg.seed

    # Create experiment and run it
    if cfg.parallel:
        ray.init(
            _temp_dir=f"{os.environ['HOME']}/tmp",
            num_cpus=cfg.num_envs + 1,
            object_store_memory=int(cfg["max_memory"]),
        )
    robot = RobotFleet(cfg)
    save_args_hydra(robot.logdir, cfg)

    start_time = time.time()
    run_episode(robot)
    print("rollout total time: {:.3f}".format(time.time() - start_time))


if __name__ == "__main__":
    main()

core/run.py
liruiw/Fleet-Tools
# --------------------------------------------------------
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------

import numpy as np
import os
import sys
import random

import scipy.io as sio
import IPython
import time

import cv2

import matplotlib.pyplot as plt
import tabulate
import yaml


import copy
import math

# transform stuff
from transforms3d.quaternions import *
from transforms3d.euler import *
from transforms3d.axangles import *

# from core.utils.torch_utils import *


def normalize(v, axis=None, eps=1e-10):
    """L2 Normalize along specified axes."""
    return v / max(np.linalg.norm(v, axis=axis, keepdims=True), eps)


def inv_lookat(eye, target=[0, 0, 0], up=[0, 1, 0]):
    """Generate LookAt matrix."""
    eye = np.float32(eye)
    forward = normalize(target - eye)
    side = normalize(np.cross(forward, up))
    up = np.cross(side, forward)
    R = np.stack([side, up, -forward], axis=-1)
    return R


def rotZ(rotz):
    RotZ = np.array(
        [
            [np.cos(rotz), -np.sin(rotz), 0, 0],
            [np.sin(rotz), np.cos(rotz), 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1],
        ]
    )
    return RotZ


def rotY(roty):
    RotY = np.array(
        [
            [np.cos(roty), 0, np.sin(roty), 0],
            [0, 1, 0, 0],
            [-np.sin(roty), 0, np.cos(roty), 0],
            [0, 0, 0, 1],
        ]
    )
    return RotY


def rotX(rotx):
    RotX = np.array(
        [
            [1, 0, 0, 0],
            [0, np.cos(rotx), -np.sin(rotx), 0],
            [0, np.sin(rotx), np.cos(rotx), 0],
            [0, 0, 0, 1],
        ]
    )
    return RotX


def ros_quat(tf_quat):  # wxyz -> xyzw
    quat = np.zeros(4)
    quat[-1] = tf_quat[0]
    quat[:-1] = tf_quat[1:]
    return quat


def tf_quat(ros_quat):  # xyzw -> wxyz
    quat = np.zeros(4)
    quat[0] = ros_quat[-1]
    quat[1:] = ros_quat[:-1]
    return quat


def se3_inverse(RT):
    RT = RT.reshape(4, 4)
    R = RT[:3, :3]
    T = RT[:3, 3].reshape((3, 1))
    RT_new = np.eye(4, dtype=np.float32)
    RT_new[:3, :3] = R.transpose()
    RT_new[:3, 3] = -1 * np.dot(R.transpose(), T).reshape((3))
    return RT_new


def se3_inverse_tensor(RT):
    R = RT[:, :3, :3]
    T = RT[:, :3, 3].view((-1, 3, 1))
    RT_new = torch.eye(4).repeat((len(R), 1, 1))
    RT_new[:, :3, :3] = R.permute(0, 2, 1)
    RT_new[:, :3, 3] = -1 * torch.matmul(R.permute(0, 2, 1), T).view(-1, 3)
    return RT_new


def unpack_action_batch(actions):
    pose_deltas = []
    for action in actions:
        pose_delta = np.eye(4)
        pose_delta[:3, :3] = euler2mat(action[3], action[4], action[5])
        pose_delta[:3, 3] = action[:3]
        pose_deltas.append(pose_delta)
    return np.stack(pose_deltas, axis=0)


def unpack_action(action, rot_type="euler"):
    if rot_type == "euler":
        rot_func = euler2mat_
    elif rot_type == "axangle":
        rot_func = axangle2mat_

    pose_delta = np.eye(4)
    pose_delta[:3, :3] = rot_func(action[3:])
    pose_delta[:3, 3] = action[:3]

    return pose_delta


def invert_action_batch(actions):
    pose_deltas = []
    for action in actions:
        pose_delta = np.eye(4)
        pose_delta[:3, :3] = euler2mat(action[3], action[4], action[5])
        pose_delta[:3, 3] = action[:3]
        pose_deltas.append(pack_pose(se3_inverse(pose_delta)))
    return np.stack(pose_deltas, axis=0)


def invert_action(action):
    pose_delta = np.eye(4)
    pose_delta[:3, :3] = euler2mat(action[3], action[4], action[5])
    pose_delta[:3, 3] = action[:3]
    return pack_pose(se3_inverse(pose_delta))


def invert_action_pose(action):
    pose_delta = np.eye(4)
    pose_delta[:3, :3] = euler2mat(action[3], action[4], action[5])
    pose_delta[:3, 3] = action[:3]
    return se3_inverse(pose_delta)


def unpack_poses_batch(pose):
    pose_mat = torch.eye(4, device=pose.device).repeat(len(pose), 1, 1)
    pose_mat[:, :3, 3] = pose[:, :3]
    pose_mat[:, :3, :3] = quat2mat_batch(pose[:, 3:7])
    return pose_mat


def relative_pose_between_bodies(pose1, pose2):
    """
    get the relative pose from pose2 to pose1
    """
    pose_mat1 = unpack_poses_batch(pose1)
    pose_mat2 = unpack_poses_batch(pose2)
    return se3_inverse_tensor(pose_mat1) @ pose_mat2


def unpack_pose_pytorch3d(pose):
    pose_mat = torch.eye(4, device=pose.device).repeat(len(pose), 1, 1)
    pose_mat[:, :3, 3] = pose[:, :3]
    pose_mat[:, :3, :3] = quat2mat_batch(pose[:, 3:7])
    return pose_mat


def se3_inverse_batch(RT):
    R = RT[:, :3, :3]
    T = RT[:, :3, 3].reshape((-1, 3, 1))
    RT_new = np.tile(np.eye(4, dtype=np.float32), (len(R), 1, 1))
    RT_new[:, :3, :3] = R.transpose(0, 2, 1)
    RT_new[:, :3, 3] = -1 * np.matmul(R.transpose(0, 2, 1), T).reshape(-1, 3)
    return RT_new


def unpack_action_pytorch3d(pose):
    pose_mat = torch.eye(4, device=pose.device).repeat(len(pose), 1, 1)
    pose_mat[:, :3, 3] = pose[:, :3]
    pose_mat[:, :3, :3] = to_rotation_matrix(pose[:, 3:])
    return pose_mat


def mat2euler_torch(rot_mat, seq="xyz"):
    """
    convert rotation matrix to euler angle
    :param rot_mat: rotation matrix rx*ry*rz [B, 3, 3]
    :param seq: seq is xyz(rotate along z first) or zyx
    :return: three angles, x, y, z
    """
    r11 = rot_mat[:, 0, 0]
    r12 = rot_mat[:, 0, 1]
    r13 = rot_mat[:, 0, 2]
    r21 = rot_mat[:, 1, 0]
    r22 = rot_mat[:, 1, 1]
    r23 = rot_mat[:, 1, 2]
    r31 = rot_mat[:, 2, 0]
    r32 = rot_mat[:, 2, 1]
    r33 = rot_mat[:, 2, 2]
    if seq == "xyz":
        z = torch.atan2(-r12, r11)
        y = torch.asin(r13)
        x = torch.atan2(-r23, r33)
    else:
        y = torch.asin(-r31)
        x = torch.atan2(r32, r33)
        z = torch.atan2(r21, r11)
    return torch.stack((z, y, x), dim=1)


def pack_action_pytorch3d(pose_mat):
    pose = torch.zeros(len(pose_mat), 6).to(pose_mat.device)
    pose[:, :3] = pose_mat[:, :3, 3]
    pose[:, 3:6] = mat2euler_torch(pose_mat[:, :3, :3])
    return pose


def euler2mat_(euler):
    return euler2mat(euler[0], euler[1], euler[2])


def unpack_pose(pose, rot_first=False, rot_type="quat"):
    rot_index = 4 if rot_type == "quat" else 3

    if rot_type == "quat":
        rot_func = quat2mat
    elif rot_type == "euler":
        rot_func = euler2mat_
    elif rot_type == "axangle":
        rot_func = axangle2mat_

    unpacked = np.eye(4)
    if rot_first:
        unpacked[:3, :3] = rot_func(pose[:rot_index])
        unpacked[:3, 3] = pose[rot_index:]
    else:
        unpacked[:3, :3] = rot_func(pose[rot_index:])
        unpacked[:3, 3] = pose[:rot_index]
    return unpacked


def _hat(x):
    """
    hat operator, compute skew matrix for cross product
    """
    return np.array([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]])


def rodrigue_formula(p, q):
    """
    allocentric to egocentric
    """
    r = _hat(np.cross(p, q))
    R_a2e = np.eye(3) + r + r.dot(r) / (1 + np.dot(p, q))
    return R_a2e


def mat2axangle_(mat):
    axangle = mat2axangle(mat[:3, :3])
    return axangle[0] * axangle[1]


def axangle2mat_(axangle):
    angle = np.linalg.norm(axangle)
    axis = axangle / angle
    return axangle2mat(axis, angle)


def action_rot_repr_from_euler_to_axangle(action):
    return pack_pose(unpack_action(action), rot_type="axangle")


def action_rot_repr_from_axangle_to_euler(action):
    return pack_pose(unpack_action(action, rot_type="axangle"))


def pack_pose(pose, rot_type="euler"):  # for action
    rot_index = 4 if rot_type == "quat" else 3
    if rot_type == "quat":
        rot_func = mat2quat
    elif rot_type == "euler":
        rot_func = mat2euler
    elif rot_type == "axangle":
        rot_func = mat2axangle_
    packed = np.zeros(3 + rot_index)
    packed[3:] = rot_func(pose[:3, :3])
    packed[:3] = pose[:3, 3]
    return packed


def rotZ(rotz):
    RotZ = np.array(
        [
            [np.cos(rotz), -np.sin(rotz), 0, 0],
            [np.sin(rotz), np.cos(rotz), 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1],
        ]
    )
    return RotZ


def rotY(roty):
    RotY = np.array(
        [
            [np.cos(roty), 0, np.sin(roty), 0],
            [0, 1, 0, 0],
            [-np.sin(roty), 0, np.cos(roty), 0],
            [0, 0, 0, 1],
        ]
    )
    return RotY


def rotX(rotx):
    RotX = np.array(
        [
            [1, 0, 0, 0],
            [0, np.cos(rotx), -np.sin(rotx), 0],
            [0, np.sin(rotx), np.cos(rotx), 0],
            [0, 0, 0, 1],
        ]
    )
    return RotX


def normalize(v, axis=None, eps=1e-10):
    """L2 Normalize along specified axes."""
    return v / max(np.linalg.norm(v, axis=axis, keepdims=True), eps)


def compose_orn(curr_quat, action_rpy):
    q = quat_from_euler_xyz(action_rpy[..., 0], action_rpy[..., 1], action_rpy[..., 2])
    return quat_mul(curr_quat, q)


def quat2mat_batch(norm_quat):
    """Convert quaternion coefficients to rotation matrix."""
    w, x, y, z = norm_quat[:, 0], norm_quat[:, 1], norm_quat[:, 2], norm_quat[:, 3]
    B = norm_quat.size(0)

    w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
    wx, wy, wz = w * x, w * y, w * z
    xy, xz, yz = x * y, x * z, y * z

    rotMat = torch.stack(
        [
            w2 + x2 - y2 - z2,
            2 * xy - 2 * wz,
            2 * wy + 2 * xz,
            2 * wz + 2 * xy,
            w2 - x2 + y2 - z2,
            2 * yz - 2 * wx,
            2 * xz - 2 * wy,
            2 * wx + 2 * yz,
            w2 - x2 - y2 + z2,
        ],
        dim=1,
    ).reshape(B, 3, 3)
    return rotMat


def get_frame_spatial_velocity(frame_A, frame_B, dt):
    """ """
    rel_pose = se3_inverse(frame_A) @ frame_B
    translation_diff = rel_pose[:3, 3]
    rot_diff = mat2axangle_(rel_pose[:3, :3])
    v = np.concatenate((translation_diff, rot_diff)) / dt
    # print("velocity:", v)
    # zero velocity
    return np.zeros_like(v)

core/utils/geometry.py
liruiw/Fleet-Tools
import numpy as np
import threading
import time
import os
from ppadb.client import Client as AdbClient
import sys


class FPSCounter:
    def __init__(self):
        current_time = time.time()
        self.start_time_for_display = current_time
        self.last_time = current_time
        self.x = 5  # displays the frame rate every X second
        self.time_between_calls = []
        self.elements_for_mean = 50

    def getAndPrintFPS(self, print_fps=True):
        current_time = time.time()
        self.time_between_calls.append(1.0 / (current_time - self.last_time + 1e-9))
        if len(self.time_between_calls) > self.elements_for_mean:
            self.time_between_calls.pop(0)
        self.last_time = current_time
        frequency = np.mean(self.time_between_calls)
        if (current_time - self.start_time_for_display) > self.x and print_fps:
            print("Frequency: {}Hz".format(int(frequency)))
            self.start_time_for_display = current_time
        return frequency


def parse_buttons(text):
    split_text = text.split(",")
    buttons = {}
    if "R" in split_text:  # right hand if available
        split_text.remove("R")  # remove marker
        buttons.update(
            {
                "A": False,
                "B": False,
                "RThU": False,  # indicates that right thumb is up from the rest position
                "RJ": False,  # joystick pressed
                "RG": False,  # boolean value for trigger on the grip (delivered by SDK)
                "RTr": False,  # boolean value for trigger on the index finger (delivered by SDK)
            }
        )
        # besides following keys are provided:
        # 'rightJS' / 'leftJS' - (x, y) position of joystick. x, y both in range (-1.0, 1.0)
        # 'rightGrip' / 'leftGrip' - float value for trigger on the grip in range (0.0, 1.0)
        # 'rightTrig' / 'leftTrig' - float value for trigger on the index finger in range (0.0, 1.0)

    if "L" in split_text:  # left hand accordingly
        split_text.remove("L")  # remove marker
        buttons.update({"X": False, "Y": False, "LThU": False, "LJ": False, "LG": False, "LTr": False})
    for key in buttons.keys():
        if key in list(split_text):
            buttons[key] = True
            split_text.remove(key)
    for elem in split_text:
        split_elem = elem.split(" ")
        if len(split_elem) < 2:
            continue
        key = split_elem[0]
        value = tuple([float(x) for x in split_elem[1:]])
        buttons[key] = value
    return buttons


def eprint(*args, **kwargs):
    RED = "\033[1;31m"
    sys.stderr.write(RED)
    print(*args, file=sys.stderr, **kwargs)
    RESET = "\033[0;0m"
    sys.stderr.write(RESET)


class OculusReader:
    def __init__(self, ip_address=None, port=5555, APK_name="com.rail.oculus.teleop", print_FPS=False, run=True):
        self.running = False
        self.last_transforms = {}
        self.last_buttons = {}
        self._lock = threading.Lock()
        self.tag = "wE9ryARX"

        self.ip_address = ip_address
        self.port = port
        self.APK_name = APK_name
        self.print_FPS = print_FPS
        if self.print_FPS:
            self.fps_counter = FPSCounter()

        self.device = self.get_device()
        self.install(verbose=False)
        if run:
            self.run()

    def __del__(self):
        self.stop()

    def run(self):
        self.running = True
        self.device.shell(
            'am start -n "com.rail.oculus.teleop/com.rail.oculus.teleop.MainActivity" -a android.intent.action.MAIN -c android.intent.category.LAUNCHER'
        )
        self.thread = threading.Thread(target=self.device.shell, args=("logcat -T 0", self.read_logcat_by_line))
        self.thread.start()

    def stop(self):
        self.running = False
        if hasattr(self, "thread"):
            self.thread.join()

    def get_network_device(self, client, retry=0):
        try:
            client.remote_connect(self.ip_address, self.port)
        except RuntimeError:
            os.system("adb devices")
            client.remote_connect(self.ip_address, self.port)
        device = client.device(self.ip_address + ":" + str(self.port))

        if device is None:
            if retry == 1:
                os.system("adb tcpip " + str(self.port))
            if retry == 2:
                eprint(
                    "Make sure that device is running and is available at the IP address specified as the OculusReader argument `ip_address`."
                )
                eprint("Currently provided IP address:", self.ip_address)
                eprint("Run `adb shell ip route` to verify the IP address.")
                exit(1)
            else:
                self.get_device(client=client, retry=retry + 1)
        return device

    def get_usb_device(self, client):
        try:
            devices = client.devices()
        except RuntimeError:
            os.system("adb devices")
            devices = client.devices()
        for device in devices:
            if device.serial.count(".") < 3:
                return device
        eprint("Device not found. Make sure that device is running and is connected over USB")
        eprint("Run `adb devices` to verify that the device is visible.")
        exit(1)

    def get_device(self):
        # Default is "127.0.0.1" and 5037
        client = AdbClient(host="127.0.0.1", port=5037)
        if self.ip_address is not None:
            return self.get_network_device(client)
        else:
            return self.get_usb_device(client)

    def install(self, APK_path=None, verbose=True, reinstall=False):
        try:
            installed = self.device.is_installed(self.APK_name)
            if not installed or reinstall:
                if APK_path is None:
                    APK_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "APK", "teleop-debug.apk")
                success = self.device.install(APK_path, test=True, reinstall=reinstall)
                installed = self.device.is_installed(self.APK_name)
                if installed and success:
                    print("APK installed successfully.")
                else:
                    eprint("APK install failed.")
            elif verbose:
                print("APK is already installed.")
        except RuntimeError:
            eprint("Device is visible but could not be accessed.")
            eprint("Run `adb devices` to verify that the device is visible and accessible.")
            eprint(
                'If you see "no permissions" next to the device serial, please put on the Oculus Quest and allow the access.'
            )
            exit(1)

    def uninstall(self, verbose=True):
        try:
            installed = self.device.is_installed(self.APK_name)
            if installed:
                success = self.device.uninstall(self.APK_name)
                installed = self.device.is_installed(self.APK_name)
                if not installed and success:
                    print("APK uninstall finished.")
                    print('Please verify if the app disappeared from the list as described in "UNINSTALL.md".')
                    print(
                        "For the resolution of this issue, please follow https://github.com/Swind/pure-python-adb/issues/71."
                    )
                else:
                    eprint("APK uninstall failed")
            elif verbose:
                print("APK is not installed.")
        except RuntimeError:
            eprint("Device is visible but could not be accessed.")
            eprint("Run `adb devices` to verify that the device is visible and accessible.")
            eprint(
                'If you see "no permissions" next to the device serial, please put on the Oculus Quest and allow the access.'
            )
            exit(1)

    @staticmethod
    def process_data(string):
        try:
            transforms_string, buttons_string = string.split("&")
        except ValueError:
            return None, None
        split_transform_strings = transforms_string.split("|")
        transforms = {}
        for pair_string in split_transform_strings:
            transform = np.empty((4, 4))
            pair = pair_string.split(":")
            if len(pair) != 2:
                continue
            left_right_char = pair[0]  # is r or l
            transform_string = pair[1]
            values = transform_string.split(" ")
            c = 0
            r = 0
            count = 0
            for value in values:
                if not value:
                    continue
                transform[r][c] = float(value)
                c += 1
                if c >= 4:
                    c = 0
                    r += 1
                count += 1
            if count == 16:
                transforms[left_right_char] = transform
        buttons = parse_buttons(buttons_string)
        return transforms, buttons

    def extract_data(self, line):
        output = ""
        if self.tag in line:
            try:
                output += line.split(self.tag + ": ")[1]
            except ValueError:
                pass
        return output

    def get_transformations_and_buttons(self):
        with self._lock:
            return self.last_transforms, self.last_buttons

    def read_logcat_by_line(self, connection):
        file_obj = connection.socket.makefile()
        while self.running:
            try:
                line = file_obj.readline().strip()
                data = self.extract_data(line)
                if data:
                    transforms, buttons = OculusReader.process_data(data)
                    with self._lock:
                        self.last_transforms, self.last_buttons = transforms, buttons
                    if self.print_FPS:
                        self.fps_counter.getAndPrintFPS()
            except UnicodeDecodeError:
                pass
        file_obj.close()
        connection.close()


def main():
    oculus_reader = OculusReader()

    while True:
        time.sleep(0.3)
        print(oculus_reader.get_transformations_and_buttons())


if __name__ == "__main__":
    main()
core/utils/oculus_reader.py
liruiw/Fleet-Tools
# https://github.com/pangtao22/robotics_utilities/blob/master/iiwa_controller/robot_internal_controller.py

import numpy as np

from pydrake.systems.framework import LeafSystem, BasicVector, PortDataType
from pydrake.multibody.tree import MultibodyForces, JacobianWrtVariable
from core.utils import *


class RobotInternalController(LeafSystem):
    def __init__(
        self,
        plant_robot,
        joint_stiffness,
        damping_ratio=1,
        kd=None,
        controller_mode="inverse_dynamics",
        name="robot_internal_controller",
    ):
        """
        Impedance controller implements the controller in Ott2008 for IIWA.
        Inverse dynamics controller makes use of drake's
            InverseDynamicsController.
        :param plant_robot:
        :param joint_stiffness: (nq,) numpy array which defines the stiffness
            of all joints.
        """
        LeafSystem.__init__(self)
        self.set_name(name)
        self.plant = plant_robot
        self.context = plant_robot.CreateDefaultContext()

        self.nq = plant_robot.num_positions()
        self.nv = plant_robot.num_velocities()
        self.robot_state_input_port = self.DeclareInputPort(
            "robot_state", PortDataType.kVectorValued, self.nq + self.nv
        )
        self.tau_feedforward_input_port = self.DeclareInputPort("tau_feedforward", PortDataType.kVectorValued, self.nq)
        self.joint_angle_commanded_input_port = self.DeclareInputPort(
            "q_robot_commanded", PortDataType.kVectorValued, self.nq
        )
        self.joint_torque_output_port = self.DeclareVectorOutputPort(
            "joint_torques", BasicVector(self.nv), self.CalcJointTorques
        )

        # control rate
        self.t = 0
        self.control_period = 1e-3  # 1000Hz.
        self.DeclareDiscreteState(self.nv)
        self.DeclarePeriodicDiscreteUpdate(period_sec=self.control_period)

        # joint velocity estimator
        self.q_prev = None
        self.w_cutoff = 2 * np.pi * 4000
        self.velocity_estimator = LowPassFilter(self.nv, self.control_period, self.w_cutoff)
        self.acceleration_estimator = LowPassFilter(self.nv, self.control_period, self.w_cutoff)

        # damping coefficient filter
        self.Kv_filter = LowPassFilter(self.nq, self.control_period, 2 * np.pi)

        # controller gains
        self.Kp = joint_stiffness
        self.damping_ratio = damping_ratio
        self.Kv = 2 * self.damping_ratio * np.sqrt(self.Kp)
        if kd is not None:
            self.Kv = np.array(kd)
        self.controller_mode = controller_mode

        # logs
        self.Kv_log = []
        self.tau_stiffness_log = []
        self.tau_damping_log = []
        self.sample_times = []
        self.cartesian_impedance_values = np.array([20, 20, 20, 200, 200, 200])
        self.joint_impedance_values = np.array([120, 100, 100, 80, 20, 50, 20, 20])
        self.KI = np.ones(9)
        self.integral_errors = np.zeros(9)

        # https://github.com/rachelholladay/franka_ros_interface/blob/master/franka_ros_controllers/config/ros_controllers.yaml
        self.joint_min_limit = np.array([-2.8973, -1.7628, -2.8973, -3.0718, -2.8973, -0.0175, -2.8973, 0, 0])
        self.joint_max_limit = np.array([2.8973, 1.7628, 2.8973, -0.0698, 2.8973, 3.7525, 2.8973, 0.04, 0.04])
        self.joint_null_space = (self.joint_min_limit + self.joint_max_limit) / 2
        self.effort_limit = np.array([87, 87, 87, 87, 12, 12, 12, 100, 100])

    def get_input_port_estimated_state(self):
        return self.robot_state_input_port

    def get_input_port_desired_state(self):
        return self.joint_angle_commanded_input_port

    def get_output_port_control(self):
        return self.joint_torque_output_port

    def get_torque_feedforward_port_control(self):
        return self.tau_feedforward_input_port

    # adding hacks for cartesian force control
    def set_desired_cartesian_force(self):
        pass

    def set_cartesian_controller_parameters(self, set_point, stiffness=None):
        if stiffness is not None:
            self.Kp = stiffness
        self.Kv = 2 * self.damping_ratio * np.sqrt(self.Kp)
        self.cartesian_setpoint = set_point  # check frame

    def compute_cartesian_impedance_control(self, context, v, tau):
        # https://github.com/NVIDIA-Omniverse/IsaacGymEnvs/blob/ca7a4fb762f9581e39cc2aab644f18a83d6ab0ba/isaacgymenvs/tasks/factory/factory_control.py#L270
        # https://github.com/rachelholladay/franka_ros_interface/blob/master/franka_ros_controllers/src/cartesian_impedance_controller.cpp
        # could be tricky to choose the coordinate frame
        x = self.robot_state_input_port.Eval(context)
        q = x[: self.nq]
        v = x[self.nq :]

        self.joint_jacobian = self.plant.CalcJacobianSpatialVelocity(
            self.context,
            JacobianWrtVariable.kQDot,
            self.plant.GetBodyByName("panda_hand").body_frame(),
            [0, 0, 0],
            self.plant.world_frame(),
            self.plant.world_frame(),
        )[:, :7]

        # compute error
        self.curr_ee = np.array(self.plant.GetFrameByName("panda_hand").CalcPoseInWorld(self.context).GetAsMatrix4())
        self.cartesian_error = se3_inverse(self.curr_ee) @ self.cartesian_setpoint
        # print("self.cartesian_error:", self.cartesian_error)
        axangle_error = mat2axangle(self.cartesian_error[:3, :3])
        axangle_error = axangle_error[0] * axangle_error[1]
        cartesian_error = np.concatenate((axangle_error, self.cartesian_error[:3, 3]))
        Kp = self.cartesian_impedance_values
        tau_stiffness = -(Kp * cartesian_error)
        Kv = 2 * self.damping_ratio * np.sqrt(Kp)
        tau_damping = -Kv * (self.joint_jacobian @ v[:7])

        # also add null space
        tau_task = self.joint_jacobian.T @ (tau_stiffness + tau_damping)

        # tau_nullspace = ( np.eye(7) -  self.joint_jacobian.T * np.linalg.pinv(self.joint_jacobian.T)) *
        #                (self.Kp * (self.joint_null_space - x) -
        #                 (2.0 * sqrt(self.Kp)) * v);

        tau_task = np.concatenate((tau_task, np.zeros(2)))

        # jacobian inverse gives
        tau += tau_task  #
        self.tau_stiffness_log.append(tau_stiffness.copy())
        self.tau_damping_log.append(tau_damping.copy())
        return tau

    def switch_controller(self, controller_mode):
        self.controller_mode = controller_mode

    # TODO(liruiw) consider cartesian impedance control and joint velocity control
    def DoCalcDiscreteVariableUpdates(self, context, events, discrete_state):
        LeafSystem.DoCalcDiscreteVariableUpdates(self, context, events, discrete_state)

        # read input ports
        x = self.robot_state_input_port.Eval(context)
        q_cmd = self.joint_angle_commanded_input_port.Eval(context)
        tau_ff = self.tau_feedforward_input_port.Eval(context)

        # TODO: could even implement a hybrid policy in here.
        q = x[: self.nq]
        v = x[self.nq :]

        # estimate velocity
        if self.q_prev is None:
            self.velocity_estimator.update(np.zeros(self.nv))
            self.acceleration_estimator.update(np.zeros(self.nv))
            v_diff = np.zeros(self.nv)
        else:
            # low pass filter velocity.
            v_diff = (q - self.q_prev) / self.control_period
            print(f"v_diff: {np.linalg.norm(v_diff):.3f}")
            self.velocity_estimator.update(v_diff)
            self.acceleration_estimator.update((v_diff - self.v_prev) / self.control_period)

        self.q_prev = q
        self.v_prev = v_diff
        v_est = self.velocity_estimator.get_current_state()
        a_est = self.acceleration_estimator.get_current_state()

        # log the P and D parts of desired acceleration
        self.sample_times.append(context.get_time())

        # update plant context
        self.plant.SetPositions(self.context, q)  # why is this called ?
        # self.plant.SetVelocities(self.context, v_est)

        # gravity compenstation
        tau_g = self.plant.CalcGravityGeneralizedForces(self.context)  # compensate for gravity
        tau = -tau_g

        # TODO(liruiw) think about controller switch in the context of drake simulation
        if self.controller_mode == "impedance":
            M = self.plant.CalcMassMatrixViaInverseDynamics(self.context)

            # m = np.sort(np.linalg.eig(M)[0])[::-1]
            m = M.diagonal()
            Kv = 2 * self.damping_ratio * np.sqrt(self.Kp * m)
            self.Kv_filter.update(Kv)
            Kv = self.Kv_filter.get_current_state()
            tau_stiffness = self.Kp * (q_cmd - q)
            tau_damping = -Kv * v
            tau += tau_damping + tau_stiffness

            self.Kv_log.append(Kv)
            self.tau_stiffness_log.append(tau_stiffness.copy())
            self.tau_damping_log.append(tau_damping.copy())

        elif self.controller_mode == "inverse_dynamics":
            qDDt_d = self.Kp * (q_cmd - q) + self.Kv * (-v_est)
            # self.integral_errors  += (q_cmd - q) * self.control_period
            # qDDt_d = qDDt_d + self.integral_errors * self.KI
            # print("ki term:", (self.integral_errors * self.KI).sum())
            tau += self.plant.CalcInverseDynamics(
                context=self.context,
                known_vdot=qDDt_d,
                external_forces=MultibodyForces(self.plant),
            )

        elif self.controller_mode == "cartesian_impedance":
            # self.plant_controller.erset_cartesian_controller_parameters()
            qDDt_d = self.Kp * (q_cmd - q) + self.Kv * (-v_est)
            tau_inv_dyn = self.plant.CalcInverseDynamics(
                context=self.context,
                known_vdot=qDDt_d,
                external_forces=MultibodyForces(self.plant),
            )
            tau_old = tau.copy()
            tau = self.compute_cartesian_impedance_control(context, v, tau)
            # https://drake.mit.edu/doxygen_cxx/classdrake_1_1multibody_1_1_multibody_plant.html#afcb4de9b58cb375d056da7e8ec30937a
            C = self.plant.CalcBiasTerm(self.context)
            M = self.plant.CalcMassMatrixViaInverseDynamics(self.context)
            tau += C @ v + M @ a_est
            tau -= MultibodyForces(self.plant).generalized_forces()
            # if  np.linalg.norm(v) > 2:
            #     print("inverse dynamics control:", tau_inv_dyn)
            #     print("cartesian_impedance control:", tau - tau_inv_dyn)  #100x greater
            #     print("tau update:", tau - self.prev_tau)  #100x greater

        output = discrete_state.get_mutable_vector().get_mutable_value()
        output[:] = tau + tau_ff
        self.t += self.control_period  # debug
        self.prev_tau = tau
        output[:] = np.clip(output, -self.effort_limit, self.effort_limit)
        print(f"v: {np.linalg.norm(v):.3f}")

    def CalcJointTorques(self, context, y_data):
        state = context.get_discrete_state_vector().get_value()
        y = y_data.get_mutable_value()
        y[:] = state


class LowPassFilter:
    def __init__(self, dimension: int, h: float, w_cutoff: float):
        if w_cutoff == np.inf:
            self.a = 1.0
        else:
            self.a = h * w_cutoff / (1 + h * w_cutoff)

        self.n = dimension
        self.x = None

    def update(self, u: np.array):
        assert u.size == self.n

        if self.x is None:
            self.x = u
        else:
            self.x = (1 - self.a) * self.x + self.a * u

    def reset_state(self):
        self.x = None

    def has_valid_state(self):
        return self.x is not None

    def get_current_state(self):
        assert self.x is not None
        return self.x.copy()

core/utils/robot_internal_controller.py
liruiw/Fleet-Tools
import numpy as np
import time
import numpy as np
from numpy.linalg import inv

from pydrake.common.cpp_param import List
from pydrake.common.value import Value
from pydrake.math import RigidTransform
from pydrake.multibody.math import SpatialForce, SpatialVelocity
from pydrake.multibody.plant import (
    ExternallyAppliedSpatialForce,
    MultibodyPlant,
)
from pydrake.multibody.tree import (
    Frame,
    JacobianWrtVariable,
    ModelInstanceIndex,
    SpatialInertia,
)
from pydrake.systems.framework import (
    BasicVector,
    DiagramBuilder,
    FixedInputPortValue,
    InputPort,
    LeafSystem,
    OutputPort,
)
from pydrake.systems.primitives import FirstOrderLowPassFilter
from core.utils.utils import *
import dataclasses as dc


def rotation_matrix_to_axang3(rot_mat):
    axangle = mat2axangle(rot_mat.matrix()[:3, :3])
    return axangle[0] * axangle[1]


@dc.dataclass
class PoseFrames:
    frame_T: Frame
    frame_P: Frame

    def __eq__(self, rhs):
        return self.frame_T is rhs.frame_T and self.frame_P is rhs.frame_P

    def AssertValid(self):
        assert self.frame_T is not None
        assert self.frame_P is not None


@dc.dataclass
class PoseFeedback:
    kp_xyz: np.ndarray
    kd_xyz: np.ndarray
    kp_rot: np.ndarray
    kd_rot: np.ndarray

    def __call__(self, X_WP, V_WP, X_WPdes, V_WPdes):
        # Transform to "negative error": desired w.r.t. actual,
        # expressed in world frame (for applying the force).
        # TODO(eric.cousineau): Use ComputePoseDiffInCommonFrame.
        p_PPdes_W = X_WPdes.translation() - X_WP.translation()
        R_WP = X_WP.rotation()
        R_PPdes = R_WP.inverse() @ X_WPdes.rotation()
        axang3_PPdes = rotation_matrix_to_axang3(R_PPdes)
        axang3_PPdes_W = R_WP @ axang3_PPdes
        V_PPdes_W = V_WPdes - V_WP
        v_PPdes_W = V_PPdes_W.translational()
        w_PPdes_W = V_PPdes_W.rotational()
        # Compute wrench components.

        f_P_W = self.kp_xyz * p_PPdes_W + self.kd_xyz * v_PPdes_W
        tau_P_W = self.kp_rot * axang3_PPdes_W + self.kd_rot * w_PPdes_W
        F_P_W_feedback = SpatialForce(tau=tau_P_W, f=f_P_W)
        return F_P_W_feedback


@dc.dataclass
class PoseReferenceInputPorts:
    frames: PoseFrames
    X_TPdes: InputPort
    V_TPdes: InputPort

    def fix(self, plant, context, plant_context=None):
        if plant_context is None:
            plant_context = plant.GetMyContextFromRoot(context)
        X_TP_init = plant.CalcRelativeTransform(
            plant_context,
            self.frames.frame_T,
            self.frames.frame_P,
        )
        V_TP_init = SpatialVelocity.Zero()
        return PosePortValues(
            X_TPdes=fix_port(self.X_TPdes, context, X_TP_init),
            V_TPdes=fix_port(self.V_TPdes, context, V_TP_init),
        )


@dc.dataclass
class PosePortValues:
    X_TPdes: FixedInputPortValue
    V_TPdes: FixedInputPortValue

    def update(self, X_TPdes, V_TPdes=None):
        if V_TPdes is None:
            V_TPdes = SpatialVelocity.Zero()
        self.X_TPdes.GetMutableData().set_value(X_TPdes)
        self.V_TPdes.GetMutableData().set_value(V_TPdes)


@dc.dataclass
class JointFeeback:
    Kp: np.ndarray
    Kd: np.ndarray

    def __call__(self, q, v, q_des, v_des):
        return self.Kp * (q_des - q) + self.Kd * (v_des - v)


@dc.dataclass
class DofSetForDynamics:
    """
    Like DofSet, but for dynamics where we are not constrained to nq = nv = nu.
    """

    q: np.ndarray  # bool
    v: np.ndarray  # bool
    u: np.ndarray  # bool


def get_frame_spatial_velocity(plant, context, frame_T, frame_F, frame_E=None):
    """
    Returns:
        SpatialVelocity of frame F's origin w.r.t. frame T, expressed in E
        (which is frame T if unspecified).
    """
    if frame_E is None:
        frame_E = frame_T
    Jv_TF_E = plant.CalcJacobianSpatialVelocity(
        context,
        with_respect_to=JacobianWrtVariable.kV,
        frame_B=frame_F,
        p_BoBp_B=[0, 0, 0],
        frame_A=frame_T,
        frame_E=frame_E,
    )
    v = plant.GetVelocities(context)
    V_TF_E = SpatialVelocity(Jv_TF_E @ v)
    return V_TF_E


def make_empty_dofset(plant):
    return DofSetForDynamics(
        q=np.zeros(plant.num_positions(), dtype=bool),
        v=np.zeros(plant.num_velocities(), dtype=bool),
        u=np.zeros(plant.num_actuated_dofs(), dtype=bool),
    )


def articulated_model_instance_dofset(plant, instance):
    """
    Returns DofSetForDynamics given a particular model instance.
    """
    dofs = make_empty_dofset(plant)
    joint_indices = plant.GetJointIndices(instance)
    for joint_index in joint_indices:
        joint = plant.get_joint(joint_index)
        start_in_q = joint.position_start()
        end_in_q = start_in_q + joint.num_positions()
        dofs.q[start_in_q:end_in_q] = True
        start_in_v = joint.velocity_start()
        end_in_v = start_in_v + joint.num_velocities()
        dofs.v[start_in_v:end_in_v] = True

    B_full = plant.MakeActuationMatrix()
    B = B_full[dofs.v]
    assert set(B.flat) <= {0.0, 1.0}
    B_rows = B.sum(axis=0)
    assert set(B_rows.flat) <= {0.0, 1.0}
    dofs.u[:] = B_rows.astype(bool)

    nq = plant.num_positions(instance)
    nv = plant.num_velocities(instance)
    nu = plant.num_actuated_dofs(instance)
    assert dofs.q.sum() == nq
    assert dofs.v.sum() == nv
    assert dofs.u.sum() == nu
    return dofs


class PoseController(LeafSystem):
    """
    Leveraging notation in line with Drake and Russ's Manipulation course
    notes:
    https://manipulation.csail.mit.edu/force.html#section3

    For additional operationaln space refs, see:

    - A. D. Ames and M. Powell, “Towards the Unification of Locomotion and
      Manipulation through Control Lyapunov Functions and Quadratic Programs,”
      in Control of Cyber-Physical Systems, 2013.
      http://link.springer.com/10.1007/978-3-319-01159-2_12
      - Easier to digest this notation.

    - O. Khatib, L. Sentis, J. Park, and J. Warren, “Whole-Body Dynamic
      Behavior and Control of Human-Like Robots,” Int. J. Human. Robot. 2004.
      https://www.worldscientific.com/doi/abs/10.1142/S0219843604000058
      https://khatib.stanford.edu/publications/pdfs/Khatib_2004_IJHR.pdf
      - Good for main reference of above paper.

    - O. Khatib, “A unified approach for motion and force control of robot
      manipulators: The operational space formulation,” IEEE Journal on
      Robotics and Automation, 1987.
      https://khatib.stanford.edu/publications/pdfs/Khatib_1987_RA.pdf

    - Robotics Handbook, p.236, Sec. 10.6, Second-Order Redundancy Resolution.

    Derivation:
    https://toyotaresearchinstitute.sharepoint.com/:o:/r/sites/ToyotaResearchInstitute/Shared%20Documents/Dexterous%20Manipulation/Issues/Anzu%20Issue%208783%20-%20Pose%20Tracking/2022-09-14%20Pose%20Tracking?d=w94c38e6f0ac547a78e188b2d23a734b5&csf=1&web=1&e=HkxogD
    """

    def __init__(
        self,
        plant,
        model_instance,
        frame_T,
        frame_P,
    ):
        """frame_T is the task frame (inertial frame) and frame_P is the end effector frame"""
        super().__init__()

        context = plant.CreateDefaultContext()
        ndof = plant.num_positions(model_instance)
        nv = plant.num_velocities(model_instance)
        assert ndof == nv
        nx = 2 * ndof
        nu = plant.num_actuated_dofs(model_instance)
        assert ndof == nu
        dofs = articulated_model_instance_dofset(plant, model_instance)
        assert np.sum(dofs.q) == ndof
        assert np.sum(dofs.v) == ndof
        FIX_FINGER = True  # always command certain torques to the fingers

        kpt_xyz = 400  # 100.0
        kdt_xyz = 40  # 20.0
        kpt_rot = 400  # 100.0
        kdt_rot = 40  # 20.0 # 400, 40, 400, 40
        self.pose_feedback = PoseFeedback(
            kp_xyz=kpt_xyz,
            kd_xyz=kdt_xyz,  # lambda R_TP:
            kp_rot=kpt_rot,
            kd_rot=kdt_rot,  # lambda R_TP:
        )

        # Very stiff.
        kpp = 200.0  # 100.0
        kdp = 20.0
        # # Less stiff.
        # kpp = 50.0
        # kdp = 12.0
        # # Wiggly.
        # kpp = 5.0
        # kdp = 1.0
        self.scaled_joint_feedback = JointFeeback(
            Kp=kpp,
            Kd=kdp,
        )

        self.direct_joint_feedback = JointFeeback(
            Kp=np.diag([160.0, 200.0, 200.0, 160.0, 25.0, 70.0, 16.0]) * 0.4,
            Kd=np.diag([37.5, 50.0, 37.5, 25.0, 5.0, 3.75, 2.5]) * 0.6,
        )

        # Should use initialization event.
        q0 = plant.GetPositions(context, model_instance)
        v0 = np.zeros(nv)

        q_lower = plant.GetPositionLowerLimits()[dofs.q]
        q_upper = plant.GetPositionUpperLimits()[dofs.q]
        v_lower = plant.GetVelocityLowerLimits()[dofs.v]
        v_upper = plant.GetVelocityUpperLimits()[dofs.v]
        u_lower = plant.GetEffortLowerLimits()[dofs.u]
        u_upper = plant.GetEffortUpperLimits()[dofs.u]

        def assert_within_limits(value, *, lower, upper):
            too_low = value < lower
            too_high = value > upper
            assert not (too_low | too_high).any()

        def assert_values_within_limits(q, v, u_des):
            assert_within_limits(q, lower=q_lower, upper=q_upper)
            assert_within_limits(v, lower=v_lower, upper=v_upper)
            assert_within_limits(u_des, lower=u_lower, upper=u_upper)

        def control_math(X_TPdes, V_TPdes):
            q = plant.GetPositions(context, model_instance)
            v = plant.GetVelocities(context, model_instance)
            X_TP = plant.CalcRelativeTransform(context, frame_T, frame_P)
            V_TP = get_frame_spatial_velocity(plant, context, frame_T, frame_P)

            Jv_TP_full = plant.CalcJacobianSpatialVelocity(
                context,
                with_respect_to=JacobianWrtVariable.kV,
                frame_B=frame_P,
                p_BoBp_B=[0, 0, 0],
                frame_A=frame_T,
                frame_E=frame_T,
            )  # J
            Jv_TP = Jv_TP_full[:, dofs.v]
            Jvdot_v_TP = plant.CalcBiasSpatialAcceleration(
                context,
                with_respect_to=JacobianWrtVariable.kV,
                frame_B=frame_P,
                p_BoBp_B=[0, 0, 0],
                frame_A=frame_T,
                frame_E=frame_T,
            )  #

            # Full-plant dynamics.
            M_full = plant.CalcMassMatrix(context)
            C_full = plant.CalcBiasTerm(context)
            tau_g_full = plant.CalcGravityGeneralizedForces(context)

            # Full-order dynamics, but selecting controlled DoFs.
            M = M_full[dofs.v, :][:, dofs.v]
            C = C_full[dofs.v]
            tau_g = tau_g_full[dofs.v]

            # Task-space dynamics.
            Jt = Jv_TP
            Jtdot_v = Jvdot_v_TP.get_coeffs()  # J_dot q_dot
            # Purely kinematic.
            Nt_T = np.eye(ndof) - np.linalg.pinv(Jt) @ Jt  # null space

            # For singularities:
            # WARNING: Jv_TP dropping rank (beyond redundancy) will cause this
            # stuff to explode. Should look at nasa jsc code or Khatib '87 for
            # pointers on how to more intelligently avoid, maybe?
            # Below can fudge things, but may induce a loss of tracking.
            nt = 6
            t_fudge = 0.01 * np.eye(nt)

            # TODO(eric.cousineau): Rewrite w/ more friendly inversions.
            inv = np.linalg.inv
            Mtinv = Jt @ inv(M) @ Jt.T  # eq 18
            Mt = inv(Mtinv + t_fudge)  # regularization
            Jtbar = inv(M) @ Jt.T @ Mt
            Nt_T = np.eye(ndof) - Jt.T @ Jtbar.T
            Ct = Jtbar.T @ C - Mt @ Jtdot_v  # eq 24
            tau_gt = Jtbar.T @ tau_g  # virtual work

            # N.B. This do not currently add M*vdot for task or joint
            # centering.
            # TODO(eric.cousineau): Add in inertial term when computing finite
            # diff of V_TPdes.

            # # Task-space control.
            Ft_feedback = (  # desired behavior when being exerted force
                Mt @ self.pose_feedback(X_TP, V_TP, X_TPdes, V_TPdes).get_coeffs()
            )
            Ft_feedforward = Ct - tau_gt  # cancel out dynamics
            # - Use task-space inverse dynamics.
            # N.B. Simulation seems to go unstable if not using task-space
            # inertial scaling.
            Ft = Ft_feedback + Ft_feedforward

            # Posture regularization.
            # N.B. Because Jp = I, we don't need any projected dynamics.
            Jp = np.eye(ndof)
            Mp = M
            Cp = C
            tau_gp = tau_g

            # - Use task-space inverse dynamics.
            Fp_feedback = Mp @ self.scaled_joint_feedback(q, v, q0, v0)
            Fp_feedforward = Cp - tau_gp
            Fp = Fp_feedback + Fp_feedforward

            # Combined joint-space control.
            ut = Jt.T @ Ft
            up = Nt_T @ Jp.T @ Fp  # null space part in the joint space eq 55
            # up = Jp.T @ Fp  # joint-space alone.

            u = ut + up
            # https://github.com/rachelholladay/franka_ros_interface/blob/fdbc28bf78e1f7eaa1ff627e1a34be72fafbdedc/franka_ros_controllers/src/cartesian_impedance_controller.cpp#L192

            # # Inverse Dynamics. (Feedback linearization?)
            # vd = scaled_joint_feedback(q, v, q0, v0)
            # u = M @ vd + C - tau_g

            # # Very simple joint feedback.
            # u = direct_joint_feedback(q, v, q0, v0) - tau_g

            # TODO(eric.cousineau): Implement as QP for joint position,
            # velocity, accel, and torque limits. May need CBF?
            # assert_values_within_limits(q, v, u)

            if FIX_FINGER:
                # print("finger torque:", u[-2:])
                u[-2:] = -100
            return u

        self.plant_state_input = self.DeclareVectorInputPort("plant_state", nx)
        X_TPdes_input = self.DeclareAbstractInputPort("X_TPdes", Value[RigidTransform]())
        V_TPdes_input = self.DeclareAbstractInputPort("V_TPdes", Value[SpatialVelocity]())
        self.pose_inputs = PoseReferenceInputPorts(
            frames=PoseFrames(frame_T=frame_T, frame_P=frame_P),
            X_TPdes=X_TPdes_input,
            V_TPdes=V_TPdes_input,
        )

        def control_calc(sys_context, output):
            x = self.plant_state_input.Eval(sys_context)
            plant.SetPositionsAndVelocities(context, model_instance, x)
            X_TPdes = self.pose_inputs.X_TPdes.Eval(sys_context)
            V_TPdes = self.pose_inputs.V_TPdes.Eval(sys_context)
            u = control_math(X_TPdes, V_TPdes)
            output.set_value(u)

        self.torques_output = self.DeclareVectorOutputPort("torques_output", size=nu, calc=control_calc)

        self._plant = plant
        self._model_instance = model_instance
        self.V_TPdes_input = V_TPdes_input
        self.X_TPdes_input = X_TPdes_input

    def set_joint_stiffness(self, mode="stiff"):
        if mode == "very_stiff":
            kpp = 400.0
            kdp = 40.0

        if mode == "stiff":
            kpp = 100.0
            kdp = 20.0
        # # Less stiff.

        if mode == "normal":
            kpp = 50.0
            kdp = 12.0

        # # Wiggly.
        if mode == "compliant":
            kpp = 5.0
            kdp = 1.0

        self.scaled_joint_feedback = JointFeeback(Kp=kpp, Kd=kdp)

    def makeJointFeeback(self, kpp, kdp):
        self.scaled_joint_feedback = JointFeeback(Kp=kpp, Kd=kdp)

    def set_pose_stiffness(self, parameters=[100, 20, 100, 20]):
        (kpt_xyz, kdt_xyz, kpt_rot, kdt_rot) = parameters
        self.pose_feedback = PoseFeedback(
            kp_xyz=kpt_xyz,
            kd_xyz=kdt_xyz,  # lambda R_TP:
            kp_rot=kpt_rot,
            kd_rot=kdt_rot,  # lambda R_TP:
        )

    def get_plant_state_input_port(self):
        return self.plant_state_input

    def get_velocity_desired_port(self):
        return self.V_TPdes_input

    def get_state_desired_port(self):
        return self.X_TPdes_input

    @staticmethod
    def AddToBuilder(builder, controller):
        plant = controller._plant
        model_instance = controller._model_instance

        builder.AddSystem(controller)
        builder.Connect(
            plant.get_state_output_port(model_instance),
            controller.plant_state_input,
        )
        builder.Connect(
            controller.torques_output,
            plant.get_actuation_input_port(model_instance),
        )

core/utils/spatialstiffness_controller.py
liruiw/Fleet-Tools
# --------------------------------------------------------
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------


import numpy as np
import time

import cv2
import matplotlib.pyplot as plt
import tabulate
import yaml
from collections import deque
import shutil

from .geometry import *
from transforms3d.quaternions import *
from transforms3d.euler import *
from transforms3d.axangles import *

from colored import fg

joystick_lowpass_filter = deque([np.zeros(6)], maxlen=30)
vr_has_init = False
rec_action_pose = None
vr_frame_origin = None
oculus_reader = None


def get_keyboard_input(action, robot_id):
    finger_state = ""
    cv2.imshow(
        "robot_{}_teleop_wsadqe_translation_ikjluo_rotation".format(robot_id), np.zeros((64, 64, 3), dtype=np.uint8)
    )
    k = cv2.waitKey(1)
    action = np.zeros_like(action)
    trans_scale = 0.01
    rot_scale = 0.05

    if k == ord("w"):
        action[0] = trans_scale
    elif k == ord("s"):
        action[0] = -trans_scale
    elif k == ord("a"):
        action[1] = trans_scale
    elif k == ord("d"):
        action[1] = -trans_scale
    elif k == ord("q"):
        action[2] = trans_scale
    elif k == ord("e"):
        action[2] = -trans_scale

    elif k == ord("i"):
        action[3] = rot_scale
    elif k == ord("k"):
        action[3] = -rot_scale
    elif k == ord("j"):
        action[4] = rot_scale
    elif k == ord("l"):
        action[4] = -rot_scale
    elif k == ord("u"):
        action[5] = rot_scale
    elif k == ord("o"):
        action[5] = -rot_scale

    elif k == ord("n"):
        finger_state = "open"
    elif k == ord("m"):
        finger_state = "close"
    return action, finger_state


def get_mouse_input(action, robot_id):
    # Space mouse
    global joystick_lowpass_filter
    event = spacenav.poll()
    finger_state = ""
    action = np.zeros_like(action)
    trans_scale = 0.02
    rot_scale = 0.05
    normalize_scale = 512
    deadband = 0.1

    if event is spacenav.MotionEvent:
        if event.x != -2147483648:
            action[0] = event.z / normalize_scale
            action[1] = -event.x / normalize_scale
            action[2] = event.y / normalize_scale
            action[3] = event.rz / normalize_scale
            action[4] = -event.rx / normalize_scale
            action[5] = event.ry / normalize_scale

            if (np.abs(action[:3]) > deadband).sum() == 0:
                action[:3] = 0
            if (np.abs(action[3:]) > deadband).sum() == 0:
                action[3:] = 0

            action[:3] *= trans_scale
            action[3:] *= rot_scale

            spacenav_prev_action[:] = action
            # print("action:", action)

        if hasattr(event, "pressed"):
            if event.pressed:
                finger_state = "close"
    print("action: ", joystick_lowpass_filter)
    joystick_lowpass_filter.append(action)
    action = np.mean(list(joystick_lowpass_filter), axis=0)
    return action, finger_state


class LowPassFilter:
    def __init__(self, dimension: int, h: float, w_cutoff: float):
        if w_cutoff == np.inf:
            self.a = 1.0
        else:
            self.a = h * w_cutoff / (1 + h * w_cutoff)

        self.n = dimension
        self.x = None

    def update(self, u: np.array):
        assert u.size == self.n

        if self.x is None:
            self.x = u
        else:
            self.x = (1 - self.a) * self.x + self.a * u

    def reset_state(self):
        self.x = None

    def has_valid_state(self):
        return self.x is not None

    def get_current_state(self):
        assert self.x is not None
        return self.x.copy()


class PoseFilter:
    def __init__(self, control_period=1000, translation_ratio=1, rotation_ratio=1):
        self.control_period = 1.0 / control_period
        self.w_cutoff_r = 1 * np.pi * control_period  # tune this smaller to get more smoothed behavior
        self.w_cutoff_t = 2 * np.pi * control_period * 2  # tune this smaller to get more smoothed behavior
        self.translation_ratio = translation_ratio
        self.rotation_ratio = rotation_ratio

        n = 3
        self.rotation_position_estimator = LowPassFilter(n, self.control_period, self.w_cutoff_r)
        self.rotation_velocity_estimator = LowPassFilter(n, self.control_period, self.w_cutoff_r)
        self.rotation_acceleration_estimator = LowPassFilter(n, self.control_period, self.w_cutoff_r)

        self.translation_position_estimator = LowPassFilter(n, self.control_period, self.w_cutoff_t)
        self.translation_velocity_estimator = LowPassFilter(n, self.control_period, self.w_cutoff_t)
        self.translation_acceleration_estimator = LowPassFilter(n, self.control_period, self.w_cutoff_t)

        self.translation_prev = np.zeros(n)
        self.translation_v_prev = np.zeros(n)
        self.translation_a_prev = np.zeros(n)

        self.rotation_prev = np.eye(n)
        self.rotation_vec_prev = np.zeros(n)
        self.rotation_v_prev = np.zeros(n)
        self.rotation_a_prev = np.zeros(n)
        self.curr_step = 0

    def update(self, pose):
        # self.saturate_pose(pose)
        self.update_translation(pose[:3, 3])
        self.update_rotation(pose[:3, :3])
        self.curr_step += 1
        # if self.curr_step % 100 == 0:
        #     self.rotation_position_estimator.reset_state()
        #     self.rotation_velocity_estimator.reset_state()
        #     self.rotation_acceleration_estimator.reset_state()

        #     self.translation_position_estimator.reset_state()
        #     self.translation_velocity_estimator.reset_state()
        #     self.translation_acceleration_estimator.reset_state()

    def update_rotation(self, rotation_mat):
        # filter based on the delta otherwise it has the wrapped around issues.
        if hasattr(self, "rotation_prev"):
            axis, angle = mat2axangle(self.rotation_prev.T @ rotation_mat)
            rotation = angle * axis * self.rotation_ratio

            # print("update rotation:", axis * angle)
            self.rotation_position_estimator.update(rotation)

            # the control period is hard-coded for now.
            v_diff = (rotation - self.rotation_vec_prev) / self.control_period
            a_diff = (v_diff - self.rotation_v_prev) / self.control_period

            self.rotation_velocity_estimator.update(v_diff)
            self.rotation_acceleration_estimator.update(a_diff)
            self.rotation_v_prev = v_diff
            self.rotation_a_prev = a_diff

        else:
            # first time
            rotation = np.zeros(3)
            self.rotation_position_estimator.update(rotation)
        self.rotation_prev = rotation_mat
        self.rotation_vec_prev = rotation

    def update_translation(self, translation):
        self.translation_position_estimator.update(translation * self.translation_ratio)
        if hasattr(self, "translation_prev"):
            v_diff = (translation - self.translation_prev) / self.control_period
            a_diff = (v_diff - self.translation_v_prev) / self.control_period

            self.translation_velocity_estimator.update(v_diff)
            self.translation_acceleration_estimator.update(a_diff)
            self.translation_v_prev = v_diff
            self.translation_a_prev = a_diff

        self.translation_prev = translation

    def get_estimates(self, saturate=False):
        p_est = self.translation_position_estimator.get_current_state()
        v_est = self.translation_velocity_estimator.get_current_state()
        a_est = self.translation_acceleration_estimator.get_current_state()

        if not hasattr(self, "rotation_est_prev"):
            # first step
            self.rotation_est_prev = self.rotation_prev

        # integration
        rp_est = self.rotation_position_estimator.get_current_state()
        angle = np.linalg.norm(rp_est)
        axis = rp_est / angle
        rot_mat = self.rotation_est_prev @ axangle2mat(axis, angle, is_normalized=True)
        self.rotation_est_prev = rot_mat

        if saturate:
            p_est, rot_mat = self.saturate_pose(p_est, rot_mat)

        rv_est = self.rotation_velocity_estimator.get_current_state()
        ra_est = self.rotation_acceleration_estimator.get_current_state()
        return p_est, v_est, a_est, rot_mat, rv_est, ra_est

    def saturate_pose(self, p_est, rot_mat):
        """
        cap the maximum rotations and translation offsets from the starting pose
        """
        translation_limit = 0.5
        rotation_limit = np.pi / 4
        p_est = np.clip(p_est, -translation_limit, translation_limit)
        axis, angle = mat2axangle(rot_mat)
        angle_scaled = np.clip(angle, -rotation_limit, rotation_limit)
        rot_mat = axangle2mat(axis, angle_scaled)

        return p_est, rot_mat


class VR_Interface:
    def __init__(self, freq=1000):
        self.vr_has_init = False
        self.rec_action_pose = None
        self.vr_frame_origin = None
        self.oculus_reader = None
        self.scale_translation_ratio = 0.8 if SCALE else 1.0
        self.scale_rotation_ratio = 0.3 if SCALE else 1.0

        self.pose_filter = PoseFilter(freq, self.scale_translation_ratio, self.scale_rotation_ratio)
        self.get_vr_init()

    def set_X_WG_Init(self, X_WGinit):
        self.X_WGinit = X_WGinit

    def get_vr_init(self):
        if not self.oculus_reader:
            try:
                self.oculus_reader = OculusReader()
            except:
                print("please install oculus reader first")
                pass

        print("VR Init, click rightTrig")
        while True:
            time.sleep(0.0001)
            transformations, buttons = self.oculus_reader.get_transformations_and_buttons()
            if "rightTrig" not in buttons:
                continue

            # print("righttrig:", buttons["rightTrig"][0])
            if buttons["rightTrig"][0] > 0.8 or buttons["rightGrip"][0] > 0.5:  # trig
                self.vr_frame_origin = self.get_X_WU(transformations)
                self.vr_has_init = True
                break

        print("finish VR init")

    def get_X_WU(self, transformations):
        X_UhatU = RigidTransform(rotX(-np.pi / 2) @ rotZ(np.pi / 2))  # (rotY(-np.pi / 2) @ rotZ(np.pi/2))
        X_WU_hat = RigidTransform(transformations["r"])
        return X_WU_hat  # @ X_UhatU

    def get_curr_filtered_pose(self, limit=True):
        p_est, v_est, a_est, rot_mat, rv_est, ra_est = self.pose_filter.get_estimates()
        dX_GinitG_W_filtered = np.eye(4)

        dX_GinitG_W_filtered[:3, 3] = p_est
        dX_GinitG_W_filtered[:3, :3] = rot_mat if not NOROTATION else np.eye(3)

        velocity_estimates = np.concatenate((v_est, rv_est))
        acceleration_estimates = np.concatenate((a_est, ra_est))

        return dX_GinitG_W_filtered, velocity_estimates, acceleration_estimates

    def pose_postprocess(self, dX_UinitU_What):
        # 1. world frame velocity. push forward: x, push up: z, push left: y
        X_What_W = RigidTransform(rotZ(np.pi / 2) @ rotX(np.pi / 2))
        # 2. end effector frame. forward z. up x. and right y.
        X_U_Ustar = RigidTransform(rotZ(np.pi / 2))
        dX_UinitU_W = X_What_W @ dX_UinitU_What @ X_U_Ustar

        X_U_G = RigidTransform(rotZ(np.pi / 2)) @ RigidTransform(rotZ(np.pi / 2) @ rotX(np.pi / 2) @ rotZ(np.pi / 2))
        dX_GinitG_W = X_What_W @ dX_UinitU_What @ X_U_G
        self.pose_filter.update(dX_GinitG_W.GetAsMatrix4())
        dX_GinitG_W_filtered, V_GinitG_W_filtered, A_GinitG_W_filtered = self.get_curr_filtered_pose()

        if DEBUG:
            curr = pack_action(dX_GinitG_W.GetAsMatrix4())
            filtered = pack_action(dX_GinitG_W_filtered)
            res = curr  # np.concatenate((curr, filtered))
            update_line(tracking_online_logger, [curr[3], filtered[3]], rospy.get_time() - initial_time)
            # update_line(tracking_online_logger_filtered, pack_action(dX_GinitG_W_filtered), rospy.get_time() - initial_time)

        if FILTER:
            dX_GinitG_W = RigidTransform(dX_GinitG_W_filtered)

        X_WGdesired = apply_premultiply_pose_delta(RigidTransform(self.X_WGinit), dX_GinitG_W)
        return dX_UinitU_W.GetAsMatrix4(), X_WGdesired.GetAsMatrix4()

    def get_vr_input_inner(self):
        """Transform oculus motion to actual end effector motion"""
        transformations, buttons = self.oculus_reader.get_transformations_and_buttons()

        if "r" not in transformations:
            return None

        X_WUinit = self.vr_frame_origin
        X_WU = self.get_X_WU(transformations)
        dX_UinitU_W = compute_premultiply_pose_delta(X_WUinit, X_WU)
        dX_UinitU_W, X_WGdesired = self.pose_postprocess(dX_UinitU_W)

        exit_loop = "rightGrip" in buttons and buttons["rightGrip"][0] > 0.5
        reset = buttons["B"]
        finger_state = "close" if buttons["A"] else "open"
        teleop_hold = buttons["rightTrig"][0] < 0.2

        if (
            teleop_hold and not reset and not exit_loop
        ):  # clutching. move the vr to a more comfortable place to restart.
            print("================== clutching. ======================")
            self.set_X_WG_Init(X_WGdesired)
            self.get_vr_init()

        return (
            dX_UinitU_W,
            finger_state,  # Right Grip
            reset,  # B
            teleop_hold,  # Right Trig
            X_WGdesired,  # A
            exit_loop,
        )

    def get_vr_input():
        global vr_interface
        abs_pose = X_WGdesired_default
        (
            action_pose,
            finger_state,  # A
            reset,  # B
            teleop_hold,  # not pressing Right Trig
            X_WGdesired,
            exit_loop,  # Right Grip
        ) = vr_interface.get_vr_input_inner()
        return (X_WGdesired, teleop_hold, reset)


X_WGdesired_default = np.array(
    [
        [0.924, -0.006, 0.382, 0.480],
        [0.043, -0.992, -0.118, -0.003],
        [0.379, 0.126, -0.917, 0.53238],
        [0.0, 0.0, 0.0, 1.0],
    ]
)
FILTER = True  # use filtered pose
DEBUG = False  # draw visualization would be slower
SCALE = True  # scale translation or not
init_joints = (0.0, -1.285, 0, -2.356, 0.0, 1.571, 0.785)
VR_MOVE = True  # use VR to control end effector
NOROTATION = False  # set rotation offset to be identity

try:
    from .oculus_reader import OculusReader

    vr_interface = VR_Interface()
    vr_interface.set_X_WG_Init(abs_pose)
except:
    pass

core/utils/teleop_utils.py
liruiw/Fleet-Tools
# ---------------------------------------------------------------
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# This work is licensed under the NVIDIA Source Code License
# for OSCAR. To view a copy of this license, see the LICENSE file.
# ---------------------------------------------------------------

# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

import torch
import numpy as np


def to_torch(x, dtype=torch.float, device="cuda:0", requires_grad=False):
    return torch.tensor(x, dtype=dtype, device=device, requires_grad=requires_grad)


@torch.jit.script
def quat_mul(a, b):
    assert a.shape == b.shape
    shape = a.shape
    a = a.reshape(-1, 4)
    b = b.reshape(-1, 4)

    x1, y1, z1, w1 = a[:, 0], a[:, 1], a[:, 2], a[:, 3]
    x2, y2, z2, w2 = b[:, 0], b[:, 1], b[:, 2], b[:, 3]
    ww = (z1 + x1) * (x2 + y2)
    yy = (w1 - y1) * (w2 + z2)
    zz = (w1 + y1) * (w2 - z2)
    xx = ww + yy + zz
    qq = 0.5 * (xx + (z1 - x1) * (x2 - y2))
    w = qq - ww + (z1 - y1) * (y2 - z2)
    x = qq - xx + (x1 + w1) * (x2 + w2)
    y = qq - yy + (w1 - x1) * (y2 + z2)
    z = qq - zz + (z1 + y1) * (w2 - x2)

    quat = torch.stack([x, y, z, w], dim=-1).view(shape)

    return quat


@torch.jit.script
def unit_vec(vec, dim=-1):
    """
    Returns a normalized unit vector along the provided dimension

    Args:
        vec (tensor): (..., N, ...) arbitrary vector
        dim (int): Dimension to normalize

    Returns:
        tensor: (..., N, ...) vector normalized along @dim
    """
    # type: (Tensor, int) -> Tensor
    return vec / torch.norm(vec, dim=dim, keepdim=True)


@torch.jit.script
def dot(v1, v2, dim=-1, keepdim=False):
    """
    Computes dot product between two vectors along the provided dim, optionally keeping the dimension

    Args:
        v1 (tensor): (..., N, ...) arbitrary vector
        v2 (tensor): (..., N, ...) arbitrary vector
        dim (int): Dimension to sum over for dot product
        keepdim (bool): Whether to keep dimension over which dot product is calculated

    Returns:
        tensor: (..., [1,] ...) dot product of vectors, with optional dimension kept if @keepdim is True
    """
    # type: (Tensor, Tensor, int, bool) -> Tensor
    return torch.sum(v1 * v2, dim=dim, keepdim=keepdim)


@torch.jit.script
def quat_slerp(quat0, quat1, frac, shortestpath=True, eps=1.0e-15):
    """
    Return spherical linear interpolation between two quaternions.

    Args:
        quat0 (tensor): (..., 4) tensor where the final dim is (x,y,z,w) initial quaternion
        quat1 (tensor): (..., 4) tensor where the final dim is (x,y,z,w) final quaternion
        frac (tensor): Values in [0.0, 1.0] representing fraction of interpolation
        shortestpath (bool): If True, will calculate shortest path
        eps (float): Value to check for singularities
    Returns:
        tensor: (..., 4) Interpolated
    """
    # type: (Tensor, Tensor, Tensor, bool, float) -> Tensor
    # reshape quaternion
    quat_shape = quat0.shape
    quat0 = unit_vec(quat0.reshape(-1, 4), dim=-1)
    quat1 = unit_vec(quat1.reshape(-1, 4), dim=-1)

    # Check for endpoint cases
    where_start = frac <= 0.0
    where_end = frac >= 1.0

    d = dot(quat0, quat1, dim=-1, keepdim=True)
    if shortestpath:
        quat1 = torch.where(d < 0.0, -quat1, quat1)
        d = torch.abs(d)
    angle = torch.acos(torch.clip(d, -1.0, 1.0))

    # Check for small quantities (i.e.: q0 = q1)
    where_small_diff = torch.abs(torch.abs(d) - 1.0) < eps
    where_small_angle = abs(angle) < eps

    isin = 1.0 / torch.sin(angle)
    val = (
        quat0 * torch.sin((1.0 - frac) * angle) * isin
        + quat1 * torch.sin(frac * angle) * isin
    )

    # Filter edge cases
    val = torch.where(
        where_small_diff | where_small_angle | where_start,
        quat0,
        torch.where(
            where_end,
            quat1,
            val,
        ),
    )

    # Reshape and return values
    return val.reshape(list(quat_shape))


@torch.jit.script
def axisangle_between_vec(vec1, vec2, eps=1e-6):
    """
    Calculates the rotation between two 3D vectors from @vec1 to @vec2 in axis angle representation
    See https://www.euclideanspace.com/maths/algebra/vectors/angleBetween/ for formula

    Args:
        vec1 (tensor): (..., 3) where final dim represents (x, y, z) 3D vector
        vec2 (tensor): (..., 3) where final dim represents (x, y, z) 3D vector
        eps (float): Stability value below which small values will be mapped to 0

    Returns:
        tensor: (..., 3) where final dim represents (ax, ay, az) axis-angle rotation from @vec1 to @vec2.
    """
    # type: (Tensor, Tensor, float) -> Tensor
    where_small = torch.abs(torch.sum(vec2 - vec1, dim=-1)) < eps
    # Normalize vectors
    vec1 = vec1 / torch.norm(vec1, dim=-1, keepdim=True)
    vec2 = vec2 / torch.norm(vec2, dim=-1, keepdim=True)
    # Calculate angle and axis
    angle = torch.acos(torch.sum(vec1 * vec2, dim=-1, keepdim=True))
    axis = torch.cross(vec1, vec2)
    axis_norm = torch.norm(axis, dim=-1)
    where_small = where_small | (axis_norm < eps)
    axis[where_small] = 0.0
    axis[~where_small] = axis[~where_small] / axis_norm[~where_small].unsqueeze(-1)
    # Combine and return value
    return angle * axis


@torch.jit.script
def rotate_vec_by_quat(vec, quat):
    """
    This function rotates a 3D vector @vec by the direction and angle represented by @quat.

    See https://gamedev.stackexchange.com/questions/28395/rotating-vector3-by-a-quaternion
    for more information.

    Args:
        vec (tensor): (..., 3) where final dim represents (x, y, z) vectors
        quat (tensor): (..., 4) where final dim represents desired (x, y, z, w) quaternion rotations
    Returns:
        tensor: (..., 3) where final dim is newly rotated (x, y, z) vectors
    """
    # Parse quaternion into vector and scalar components
    q_vec, q_s = quat[..., :3], quat[..., 3].unsqueeze(-1)
    # Calculate new vector and return
    return (
        2.0 * dot(q_vec, vec, keepdim=True) * q_vec
        + (q_s * q_s - dot(q_vec, q_vec, keepdim=True)) * vec
        + 2.0 * q_s * torch.cross(q_vec, vec)
    )


@torch.jit.script
def rotate_vec_by_axisangle(vec, aa_vec):
    """
    This function rotates a 3D vector @vec by the direction and angle represented by @aa_vec.

    See https://stackoverflow.com/questions/32485772/how-do-axis-angle-rotation-vectors-work-and-how-do-they-compare-to-rotation-matr
    for more information.

    Args:
        vec (tensor): (..., 3) where final dim represents (x, y, z) vectors
        aa_vec (tensor): (..., 3) where final dim represents desired (ax, ay, az) axis-angle rotations
    Returns:
        tensor: (..., 3) where final dim is newly rotated (x, y, z) vectors
    """
    # Extract angle and unit vector from axis-angle vectors
    angle = torch.norm(aa_vec, dim=-1, keepdim=True)
    aa_v = aa_vec / angle
    # Map all NaNs to 0
    aa_v[aa_v != aa_v] = 0.0

    # # Record angles that are zero so we don't get nan's
    # idx = torch.nonzero(angle.squeeze(dim=-1)).tolist()
    # aa_v = torch.zeros_like(aa_vec)
    # aa_v[idx] = aa_vec[idx] / angle[idx]

    # Rotate the vector using the formula (see link above)
    c, s = torch.cos(angle), torch.sin(angle)
    vec_rot = (
        vec * c
        + torch.cross(aa_v, vec) * s
        + aa_v * (torch.sum(aa_v * vec, dim=-1, keepdim=True)) * (1 - c)
    )

    return vec_rot


@torch.jit.script
def quat_conjugate(q):
    input_shape = q.shape
    q = (-q).reshape(-1, input_shape[-1])
    q[:, 3] = -q[:, 3]
    return q.reshape(list(input_shape))


@torch.jit.script
def orientation_error(desired, current):
    """
    This function calculates a 3-dimensional orientation error vector for use in the
    impedance controller. It does this by computing the delta rotation between the
    inputs and converting that rotation to exponential coordinates (axis-angle
    representation, where the 3d vector is axis * angle).
    See https://en.wikipedia.org/wiki/Axis%E2%80%93angle_representation for more information.
    Optimized function to determine orientation error from matrices

    Args:
        desired (tensor): (..., 3, 3) where final two dims are 2d array representing target orientation matrix
        current (tensor): (..., 3, 3) where final two dims are 2d array representing current orientation matrix
    Returns:
        tensor: (..., 3) where final dim is (ax, ay, az) axis-angle representing orientation error
    """
    # convert input shapes
    input_shape = desired.shape[:-2]
    desired = desired.reshape(-1, 3, 3)
    current = current.reshape(-1, 3, 3)

    # grab relevant info
    rc1 = current[:, :, 0]
    rc2 = current[:, :, 1]
    rc3 = current[:, :, 2]
    rd1 = desired[:, :, 0]
    rd2 = desired[:, :, 1]
    rd3 = desired[:, :, 2]

    error = 0.5 * (
        torch.cross(rc1, rd1, dim=-1)
        + torch.cross(rc2, rd2, dim=-1)
        + torch.cross(rc3, rd3, dim=-1)
    )

    # Reshape
    error = error.reshape(list(input_shape) + [3])

    return error


@torch.jit.script
def orientation_error_from_quat(desired, current):
    """
    This function calculates a 3-dimensional orientation error vector, where inputs are quaternions

    Args:
        desired (tensor): (..., 4) where final dim is (x,y,z,w) quaternion
        current (tensor): (..., 4) where final dim is (x,y,z,w) quaternion
    Returns:
        tensor: (..., 3) where final dim is (ax, ay, az) axis-angle representing orientation error
    """
    # convert input shapes
    input_shape = desired.shape[:-1]
    desired = desired.reshape(-1, 4)
    current = current.reshape(-1, 4)

    cc = quat_conjugate(current)
    q_r = quat_mul(desired, cc)
    return (q_r[:, 0:3] * torch.sign(q_r[:, 3]).unsqueeze(-1)).reshape(
        list(input_shape) + [3]
    )


@torch.jit.script
def quat2mat(quaternion):
    """
    Converts given quaternion to matrix.
    Args:
        quaternion (tensor): (..., 4) tensor where the final dim is (x,y,z,w) quaternion
    Returns:
        tensor: (..., 3, 3) tensor whose final two dimensions are 3x3 rotation matrices
    """
    # convert quat convention
    inds = torch.tensor([3, 0, 1, 2])
    input_shape = quaternion.shape[:-1]
    q = quaternion.reshape(-1, 4)[:, inds]

    # Conduct dot product
    n = torch.bmm(q.unsqueeze(1), q.unsqueeze(-1)).squeeze(-1).squeeze(-1)  # shape (-1)
    idx = torch.nonzero(n).reshape(-1)
    q_ = q.clone()  # Copy so we don't have inplace operations that fail to backprop
    q_[idx, :] = q[idx, :] * torch.sqrt(2.0 / n[idx].unsqueeze(-1))

    # Conduct outer product
    q2 = (
        torch.bmm(q_.unsqueeze(-1), q_.unsqueeze(1)).squeeze(-1).squeeze(-1)
    )  # shape (-1, 4 ,4)

    # Create return array
    ret = (
        torch.eye(3, 3, device=q.device)
        .reshape(1, 3, 3)
        .repeat(torch.prod(torch.tensor(input_shape)), 1, 1)
    )
    ret[idx, :, :] = torch.stack(
        [
            torch.stack(
                [
                    1.0 - q2[idx, 2, 2] - q2[idx, 3, 3],
                    q2[idx, 1, 2] - q2[idx, 3, 0],
                    q2[idx, 1, 3] + q2[idx, 2, 0],
                ],
                dim=-1,
            ),
            torch.stack(
                [
                    q2[idx, 1, 2] + q2[idx, 3, 0],
                    1.0 - q2[idx, 1, 1] - q2[idx, 3, 3],
                    q2[idx, 2, 3] - q2[idx, 1, 0],
                ],
                dim=-1,
            ),
            torch.stack(
                [
                    q2[idx, 1, 3] - q2[idx, 2, 0],
                    q2[idx, 2, 3] + q2[idx, 1, 0],
                    1.0 - q2[idx, 1, 1] - q2[idx, 2, 2],
                ],
                dim=-1,
            ),
        ],
        dim=1,
    )

    # Reshape and return output
    ret = ret.reshape(list(input_shape) + [3, 3])
    return ret


@torch.jit.script
def quat2axisangle(quat):
    """
    Converts quaternion to axis-angle format.
    Returns a unit vector direction scaled by its angle in radians.
    Args:
        quat (tensor): (..., 4) tensor where the final dim is (x,y,z,w) quaternion
    Returns:
        tensor: (..., 3) axis-angle exponential coordinates
    """
    # reshape quaternion
    quat_shape = quat.shape[:-1]  # ignore last dim
    quat = quat.reshape(-1, 4)
    # clip quaternion
    quat[:, 3] = torch.clamp(quat[:, 3], -1.0, 1.0)
    # Calculate denominator
    den = torch.sqrt(1.0 - quat[:, 3] * quat[:, 3])
    # Map this into a mask

    # Create return array
    ret = torch.zeros_like(quat)[:, :3]
    idx = torch.nonzero(den).reshape(-1)
    ret[idx, :] = (quat[idx, :3] * 2.0 * torch.acos(quat[idx, 3]).unsqueeze(-1)) / den[
        idx
    ].unsqueeze(-1)

    # Reshape and return output
    ret = ret.reshape(
        list(quat_shape)
        + [
            3,
        ]
    )
    return ret


@torch.jit.script
def axisangle2quat(vec, eps=1e-6):
    """
    Converts scaled axis-angle to quat.
    Args:
        vec (tensor): (..., 3) tensor where final dim is (ax,ay,az) axis-angle exponential coordinates
        eps (float): Stability value below which small values will be mapped to 0

    Returns:
        tensor: (..., 4) tensor where final dim is (x,y,z,w) vec4 float quaternion
    """
    # type: (Tensor, float) -> Tensor
    # store input shape and reshape
    input_shape = vec.shape[:-1]
    vec = vec.reshape(-1, 3)

    # Grab angle
    angle = torch.norm(vec, dim=-1, keepdim=True)

    # Create return array
    quat = torch.zeros(torch.prod(torch.tensor(input_shape)), 4, device=vec.device)
    quat[:, 3] = 1.0

    # Grab indexes where angle is not zero an convert the input to its quaternion form
    idx = angle.reshape(-1) > eps  # torch.nonzero(angle).reshape(-1)
    quat[idx, :] = torch.cat(
        [
            vec[idx, :] * torch.sin(angle[idx, :] / 2.0) / angle[idx, :],
            torch.cos(angle[idx, :] / 2.0),
        ],
        dim=-1,
    )

    # Reshape and return output
    quat = quat.reshape(
        list(input_shape)
        + [
            4,
        ]
    )
    return quat


@torch.jit.script
def normalize(x, eps: float = 1e-9):
    return x / x.norm(p=2, dim=-1).clamp(min=eps, max=None).unsqueeze(-1)


@torch.jit.script
def quat_apply(a, b):
    shape = b.shape
    a = a.reshape(-1, 4)
    b = b.reshape(-1, 3)
    xyz = a[:, :3]
    t = xyz.cross(b, dim=-1) * 2
    return (b + a[:, 3:] * t + xyz.cross(t, dim=-1)).view(shape)


@torch.jit.script
def quat_rotate(q, v):
    shape = q.shape
    q_w = q[:, -1]
    q_vec = q[:, :3]
    a = v * (2.0 * q_w**2 - 1.0).unsqueeze(-1)
    b = torch.cross(q_vec, v) * q_w.unsqueeze(-1) * 2.0
    c = (
        q_vec
        * torch.bmm(q_vec.view(shape[0], 1, 3), v.view(shape[0], 3, 1)).squeeze(-1)
        * 2.0
    )
    return a + b + c


@torch.jit.script
def quat_rotate_inverse(q, v):
    shape = q.shape
    q_w = q[:, -1]
    q_vec = q[:, :3]
    a = v * (2.0 * q_w**2 - 1.0).unsqueeze(-1)
    b = torch.cross(q_vec, v) * q_w.unsqueeze(-1) * 2.0
    c = (
        q_vec
        * torch.bmm(q_vec.view(shape[0], 1, 3), v.view(shape[0], 3, 1)).squeeze(-1)
        * 2.0
    )
    return a - b + c


@torch.jit.script
def quat_conjugate(a):
    shape = a.shape
    a = a.reshape(-1, 4)
    return torch.cat((-a[:, :3], a[:, -1:]), dim=-1).view(shape)


@torch.jit.script
def quat_unit(a):
    return normalize(a)


@torch.jit.script
def quat_from_angle_axis(angle, axis):
    theta = (angle / 2).unsqueeze(-1)
    xyz = normalize(axis) * theta.sin()
    w = theta.cos()
    return quat_unit(torch.cat([xyz, w], dim=-1))


@torch.jit.script
def get_identity_transform(x):
    x_ = torch.zeros_like(x)
    x_[:, 3] = 1.0
    return x_


@torch.jit.script
def normalize_angle(x):
    return torch.atan2(torch.sin(x), torch.cos(x))


@torch.jit.script
def tf_inverse(q, t):
    q_inv = quat_conjugate(q)
    return q_inv, quat_apply(q_inv, t)


@torch.jit.script
def tf_apply(q, t, v):
    return quat_apply(q, v) + t


@torch.jit.script
def tf_vector(q, v):
    return quat_apply(q, v)


@torch.jit.script
def tf_combine(q1, t1, q2, t2):
    return quat_mul(q1, q2), quat_apply(q1, t2) + t1


@torch.jit.script
def get_basis_vector(q, v):
    return quat_rotate(q, v)


def get_axis_params(value, axis_idx, x_value=0.0, dtype=np.float32, n_dims=3):
    """construct arguments to `Vec` according to axis index."""
    zs = np.zeros((n_dims,))
    assert axis_idx < n_dims, "the axis dim should be within the vector dimensions"
    zs[axis_idx] = 1.0
    params = np.where(zs == 1.0, value, zs)
    params[0] = x_value
    return list(params.astype(dtype))


@torch.jit.script
def copysign(a, b):
    # type: (float, Tensor) -> Tensor
    a = torch.tensor(a, device=b.device, dtype=torch.float).repeat(b.shape[0])
    return torch.abs(a) * torch.sign(b)


@torch.jit.script
def get_euler_xyz(q):
    qx, qy, qz, qw = 0, 1, 2, 3
    # roll (x-axis rotation)
    sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz])
    cosr_cosp = (
        q[:, qw] * q[:, qw]
        - q[:, qx] * q[:, qx]
        - q[:, qy] * q[:, qy]
        + q[:, qz] * q[:, qz]
    )
    roll = torch.atan2(sinr_cosp, cosr_cosp)

    # pitch (y-axis rotation)
    sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx])
    pitch = torch.where(
        torch.abs(sinp) >= 1, copysign(np.pi / 2.0, sinp), torch.asin(sinp)
    )

    # yaw (z-axis rotation)
    siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy])
    cosy_cosp = (
        q[:, qw] * q[:, qw]
        + q[:, qx] * q[:, qx]
        - q[:, qy] * q[:, qy]
        - q[:, qz] * q[:, qz]
    )
    yaw = torch.atan2(siny_cosp, cosy_cosp)

    return roll % (2 * np.pi), pitch % (2 * np.pi), yaw % (2 * np.pi)


@torch.jit.script
def quat_from_euler_xyz(roll, pitch, yaw):
    cy = torch.cos(yaw * 0.5)
    sy = torch.sin(yaw * 0.5)
    cr = torch.cos(roll * 0.5)
    sr = torch.sin(roll * 0.5)
    cp = torch.cos(pitch * 0.5)
    sp = torch.sin(pitch * 0.5)

    qw = cy * cr * cp + sy * sr * sp
    qx = cy * sr * cp - sy * cr * sp
    qy = cy * cr * sp + sy * sr * cp
    qz = sy * cr * cp - cy * sr * sp

    return torch.stack([qx, qy, qz, qw], dim=-1)


@torch.jit.script
def torch_rand_float(lower, upper, shape, device):
    # type: (float, float, Tuple[int, int], str) -> Tensor
    return (upper - lower) * torch.rand(*shape, device=device) + lower


@torch.jit.script
def torch_random_dir_2(shape, device):
    # type: (Tuple[int, int], str) -> Tensor
    angle = torch_rand_float(-np.pi, np.pi, shape, device).squeeze(-1)
    return torch.stack([torch.cos(angle), torch.sin(angle)], dim=-1)


@torch.jit.script
def tensor_clamp(t, min_t, max_t):
    return torch.max(torch.min(t, max_t), min_t)


@torch.jit.script
def scale(x, lower, upper):
    return 0.5 * (x + 1.0) * (upper - lower) + lower


@torch.jit.script
def unscale(x, lower, upper):
    return (2.0 * x - upper - lower) / (upper - lower)


def unscale_np(x, lower, upper):
    return (2.0 * x - upper - lower) / (upper - lower)


@torch.jit.script
def compute_heading_and_up(
    torso_rotation, inv_start_rot, to_target, vec0, vec1, up_idx
):
    # type: (Tensor, Tensor, Tensor, Tensor, Tensor, int) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]
    num_envs = torso_rotation.shape[0]
    target_dirs = normalize(to_target)

    torso_quat = quat_mul(torso_rotation, inv_start_rot)
    up_vec = get_basis_vector(torso_quat, vec1).view(num_envs, 3)
    heading_vec = get_basis_vector(torso_quat, vec0).view(num_envs, 3)
    up_proj = up_vec[:, up_idx]
    heading_proj = torch.bmm(
        heading_vec.view(num_envs, 1, 3), target_dirs.view(num_envs, 3, 1)
    ).view(num_envs)

    return torso_quat, up_proj, heading_proj, up_vec, heading_vec


@torch.jit.script
def compute_rot(torso_quat, velocity, ang_velocity, targets, torso_positions):
    vel_loc = quat_rotate_inverse(torso_quat, velocity)
    angvel_loc = quat_rotate_inverse(torso_quat, ang_velocity)

    roll, pitch, yaw = get_euler_xyz(torso_quat)

    walk_target_angle = torch.atan2(
        targets[:, 2] - torso_positions[:, 2], targets[:, 0] - torso_positions[:, 0]
    )
    angle_to_target = walk_target_angle - yaw

    return vel_loc, angvel_loc, roll, pitch, yaw, angle_to_target


def qrot(q, v):
    """
    Rotate vector(s) v about the rotation described by quaternion(s) q.
    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
    where * denotes any number of dimensions.
    Returns a tensor of shape (*, 3).
    """
    assert q.shape[-1] == 4
    assert v.shape[-1] == 3

    original_shape = list(v.shape)
    q = q.view(-1, 4)
    v = v.view(-1, 3)

    qvec = q[:, 1:]
    uv = torch.cross(qvec, v, dim=1)
    uuv = torch.cross(qvec, uv, dim=1)
    return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)


def qrot2(q, v):
    """
    Rotate vector(s) v about the rotation described by quaternion(s) q.
    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
    where * denotes any number of dimensions.
    Returns a tensor of shape (*, 3).
    """
    assert q.shape[1] == 4
    assert v.shape[1] == 3

    # original_shape = list(v.shape)

    qvec = q[:, 1:]
    uv = torch.cross(qvec, v, dim=1)
    uuv = torch.cross(qvec, uv, dim=1)
    return v + 2 * (q[..., :1, :] * uv + uuv)  # .view(original_shape)


def to_rotation_matrix(angles):
    az, el, th = angles[..., 0], angles[..., 1], angles[..., 2]
    batch_dims = list(az.shape)

    cx = torch.cos(az)
    cy = torch.cos(el)
    cz = torch.cos(th)
    sx = torch.sin(az)
    sy = torch.sin(el)
    sz = torch.sin(th)

    ones = torch.ones_like(cx)
    zeros = torch.zeros_like(cx)

    rx = torch.stack([ones, zeros, zeros, zeros, cx, -sx, zeros, sx, cx], dim=-1)
    ry = torch.stack([cy, zeros, sy, zeros, ones, zeros, -sy, zeros, cy], dim=-1)
    rz = torch.stack([cz, -sz, zeros, sz, cz, zeros, zeros, zeros, ones], dim=-1)

    rot_shape = batch_dims + [3, 3]
    rx = torch.reshape(rx, rot_shape)
    ry = torch.reshape(ry, rot_shape)
    rz = torch.reshape(rz, rot_shape)

    return torch.matmul(rz, torch.matmul(ry, rx))

core/utils/torch_utils.py
liruiw/Fleet-Tools
# --------------------------------------------------------
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
try:
    import open3d as o3d
except:
    pass
import numpy as np
import os
import sys
import random

import scipy.io as sio
import IPython
import time

import cv2
import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import tabulate

# import torch
# import torch.nn.functional as F
import shutil

from .geometry import *
from transforms3d.quaternions import *
from transforms3d.euler import *
from transforms3d.axangles import *
import sys
import gym
import json

from .teleop_utils import *
from colored import fg
import GPUtil
import psutil

np.set_printoptions(precision=3)
import warnings
from pydrake.common.deprecation import DrakeDeprecationWarning

warnings.simplefilter("always", DrakeDeprecationWarning)

# anchor seeds for franka panda IK solutions
anchor_seeds = np.array(
    [
        [0.0, -1.285, 0, -2.356, 0.0, 1.571, 0.785, 0, 0],
        [2.5, 0.23, -2.89, -1.69, 0.056, 1.46, -1.27, 0, 0],
        [2.8, 0.23, -2.89, -1.69, 0.056, 1.46, -1.27, 0, 0],
        [2, 0.23, -2.89, -1.69, 0.056, 1.46, -1.27, 0, 0],
        [2.5, 0.83, -2.89, -1.69, 0.056, 1.46, -1.27, 0, 0],
        [0.049, 1.22, -1.87, -0.67, 2.12, 0.99, -0.85, 0, 0],
        [-2.28, -0.43, 2.47, -1.35, 0.62, 2.28, -0.27, 0, 0],
        [-2.02, -1.29, 2.20, -0.83, 0.22, 1.18, 0.74, 0, 0],
        [-2.2, 0.03, -2.89, -1.69, 0.056, 1.46, -1.27, 0, 0],
        [-2.5, -0.71, -2.73, -0.82, -0.7, 0.62, -0.56, 0, 0],
        [-2, -0.71, -2.73, -0.82, -0.7, 0.62, -0.56, 0, 0],
        [-2.66, -0.55, 2.06, -1.77, 0.96, 1.77, -1.35, 0, 0],
        [1.51, -1.48, -1.12, -1.55, -1.57, 1.15, 0.24, 0, 0],
        [-2.61, -0.98, 2.26, -0.85, 0.61, 1.64, 0.23, 0, 0],
    ]
)

# predefined lists for replay buffer, environment, and policy
key_list = [
    "state",
    "action",
    "reward",
    "next_state",
    "done",
    "obs",
    "timestep",
    "joint_pos",
    "hand_camera_pose",
    "base_camera_pose",
    "ee_pose",
    "cam_intr",
    "curr_keypoint",
    "goal_keypoint",
    "goal_pose",
    "keypoint_wrench",
    "external_torque",
    "endeffector_wrench",
    "history_endeffector_wrench",
    "point_cloud",
    "tool_rel_pose",
    "img_feature",
]

key_shape_list = [
    (27,),
    (6,),
    (1,),
    (27,),
    (1,),
    (480, 640, 10),
    (1,),
    (9,),
    (4, 4),
    (4, 4),
    (16,),
    (3, 3),
    (4, 3),
    (3, 3),
    (4, 4),
    (3, 6),
    (9,),
    (6,),
    (6 * 32,),
    (4, 1024),
    (4, 4),
    (384,),
]


def color_print(str, color="white"):
    print(str)  #


def flatten_list(l):
    return [item for sublist in l for item in sublist]


def mkdir_if_missing(dst_dir, cleanup=False):
    if os.path.exists(dst_dir) and cleanup:
        shutil.rmtree(dst_dir)
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)


def print_and_write(file_handle, text):
    print(text)
    if file_handle is not None:
        file_handle.write(text + "\n")
    return text


def cycle(iterable):
    while True:
        for x in iterable:
            yield x


def normalize(v, axis=None, eps=1e-10):
    """L2 Normalize along specified axes."""
    return v / max(np.linalg.norm(v, axis=axis, keepdims=True), eps)


def merge_two_dicts(x, y):
    z = x.copy()
    z.update(y)
    return z


def get_usage():
    GPUs = GPUtil.getGPUs()
    memory_usage = psutil.virtual_memory().percent
    try:
        gpu_usage = max([GPU.memoryUsed for GPU in GPUs])
    except:
        gpu_usage = 1e4  # why?
    return gpu_usage, memory_usage


def visualize_image_observation(state):
    if type(state) is not np.ndarray:
        state = state.detach().cpu().numpy()

    for idx in range(0, len(state), 2):
        fig = plt.figure(figsize=(25.6, 9.6))
        ax = fig.add_subplot(1, 3, 1)
        plt.imshow((state[idx][:3].transpose(1, 2, 0) * 255).astype(np.uint8))

        if state[idx].shape[0] == 4:
            ax = fig.add_subplot(1, 3, 2)
            # IPython.embed()
            plt.imshow(state[idx][-1])

        if state[idx].shape[0] >= 5:
            ax = fig.add_subplot(1, 3, 2)
            plt.imshow(state[idx][-2])

            ax = fig.add_subplot(1, 3, 3)
            mask = (state[idx][-1]).astype(np.uint8)
            plt.imshow(mask)
        plt.show()
        # break


def randomize_sphere_lookat(target, near, far, x_near, x_far, theta_low=0, theta_high=np.pi / 2):
    """
    randomly initialize pose to look at some center.
    useful for initializing wrist or overhead cameras
    """

    res = None
    # the object is not on the table in the beginning
    theta = np.random.uniform(low=theta_low, high=theta_high)
    phi = np.random.uniform(low=np.pi / 2, high=3 * np.pi / 2)
    # top sphere
    r = np.random.uniform(low=near, high=far)
    # sphere radius
    pos = np.array(
        [
            r * np.sin(theta) * np.cos(phi),
            r * np.sin(theta) * np.sin(phi),
            r * np.cos(theta),
        ]
    )

    # noise and clip
    trans = pos + target + np.random.uniform(-0.03, 0.03, 3)
    trans[2] = np.clip(trans[2], 0.3, 0.7)
    trans[1] = np.clip(trans[1], -0.4, 0.4)
    trans[0] = np.clip(trans[0], x_near, x_far)

    rand_vec = [0, 0, -1]
    pos = trans - target

    # open gl coordinate transform
    R = inv_lookat(pos, 2 * pos, rand_vec).dot(rotZ(np.pi / 2)[:3, :3])

    # from end effector to the world
    target_pose = np.eye(4)
    target_pose[:3, 3] = trans
    target_pose[:3, :3] = R

    eff_up_vec = R.dot(np.array([1, 0, 0]))
    eff_forward_vec = R.dot(np.array([0, 0, 1]))

    return target_pose


def projection_to_intrinsics(mat, width=224, height=224):
    intrinsic_matrix = np.eye(3)
    mat = np.array(mat).reshape([4, 4]).T
    fv = width / 2 * mat[0, 0]
    fu = height / 2 * mat[1, 1]
    u0 = width / 2
    v0 = height / 2

    intrinsic_matrix[0, 0] = fu
    intrinsic_matrix[1, 1] = fv
    intrinsic_matrix[0, 2] = u0
    intrinsic_matrix[1, 2] = v0
    return intrinsic_matrix


def backproject_camera_target(im_depth, K, target_mask=None):
    Kinv = np.linalg.inv(K)

    width = im_depth.shape[1]
    height = im_depth.shape[0]
    depth = im_depth.astype(np.float32, copy=True).flatten()

    x, y = np.meshgrid(np.arange(width), np.arange(height))
    ones = np.ones((height, width), dtype=np.float32)
    x2d = np.stack((x, y, ones), axis=2).reshape(width * height, 3)  # each pixel

    # backprojection
    R = Kinv.dot(x2d.transpose())  #
    X = np.multiply(np.tile(depth.reshape(1, width * height), (3, 1)), R)
    # X[1] *= -1 # not sure if this is needed

    if target_mask is not None:
        target_mask = (depth != 0) * (target_mask.flatten() == 0)
        return X[:, target_mask]
    return X


def inv_lookat(eye, target=[0, 0, 0], up=[0, 1, 0]):
    """Generate LookAt matrix."""
    up = np.random.uniform(size=(3,))
    eye = np.float32(eye)  #  trans - target
    forward = normalize(target - eye)  #  trans - target
    side = normalize(np.cross(forward, up))
    up = np.cross(side, forward)  # redo

    R = np.stack([side, up, -forward], axis=-1)
    return R


def safe_div(a, b):
    return a / (b + 1e-10)


def save_args_json(path, args):
    mkdir_if_missing(path)
    arg_json = os.path.join(path, "config.json")
    with open(arg_json, "w") as f:
        # args = vars(args)
        json.dump(args, f, indent=4, sort_keys=True)


def make_video_writer(name, window_width, window_height):
    fourcc = cv2.VideoWriter_fourcc(*"MJPG")
    return cv2.VideoWriter(name, fourcc, 10.0, (window_width, window_height))


def write_video(
    traj,
    scene_file,
    expert_traj=None,
    IMG_SIZE=(112, 112),
    output_dir="output_misc/",
    target_name="",
    extra_text=None,
):
    ratio = 1 if expert_traj is None else 2
    video_writer = make_video_writer(
        os.path.join(
            output_dir + "_{}_rollout.avi".format(scene_file),
        ),
        int(ratio * IMG_SIZE[1]),
        int(IMG_SIZE[0]),
    )

    for i in range(len(traj)):
        img = traj[i][..., :3] * 255
        # cv2.resize(, (IMG_SIZE[:2]), interpolation = cv2.INTER_AREA)[..., [2, 1, 0]]
        if expert_traj is not None:
            idx = min(len(expert_traj) - 1, i)
            img = np.concatenate((img, expert_traj[idx][..., [2, 1, 0]]), axis=1)

        img = img.astype(np.uint8)
        if extra_text is not None:
            img = add_extra_text(img, extra_text)
        video_writer.write(img)


def merge_two_dicts(x, y):
    z = x.copy()
    z.update(y)
    return z


def PandaTaskSpace6D(trans_scale=0.04, rotation_scale=np.pi / 8):
    """end effector action space for the franka environment"""
    high = np.array(
        [
            trans_scale,
            trans_scale,
            trans_scale,
            rotation_scale,
            rotation_scale,
            rotation_scale,
        ]
    )  # np.pi/10
    low = np.array(
        [
            -trans_scale,
            -trans_scale,
            -trans_scale,
            -rotation_scale,
            -rotation_scale,
            -rotation_scale,
        ]
    )  # -np.pi/3
    shape = [6]
    # self.bounds = np.vstack([self.low, self.high])
    return gym.spaces.Box(low=low, high=high, shape=shape)


def get_interp_time(curr_time, finish_time, ratio):
    """get interpolated time between curr and finish"""
    return (finish_time - curr_time) * ratio + curr_time


def save_args_hydra(path, cfg):
    from omegaconf import OmegaConf

    cfg_dict = OmegaConf.to_container(cfg, resolve=True)
    save_args_json(path, cfg_dict)


def tabulate_print_state(state_dict):
    """
    Print ego car state dict
    """
    state_dict = sorted(state_dict.items())
    headers = [kv[0] for kv in state_dict]
    data = [[kv[1] for kv in state_dict]]
    print(tabulate.tabulate(data, headers, tablefmt="psql", floatfmt=".2f"))


def proj_point_img(
    img,
    K,
    offset_pose,
    points,
    color=(255, 0, 0),
    vis=False,
    neg_y=False,
    real_world=False,
):
    xyz_points = offset_pose[:3, :3].dot(points) + offset_pose[:3, [3]]
    if neg_y:
        xyz_points[:2] *= -1
    p_xyz = K.dot(xyz_points)
    p_xyz = p_xyz[:, p_xyz[2] > 0.03]
    x, y = (p_xyz[0] / p_xyz[2]).astype(np.int), (p_xyz[1] / p_xyz[2]).astype(np.int)
    valid_idx_mask = (x > 0) * (x < img.shape[1] - 1) * (y > 0) * (y < img.shape[0] - 1)
    img[y[valid_idx_mask], x[valid_idx_mask], :3] = [0, 1, 1]
    return img

core/utils/utils.py
liruiw/Fleet-Tools
from .frankadrakeenv import FrankaDrakeEnv
from .frankadrakeknifeenv import FrankaDrakeKnifeEnv
from .frankadrakewrenchenv import FrankaDrakeWrenchEnv
from .frankadrakespatulaenv import FrankaDrakeSpatulaEnv
from .frankadrakehammerenv import FrankaDrakeHammerEnv

from gym.envs.registration import register
import gym

ENV_ID = {
    "FrankaDrakeEnv": "FrankaDrakeEnv-v0",
    "FrankaDrakeKnifeEnv": "FrankaDrakeKnifeEnv-v0",
    "FrankaDrakeSpatulaEnv": "FrankaDrakeSpatulaEnv-v0",
    "FrankaDrakeWrenchEnv": "FrankaDrakeWrenchEnv-v0",
    "FrankaDrakeHammerEnv": "FrankaDrakeHammerEnv-v0",
}

ENV_CLASS = {
    "FrankaDrakeEnv": "FrankaDrakeEnv",
    "FrankaDrakeKnifeEnv": "FrankaDrakeKnifeEnv",
    "FrankaDrakeSpatulaEnv": "FrankaDrakeSpatulaEnv",
    "FrankaDrakeWrenchEnv": "FrankaDrakeWrenchEnv",
    "FrankaDrakeHammerEnv": "FrankaDrakeHammerEnv",
}


def make_env(env_name, **kwargs):
    env_id = ENV_ID[env_name]
    return gym.make(env_id, **kwargs)  #  eval(env_name)(**kwargs) #


def register_env(env_name):
    assert env_name in ENV_ID, "unknown environment"
    env_id = ENV_ID[env_name]
    env_class = ENV_CLASS[env_name]
    register(id=env_id, entry_point="env." + env_name.lower() + ":" + env_class)

env/__init__.py
liruiw/Fleet-Tools
import gym
import numpy as np

import warnings
from gym.utils import seeding
from pydrake.all import Simulator, SimulatorStatus, RandomGenerator
from pydrake.math import RollPitchYaw, RigidTransform

try:
    import open3d as o3d
except:
    pass

from pydrake.common import RandomGenerator
from env.env_util import *
import IPython
from core.utils import *
from colored import fg
import matplotlib.pyplot as plt
import time


class BaseDrakeRobotEnv(gym.Env):
    """
    Base Drake Environment Class for Manipulation
    """

    def __init__(self, task_config):
        self.reset_from_scene_file = False

    def cleanup_meshcat(self):
        # cleanup meshcat
        from env.env_util import meshcat

        if meshcat is not None:
            if not self.task_config.vis_collision_geometry:
                meshcat.SetProperty("collision", "visible", False)
            for index in range(self.plant.num_model_instances()):
                model_name = self.plant.GetModelInstanceName(ModelInstanceIndex(index))
                if "body_link" in model_name:  # cleanup objects and tools
                    meshcat.Delete("visualizer/" + model_name)
            time.sleep(0.1)

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        self.drake_random = RandomGenerator(seed)
        return [seed]

    def advance_time(self, context):
        self.time = context.get_time()
        catch = False
        s = time.time()
        repeat_num = self.task_config.sim.substep_num

        try:
            for _ in range(repeat_num):  # make the simulation smooth and visulizable
                self.time = context.get_time()
                status = self.simulator.AdvanceTo(self.time + self.sim_dt)
                self._get_robot_freq_data()  # 0.02s slower per step

        except RuntimeError as e:
            warnings.warn("Calling Done after catching RuntimeError:")
            warnings.warn(e.args[0])
            catch = True
            self.info["sim_failure"] = True  # emergent reset
            color_print("simulation failed, reset!", "light_slate_blue")

        sim_done = catch or status.reason() == SimulatorStatus.ReturnReason.kReachedTerminationCondition
        return sim_done

    def get_ee_pos(self):
        """get end effector poses"""
        plant_context = self.system.GetSubsystemContext(self.plant, self.context)
        return self.plant.GetFrameByName("panda_hand").CalcPoseInWorld(plant_context)

    def _get_reward(self, context, action):
        return 0

    def _get_progress(self, context):
        # determine which stage is at for the expert
        return None

    def control_finger(self, state=None):
        """extra function to control the finger"""
        if state == "open":
            self.finger_state = [0, 0]
        if state == "close":
            self.finger_state = [0.04, 0.04]

    def get_joint_angles(self):
        """get the joint angle of the robot"""

        # q0 = self.joint_measurement_port.Eval(self.context)
        q0 = self.plant.GetPositions(self.plant_context, self.plant.GetModelInstanceByName("panda"))
        self.joint_positions = np.array(q0).copy()
        return self.joint_positions

    def get_actual_joint_infos(self):
        joint_pos = self.plant.GetPositions(self.plant_context, self.plant.GetModelInstanceByName("panda"))  #
        joint_vel = self.plant.GetVelocities(self.plant_context, self.plant.GetModelInstanceByName("panda"))
        return joint_pos, joint_vel

    def get_obs(self, render=True):
        """get the environment observation"""
        return self._get_obs(self.context, render=render)

    def check_collision_name(self, body_a_name, body_b_name):
        """only return collisions if two body parts collide or anything collide with anything but tool and object"""
        self_collision = "panda_link" in body_a_name and "panda_link" in body_b_name
        external_collision = ("model" not in body_a_name and "panda" not in body_a_name and "panda" in body_b_name) or (
            "model" not in body_b_name and "panda" not in body_b_name and "panda" in body_a_name
        )
        return self_collision or external_collision  # or tool_collision

    def export_context(self):
        context = self.simulator.get_context()
        return context.Clone()

    def export_all_context(self):
        context = self.simulator.get_context()
        self.controller_context = self.controller_plant_system.GetSubsystemContext(self.controller_plant, self.context)

        return context.Clone(), self.controller_context.Clone()

    def load_all_context(self, context, controller_context):
        self.context.SetTimeStateAndParametersFrom(context)
        self.controller_context.SetTimeStateAndParametersFrom(controller_context)

        self.simulator.Initialize()

    def load_context(self, context):
        self.context.SetTimeStateAndParametersFrom(context)
        self.simulator.Initialize()

    def seed(self, seed=None):
        """Implements gym.Env.seed using Drake's RandomGenerator."""
        if seed:
            self.generator = RandomGenerator(seed)
        else:
            seed = self.generator()

        return [seed]

    def set_plant_positions(self, joints, joint_vel=None):
        # sync all plants that have the robot states
        self.plant_context = self.system.GetSubsystemContext(self.plant, self.context)
        self.plant.SetPositions(self.plant_context, self.plant.GetModelInstanceByName("panda"), joints)  #
        if joint_vel is None:
            self.plant.SetVelocities(self.plant_context, self.plant.GetModelInstanceByName("panda"), [0] * 9)
        else:
            self.plant.SetVelocities(self.plant_context, self.plant.GetModelInstanceByName("panda"), joint_vel)

        ee_pose = self.set_diffik_state(self.context, reinit=True, set_joints=joints)
        controller_context = self.controller_plant.CreateDefaultContext()

        self.controller_plant.SetPositions(
            controller_context,
            self.controller_plant.GetModelInstanceByName("panda"),
            joints,
        )
        self.controller_plant.SetVelocities(
            controller_context,
            self.controller_plant.GetModelInstanceByName("panda"),
            [0] * 9,
        )
        self.joint_positions = joints
        self.diffik_port.FixValue(self.context, RigidTransform(ee_pose))
        if self.use_joint_controller():
            self.joint_action_port.FixValue(self.context, self.joint_positions)
        self.ee_pose = np.array(ee_pose.GetAsMatrix4()).copy().flatten()

    def reset_object(self, context, curr_object_id=-1):
        """sample a random object to drop on the table"""
        self.plant_context = self.system.GetSubsystemContext(self.plant, self.context)
        x_range = self.task_config.env.init_object_x_range
        y_range = self.task_config.env.init_object_y_range

        # RollPitchYaw(0, 0, np.random.uniform(-np.pi, np.pi)
        initial_rot = RollPitchYaw(0, 0, 0)  # np.random.uniform(-np.pi, np.pi)

        tf = RigidTransform(
            initial_rot,
            [
                np.random.uniform(x_range[0], x_range[1]),
                np.random.uniform(y_range[0], y_range[1]),
                self.table_init_height,
            ],
        )

        self.set_object_pose(tf)
        if self.task_config.randomize_mass:
            random_mass_value = np.random.uniform(0.1, 3)
            self.plant.GetBodyByName(self.object_name).SetMass(self.context, mass=1)

    def set_object_pose(self, tf):
        """setting object pose
        :param tf: Rigidtransform tf for the pose
        """
        all_object_bodies = self.plant.GetBodyIndices(self.object_asset)

        for idx, obj_id in enumerate(all_object_bodies):
            body = self.plant.get_body(obj_id)
            self.plant.SetFreeBodyPose(self.plant_context, body, tf)
            self.plant.SetFreeBodySpatialVelocity(body, SpatialVelocity(np.zeros((6, 1))), self.plant_context)

    def _get_object_pose(self, context=None):
        if context is None:
            context = self.context
        plant_context = self.system.GetSubsystemContext(self.plant, self.context)
        self.object_pose = self.plant.GetFrameByName(self.object_name).CalcPoseInWorld(plant_context).GetAsMatrix4()
        return self.object_pose

    def _get_tool_pose(self, context=None):
        if context is None:
            context = self.context
        plant_context = self.system.GetSubsystemContext(self.plant, self.context)
        self.tool_pose = self.plant.GetFrameByName(self.tool_name).CalcPoseInWorld(plant_context).GetAsMatrix4()
        return self.tool_pose

    def solve_ik(self, target_pose, seed, rot_tol=0.01, filter_manipulability=False, table_col=True):
        s = time.time()
        diagram_context = self.controller_plant_system.CreateDefaultContext()

        res = solve_ik(
            self.controller_plant,
            self.controller_plant.GetFrameByName("panda_hand"),
            RigidTransform(target_pose),
            seed.reshape(-1, 1),
            rot_tol=rot_tol,
            add_table_col=table_col,
            timeout=True,
            # add_gripper_faceup=False,
            filter_manipulability=filter_manipulability,
        )
        return res

    ### function for panda with tools
    def reset_robot(self, context, ik_consider_collision=False):
        """sample a random half sphere view and solve for inverse kinematics"""
        self.plant_context = self.system.GetSubsystemContext(self.plant, self.context)
        self.joint_positions = np.array(self.task_config.env.init_joints)
        self.plant.SetPositions(
            self.plant_context,
            self.plant.GetModelInstanceByName("panda"),
            self.joint_positions,
        )  #
        self.plant.SetVelocities(self.plant_context, self.plant.GetModelInstanceByName("panda"), [0] * 9)

        if not self.task_config.env.randomize_init_eff:
            return self.get_ee_pos()

        target_obj_pose = self.plant.GetFrameByName(self.object_name).CalcPoseInWorld(self.plant_context)
        target = target_obj_pose.translation()[:3]
        target = np.array(target)
        target[-1] = 0
        ik = None
        loop_num = 50
        initial_near = self.task_config.env.init_endeffector_range[0]
        initial_far = self.task_config.env.init_endeffector_range[1]
        inner_loop_num = 5
        res = None

        for i in range(loop_num):
            # the object is not on the table in the beginning
            target_pose = randomize_sphere_lookat(target, near=initial_near, far=initial_far, x_near=0.2, x_far=0.5)

            if target_pose[2, 3] < 0.3:
                print("end effector init too low. skip: ", target_pose[2, 3])
                continue

            idx = 0 if target_pose[1, 3] > target[1] else 1
            for j in range(inner_loop_num):

                seed = np.array(anchor_seeds[np.random.randint(len(anchor_seeds))])
                res = self.solve_ik(target_pose, seed.reshape(-1, 1), filter_manipulability=True)
                if res is not None:
                    res = res.get_x_val()[:9]
                    res[-2:] = 0.0
                    self.set_plant_positions(res[:9])
                    self.initial_ee_pose = self.get_ee_pos().GetAsMatrix4()
                    return self.initial_ee_pose

        color_print("robot end effector does not get to randomize", "light_slate_blue")
        res = self.task_config.env.init_joints
        self.set_plant_positions(res[:9])
        self.initial_ee_pose = self.get_ee_pos().GetAsMatrix4()
        return self.initial_ee_pose

    def setup(self, tool_obj_info, tool_info):
        self.object_asset, self.object_name, self.tool_obj_info, _ = tool_obj_info
        (
            self.tool_asset,
            self.tool_name,
            self.tool_info,
            _,
            self.tool_rel_pose,
        ) = tool_info

env/basedrakeenv.py
liruiw/Fleet-Tools
import numpy as np
import os
import warnings

from core.utils.robot_internal_controller import RobotInternalController
from pydrake.all import *
import warnings

warnings.simplefilter("ignore")

import IPython
from core.utils import *
from env.meshcat_cpp_utils import *
import time
import json
from colored import fg
import gym

screw_joint_ = None
bolt_joint = None
extra_force_system = None
TRAIN_TEST_SPLIT = 0.5


class ExternalForceSystem(LeafSystem):
    """
    This class provides external force for rigid bodies in the MBP to change the force behavior.
    """

    def __init__(self, plant, target_idx):
        LeafSystem.__init__(self)

        self.nv = 0  # plant.num_velocities()
        self.target_body_index = target_idx
        forces_cls = Value[List[ExternallyAppliedSpatialForce]]

        self.DeclareAbstractOutputPort("spatial_forces_vector", lambda: forces_cls(), self.DoCalcAbstractOutput)

        self.DeclareVectorOutputPort("generalized_forces", self.nv, self.DoCalcVectorOutput)

        self.wrench = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0])

    def DoCalcAbstractOutput(self, context, y_data):
        test_force = ExternallyAppliedSpatialForce()
        test_force.body_index = self.target_body_index
        test_force.p_BoBq_B = np.array([0, 0, 1e-3])  # the origin of the torque  np.zeros(3) #
        test_force.F_Bq_W = SpatialForce(tau=self.wrench[3:], f=self.wrench[:3])
        y_data.set_value([test_force])

    def DoCalcVectorOutput(self, context, y_data):
        y_data.SetFromVector([0] * self.nv)


def AddFloatingRpyJoint(plant, frame, instance):
    inertia = UnitInertia.SolidSphere(1.0)
    x_body = plant.AddRigidBody("x", instance, SpatialInertia(mass=0, p_PScm_E=[0.0, 0.0, 0.0], G_SP_E=inertia))
    plant.AddJoint(PrismaticJoint("x", plant.world_frame(), x_body.body_frame(), [1, 0, 0]))
    y_body = plant.AddRigidBody("y", instance, SpatialInertia(mass=0, p_PScm_E=[0.0, 0.0, 0.0], G_SP_E=inertia))
    plant.AddJoint(PrismaticJoint("y", x_body.body_frame(), y_body.body_frame(), [0, 1, 0]))
    z_body = plant.AddRigidBody("z", instance, SpatialInertia(mass=0, p_PScm_E=[0.0, 0.0, 0.0], G_SP_E=inertia))
    plant.AddJoint(PrismaticJoint("z", y_body.body_frame(), z_body.body_frame(), [0, 0, 1]))
    plant.AddJoint(BallRpyJoint("ball", z_body.body_frame(), frame))


def FindResource(filename):
    return os.path.join(os.path.dirname(__file__), filename)


def AddFranka(plant, collision_model="no_collision"):
    """Add the default franka panda robot to the scene. its origin is also the origin of the scene"""

    franka_combined_path = FindResource("models/panda_arm_hand.urdf")
    parser = Parser(plant)
    franka = parser.AddModelFromFile(franka_combined_path)
    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("panda_link0"))

    # Set default positions:
    q0 = [0.0, -1.285, 0, -2.356, 0.0, 1.571, 0.785, 0.0, 0.0]
    index = 0
    for joint_index in plant.GetJointIndices(franka):
        joint = plant.get_mutable_joint(joint_index)
        if isinstance(joint, RevoluteJoint):
            joint.set_default_angle(q0[index])
            index += 1
    return franka


def AddTable(plant, blender=False):
    """Add the default table to the scene"""
    parser = Parser(plant)
    if blender:
        sdf_path = FindResource("models/table/table.sdf")
    else:
        sdf_path = FindResource("models/table.sdf")
    table_pose = RigidTransform(RollPitchYaw(0, 0, 0), [0.96, 0, 0.15])
    table = parser.AddModelFromFile(sdf_path)
    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("top_center", table), table_pose)
    return table


def AddCameraBox(plant, X_WC, name="camera0", parent_frame=None):
    if not parent_frame:
        parent_frame = plant.world_frame()

    camera = Parser(plant).AddModelFromFile(FindResource("models/camera_box.sdf"), name)
    plant.WeldFrames(parent_frame, plant.GetFrameByName("base", camera), X_WC)


def AddShape(plant, shape, name, mass=1, mu=1, color=[0.5, 0.5, 0.9, 1.0], collision=True):
    instance = plant.AddModelInstance(name)
    if isinstance(shape, Box):
        inertia = UnitInertia.SolidBox(shape.width(), shape.depth(), shape.height())
    elif isinstance(shape, Cylinder):
        inertia = UnitInertia.SolidCylinder(shape.radius(), shape.length())
    elif isinstance(shape, Sphere):
        inertia = UnitInertia.SolidSphere(shape.radius())
    else:
        raise RuntimeError(f"need to write the unit inertia for shapes of type {shape}")
    body = plant.AddRigidBody(
        name,
        instance,
        SpatialInertia(mass=mass, p_PScm_E=np.array([0.0, 0.0, 0.0]), G_SP_E=inertia),
    )

    if plant.geometry_source_is_registered():
        """add small spheres"""
        if collision:
            if isinstance(shape, Box):
                plant.RegisterCollisionGeometry(
                    body,
                    RigidTransform(),
                    Box(
                        shape.width() - 0.001,
                        shape.depth() - 0.001,
                        shape.height() - 0.001,
                    ),
                    name,
                    CoulombFriction(mu, mu),
                )
                i = 0
                for x in [-shape.width() / 2.0, shape.width() / 2.0]:
                    for y in [-shape.depth() / 2.0, shape.depth() / 2.0]:
                        for z in [-shape.height() / 2.0, shape.height() / 2.0]:
                            plant.RegisterCollisionGeometry(
                                body,
                                RigidTransform([x, y, z]),
                                Sphere(radius=1e-7),
                                f"contact_sphere{i}",
                                CoulombFriction(mu, mu),
                            )
                            i += 1
            else:
                plant.RegisterCollisionGeometry(body, RigidTransform(), shape, name, CoulombFriction(mu, mu))
        plant.RegisterVisualGeometry(body, RigidTransform(), shape, name, color)
    return instance


def LoadMeshFrom(parent_path, name, fixed_idx, info):
    object_paths = sorted(
        [os.path.join(parent_path, name, o, o + ".sdf") for o in os.listdir(os.path.join(parent_path, name))]
    )
    if fixed_idx != -1:
        object_paths = [n for n in object_paths if f"model{fixed_idx}" in n]
        if len(object_paths) == 0:
            print("no matching idx tool", fixed_idx, object_paths)
            object_paths = sorted(
                [os.path.join(parent_path, name, o, o + ".sdf") for o in os.listdir(os.path.join(parent_path, name))]
            )

    split_idx = int(TRAIN_TEST_SPLIT * len(object_paths))
    idx = np.random.randint(0, max(1, split_idx))

    path = object_paths[idx]
    idx = idx

    fixedpath = [p for p in object_paths if str(idx) in p]
    if len(fixedpath) > 0:
        path = fixedpath[0]

    if len(info) > 0:

        matched_paths = [p for p in object_paths if p.split("/")[-1][:-4] in info]
        if len(matched_paths) > 0:
            path = matched_paths[0]
        else:
            print(object_paths)
            # raise ValueError("ERROR: matching instance not found: {}".format(info))
            print("ERROR: matching instance not found: {}".format(info))
    return path


def AddTools(
    plant,
    tool_name,
    instance_info={},
    fixed_idx=-1,
    train=False,
    actuated=False,
    load_all_objs=False,
    blender_render=False,
):
    """
    Adding the tool to the robot
    """
    parser = Parser(plant)
    # obj_path = "assets/tools"  # objects_sdf
    obj_path = "assets/tools"  # objects_sdf
    info = "" if len(instance_info) == 0 else instance_info["tool_name"]
    path = LoadMeshFrom(obj_path, tool_name, fixed_idx, info)
    model_name = path.split("/")[-1][:-4]

    tool_info_path = path.replace(".sdf", "_info.json")
    tool_info = json.load(open(tool_info_path))
    object_name = tool_name + "_" + model_name + "_body_link"
    if blender_render:
        path = path.replace(".sdf", "_blender.sdf")

    tool_obj = parser.AddModelFromFile(path, object_name)
    return [tool_obj, object_name, tool_info, path]


def AddToolObjects(
    builder, plant, scene_graph, tool_name, instance_info={}, fixed_idx=-1, train=False, load_all_objs=False
):
    """
    Adding the object that the tool needs to interact with
    """
    if tool_name == "wrench":
        return AddWrenchObject(
            builder,
            plant,
            scene_graph,
            tool_name,
            instance_info,
            fixed_idx=fixed_idx,
            train=train,
        )

    if tool_name == "knife":
        return AddKnifeObject(
            builder,
            plant,
            scene_graph,
            tool_name,
            instance_info,
            fixed_idx=fixed_idx,
            train=train,
        )

    if tool_name == "hammer":
        return AddHammerObject(
            builder,
            plant,
            scene_graph,
            tool_name,
            instance_info,
            fixed_idx=fixed_idx,
            train=train,
        )

    if tool_name == "spatula":
        return AddSpatulaObject(
            builder,
            plant,
            scene_graph,
            tool_name,
            instance_info,
            fixed_idx=fixed_idx,
            train=train,
        )

    if tool_name == "scoop" or tool_name == "whisk":
        return AddBowlObject(
            builder,
            plant,
            scene_graph,
            tool_name,
            instance_info,
            fixed_idx=fixed_idx,
            train=train,
        )

    parser = Parser(plant)
    # obj_path = "assets/tool_obj"  # objects_sdf
    obj_path = "assets/tool_obj"  # objects_sdf
    info = "" if len(instance_info) == 0 else instance_info["object_name"]
    path = LoadMeshFrom(obj_path, tool_name, fixed_idx, info)
    print(f"Using object at {path}")

    model_name = path.split("/")[-1][:-4]
    tool_obj_keypoint_path = path.replace(".sdf", "_info.json")

    tool_obj_info = json.load(open(tool_obj_keypoint_path))  # ['keypoints']
    object_name = model_name + "_body_link"
    tool_obj = parser.AddModelFromFile(path, object_name)
    return [tool_obj, object_name, tool_obj_info, path]


def AddBowlObject(
    builder, plant, scene_graph, tool_name, instance_info={}, fixed_idx=-1, train=False, load_all_objs=False
):
    """add a welded bowl to the table"""
    parser = Parser(plant)
    # obj_path = "assets/tool_obj"  # objects_sdf
    obj_path = "assets/tool_obj"  # objects_sdf
    info = "" if len(instance_info) == 0 else instance_info["object_name"]
    path = LoadMeshFrom(obj_path, tool_name, fixed_idx, info)
    print(f"Using object at {path}")

    initial_rot = RollPitchYaw(0, 0, np.random.uniform(-np.pi, np.pi))  # -np.pi/2
    tf = RigidTransform(initial_rot, [np.random.uniform(0.2, 0.4), np.random.uniform(-0.2, 0.2), 0.02])

    if len(instance_info) > 0:
        object_pose = instance_info["object_pose"][0]
        tf = RigidTransform(object_pose)

    model_name = path.split("/")[-1][:-4]
    tool_obj_keypoint_path = path.replace(".sdf", "_info.json")
    tool_obj_info = json.load(open(tool_obj_keypoint_path))  # ['keypoints']
    object_name = model_name + "_body_link"
    tool_obj = parser.AddModelFromFile(path, object_name)
    tool_body_frame = plant.GetBodyByName(object_name).body_frame()

    plant.WeldFrames(plant.world_frame(), tool_body_frame, tf)
    # self.set_object_pose(plant_context, tf)
    return [tool_obj, object_name, tool_obj_info, path]


def AddSpatulaObject(
    builder, plant, scene_graph, tool_name, instance_info={}, fixed_idx=-1, train=False, load_all_objs=False
):
    """Add four small spheres under spatula for scooping"""
    parser = Parser(plant)
    obj_path = "assets/tool_obj"  # objects_sdf
    info = "" if len(instance_info) == 0 else instance_info["object_name"]
    path = LoadMeshFrom(obj_path, tool_name, fixed_idx, info)
    print(f"Using object at {path}")

    model_name = path.split("/")[-1][:-4]
    tool_obj_keypoint_path = path.replace(".sdf", "_info.json")
    tool_obj_info = json.load(open(tool_obj_keypoint_path))

    object_name = model_name + "_body_link"
    tool_obj = parser.AddModelFromFile(path, object_name)

    body_frame_id = plant.GetBodyFromFrameId(plant.GetBodyFrameIdOrThrow(plant.GetBodyIndices(tool_obj)[0]))

    # add small spheres for easy scooping
    i = 0
    for x in [-0.02, 0.02]:
        for y in [-0.02, 0.02]:
            plant.RegisterCollisionGeometry(
                body_frame_id,
                RigidTransform([x, y, -0.02]),
                Sphere(radius=1e-5),
                f"contact_sphere{i}",
                CoulombFriction(1, 1),
            )
            i += 1

    return [tool_obj, object_name, tool_obj_info, path]


def AddHammerObject(
    builder, plant, scene_graph, tool_name, instance_info={}, fixed_idx=-1, train=False, load_all_objs=False
):
    """
    Adding a bolt with prismatic joint for hammering
    """
    parser = Parser(plant)
    obj_path = "assets/tool_obj"  # objects_sdf
    info = "" if len(instance_info) == 0 else instance_info["object_name"]
    path = LoadMeshFrom(obj_path, tool_name, fixed_idx, info)
    print(f"Using object at {path}")

    model_name = path.split("/")[-1][:-4]
    tool_obj_keypoint_path = path.replace(".sdf", "_info.json")
    tool_obj_info = json.load(open(tool_obj_keypoint_path))  # ['keypoints']

    object_name = model_name + "_body_link"
    tool_obj = parser.AddModelFromFile(path, object_name)

    # build the bolt-nut combo
    # child rotates and translates on the axis of Z of the parent
    bolt = tool_obj

    # get table
    bolt_body_frame_id = plant.GetBodyFrameIdOrThrow(plant.GetBodyIndices(bolt)[0])

    try:
        table_body_frame = plant.GetBodyByName("link").body_frame()
    except:
        table_body_frame = plant.GetBodyByName("top_center").body_frame()
    bolt_body_frame = plant.GetBodyFromFrameId(bolt_body_frame_id).body_frame()
    global bolt_joint, extra_force_system

    # transparent
    box_instance = AddShape(
        plant,
        Box(0.1, 0.1, 0.03),
        "box",
        1,
        1,
        color=[0.1, 0.5, 0.1, 0.0],
        collision=True,
    )
    box_body_frame = plant.GetBodyByName("box").body_frame()
    bolt_pose = RigidTransform(RollPitchYaw(0, 0, 0), [0.56, 0.0, 0.2])  # 0.13

    plant.WeldFrames(plant.world_frame(), box_body_frame, bolt_pose)
    bolt_joint = PrismaticJoint(
        "bolt_joint",
        box_body_frame,
        bolt_body_frame,
        [0, 0, -1],
        damping=1,
        pos_upper_limit=-0.01,
        pos_lower_limit=-0.025,
    )
    bolt_joint.set_default_translation(-0.025)
    bolt_joint = plant.AddJoint(bolt_joint)
    bolt_body_id = plant.GetBodyIndices(bolt)[0]

    # maybe add a board first
    extra_force_system = builder.AddSystem(ExternalForceSystem(plant, bolt_body_id))
    return [tool_obj, object_name, tool_obj_info, path]


def AddKnifeObject(
    builder, plant, scene_graph, tool_name, instance_info={}, fixed_idx=-1, train=False, load_all_objs=False
):
    """Add two separate parts for knife splitting"""
    parser = Parser(plant)
    obj_path = "assets/tool_obj"  # objects_sdf
    info = "" if len(instance_info) == 0 else instance_info["object_name"]
    path = LoadMeshFrom(obj_path, tool_name, fixed_idx, info)
    print(f"Using object at {path}")

    model_name = path.split("/")[-1][:-4]
    tool_obj_keypoint_path = path.replace(".sdf", "_info.json")
    tool_obj_info = json.load(open(tool_obj_keypoint_path))

    object_name = model_name + "_body_link"
    object_name2 = model_name + "_body_link_2"

    tool_obj = parser.AddModelFromFile(path, object_name)  # add two parts

    copy_path = path.replace(".sdf", "_copy.sdf")
    if not os.path.exists(copy_path):
        copy_file_content = open(path).read()
        copy_file_content = copy_file_content.replace("body_link", "body_link_2")
        with open(copy_path, "w") as f:
            f.write(copy_file_content)

    # manually create the copy
    tool_obj2 = parser.AddModelFromFile(copy_path, object_name2)
    # add small spheres for easy splitting
    body_frame_id = plant.GetBodyFromFrameId(plant.GetBodyFrameIdOrThrow(plant.GetBodyIndices(tool_obj)[0]))

    i = 0
    for x in [-0.02, 0.02]:
        for y in [-0.02, 0.02]:
            plant.RegisterCollisionGeometry(
                body_frame_id,
                RigidTransform([x, y, -0.04]),
                Sphere(radius=1e-4),
                f"contact_sphere{i}",
                CoulombFriction(0.5, 0.5),
            )
            i += 1

    body_frame_id = plant.GetBodyFromFrameId(plant.GetBodyFrameIdOrThrow(plant.GetBodyIndices(tool_obj2)[0]))

    # add small spheres for easy splitting
    for x in [-0.02, 0.02]:
        for y in [-0.02, 0.02]:
            plant.RegisterCollisionGeometry(
                body_frame_id,
                RigidTransform([x, y, -0.04]),
                Sphere(radius=1e-4),
                f"contact_sphere{i}",
                CoulombFriction(0.5, 0.5),
            )
            i += 1

    return [tool_obj, object_name, tool_obj_info, path]


def AddWrenchObject(
    builder, plant, scene_graph, tool_name, instance_info={}, fixed_idx=-1, train=False, load_all_objs=False
):
    """A way to handle the screw joint special case"""
    parser = Parser(plant)
    obj_path = "assets/tool_obj"  # objects_sdf
    # obj_path = "assets/objects_scaled"  # objects_sdf
    info = "" if len(instance_info) == 0 else instance_info["object_name"]
    path = LoadMeshFrom(obj_path, tool_name, fixed_idx, info)
    print(f"Using object at {path}")

    model_name = path.split("/")[-1][:-4]
    tool_obj_keypoint_path = path.replace(".sdf", "_nut_info.json")
    tool_obj_info = json.load(open(tool_obj_keypoint_path))  # ['keypoints']

    object_name = model_name + "_body_link"
    tool_obj = parser.AddModelFromFile(path, object_name)

    # build the bolt-nut combo
    # child rotates and translates on the axis of Z of the parent
    bolt = tool_obj
    nut_name = object_name.replace("_body", "_nut_body")
    nut_path = path.replace(".sdf", "_nut.sdf")

    nut = parser.AddModelFromFile(nut_path, nut_name)
    nut_body_id = plant.GetBodyIndices(nut)[0]
    nut_body_frame_id = plant.GetBodyFrameIdOrThrow(nut_body_id)
    bolt_body_frame_id = plant.GetBodyFrameIdOrThrow(plant.GetBodyIndices(bolt)[0])

    # body_frame
    nut_body_frame = plant.GetBodyFromFrameId(nut_body_frame_id).body_frame()
    bolt_body_frame = plant.GetBodyFromFrameId(bolt_body_frame_id).body_frame()
    # the origin of the nut must be at the screw
    screw_joint = ScrewJoint(
        "nut_bolt_joint",
        bolt_body_frame,
        nut_body_frame,
        screw_pitch=5e-3,
        damping=5e-2,
    )

    # weld the bolt to the table, at a fixed pose for now.
    plant.AddJoint(screw_joint)
    x_random_range = np.random.uniform(0.35, 0.5)
    y_random_range = np.random.uniform(-0.15, 0.15)
    bolt_pose = RigidTransform(RollPitchYaw(0, 0, 0), [x_random_range, y_random_range, 0.18])

    # bolt_pose = RigidTransform(RollPitchYaw(0, 0, 0), [0.4, 0, 0.18])
    plant.WeldFrames(plant.world_frame(), bolt_body_frame, bolt_pose)
    tool_obj = nut
    global screw_joint_, extra_force_system
    screw_joint_ = screw_joint
    extra_force_system = builder.AddSystem(ExternalForceSystem(plant, nut_body_id))
    # extra_force_system = builder.AddSystem(ExternalForceSystem(plant, nut_body_id))

    return [tool_obj, object_name, tool_obj_info, path]

env/env_object_related.py
liruiw/Fleet-Tools
import numpy as np
import os
import sys
import torch

from core.utils.robot_internal_controller import RobotInternalController
from core.utils.spatialstiffness_controller import PoseController
from pydrake.all import *
import warnings

warnings.simplefilter("ignore")
import IPython
from core.utils.utils import *
from env.meshcat_cpp_utils import *
import time
from colored import fg
from env.env_object_related import *

try:
    from pointnet2_ops.pointnet2_utils import furthest_point_sample, gather_operation
except:
    pass
from scipy import interpolate

g = RandomGenerator()
#  global variables
meshcat = None
visualizer = None
blender_color_cam = None
robot_plant = None  # for computing FK and IK'

# drake camera coordinate system has z facing forward and y facing downward
OVERHEAD_CAM_EXTR = (
    np.array(
        [
            [-0.21092498, -0.05561141, -0.97591907, 1.51906],
            [-0.13466501, -0.9872077, 0.08535977, 0.161601],
            [-0.9681818, 0.14942665, 0.20073785, 0.634862],
            [0.0, 0.0, 0.0, 1.0],
        ]
    )
    .dot(rotY(np.pi / 6))
    .dot(rotZ(-np.pi / 2))
)


HAND_CAM_EXTR = (
    np.array(
        [
            [0.00266189, -0.00885279, 0.99995727, 0.1260556],
            [0.02328945, 0.99969013, 0.00878842, -0.036485],
            [-0.99972522, 0.02326506, 0.00286725, 0.057091],
            [0.0, 0.0, 0.0, 1.0],
        ]
    )
    .dot(rotY(-np.pi / 2))
    .dot(rotZ(-np.pi / 2))
    .dot(rotX(np.pi / 10))
)

CAM_INTR = np.array(
    [
        614.88671875,
        0.0,
        323.3525390625,
        0.0,
        615.4398193359375,
        239.89163208007812,
        0.0,
        0.0,
        1.0,
    ]
).reshape(3, 3)


def reset_global_variables():
    # hack to reset all global variables
    from core.utils.utils import (
        vr_has_init,
        vr_frame_origin,
        rec_action_pose,
        oculus_reader,
    )

    vr_has_init = False
    rec_action_pose = None
    coordinate_origin = None
    oculus_reader = None


def color_print(str, color="white"):
    print(fg(color) + str)


def AddRgbdSensors(
    builder,
    plant,
    scene_graph,
    also_add_point_clouds=True,
    model_instance_prefix="panda",
    depth_camera=None,
    renderer=None,
    width=640,
    height=480,
    relative_pose=None,
    attach_body_name="",
):
    """
    Adds a RgbdSensor to every body in the plant with a name starting with body_prefix.
    """
    if sys.platform == "linux" and os.getenv("DISPLAY") is None:
        from pyvirtualdisplay import Display

        virtual_display = Display(visible=0, size=(1400, 900))
        virtual_display.start()

    if not renderer:
        renderer = "my_renderer"

    if not scene_graph.HasRenderer(renderer):
        scene_graph.AddRenderer(renderer, MakeRenderEngineVtk(RenderEngineVtkParams()))

    if not depth_camera:
        depth_camera = DepthRenderCamera(
            RenderCameraCore(
                renderer,
                CameraInfo(width=width, height=height, intrinsic_matrix=CAM_INTR),
                ClippingRange(near=0.02, far=3.0),
                RigidTransform(),
            ),
            #  overlaped with the rgbd sensor pose
            DepthRange(0.2, 3.0),
        )

    for index in range(plant.num_model_instances()):
        model_instance_index = ModelInstanceIndex(index)
        model_name = plant.GetModelInstanceName(model_instance_index)

        if model_name.startswith(model_instance_prefix):
            if len(attach_body_name) == 0:
                body_index = plant.GetBodyIndices(model_instance_index)[0]
                model_name = "_".join([model_name, "base", "cam"])

            else:
                body_index = plant.GetBodyByName(attach_body_name).index()
                model_name = "_".join([attach_body_name, "cam"])

            if relative_pose is not None:
                X_PB = RigidTransform(relative_pose)
            else:
                X_PB = RigidTransform()
            rgbd = builder.AddSystem(
                RgbdSensor(
                    parent_id=plant.GetBodyFrameIdOrThrow(body_index),
                    X_PB=X_PB,
                    depth_camera=depth_camera,
                    show_window=False,
                )
            )

            rgbd.set_name(model_name)
            builder.Connect(scene_graph.get_query_output_port(), rgbd.query_object_input_port())

            # Export the camera outputs
            builder.ExportOutput(rgbd.color_image_output_port(), f"{model_name}_rgb_image")
            builder.ExportOutput(rgbd.depth_image_32F_output_port(), f"{model_name}_depth_image")
            builder.ExportOutput(rgbd.label_image_output_port(), f"{model_name}_label_image")

            if also_add_point_clouds:
                # Add a system to convert the camera output into a point cloud
                to_point_cloud = builder.AddSystem(
                    DepthImageToPointCloud(
                        camera_info=rgbd.depth_camera_info(),
                        fields=BaseField.kXYZs | BaseField.kRGBs,
                    )
                )
                builder.Connect(
                    rgbd.depth_image_32F_output_port(),
                    to_point_cloud.depth_image_input_port(),
                )
                builder.Connect(
                    rgbd.color_image_output_port(),
                    to_point_cloud.color_image_input_port(),
                )

                class ExtractBodyPose(LeafSystem):
                    def __init__(self, body_index):
                        LeafSystem.__init__(self)
                        self.body_index = body_index
                        self.DeclareAbstractInputPort("poses", plant.get_body_poses_output_port().Allocate())
                        self.DeclareAbstractOutputPort(
                            "pose",
                            lambda: AbstractValue.Make(RigidTransform()),
                            self.CalcOutput,
                        )

                    def CalcOutput(self, context, output):
                        poses = self.EvalAbstractInput(context, 0).get_value()
                        pose = poses[int(self.body_index)]
                        output.get_mutable_value().set(pose.rotation(), pose.translation())
                        #
                        new_output = RigidTransform(output.get_mutable_value().GetAsMatrix4().dot(relative_pose))
                        output.get_mutable_value().set(new_output.rotation(), new_output.translation())

                camera_pose = builder.AddSystem(ExtractBodyPose(body_index))
                builder.Connect(plant.get_body_poses_output_port(), camera_pose.get_input_port())
                builder.Connect(
                    camera_pose.get_output_port(),
                    to_point_cloud.GetInputPort("camera_pose"),
                )
                builder.ExportOutput(camera_pose.get_output_port(), f"{model_name}_camera_pose")
                # Export the point cloud output.
                builder.ExportOutput(
                    to_point_cloud.point_cloud_output_port(),
                    f"{model_name}_point_cloud",
                )


def UpdateObjectPhysicalProperties(
    scene_graph,
    inspector,
    source_id,
    oid,
    obj_name,
    range_lo=0.1,
    range_hi=0.9,
    randomize_physics=False,
    frictions_hyper_param=None,
    hydroelastic_contact=False,
    change_mass=False,
    compliant_model=False,
):
    """change physics parameters such as friction, mass, and hydroelastic contacts"""
    if not inspector.BelongsToSource(oid, source_id):
        return
    obj_gids = []
    for geometry_id in inspector.GetAllGeometryIds():
        name = inspector.GetName(geometry_id)
        if obj_name in name and "collision" in name:
            obj_gids.append(geometry_id)

    for gid in obj_gids:
        props = inspector.GetProximityProperties(gid)

        if randomize_physics:
            static_friction = np.random.uniform(range_lo, range_hi)  # simple for now
            friction = CoulombFriction(static_friction, static_friction)
            if props is not None:
                props.UpdateProperty("material", "coulomb_friction", friction)
            try:
                scene_graph.AssignRole(source_id, gid, props, RoleAssign.kReplace)
            except:
                pass

        if hydroelastic_contact:
            try:
                # low friction and less bumpy
                dissipation = 5.0
                mu_static = 0.8  # 0.6
                mu_dynamic = 0.6  # 0.2
                if frictions_hyper_param is not None:
                    dissipation = frictions_hyper_param[0]
                    mu_static = frictions_hyper_param[1]  # 0.6
                    mu_dynamic = frictions_hyper_param[2]  # 0.2

                friction = CoulombFriction(0.7 * mu_static, 0.7 * mu_dynamic)
                if props is not None:
                    if not props.HasProperty("hydroelastic", "resolution_hint"):
                        if compliant_model:
                            AddCompliantHydroelasticProperties(
                                resolution_hint=0.005,
                                hydroelastic_modulus=5e6,
                                properties=props,
                            )  # 0.005
                        else:
                            AddRigidHydroelasticProperties(resolution_hint=0.005, properties=props)

                    if not props.HasProperty("material", "coulomb_friction"):
                        AddContactMaterial(dissipation=dissipation, friction=friction, properties=props)
                scene_graph.AssignRole(source_id, gid, props, RoleAssign.kReplace)
            except:
                print("adding complaint hydroelastic_contact fails")


def UpdateObjectVisualProperties(
    scene_graph,
    inspector,
    source_id,
    oid,
    obj_name,
    range_lo=0.1,
    range_hi=0.9,
    color=None,
    alpha=-1,
    randomize_textures=False,
):
    """change color, transparency, and texture"""

    if not inspector.BelongsToSource(oid, source_id):
        return

    # visual
    obj_gids = []
    for geometry_id in inspector.GetAllGeometryIds():
        name = inspector.GetName(geometry_id)
        if obj_name in name and "visual" in name:
            obj_gids.append(geometry_id)

    for gid in obj_gids:
        if color is not None:
            props = inspector.GetIllustrationProperties(gid)
            props = SetColor(props, color)

        if alpha > 0:
            props = inspector.GetIllustrationProperties(gid)
            props = SetTransparency(props, alpha)

        if randomize_textures:
            props = inspector.GetPerceptionProperties(gid)
            props = SetTextures(props)


def SetTextures(props):
    """update texture in drake"""
    texture_dir = "assets/textures/new_textures"
    files = os.listdir(texture_dir)
    texture_path = os.path.join(texture_dir, np.random.choice(files, 1)[0])

    if props is not None and props.HasProperty("phong", "diffuse_map"):
        c = props.GetProperty("phong", "diffuse_map")
        props.UpdateProperty("phong", "diffuse_map", texture_path)
    else:
        props.AddProperty("phong", "diffuse_map", texture_path)
    return props


def SetTransparency(props, alpha):
    """set transparency on meshcat"""
    if props is not None and props.HasProperty("phong", "diffuse"):
        c = props.GetProperty("phong", "diffuse")
        new_color = Rgba(c.r(), c.g(), c.b(), alpha)
        props.UpdateProperty("phong", "diffuse", new_color)
    else:
        new_color = Rgba(1.0, 1.0, 1.0, alpha)
        props.AddProperty("phong", "diffuse", new_color)
    return props


def SetColor(props, color):
    if props is not None and props.HasProperty("phong", "diffuse"):
        new_color = Rgba(color[0], color[1], color[2], color[3])
        props.UpdateProperty("phong", "diffuse", new_color)
    return props


def AddTriad(
    source_id,
    frame_id,
    scene_graph,
    length=0.25,
    radius=0.01,
    opacity=1.0,
    X_FT=RigidTransform(),
    name="frame",
):
    """
    Adds illustration geometry representing the coordinate frame, with the
    x-axis drawn in red, the y-axis in green and the z-axis in blue. The axes
    point in +x, +y and +z directions, respectively.
    """
    # x-axis
    X_TG = RigidTransform(RotationMatrix.MakeYRotation(np.pi / 2), [length / 2.0, 0, 0])
    geom = GeometryInstance(X_FT.multiply(X_TG), Cylinder(radius, length), name + " x-axis")
    geom.set_illustration_properties(MakePhongIllustrationProperties([1, 0, 0, opacity]))
    scene_graph.RegisterGeometry(source_id, frame_id, geom)

    # y-axis
    X_TG = RigidTransform(RotationMatrix.MakeXRotation(np.pi / 2), [0, length / 2.0, 0])
    geom = GeometryInstance(X_FT.multiply(X_TG), Cylinder(radius, length), name + " y-axis")
    geom.set_illustration_properties(MakePhongIllustrationProperties([0, 1, 0, opacity]))
    scene_graph.RegisterGeometry(source_id, frame_id, geom)

    # z-axis
    X_TG = RigidTransform([0, 0, length / 2.0])
    geom = GeometryInstance(X_FT.multiply(X_TG), Cylinder(radius, length), name + " z-axis")
    geom.set_illustration_properties(MakePhongIllustrationProperties([0, 0, 1, opacity]))
    scene_graph.RegisterGeometry(source_id, frame_id, geom)


def visualize_rendered_image(state):
    """visualize the observation with pyplot"""
    print("visualize data shape:", state.shape)
    plt.switch_backend("TKAgg")

    fig = plt.figure(figsize=(25.6, 9.6))
    ax = fig.add_subplot(1, 3, 1)
    ax.set_title("color")
    plt.imshow((state[..., :3] * 255).astype(np.uint8))

    if state.shape[0] == 4:
        ax = fig.add_subplot(1, 3, 2)
        ax.set_title("depth")
        plt.imshow(state[..., -1], cmap=plt.get_cmap("magma"))

    if state.shape[0] >= 5:
        ax = fig.add_subplot(1, 3, 2)
        ax.set_title("depth")
        plt.imshow(state[..., -2], cmap=plt.get_cmap("magma"))

        ax = fig.add_subplot(1, 3, 3)
        ax.set_title("mask")
        mask = (state[..., -1]).astype(np.uint8)
        mask[mask == 255] = 0
        mask[mask == 254] = 0

        plt.imshow(mask)
    plt.show()


def AddMeshcatTriad(meshcat, path, length=0.25, radius=0.01, opacity=1.0, X_PT=RigidTransform()):
    """add meshcat visualization of a coord"""
    meshcat.SetTransform(path, X_PT)
    # x-axis
    X_TG = RigidTransform(RotationMatrix.MakeYRotation(np.pi / 2), [length / 2.0, 0, 0])
    meshcat.SetTransform(path + "/x-axis", X_TG)
    meshcat.SetObject(path + "/x-axis", Cylinder(radius, length), Rgba(1, 0, 0, opacity))

    # y-axis
    X_TG = RigidTransform(RotationMatrix.MakeXRotation(np.pi / 2), [0, length / 2.0, 0])
    meshcat.SetTransform(path + "/y-axis", X_TG)
    meshcat.SetObject(path + "/y-axis", Cylinder(radius, length), Rgba(0, 1, 0, opacity))

    # z-axis
    X_TG = RigidTransform([0, 0, length / 2.0])
    meshcat.SetTransform(path + "/z-axis", X_TG)
    meshcat.SetObject(path + "/z-axis", Cylinder(radius, length), Rgba(0, 0, 1, opacity))


def vis_tool_pose(meshcat, X_TG, path, mesh_path=None, color=[0.1, 0.1, 0.1, 0.3]):
    """render the tool at the pose tf in meshcat"""
    if type(X_TG) is not RigidTransform:
        X_TG = RigidTransform(X_TG)
    if meshcat is None:
        return

    if mesh_path is not None:
        mesh = ReadObjToTriangleSurfaceMesh(mesh_path.replace(".sdf", "_corr.obj"))
        meshcat.Delete(path)  # delete the old one
        meshcat.SetObject(path, mesh, Rgba(*color))
    meshcat.SetProperty(path, "color", color)
    meshcat.SetTransform(path, X_TG)


def vis_hand_pose(meshcat, X_TG, path, load=False, color=[0.1, 0.1, 0.1, 0.3]):
    """render the tool at the pose tf in meshcat"""
    if type(X_TG) is not RigidTransform:
        X_TG = RigidTransform(X_TG)
    if meshcat is None:
        return

    if load:
        mesh = ReadObjToTriangleSurfaceMesh("env/models/hand_finger.obj")
        meshcat.Delete(path)  # delete the old one
        meshcat.SetObject(path, mesh, Rgba(*color))
    meshcat.SetProperty(path, "color", color)
    meshcat.SetTransform(path, X_TG)


def AddWrenchVis(meshcat, path, point, wrench, length=0.03, radius=0.001, opacity=1.0):
    """compute the vector visualization of the wrench"""
    torque_line_vertices = np.array([point, point + wrench[:3]] * 100).T
    force_line_vertices = np.array([point, point + wrench[:3]] * 100).T

    meshcat.SetLine(path + "_force", force_line_vertices, radius, Rgba(1, 0, 0, opacity))
    meshcat.SetLine(path + "_wrench", torque_line_vertices, radius, Rgba(0, 1, 0, opacity))


def add_table_collision_free_constraint(ik, plant, frame, bb_size=[0.12, 0.08, 0.08], table_height=0.1):
    # apprxoimate a link with a bounding box and add as position constraints
    min_height = -0.01 + table_height
    max_num = 100
    y_bound = 1

    ik.AddPositionConstraint(
        frame,
        [0, 0, 0],
        plant.world_frame(),
        [0, -y_bound, min_height],
        [max_num, y_bound, max_num],
    )
    ik.AddPositionConstraint(
        frame,
        [0, 0, -bb_size[2]],
        plant.world_frame(),
        [-0.3, -y_bound, min_height],
        [max_num, y_bound, max_num],
    )
    ik.AddPositionConstraint(
        frame,
        [0, 0, bb_size[2]],
        plant.world_frame(),
        [-0.3, -y_bound, min_height],
        [max_num, y_bound, max_num],
    )
    ik.AddPositionConstraint(
        frame,
        [0, -bb_size[1], 0],
        plant.world_frame(),
        [-0.3, -y_bound, min_height],
        [max_num, y_bound, max_num],
    )
    ik.AddPositionConstraint(
        frame,
        [0, bb_size[1], 0],
        plant.world_frame(),
        [-0.3, -y_bound, min_height],
        [max_num, y_bound, max_num],
    )
    ik.AddPositionConstraint(
        frame,
        [bb_size[0], 0, 0],
        plant.world_frame(),
        [-0.3, -y_bound, min_height],
        [max_num, y_bound, max_num],
    )
    ik.AddPositionConstraint(
        frame,
        [-bb_size[0], 0, 0],
        plant.world_frame(),
        [-0.3, -y_bound, min_height],
        [max_num, y_bound, max_num],
    )


def compute_manipulability(plant, plant_context):
    """compute the manipulability matrix"""
    J = plant.CalcJacobianSpatialVelocity(
        plant_context,
        JacobianWrtVariable.kQDot,
        plant.GetBodyByName("panda_hand").body_frame(),
        [0, 0, 0],
        plant.world_frame(),
        plant.world_frame(),
    )[:, :7]
    return np.linalg.det(J @ J.T), np.linalg.cond(J @ J.T)


def refine_ik_traj(plant, pose_traj, q0, rot_tol=3e-2):
    """run a small trajopt on the trajectory with the solved IK from end-effector traj"""
    # make sure the beginning and the end do not get updated
    prog = MathematicalProgram()
    q = prog.NewContinuousVariables(9, len(pose_traj))
    gripper_frame = plant.GetFrameByName("panda_hand")
    plant_contexts = [plant.CreateDefaultContext() for i in range(len(pose_traj))]

    for idx, pose in enumerate(pose_traj):
        pose = RigidTransform(pose)
        prog.AddConstraint(
            PositionConstraint(
                plant,
                plant.world_frame(),
                pose.translation(),
                pose.translation(),
                gripper_frame,
                [0, 0, 0],
                plant_contexts[idx],
            ),
            q[:, idx],
        )
        prog.AddConstraint(
            PositionConstraint(
                plant,
                plant.world_frame(),
                [0.05, -1, 0],
                [1, 1, 1],
                gripper_frame,
                [0, 0, 0],
                plant_contexts[idx],
            ),
            q[:, idx],
        )
        prog.AddConstraint(
            OrientationConstraint(
                plant,
                gripper_frame,
                RotationMatrix(),
                plant.world_frame(),
                pose.rotation(),
                rot_tol,
                plant_contexts[idx],
            ),
            q[:, idx],
        )

    # add some other constraints
    prog.AddConstraint(np.sum((q[:, 0] - q0[:, 0]) ** 2) == 0)
    prog.AddConstraint(np.sum((q[:, -1] - q0[:, -1]) ** 2) == 0)

    # Add smoothness cost

    weight = np.ones((9, 1))
    weight[0] = 10.0
    weight[-1] = 10.0
    prog.AddQuadraticCost(np.sum(weight * (q[:, 1:] - q[:, :-1]) ** 2))

    # add linear constraint
    try:
        solver = SnoptSolver()
    except:
        solver = IpoptSolver()
    prog.SetInitialGuess(q, q0)
    result = solver.Solve(prog)  #
    if result.is_success():
        return result
    else:
        return None


def refine_feasible_ik_traj(plant, pose_traj, q0, rot_tol=5e-2):
    """run a small trajopt on the trajectory with the solved IK from end-effector traj"""
    # make sure the beginning and the end do not get updated
    prog = MathematicalProgram()
    q = prog.NewContinuousVariables(9, len(pose_traj))
    gripper_frame = plant.GetFrameByName("panda_hand")
    plant_contexts = [plant.CreateDefaultContext() for i in range(len(pose_traj))]

    for idx, pose in enumerate(pose_traj):
        pose = RigidTransform(pose)
        prog.AddConstraint(
            PositionConstraint(
                plant,
                plant.world_frame(),
                pose.translation(),
                pose.translation(),
                gripper_frame,
                [0, 0, 0],
                plant_contexts[idx],
            ),
            q[:, idx],
        )
        prog.AddConstraint(
            OrientationConstraint(
                plant,
                gripper_frame,
                RotationMatrix(),
                plant.world_frame(),
                pose.rotation(),
                rot_tol,
                plant_contexts[idx],
            ),
            q[:, idx],
        )

    prog.AddQuadraticCost(np.sum((q - q0) ** 2))  # np.sum((q[:, 1:] - q[:, :-1]) ** 2)
    # add linear constraint
    try:
        solver = SnoptSolver()
        prog.SetInitialGuess(q, q0)
        result = solver.Solve(prog)
    except:
        solver = IpoptSolver()
        prog.SetInitialGuess(q, q0)
        result = solver.Solve(prog)
    #
    if result.is_success():
        return result
    else:
        return None


def interpolate_waypoints(waypoints, n, m, mode="cubic"):  # linear
    """
    Interpolate the waypoints using interpolation.
    """
    data = np.zeros([n, m])
    x = np.linspace(0, 1, waypoints.shape[0])
    for i in range(waypoints.shape[1]):
        y = waypoints[:, i]

        t = np.linspace(0, 1, n + 2)
        if mode == "linear":  # clamped cubic spline
            f = interpolate.interp1d(x, y, "linear")
        if mode == "cubic":  # clamped cubic spline
            f = interpolate.CubicSpline(x, y, bc_type="clamped")
        elif mode == "quintic":  # seems overkill to me
            pass
        data[:, i] = f(t[1:-1])  #

    return data


def solve_ik_traj_with_standoff(
    plant, endpoint_pose, endpoint_joints, q_traj, waypoint_times, keytimes, keyposes, rot_tol=0.03
):
    """run a small trajopt on the trajectory with the solved IK from end-effector traj"""
    # make sure the beginning and the end do not get updated

    waypoint_num = len(waypoint_times)
    prog = MathematicalProgram()
    q = prog.NewContinuousVariables(9, waypoint_num)
    gripper_frame = robot_plant.GetFrameByName("panda_hand")
    plant_contexts = [robot_plant.CreateDefaultContext() for i in range(waypoint_num)]
    q0 = np.array([q_traj.value(t) for t in waypoint_times])

    for idx, time in enumerate(waypoint_times):
        if time in keytimes:  # standoff
            keypoint_idx = keytimes.index(time)
            pose = RigidTransform(keyposes[keypoint_idx])
            prog.AddConstraint(
                OrientationConstraint(
                    robot_plant,
                    gripper_frame,
                    RotationMatrix(),
                    robot_plant.world_frame(),
                    pose.rotation(),
                    rot_tol,
                    plant_contexts[idx],
                ),
                q[:, idx],
            )

            prog.AddConstraint(
                PositionConstraint(
                    robot_plant,
                    robot_plant.world_frame(),
                    pose.translation(),
                    pose.translation(),
                    gripper_frame,
                    [0, 0, 0],
                    plant_contexts[idx],
                ),
                q[:, idx],
            )

        # table constraint
        prog.AddConstraint(
            PositionConstraint(
                robot_plant,
                robot_plant.world_frame(),
                [0.05, -0.7, 0.06],
                [1, 0.7, 1],
                gripper_frame,
                [0, 0, 0],
                plant_contexts[idx],
            ),
            q[:, idx],
        )

    # add some other constraints
    prog.AddConstraint(np.sum((q[:, 0] - endpoint_joints[:, 0]) ** 2) == 0)

    # Add smoothness cost
    weight = np.ones((9, 1))
    weight[0] = 10.0
    weight[-1] = 10.0
    prog.AddQuadraticCost(np.sum(weight * (q[:, 1:] - q[:, :-1]) ** 2))
    prog.SetInitialGuess(q, q0.squeeze().T)

    # add linear constraint
    try:
        solver = SnoptSolver()
        result = solver.Solve(prog)
    except:
        solver = IpoptSolver()
        result = solver.Solve(prog)

    if result.is_success():
        return result
    else:
        return None


def smoothness(q):
    return np.sum((q[:, 1:] - q[:, :-1]) ** 2)


def solve_ik_kpam(
    cost_and_constraints,
    plant,
    gripper_frame,
    keypoints_in_hand,
    keypoint_obj_in_world,
    p0,
    q0,
    centering_joint,
    rot_tol=np.pi / 10,
    add_table_col=True,
    add_gripper_faceup=False,
    timeout=False,
    consider_collision=False,
    table_height=0.15,
):
    """
    the simple case for the kpam problem
    always assume the head tool keyponint match the head object keypoint
    the head and tail orthogonal to the 0,0,1
    and the tail and side orthogonal to 0,0,1
    minimize cost and both joint space and the pose space
    """

    constraint_dicts, cost_dicts = cost_and_constraints
    ik_context = robot_plant.CreateDefaultContext()
    ik = InverseKinematics(robot_plant, ik_context)

    # separate out axis
    for constraint in constraint_dicts:
        if "axis_from_keypoint_name" in constraint:
            from_idx = constraint["axis_from_keypoint_idx"]
            to_idx = constraint["axis_to_keypoint_idx"]
            target_axis = constraint["target_axis"]
            vec = keypoints_in_hand[to_idx] - keypoints_in_hand[from_idx]
            tol = constraint["tolerance"]
            tgt = np.arccos(constraint["target_inner_product"])
            lower_bound = max(tgt - tol, 0)
            upper_bound = min(tgt + tol, np.pi)

            ik.AddAngleBetweenVectorsConstraint(
                gripper_frame, vec, robot_plant.world_frame(), target_axis, lower_bound, upper_bound
            )
        else:
            idx = constraint["keypoint_idx"]
            target_position = constraint["target_position"]

            target_point = keypoints_in_hand[idx]
            tol = constraint["tolerance"]

            ik.AddPositionConstraint(
                gripper_frame, target_point, robot_plant.world_frame(), target_position - tol, target_position + tol
            )

    """solving IK to match tool head keypoint and the object keypoint"""
    # make sure the arm does not go backward
    ik.AddPositionConstraint(gripper_frame, [0, 0, 0], robot_plant.world_frame(), [0.05, -1, 0], [1, 1, 1])

    if add_gripper_faceup:
        ik.AddAngleBetweenVectorsConstraint(
            gripper_frame, [1, 0, 0], robot_plant.world_frame(), [0, 0, -1], np.pi / 12, np.pi
        )

    # not touching table constraints add elbow
    if add_table_col:
        add_table_collision_free_constraint(ik, robot_plant, gripper_frame, [0.03, 0.04, 0.08])
        add_table_collision_free_constraint(
            ik,
            robot_plant,
            robot_plant.GetFrameByName("panda_link6"),
            [0.03, 0.03, 0.03],
        )
        add_table_collision_free_constraint(
            ik,
            robot_plant,
            robot_plant.GetFrameByName("panda_link7"),
            [0.03, 0.03, 0.03],
        )

    if consider_collision:
        ik.AddMinimumDistanceConstraint(0.01)  # 0.03

    prog = ik.get_mutable_prog()
    q = ik.q()
    solver = SnoptSolver()
    # if timeout:
    #    solver.SetSolverOption(solver.id(), "Major Iterations Limit", 1000)

    # as well as pose space costs experiments
    joint_cost_mat = np.identity(len(q))
    joint_cost_mat[0, 0] = 10  # 000
    prog.AddQuadraticErrorCost(joint_cost_mat, centering_joint, q)
    ik.AddPositionCost(gripper_frame, [0, 0, 0], robot_plant.world_frame(), p0.translation(), np.eye(3))
    ik.AddOrientationCost(gripper_frame, RotationMatrix(), robot_plant.world_frame(), p0.rotation(), 1)

    prog.SetInitialGuess(q, q0)
    result = solver.Solve(ik.prog())  #
    if result.is_success():
        return result
    else:
        return None


def solve_ik(
    plant,
    gripper_frame,
    pose,
    q0,
    rot_tol=np.pi / 10,
    add_table_col=True,
    add_gripper_faceup=False,
    timeout=False,
    consider_collision=False,
    table_height=0.15,
    filter_manipulability=False,
):
    ik_context = robot_plant.CreateDefaultContext()
    ik = InverseKinematics(robot_plant, ik_context)

    """solving IK with environment and rotation constraints"""
    ik.AddPositionConstraint(
        gripper_frame,
        [0, 0, 0],
        robot_plant.world_frame(),
        pose.translation(),
        pose.translation(),
    )
    # make sure the arm does not go backward
    ik.AddPositionConstraint(gripper_frame, [0, 0, 0], robot_plant.world_frame(), [0.05, -1, 0], [1, 1, 1])
    ik.AddOrientationConstraint(gripper_frame, RotationMatrix(), robot_plant.world_frame(), pose.rotation(), rot_tol)
    if add_gripper_faceup:
        ik.AddAngleBetweenVectorsConstraint(
            gripper_frame, [1, 0, 0], robot_plant.world_frame(), [0, 0, -1], np.pi / 12, np.pi
        )

    # not touching table constraints add elbow
    if add_table_col:
        add_table_collision_free_constraint(ik, robot_plant, gripper_frame, [0.03, 0.04, 0.08])
        add_table_collision_free_constraint(
            ik,
            robot_plant,
            robot_plant.GetFrameByName("panda_link6"),
            [0.03, 0.03, 0.03],
        )
        add_table_collision_free_constraint(
            ik,
            robot_plant,
            robot_plant.GetFrameByName("panda_link7"),
            [0.03, 0.03, 0.03],
        )

    if consider_collision:
        ik.AddMinimumDistanceConstraint(0.01)  # 0.03

    prog = ik.get_mutable_prog()
    q = ik.q()

    prog.AddQuadraticErrorCost(np.identity(len(q)), q0, q)
    prog.SetInitialGuess(q, q0)
    try:
        solver = SnoptSolver()
        result = solver.Solve(ik.prog())
    except:
        solver = IpoptSolver()
        result = solver.Solve(ik.prog())

    if result.is_success():
        if result is not None and filter_manipulability:
            det_J, cond_J = compute_manipulability(robot_plant, ik_context)
            if det_J < 0.001 or cond_J > 180:
                print(f"not enough maneuverability det_J: {det_J:.3f} cond_J: {cond_J:.3f}")
                result = None

        return result
    else:
        return None


def add_tool_to_controller_plant(plant, tool_class_name, instance_info, task_config, tool_rel_pose):
    """add the tool to the controller plant"""
    tool_info = AddTools(plant, tool_class_name, instance_info, train=task_config.train)
    plant.WeldFrames(
        plant.GetFrameByName("panda_hand"),
        plant.GetFrameByName(tool_info[1], tool_info[0]),
        tool_rel_pose,
    )


def transform_point(pose, pc):
    """transform points with the given pose
    param pose: 4 x 4
    param pc: 3 x N
    """
    return pose[:3, :3] @ pc + pose[:3, [3]]


def downsample_point(pc, num_pt):
    """param pc: 3 x N
    return 3 x num_pt
    """
    if type(pc) is not torch.Tensor:
        pc = torch.from_numpy(pc).cuda().float()
    if len(pc.shape) == 2:
        pc = pc[None].cuda().float()
    if pc.shape[1] == 0:
        return torch.zeros((pc.shape[-1], 0)).float()

    try:
        # to mimic what real-world pointcloud does
        xyz = gather_operation(
            pc.transpose(1, 2).contiguous(),
            furthest_point_sample(pc[..., :3].contiguous(), num_pt),
        ).contiguous()

    except:
        # farthest point sampling is not available
        print(pc.shape)
        print("Use naive sampling. Consider installing [PointNet++](https://github.com/liruiw/Pointnet2_PyTorch)")
        pc = pc.detach().cpu().numpy()
        sample_indexes = np.random.choice(pc.shape[1], num_pt)  # ALLOW replace
        pc = pc[0, sample_indexes].T
        print(np.unique(sample_indexes))
        return torch.from_numpy(pc).float()

    return xyz[0].detach()


def finalize_builder(
    plant,
    builder,
    scene_graph,
    franka,
    time_step,
    tool_class_name,
    task_config=None,
    tool_name=None,
    object_name=None,
    tool_rel_pose=None,
    object_info=None,
):
    """add the controller / sensors and related ports to the plant"""
    from env.env_object_related import extra_force_system

    add_contact_visualizer = task_config.env.vis_contact
    use_impedance_controller = task_config.use_impedance_controller
    use_meshcat = task_config.use_meshcat

    s = time.time()
    num_positions = plant.num_positions(franka)
    demux = builder.AddSystem(Demultiplexer(2 * num_positions, num_positions))
    builder.Connect(plant.get_state_output_port(franka), demux.get_input_port())
    builder.ExportOutput(demux.get_output_port(0), "panda_position_measured")
    builder.ExportOutput(demux.get_output_port(1), "panda_velocity_estimated")
    builder.ExportOutput(plant.get_state_output_port(franka), "panda_state_estimated")

    multibody_plant_config = MultibodyPlantConfig(
        time_step=time_step,
        discrete_contact_solver=task_config.sim.discrete_contact_solver_type,
    )
    controller_builder = DiagramBuilder()
    controller_plant, controller_scene_graph = AddMultibodyPlant(multibody_plant_config, controller_builder)
    controller_panda = AddFranka(controller_plant)

    # tool is added to the controller plant in the real world. This would cause external torque.
    if tool_rel_pose is not None and "mug" not in task_config.task_name:  # hack around mug
        add_tool_to_controller_plant(
            controller_plant,
            tool_class_name,
            {"tool_name": tool_name},
            task_config,
            tool_rel_pose=tool_rel_pose,
        )

    # the controller plant needs to feel the extra torque
    controller_plant.Finalize()  # match two plants
    controller_diagram = controller_builder.Build()

    # I need a PassThrough system so that I can export the input port.
    panda_positions = builder.AddSystem(PassThrough([0] * num_positions))
    panda_kp = [task_config.env.controller_gains[0]] * num_positions  # [120, 100, 100, 80, 30, 20, 50, 50, 50]
    panda_kd = [task_config.env.controller_gains[2]] * num_positions  # critically damped otherwise
    adder = builder.AddSystem(Adder(2, num_positions))

    if use_impedance_controller:
        # https://github.com/frankaemika/franka_ros/blob/develop/franka_control/config/default_controllers.yaml#L31
        # https://github.com/NVIDIA-Omniverse/IsaacGymEnvs/blob/ca7a4fb762f9581e39cc2aab644f18a83d6ab0ba/isaacgymenvs/tasks/factory/factory_control.py
        impedance_controller = RobotInternalController(controller_plant, joint_stiffness=panda_kp, kd=panda_kd)
        franka_controller = builder.AddSystem(impedance_controller)
        builder.Connect(
            panda_positions.get_output_port(),
            franka_controller.get_input_port_desired_state(),
        )
        panda_controller_torque = builder.AddSystem(PassThrough([0] * num_positions))
        builder.Connect(
            panda_controller_torque.get_output_port(),
            franka_controller.get_torque_feedforward_port_control(),
        )

        builder.Connect(
            plant.get_state_output_port(franka),
            franka_controller.get_input_port_estimated_state(),
        )
        builder.Connect(franka_controller.get_output_port_control(), adder.get_input_port(0))

    elif task_config.pose_controller or task_config.use_controller_portswitch:
        # use eric's controller
        # input port : X_TPdes, V_TPdes
        # output port: torques_output
        franka_controller = PoseController(
            controller_plant,
            controller_panda,
            controller_plant.GetFrameByName("panda_link0"),
            controller_plant.GetFrameByName("panda_hand"),
        )
        franka_controller = builder.AddSystem(franka_controller)
        desired_eff_pose = builder.AddSystem(PassThrough(AbstractValue.Make(RigidTransform())))
        desired_eff_velocity = builder.AddSystem(PassThrough(AbstractValue.Make(SpatialVelocity())))
        builder.Connect(
            plant.get_state_output_port(franka),
            franka_controller.get_plant_state_input_port(),
        )

        builder.Connect(
            desired_eff_pose.get_output_port(),
            franka_controller.get_state_desired_port(),
        )
        builder.Connect(
            desired_eff_velocity.get_output_port(),
            franka_controller.get_velocity_desired_port(),
        )

        builder.ExportInput(desired_eff_pose.get_input_port(), "desired_eff_pose")
        builder.ExportInput(desired_eff_velocity.get_input_port(), "desired_eff_velocity")

        if task_config.use_controller_portswitch:
            # add the inverse dynamics controller together
            controller_port_switch = builder.AddSystem(PortSwitch(vector_size=9))
            osc_port = controller_port_switch.DeclareInputPort(name="operational_space_controller")
            inv_dyn_port = controller_port_switch.DeclareInputPort(name="inverse_dynamics_controller")

            inverse_dynamic_controller = InverseDynamicsController(
                controller_plant,
                kp=panda_kp,  # 100
                ki=[1] * num_positions,  # 1
                kd=panda_kd,  # 1 for less oscillating motion and 20 for reaching
                has_reference_acceleration=False,
            )

            # joint position controller. The controller block
            inv_dyn_franka_controller = builder.AddSystem(inverse_dynamic_controller)

            # Add discrete derivative to command velocities. The IK block
            desired_state_from_position = builder.AddSystem(
                StateInterpolatorWithDiscreteDerivative(num_positions, time_step, suppress_initial_transient=True)
            )
            desired_state_from_position.set_name("desired_state_from_position")
            builder.Connect(
                desired_state_from_position.get_output_port(),
                inv_dyn_franka_controller.get_input_port_desired_state(),
            )

            builder.Connect(
                panda_positions.get_output_port(),
                desired_state_from_position.get_input_port(),
            )
            builder.Connect(
                plant.get_state_output_port(franka),
                inv_dyn_franka_controller.get_input_port_estimated_state(),
            )
            builder.Connect(inv_dyn_franka_controller.get_output_port(), inv_dyn_port)  # 2
            builder.Connect(franka_controller.get_output_port(), osc_port)  # 1
            builder.Connect(controller_port_switch.get_output_port(), adder.get_input_port(0))

            builder.ExportInput(
                controller_port_switch.get_port_selector_input_port(),
                "controller_switch_port",
            )
        else:
            builder.Connect(franka_controller.get_output_port(), adder.get_input_port(0))
    else:
        # too much oscillation for learner if high kd, too much oscillation for the expert if high kp tuning required
        inverse_dynamic_controller = InverseDynamicsController(
            controller_plant,
            kp=panda_kp,  # 100
            ki=[1] * num_positions,  # 1
            kd=panda_kd,  # 1 for less oscillating motion and 20 for reaching
            has_reference_acceleration=False,
        )

        # joint position controller. The controller block
        franka_controller = builder.AddSystem(inverse_dynamic_controller)

        # Add discrete derivative to command velocities. The IK block
        desired_state_from_position = builder.AddSystem(
            StateInterpolatorWithDiscreteDerivative(num_positions, time_step, suppress_initial_transient=True)
        )
        desired_state_from_position.set_name("desired_state_from_position")

        builder.Connect(
            desired_state_from_position.get_output_port(),
            franka_controller.get_input_port_desired_state(),
        )
        builder.Connect(
            panda_positions.get_output_port(),
            desired_state_from_position.get_input_port(),
        )
        builder.Connect(
            plant.get_state_output_port(franka),
            franka_controller.get_input_port_estimated_state(),
        )
        builder.Connect(franka_controller.get_output_port_control(), adder.get_input_port(0))

    franka_controller.set_name("panda_controller")
    # Add in the feed-forward torque
    # Use a PassThrough to make the port optional (it will provide zero values if not connected).

    torque_passthrough = builder.AddSystem(PassThrough([0] * num_positions))
    builder.Connect(torque_passthrough.get_output_port(), adder.get_input_port(1))
    builder.ExportInput(torque_passthrough.get_input_port(), "feedforward_torque")
    builder.Connect(adder.get_output_port(), plant.get_actuation_input_port(franka))
    builder.ExportInput(panda_positions.get_input_port(), "panda_joint_commanded")

    # Export commanded torques.
    builder.ExportOutput(adder.get_output_port(), "panda_torque_commanded")
    builder.ExportOutput(adder.get_output_port(), "panda_torque_measured")
    builder.ExportOutput(
        plant.get_generalized_contact_forces_output_port(franka),
        "panda_torque_external",
    )

    # Inverse Kinematics
    params = DifferentialInverseKinematicsParameters(num_positions, num_positions)
    ik_timestep = time_step * 160  # differentialIK...
    velocity_limit_scale = 2.1750  # from ROS
    franka_velocity_limits = np.ones(9) * velocity_limit_scale

    # differential IK is now separated by another plant
    global robot_plant
    robot_builder = DiagramBuilder()
    robot_plant, _ = AddMultibodyPlant(multibody_plant_config, robot_builder)
    robot_panda = AddFranka(robot_plant)
    robot_plant.Finalize()

    differential_ik = builder.AddSystem(
        DifferentialInverseKinematicsIntegrator(
            robot_plant,
            robot_plant.GetFrameByName("panda_hand"),
            ik_timestep,
            params,
        )
    )
    differential_ik.get_mutable_parameters().set_nominal_joint_position(
        [0.0, -1.285, 0, -2.356, 0.0, 1.571, 0.785, 0.04, 0.04]
    )
    builder.ExportInput(differential_ik.get_input_port(0), "panda_hand_pose")
    builder.ExportOutput(differential_ik.get_output_port(), "differential_ik_output")

    if task_config.task_name == "wrench" or task_config.task_name == "hammer":
        builder.Connect(
            extra_force_system.get_output_port(0),
            plant.get_applied_spatial_force_input_port(),
        )
        builder.Connect(
            extra_force_system.get_output_port(1),
            plant.get_actuation_input_port(object_info[0]),
        )

    wrist_cam_view = HAND_CAM_EXTR.copy()
    overhead_cam_view = OVERHEAD_CAM_EXTR.copy()

    if (
        hasattr(task_config.env, "randomize_camera_extrinsics")
        and task_config.env.randomize_camera_extrinsics
        and task_config.train
    ):
        translation_noise = 0.02
        rotation_noise = 0.05
        augment_noise = np.concatenate(
            (
                np.random.uniform(-1, 1, size=(3,)) * translation_noise,
                np.random.uniform(-1, 1, size=(3,)) * rotation_noise,
            ),
            axis=-1,
        )
        augment_noise_pose = unpack_action(augment_noise)
        wrist_cam_view = wrist_cam_view @ augment_noise_pose

        # one way to choose the overhead camera distribution is to shift it to some sphere that is x axis some distance away from the table center
        rotation_offset = 0.1
        target_center = np.array([0.45, 0.0, 0.3])
        overhead_cam_view = randomize_sphere_lookat(
            target_center, near=1.5, far=2.5, x_near=0.7, x_far=1.5, theta_low=np.pi / 8, theta_high=np.pi / 4
        )

    # Overhead Cameras
    AddRgbdSensors(
        builder,
        plant,
        scene_graph,
        model_instance_prefix="panda",
        relative_pose=overhead_cam_view,
        attach_body_name="panda_link0",
    )

    # Wrist Cameras
    AddRgbdSensors(
        builder,
        plant,
        scene_graph,
        model_instance_prefix="panda",
        relative_pose=wrist_cam_view,
        attach_body_name="panda_hand",
    )
    if not task_config.blender_render:
        AddTriad(
            plant.get_source_id(),
            plant.GetBodyFrameIdOrThrow(plant.GetBodyByName("panda_link0").index()),
            scene_graph,
            length=0.05,
            radius=0.005,
            X_FT=RigidTransform(overhead_cam_view),
            name="overhead_cam_coord",
        )

        AddTriad(
            plant.get_source_id(),
            plant.GetBodyFrameIdOrThrow(plant.GetBodyByName("panda_hand").index()),
            scene_graph,
            length=0.05,
            radius=0.005,
            X_FT=RigidTransform(HAND_CAM_EXTR),
            name="hand_cam_coord",
        )

    # create a dummy body and attach a sensor to it.
    AddRgbdSensors(
        builder,
        plant,
        scene_graph,
        model_instance_prefix="dummy",
        relative_pose=np.eye(4),
        attach_body_name="dummy_cam_body",
    )
    ##########################################

    # add the tool frame and the object frame
    if tool_name is not None and not task_config.blender_render:
        AddTriad(
            plant.get_source_id(),
            plant.GetBodyFrameIdOrThrow(plant.GetBodyByName(tool_name).index()),
            scene_graph,
            length=0.05,
            radius=0.005,
            name="tool_coord",
        )

    builder.ExportOutput(scene_graph.get_query_output_port(), "query_object")
    builder.ExportOutput(plant.get_contact_results_output_port(), "contact_results")
    builder.ExportOutput(plant.get_state_output_port(), "plant_continuous_state")

    global meshcat, visualizer
    if use_meshcat:
        if not meshcat:
            meshcat = StartMeshcat()
            if task_config.verbose:
                print("meshcat url: ", meshcat.web_url())
            meshcat.ResetRenderMode()
            meshcat.SetProperty("/Background", "visible", False)
            meshcat.SetProperty("/Grid", "visible", False)
            meshcat.SetProperty("/Lights/AmbientLight/<object>", "intensity", 0.5)
            meshcat.SetRealtimeRate(task_config.sim.realtime_rate)

        visualizer = MeshcatVisualizer.AddToBuilder(
            builder,
            scene_graph,
            meshcat,
            MeshcatVisualizerParams(role=Role.kIllustration, delete_on_initialization_event=False),
        )
        collision_visualizer = MeshcatVisualizer.AddToBuilder(
            builder,
            scene_graph,
            meshcat,
            MeshcatVisualizerParams(prefix="collision", role=Role.kProximity),
        )

        if not task_config.vis_collision_geometry:
            meshcat.SetProperty("collision", "visible", False)

        if add_contact_visualizer:
            cparams = ContactVisualizerParams()
            cparams.force_threshold = 1e-2
            cparams.newtons_per_meter = 1e4  # 1e-6 .0 # 1.0
            cparams.newton_meters_per_meter = 1e1
            cparams.radius = 0.005
            contact_visualizer = ContactVisualizer.AddToBuilder(builder, plant, meshcat, cparams)

    state_logger = LogVectorOutput(differential_ik.get_output_port(), builder, 10.0)
    diagram = builder.Build()
    # to draw the diagram
    # pydot.graph_from_dot_data(diagram.GetGraphvizString())[0].write_svg("diagram.svg")
    # pydot.graph_from_dot_data(controller_diagram.GetGraphvizString())[0].write_svg("controller_diagram.svg")

    diagram.set_name("ManipulationStation")
    return [
        builder,
        state_logger,
        diagram,
        plant,
        franka_controller,
        controller_plant,
        controller_diagram,
        differential_ik,
        scene_graph,
    ]


def domain_randomization(scene_graph, inspector, plant, tool_info, tool_obj_info, task_config):
    """domain randomzation for the textures and physics of the object and tools"""
    tool_name, object_name = tool_info[1], tool_obj_info[1]
    tool_id, object_id = tool_info[0], tool_obj_info[0]

    randomize_physics = task_config.env.randomize_physics
    hydroelastic_contact = task_config.sim.hydroelastic_contact
    randomize_textures = task_config.env.randomize_textures
    compliant_hydroelastic_model = task_config.sim.compliant_hydroelastic_model
    frictions_hyper_param = [
        task_config.sim.friction_dissipation,
        task_config.sim.friction_mu_static,
        task_config.sim.friction_mu_dynamic,
    ]

    UpdateObjectPhysicalProperties(
        scene_graph,
        inspector,
        plant.get_source_id(),
        plant.GetBodyFrameIdOrThrow(plant.GetBodyByName(tool_name).index()),
        tool_name,
        randomize_physics=randomize_physics,
        hydroelastic_contact=hydroelastic_contact,
        compliant_model=compliant_hydroelastic_model,
        frictions_hyper_param=frictions_hyper_param,
    )

    UpdateObjectVisualProperties(
        scene_graph,
        inspector,
        plant.get_source_id(),
        plant.GetBodyFrameIdOrThrow(plant.GetBodyByName(tool_name).index()),
        tool_name,
        color=(1, 0, 0, 1),
        randomize_textures=randomize_textures,
    )

    UpdateObjectPhysicalProperties(
        scene_graph,
        inspector,
        plant.get_source_id(),
        plant.GetBodyFrameIdOrThrow(plant.GetBodyByName(object_name).index()),
        object_name,
        randomize_physics=randomize_physics,
        hydroelastic_contact=hydroelastic_contact,
        compliant_model=compliant_hydroelastic_model,
        frictions_hyper_param=frictions_hyper_param,
    )

    UpdateObjectVisualProperties(
        scene_graph,
        inspector,
        plant.get_source_id(),
        plant.GetBodyFrameIdOrThrow(plant.GetBodyByName(object_name).index()),
        object_name,
        randomize_textures=randomize_textures,
        color=(1, 0, 0, 1),
        alpha=0.3,
    )

    if task_config.task_name == "wrench":
        nut_object_name = object_name.replace("body", "nut_body")
        UpdateObjectPhysicalProperties(
            scene_graph,
            inspector,
            plant.get_source_id(),
            plant.GetBodyFrameIdOrThrow(plant.GetBodyByName(nut_object_name).index()),
            nut_object_name,
            randomize_physics=randomize_physics,
            hydroelastic_contact=hydroelastic_contact,
            compliant_model=compliant_hydroelastic_model,
            frictions_hyper_param=frictions_hyper_param,
        )

        UpdateObjectVisualProperties(
            scene_graph,
            inspector,
            plant.get_source_id(),
            plant.GetBodyFrameIdOrThrow(plant.GetBodyByName(nut_object_name).index()),
            nut_object_name,
            randomize_textures=randomize_textures,
            color=(1, 0, 0, 1),
            alpha=0.3,
        )

    # table
    UpdateObjectPhysicalProperties(
        scene_graph,
        inspector,
        plant.get_source_id(),
        plant.GetBodyFrameIdOrThrow(plant.GetBodyByName("link").index()),
        "::surface",
        randomize_physics=randomize_physics,
        hydroelastic_contact=hydroelastic_contact,
        compliant_model=compliant_hydroelastic_model,
        frictions_hyper_param=[0.5, 0.1, 0.1],  # low friction for the table
    )

    UpdateObjectVisualProperties(
        scene_graph,
        inspector,
        plant.get_source_id(),
        plant.GetBodyFrameIdOrThrow(plant.GetBodyByName("link").index()),
        "::visual1",
        randomize_textures=randomize_textures,
    )


def update_tool_pose(tool_class_name, tool_info, randomize_tool_pose, task_config):
    """domain randomze the tool pose"""
    tool_rel_pose = [[0, 0, 0], [0, 0, 0.15]]
    # domain randoization
    tool_name = tool_info[1]

    if randomize_tool_pose:
        # pitch augmentation
        # roll pitch yaw
        tool_rel_pose[0][1] = np.random.uniform(
            task_config.env.tool_rotation_y_range[0],
            task_config.env.tool_rotation_y_range[1],
        )
        tool_rel_pose[0][0] = np.random.uniform(
            task_config.env.tool_rotation_z_range[0],
            task_config.env.tool_rotation_z_range[1],
        )
        tool_rel_pose[0][2] = np.random.uniform(
            task_config.env.tool_rotation_x_range[0],
            task_config.env.tool_rotation_x_range[1],
        )

        # add perturbation / to make the tool tilted down
        # translation in z
        tool_rel_pose[1][2] = np.random.uniform(
            task_config.env.tool_translation_z_range[0],
            task_config.env.tool_translation_z_range[1],
        )
        tool_rel_pose[1][0] = np.random.uniform(
            task_config.env.tool_translation_x_range[0],
            task_config.env.tool_translation_x_range[1],
        )
    return RigidTransform(RollPitchYaw(*tool_rel_pose[0]), tool_rel_pose[1])


def MakeManipulationStation(time_step=0.005, task_config=None, instance_info={}):
    builder = DiagramBuilder()
    multibody_plant_config = MultibodyPlantConfig(time_step=time_step)
    plant, scene_graph = AddMultibodyPlant(multibody_plant_config, builder)
    franka = AddFranka(plant)
    AddTable(plant, task_config.blender_render)
    plant.Finalize()
    res = finalize_builder(plant, builder, scene_graph, franka, time_step, "", task_config)
    inspector = scene_graph.model_inspector()
    res.append(inspector)

    global meshcat
    vis_hand_pose(meshcat, np.eye(4), "hand_goal_pose", load=True)
    return res


def MakeManipulationStationTools(time_step=0.005, tool_class_name="pole", task_config=None, instance_info={}):
    """make the manipulation station with tools"""
    reset_global_variables()
    s = time.time()
    builder = DiagramBuilder()
    multibody_plant_config = MultibodyPlantConfig(
        time_step=time_step,
        contact_surface_representation=task_config.sim.contact_surface_representation,
        contact_model=task_config.sim.contact_model,
        discrete_contact_solver=task_config.sim.discrete_contact_solver_type,
    )
    plant, scene_graph = AddMultibodyPlant(multibody_plant_config, builder)
    franka = AddFranka(plant)
    AddTable(plant, task_config.blender_render)

    # tools
    tool_obj_info = AddToolObjects(
        builder,
        plant,
        scene_graph,
        task_config.task_name,  # actually the task determines which object to use
        instance_info,
        train=task_config.train,
        fixed_idx=task_config.tool_fix_idx,
    )  #
    tool_info = AddTools(
        plant, tool_class_name, instance_info, train=task_config.train, fixed_idx=task_config.tool_fix_idx
    )
    tool_rel_pose = update_tool_pose(tool_class_name, tool_info, task_config.env.randomize_tool_pose, task_config)
    plant.WeldFrames(
        plant.GetFrameByName("panda_hand"),
        plant.GetFrameByName(tool_info[1], tool_info[0]),
        tool_rel_pose,
    )
    if not task_config.blender_render:
        AddCameraBox(
            plant,
            RigidTransform(OVERHEAD_CAM_EXTR),
            name="overhead_cam",
            parent_frame=plant.world_frame(),
        )

    # add a dummy camera
    AddShape(
        plant,
        Box(1e-3, 1e-3, 1e-3),
        "dummy_cam_body",
        1,
        0,  # mass to be 0
        color=[0.1, 0.5, 0.1, 0.0],
        collision=False,
    )

    plant.Finalize()
    tool_name, object_name = tool_info[1], tool_obj_info[1]
    inspector = scene_graph.model_inspector()
    domain_randomization(scene_graph, inspector, plant, tool_info, tool_obj_info, task_config)
    res = finalize_builder(
        plant,
        builder,
        scene_graph,
        franka,
        time_step,
        tool_class_name,
        task_config,
        tool_name=tool_name,
        object_name=object_name,
        tool_rel_pose=tool_rel_pose,
        object_info=tool_obj_info,
    )
    res.append(inspector)

    global meshcat
    if not task_config.blender_render:
        vis_tool_pose(meshcat, np.eye(4), "tool_goal_pose", mesh_path=tool_info[-1])
        vis_hand_pose(meshcat, np.eye(4), "hand_goal_pose", load=True)

    if task_config.verbose:
        print(f"object name: {object_name} tool name: {tool_name}")
        print(f"make manipulation time: {(time.time() - s):.3f}")
    tool_info.append(tool_rel_pose.GetAsMatrix4())

    return res + [tool_obj_info, tool_info]

env/env_util.py
liruiw/Fleet-Tools
import gym
import numpy as np

import warnings
from gym.utils import seeding
from pydrake.all import Simulator, SimulatorStatus
from pydrake.math import RollPitchYaw, RigidTransform

try:
    import open3d as o3d
except:
    pass

from pydrake.common import RandomGenerator
from env.env_util import *
import IPython
from core.utils import *
from colored import fg
import matplotlib.pyplot as plt
from env.basedrakeenv import BaseDrakeRobotEnv
import torch


class FrankaDrakeEnv(BaseDrakeRobotEnv):
    def __init__(self, task_config):
        self.reset_from_scene_file = False
        self.init(task_config)

    def init(self, task_config, info={}):
        self.task_config = task_config
        self.sim_dt = self.task_config.sim.dt
        self.at_contact = False
        self.fix_pose_controller = False
        self.target_obj_idx = 0

        if hasattr(self, "tool_class_name"):
            (
                builder,
                logger,
                self.system,
                self.plant,
                self.robot_controller,
                self.controller_plant,
                self.controller_plant_system,
                self.differential_ik,
                self.scene_graph,
                self.scene_inspector,
                tool_obj_info,
                tool_info,
            ) = MakeManipulationStationTools(
                time_step=self.sim_dt,
                tool_class_name=self.tool_class_name,
                task_config=self.task_config,
                instance_info=info,
            )
        else:
            (
                builder,
                logger,
                self.system,
                self.plant,
                self.robot_controller,
                self.controller_plant,
                self.controller_plant_system,
                self.differential_ik,
                self.scene_graph,
                self.scene_inspector,
            ) = MakeManipulationStation(time_step=self.sim_dt, task_config=self.task_config)

        self.simulator = Simulator(self.system)
        self.env_dt = self.sim_dt * self.task_config.sim.substep_num
        self.diffik_port = self.system.GetInputPort("panda_hand_pose")
        self.diffik_output_port = self.system.GetOutputPort("differential_ik_output")
        if self.task_config.use_controller_portswitch:
            self.controller_switch_port = self.system.GetInputPort("controller_switch_port")

        if self.task_config.use_controller_portswitch:
            self.action_port = self.system.GetInputPort("desired_eff_pose")
            self.action_vel_port = self.system.GetInputPort("desired_eff_velocity")
            self.joint_action_port = self.system.GetInputPort("panda_joint_commanded")

        elif not self.task_config.pose_controller:
            self.joint_action_port = self.system.GetInputPort("panda_joint_commanded")

        else:
            self.action_port = self.system.GetInputPort("desired_eff_pose")
            self.action_vel_port = self.system.GetInputPort("desired_eff_velocity")

        self.torque_input_port = self.system.GetInputPort("feedforward_torque")
        self.joint_measurement_port = self.system.GetOutputPort("panda_position_measured")
        self.contact_measurement_port = self.system.GetOutputPort("contact_results")
        self.external_torque_port = self.system.GetOutputPort("panda_torque_external")
        self.object_query_port = self.system.GetOutputPort("query_object")

        # after gravity compensated
        self.joint_min_limit = np.array([-2.8973, -1.7628, -2.8973, -3.0718, -2.8973, -0.0175, -2.8973, 0, 0])
        self.joint_max_limit = np.array([2.8973, 1.7628, 2.8973, -0.0698, 2.8973, 3.7525, 2.8973, 0.04, 0.04])
        self.joint_null_space = (self.joint_min_limit + self.joint_max_limit) / 2

        # initialize info
        self.info = {k: np.zeros(1) for k in key_list}
        self.action_num = self.task_config.env.numActions
        self.observation_num = self.task_config.env.numObservations
        self.action_worldframe = self.task_config.env.action_worldframe
        self.generator = RandomGenerator()
        self.seed(0)

        self.history_wrench_obs = deque(
            [],
            maxlen=self.task_config.env.history_force_buffer_length,
        )
        self.table_height = 0.15
        self.table_init_height = 0.3
        self.finger_state = [0, 0]
        self.time_step = 0
        self.need_termination = False
        self.goal_pose = None
        self.time = 0

        if hasattr(self, "tool_class_name"):
            self.setup(tool_obj_info, tool_info)

    def use_pose_controller(self):
        contact_and_hybrid = self.task_config.use_controller_portswitch and self.at_contact
        return (self.task_config.pose_controller) or contact_and_hybrid

    def use_joint_controller(self):
        nocontact_and_hybrid = self.task_config.use_controller_portswitch and not self.at_contact
        return (not self.task_config.pose_controller or nocontact_and_hybrid) and hasattr(self, "joint_action_port")

    @property
    def action_space(self):
        """action space of the simulation environment"""
        return PandaTaskSpace6D(*self.task_config.env.action_scale)

    @property
    def observation_space(self):
        """state space of the simulation environment"""
        return gym.spaces.Box(low=-1.0, high=1.0, shape=self.observation_num)

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        self.drake_random = RandomGenerator(seed)
        return [seed]

    def advance_time(self, context):
        """
        simulate forward
        """
        self.time = context.get_time()
        catch = False
        s = time.time()
        repeat_num = self.task_config.sim.substep_num

        try:
            for _ in range(repeat_num):  # make the simulation smooth and visulizable
                self.time = context.get_time()
                status = self.simulator.AdvanceTo(self.time + self.sim_dt)
                self._get_robot_freq_data()  # 0.02s slower per step

        except RuntimeError as e:
            warnings.warn("Calling Done after catching RuntimeError:")
            warnings.warn(e.args[0])
            catch = True
            self.info["sim_failure"] = True  # emergent reset
            color_print("simulation failed, reset!", "light_slate_blue")

        sim_done = catch or status.reason() == SimulatorStatus.ReturnReason.kReachedTerminationCondition
        return sim_done

    def process_action(self, action):
        """process the action for the actuation ports
        if the action is a concatenation of joint and pose, then it would use joint unless
        the environment has been configured to use pose.
        """
        self.action = action
        use_expert_traj_joints = len(action) == 15

        if use_expert_traj_joints:
            # use the expert joint trajectories
            target_joint_positions = action[:9]
            action = action[9:]

            # compute delta action with forward kinematics
            target_pose = self.set_diffik_state(self.context, set_joints=target_joint_positions).GetAsMatrix4()
            rel_pose = se3_inverse(self.ee_pose.reshape(4, 4)).dot(target_pose)
            self.action = pack_pose(rel_pose)
            self.target_joint = target_joint_positions

        action_delta_pose = unpack_pose(action, rot_type="euler")
        cur_pose = self.get_ee_pos().GetAsMatrix4()

        self.target_pose = cur_pose.dot(action_delta_pose)
        self.commanded_pose = self.target_pose.copy()

        # use OSC if (1) configured to use pose controller (2) hybrid and has contact
        if self.use_pose_controller():  # and not use_expert_traj_joints
            self.action_port.FixValue(self.context, RigidTransform(self.target_pose))
            self.target_pose_vel = get_frame_spatial_velocity(cur_pose, self.target_pose, self.env_dt)
            self.action_vel_port.FixValue(self.context, SpatialVelocity(self.target_pose_vel))

        # (1) expert demonstration (2) hybrid and no-contact
        elif self.use_joint_controller() or use_expert_traj_joints:
            if not use_expert_traj_joints:
                res = self.solve_ik(self.target_pose, self.joint_positions)
                if res is not None:
                    next_joint = res.get_x_val()[:9]

                else:
                    # use differential IK
                    self.set_diffik_state(self.context, set_joints=self.joint_positions)
                    self.diffik_port.FixValue(self.context, RigidTransform(self.target_pose))
                    next_joint = self.diffik_output_port.Eval(self.context)
                target_joint_positions = next_joint

            target_joint_positions[-2:] = 0.0  # no need to control the finger
            if hasattr(self, "joint_action_port"):
                self.joint_action_port.FixValue(self.context, target_joint_positions)

        return action

    def forward_kinematics(self, joint):
        return self.set_diffik_state(self.context, set_joints=joint).GetAsMatrix4()

    def set_diffik_state(self, context, reinit=False, set_joints=None):
        """set the differential ik. Could also be used as forward kinematics function"""
        diff_ik_context = self.differential_ik.GetMyMutableContextFromRoot(context)
        if set_joints is None:
            set_joints = self.joint_measurement_port.Eval(context)  # use measured joint
        try:
            self.differential_ik.SetPositions(diff_ik_context, set_joints)
        except:
            print("set joint failed")
        return self.differential_ik.ForwardKinematics(diff_ik_context)

    def step(self, action):
        """environment step"""
        self.time_step += 1
        self.context = self.simulator.get_context()

        sim_done = self.advance_time(self.context)
        if sim_done:  # terminate because of simulation error
            return [], 0, True, self.info

        self._get_obs(self.context)
        obs = self.select_obs()
        reward = self._get_reward(self.context, action)
        done = sim_done

        return obs, reward, done, self.info

    def select_obs(self):
        """select different observation space"""
        if self.task_config.env.use_image:
            return self.img_obs
        elif self.task_config.env.use_keypoint:
            return self.info["curr_keypoint"].reshape(-1)
        else:
            return self.state_obs

    def init_context(self):
        """higher-level scene and robot reset function for the environment"""
        self.need_termination = False
        self.at_contact = False

        self.time_step = 0
        self.simulator = Simulator(self.system)
        self.context = self.simulator.get_context()
        self.plant_context = self.system.GetSubsystemContext(self.plant, self.context)
        self.get_joint_angles()

        if hasattr(self, "tool_name"):
            if not self.reset_from_scene_file:
                self.reset_object(self.context)
                self.reset_robot(self.context)
            else:
                # object pose
                if self.task_config.verbose:
                    print(
                        "reset from scene", self.scene_descriptions["tool_name"], self.scene_descriptions["object_name"]
                    )
                try:
                    if "object_pose" in self.scene_descriptions:
                        pose = np.array(self.scene_descriptions["object_pose"][0])
                        pose[2, 3] += 0.03
                        tf = RigidTransform(pose)
                        self.set_object_pose(tf)
                    else:
                        self.reset_object(self.context)

                    self.joint_positions = np.array(self.scene_descriptions["robot_joints"][0])
                    self.plant.SetPositions(
                        self.plant_context,
                        self.plant.GetModelInstanceByName("panda"),
                        self.joint_positions.reshape(-1, 1),
                    )  #
                    self.plant.SetVelocities(
                        self.plant_context,
                        self.plant.GetModelInstanceByName("panda"),
                        [0] * 9,
                    )
                    self.set_diffik_state(self.context)
                except:
                    pass

        self.initial_ee_pose = self.get_ee_pos().GetAsMatrix4()
        self.context.SetTime(0)

        if self.use_pose_controller():
            self.action_port.FixValue(self.context, RigidTransform(self.initial_ee_pose))
            self.target_pose_vel = SpatialVelocity(np.zeros(6))
            self.action_vel_port.FixValue(self.context, self.target_pose_vel)

        if self.use_joint_controller():
            self.joint_action_port.FixValue(self.context, self.joint_positions)

        if self.task_config.use_controller_portswitch:
            self.controller_switch_port.FixValue(self.context, InputPortIndex(2))

        try:
            if self.task_config.blender_render:
                self.simulator.set_publish_every_time_step(False)
            else:
                self.simulator.Initialize()
        except:
            print("init context exception")
            res = self.task_config.env.init_joints
            self.plant.SetPositions(self.plant_context, self.plant.GetModelInstanceByName("panda"), res[:9])  #
            self.plant.SetVelocities(self.plant_context, self.plant.GetModelInstanceByName("panda"), [0] * 9)
            self.set_diffik_state(self.context, reinit=True, set_joints=res[:9])
            self.joint_positions = res[:9]
            self.context.SetTime(0)
            self.simulator.Initialize()

        ee_pose = self.get_ee_pos()
        self.simulator.set_target_realtime_rate(self.task_config.sim.realtime_rate)  #

        self.get_joint_angles()
        self.target_pose = np.array(ee_pose.GetAsMatrix4())
        self.ee_pose = np.array(ee_pose.GetAsMatrix4()).copy().flatten()
        self.commanded_pose = np.array(ee_pose.GetAsMatrix4())  # prev target pose

        self.set_diffik_state(self.context)
        self.diffik_port.FixValue(self.context, RigidTransform(ee_pose))
        self.target_joint_positions = self.joint_positions

        # reset some states
        self.fused_tool_points = torch.zeros([4, 0])  #
        self.fused_object_points = torch.zeros([4, 0])
        self.history_wrench_obs = deque([], maxlen=self.task_config.env.history_force_buffer_length)

        self.goal_ee_pose = self.ee_pose  # used for demonstration
        self.finger_state = [0, 0]
        self.endeffector_wrench = np.zeros(6)

        self.time_step = 0
        self.controller_context = self.controller_plant.CreateDefaultContext()
        self.goal_ik = InverseKinematics(self.controller_plant, self.controller_context)
        self.at_contact = False
        self.empty_pointcloud_step = False

    def reset(self, done=False):
        self.init_context()
        self.object_pose = np.eye(4)
        self.tool_pose = np.eye(4)
        self.tool_rel_pose = np.eye(4)

        for _ in range(self.task_config.sim.warmstart_iter):
            sim_done = self.advance_time(self.context)  #
            if sim_done:  # terminate because of simulation error
                return [], 0, True, self.info
        self.action = np.zeros(self.action_num)
        self._get_obs(self.context)
        obs = self.select_obs()
        reward = self._get_reward(self.context, self.action)
        return obs, reward, done, self.info

    def compute_state_obs(self):
        return np.concatenate(
            [
                pack_pose(self.ee_pose.reshape(4, 4)),
                self.joint_positions,
                pack_pose(self.object_pose),
                pack_pose(self.tool_pose),
            ]
        )

    def reset_from_scene(self, scene_file=None):
        """load a scene including object poses and meshes from a saved file"""
        if scene_file is not None:
            scene_descriptions = json.load(open(os.path.join(scene_file, "scene_descriptions.json")))
            if "object_name" not in scene_descriptions:
                print("filling in real world object details")
                scene_descriptions["object_name"] = scene_descriptions["tool_name"] + "_model0_body_link"
                scene_descriptions["tool_name"] = scene_descriptions["tool_name"] + "_model0_body_link"

            self.scene_descriptions = scene_descriptions
            if self.task_config.verbose:
                print("init from scene file:", scene_file)
        self.init(self.task_config, self.scene_descriptions)

    def get_reward(self, *args):
        return self._get_reward(*args)

    def _get_reward(self, context, action):
        return 0

    def _get_progress(self, context):
        # determine which stage is at for the expert
        return None

    def get_keypoint_information(self):
        """get the keypoint observation"""
        self.tool_keypoints = self.tool_info["keypoints"]
        self.tool_obj_keypoints = self.tool_obj_info["keypoints"]
        self.curr_tool_keypoints = self.tool_keypoints.copy()
        self.curr_tool_obj_keypoints = self.tool_obj_keypoints.copy()
        self.tool_keypoint_head = self.tool_keypoints[0]
        self.tool_keypoint_tail = self.tool_keypoints[1]
        self.tool_keypoint_side = self.tool_keypoints[2]

    def transform_keypoints_env(self):
        self.curr_tool_keypoint_head = self.tool_pose[:3, :3].dot(self.tool_keypoint_head) + self.tool_pose[:3, 3]
        self.curr_tool_keypoint_tail = self.tool_pose[:3, :3].dot(self.tool_keypoint_tail) + self.tool_pose[:3, 3]
        self.curr_tool_keypoint_side = self.tool_pose[:3, :3].dot(self.tool_keypoint_side) + self.tool_pose[:3, 3]
        self.curr_tool_obj_keypoints = (
            self.object_pose[:3, :3].dot(self.tool_obj_keypoints[0]) + self.object_pose[:3, 3]
        )

    def update_keypoint_information(self):
        """get two keypoint positions and the direction between these two points"""
        self.transform_keypoints_env()
        if len(self.tool_obj_keypoints) > 1:
            self.curr_tool_obj_keypoint_rest = (
                self.object_pose[:3, :3].dot(self.tool_obj_keypoints[1:].T).T + self.object_pose[:3, 3]
            )

        # compute the keypoint wrench in the object keypoint coordinate
        self.tool_keypoint_head_in_hand = (
            self.tool_rel_pose[:3, :3].dot(self.tool_keypoint_head) + self.tool_rel_pose[:3, 3]
        )
        self.tool_keypoint_tail_in_hand = (
            self.tool_rel_pose[:3, :3].dot(self.tool_keypoint_tail) + self.tool_rel_pose[:3, 3]
        )
        self.tool_keypoint_side_in_hand = (
            self.tool_rel_pose[:3, :3].dot(self.tool_keypoint_side) + self.tool_rel_pose[:3, 3]
        )

        # get the jacobian that measures speed in the world frame, and expressed in the world frame
        self.head_keypoint_joint_jacobian = self.plant.CalcJacobianSpatialVelocity(
            self.plant_context,
            JacobianWrtVariable.kQDot,
            self.plant.GetBodyByName("panda_hand").body_frame(),
            self.tool_keypoint_head_in_hand,
            self.plant.world_frame(),
            self.plant.world_frame(),
        )[:, :7]
        self.tail_keypoint_joint_jacobian = self.plant.CalcJacobianSpatialVelocity(
            self.plant_context,
            JacobianWrtVariable.kQDot,
            self.plant.GetBodyByName("panda_hand").body_frame(),
            self.tool_keypoint_tail_in_hand,
            self.plant.world_frame(),
            self.plant.world_frame(),
        )[:, :7]
        self.side_keypoint_joint_jacobian = self.plant.CalcJacobianSpatialVelocity(
            self.plant_context,
            JacobianWrtVariable.kQDot,
            self.plant.GetBodyByName("panda_hand").body_frame(),
            self.tool_keypoint_side_in_hand,
            self.plant.world_frame(),
            self.plant.world_frame(),
        )[:, :7]

        # visualize the force and torque at the keypoint locations. Maybe the wrench is more useful in the local frames.
        # rotation first and then position
        torque_sensing = self.external_torque[:7]
        self.head_keypoint_wrench = np.linalg.pinv(self.head_keypoint_joint_jacobian.T).dot(torque_sensing)
        self.tail_keypoint_wrench = np.linalg.pinv(self.tail_keypoint_joint_jacobian.T).dot(torque_sensing)
        self.side_keypoint_wrench = np.linalg.pinv(self.side_keypoint_joint_jacobian.T).dot(torque_sensing)

        keypoint_loc = np.stack(
            (
                self.curr_tool_keypoint_head,
                self.curr_tool_keypoint_tail,
                self.curr_tool_keypoint_side,
                self.curr_tool_obj_keypoints,
            ),
            axis=0,
        )
        self.info["curr_keypoint"] = keypoint_loc
        self.info["keypoint_wrench"] = np.stack(
            (
                self.head_keypoint_wrench,
                self.tail_keypoint_wrench,
                self.side_keypoint_wrench,
            ),
            axis=0,
        )

        if self.task_config.use_meshcat:
            from env.env_util import meshcat

            self.visualize_task_keypoints(
                [
                    self.curr_tool_keypoint_head,
                    self.curr_tool_keypoint_tail,
                    self.curr_tool_keypoint_side,
                    self.curr_tool_obj_keypoints,
                ],
                "task_points",
            )
            # visualize force and torque
            AddWrenchVis(
                meshcat,
                "keypoint_head",
                self.curr_tool_keypoint_head,
                self.head_keypoint_wrench,
            )
            if not self.task_config.vis_collision_geometry:
                meshcat.SetProperty("collision", "visible", False)

    def _get_robot_freq_data(self):
        """get the higher-freq robot state"""
        s = time.time()
        # expressed in the world frame
        self.joint_jacobian = self.plant.CalcJacobianSpatialVelocity(
            self.plant_context,
            JacobianWrtVariable.kQDot,
            self.plant.GetBodyByName("panda_hand").body_frame(),
            [0, 0, 0],
            self.plant.world_frame(),
            self.plant.GetBodyByName("panda_hand").body_frame(),
        )[:, :7]

        try:
            self.external_torque = self.external_torque_port.Eval(self.context)
            torque_sensing = self.external_torque[:7]
            self.get_joint_angles()
        except Exception as e:
            print(e)
            print("history wrench failed")
            torque_sensing = np.zeros(7)

        try:
            if not hasattr(self, "endeffector_wrench"):
                self.endeffector_wrench = np.zeros(6)

            # current
            prev_endeffector_wrench = self.endeffector_wrench.copy()
            self.endeffector_wrench = np.linalg.pinv(self.joint_jacobian.T).dot(torque_sensing)
            self.endeffector_wrench_diff = self.endeffector_wrench - prev_endeffector_wrench
            self.endeffector_wrench_diff = self.endeffector_wrench_diff[:3]
            self.endeffector_wrench_diff_norm = np.linalg.norm(self.endeffector_wrench_diff)
            if self.endeffector_wrench_diff_norm > 1e-4:
                self.endeffector_wrench_diff = self.endeffector_wrench_diff / np.linalg.norm(
                    self.endeffector_wrench_diff
                )
            else:
                self.endeffector_wrench_diff = np.zeros(3)

            self.history_wrench_obs.append(self.endeffector_wrench)
        except Exception as e:
            print(e)
            print("history wrench failed")

    def compute_pointcloud_observation(self, obs, separate_mask, vis=False):
        """compute the torque at each point on the tool
        farthest point sampling to downsample to 1024 points point fusion of two views
        """
        s = time.time()
        pc_img_transforms = []

        # fuse all point cloud from two cameras in world coordinate, partial view.
        pose_names = ["base_camera_pose", "hand_camera_pose"]
        drop_camera = (
            hasattr(self.task_config, "randomize_drop_camera")
            and self.task_config.randomize_drop_camera
            and np.random.uniform() < 0.05
            and self.task_config.train
        )
        drop_cam_index = int(np.random.uniform() < 0.5)
        for idx, (obs_i, mask_i, cam_name) in enumerate(zip(obs, separate_mask, pose_names)):

            depth_image = obs_i[..., 3]
            mask_image = mask_i.flatten().reshape(1, -1)
            K = CAM_INTR
            extrinsics = self.info[cam_name].reshape(4, 4)

            pc_img = backproject_camera_target(depth_image, K)
            pc_img_world = transform_point(extrinsics, pc_img)

            pc_img_world = np.concatenate((pc_img_world, mask_image), axis=0)

            if drop_camera and drop_cam_index == idx:
                print("drop camera idx:", idx)
                continue

            pc_img_transforms.append(pc_img_world)

        # fuse pointclouds
        pc_img_transforms = np.concatenate(pc_img_transforms, axis=1)

        # world coordinate and use mask
        object_points = pc_img_transforms[:, pc_img_transforms[-1] == 1]
        tool_points = pc_img_transforms[:, pc_img_transforms[-1] == 2]
        tool_pc = torch.from_numpy(tool_points.T).cuda().float()[None]

        obj_pc = torch.from_numpy(object_points.T).cuda().float()[None]
        ee_pose = torch.from_numpy(self.ee_pose.reshape(4, 4)).cuda().float()
        ee_inv_pose = torch.from_numpy(se3_inverse(self.ee_pose.reshape(4, 4))).cuda().float()

        # tool points fused in local frame
        self.fused_tool_points = self.fused_tool_points.cuda()
        self.fused_object_points = self.fused_object_points.cuda()
        self.empty_pointcloud_step = False

        if tool_points.shape[1] != 0:
            tool_xyz = downsample_point(tool_pc, self.task_config.env.num_points // 2).cuda()
            # if any of the pointcloud is empty. Need to return or gives some zero paddings
            # if there is a contact, we need to reset the pointcloud
            if self.external_torque[:7].sum() > 0:
                self.fused_tool_points = torch.zeros([4, 0]).cuda()  #

            tool_xyz[-1] = 0  # mask 0 for tool
            tool_xyz[:3] = transform_point(ee_inv_pose, tool_xyz[:3])
            self.fused_tool_points = tool_xyz
            if hasattr(self.task_config, "use_fused_pointcloud") and self.task_config.use_fused_pointcloud:
                self.fused_tool_points = torch.cat((self.fused_tool_points, tool_xyz), axis=1)
            tool_points = tool_xyz

        else:  # fused points could be used as backup
            tool_points = self.fused_tool_points
            print("empty tool points at the current timestep")
            if self.time < 10:
                self.empty_pointcloud_step = True

        # object points fused in world frame
        if object_points.shape[1] != 0:
            obj_xyz = downsample_point(obj_pc, self.task_config.env.num_points // 2).cuda()

            if self.external_torque[:7].sum() > 0:  # reset points when there is contact
                self.fused_object_points = torch.zeros([4, 0]).cuda()

            obj_xyz[-1] = 1  # mask 1 for obj
            self.fused_object_points = obj_xyz
            if hasattr(self.task_config, "use_fused_pointcloud") and self.task_config.use_fused_pointcloud:
                self.fused_object_points = torch.cat((self.fused_object_points.cuda(), obj_xyz.cuda()), axis=1)
            object_points = obj_xyz

        else:  # fused points could be used as backup
            object_points = self.fused_object_points
            print("empty object points at the current timestep")
            if self.time < 10:
                self.empty_pointcloud_step = True

        if hasattr(self.task_config, "use_fused_pointcloud") and self.task_config.use_fused_pointcloud:
            self.fused_tool_points = downsample_point(
                self.fused_tool_points.T, self.task_config.env.fused_num_points // 2
            )
            self.fused_object_points = downsample_point(
                self.fused_object_points.T, self.task_config.env.fused_num_points // 2
            )

        curr_frame_tool_xyz = downsample_point(tool_points.T, self.task_config.env.num_points // 2).cuda()
        curr_frame_obj_xyz = downsample_point(object_points.T, self.task_config.env.num_points // 2).cuda()

        # object point transform from world frame to the ee frame
        curr_frame_obj_xyz[:3] = transform_point(ee_inv_pose, curr_frame_obj_xyz[:3])
        combined_pc = torch.cat((curr_frame_tool_xyz, curr_frame_obj_xyz), axis=-1)
        combined_pc = combined_pc.detach().cpu().numpy()

        # convert to hand coordinate frame
        self.info["point_cloud"] = combined_pc

        if vis and self.time_step % 5 == 0 and self.task_config.use_meshcat:
            from env.env_util import meshcat, visualizer

            pcd = o3d.geometry.PointCloud()
            point_xyzs = transform_point(self.ee_pose.reshape(4, 4), combined_pc[:3]).T
            mask = self.info["point_cloud"][-1]
            point_colors = np.zeros_like(point_xyzs)

            point_colors[mask == 1] = [1.0, 0, 0]
            point_colors[mask == 0] = [0, 1.0, 0]
            pcd.points = o3d.utility.Vector3dVector(point_xyzs)
            pcd.colors = o3d.utility.Vector3dVector(point_colors)
            draw_open3d_point_cloud(meshcat, "observed_pc", pcd, point_size=0.001)

    def _get_obs(self, context, render=True):
        """compute the observation of the environment"""
        self.get_joint_angles()
        self.ee_pose = np.array(self.get_ee_pos().GetAsMatrix4()).copy().flatten()
        self.info["action"] = self.action  # the action that brings to this observation
        self.img_obs = []
        self.full_mask_data = []
        render = self.task_config.env.use_image or self.task_config.env.use_point_cloud_model

        # get camera poses
        if render:
            cams = self.task_config.env.cam.split("+")
            if len(cams) == 1:
                self.img_obs = self._render(cams[0])
                camera_pose_port = self.system.GetOutputPort(cams[0] + "_camera_pose")
                cam_pose = camera_pose_port.Eval(context).GetAsMatrix4()
                self.info[cams[0] + "_pose"] = cam_pose.copy()

            else:  # multiple cameras
                for cam in cams:
                    self.img_obs.append(self._render(cam))
                    camera_pose_port = self.system.GetOutputPort(cam + "_camera_pose")
                    cam_pose = camera_pose_port.Eval(context).GetAsMatrix4()
                    self.info[cam + "_pose"] = cam_pose.copy()

            self.info["base_camera_pose"] = self.info["panda_link0_cam_pose"]
            self.info["hand_camera_pose"] = self.info["panda_hand_cam_pose"]

            if self.task_config.env.use_point_cloud_model:
                self.compute_pointcloud_observation(self.img_obs, self.full_mask_data, vis=True)

                if len(self.info["point_cloud"].shape) == 1 or self.empty_pointcloud_step:
                    print("empty pointcloud.")
                    if hasattr(self.task_config, "data_collection"):
                        print("reset!")
                        return self.reset()  # will count as failures

        if hasattr(self.task_config, "save_dino_features") and self.task_config.save_dino_features:
            # compute features
            from misc.process_dino_features import extract_dino_v2_features

            # overhead first
            descriptors, features = extract_dino_v2_features(self.img_obs[0][..., :3], self.img_obs[0][..., [-1]])
            self.info["img_feature"] = features[0]

        self.info["joint_pos"] = self.joint_positions
        self.info["ee_pose"] = self.ee_pose
        self.info["cam_intr"] = CAM_INTR
        self.info["timestep"] = self.time_step
        self.info["task_name"] = self.task_config.task_name
        self.info["success"] = self._get_reward(self.context, self.action) > 0

        if hasattr(self, "tool_class_name"):
            self.info["tool_class_name"] = self.tool_class_name

        if hasattr(self, "tool_name"):
            self.info["tool_name"] = self.tool_name
            self.info["tool_pose"] = self._get_tool_pose(context)

        if hasattr(self, "object_name"):
            self.info["tool_object_name"] = self.object_name
            self.info["object_pose"] = self._get_object_pose(context)

        self.state_obs = self.compute_state_obs()
        self.info["state"] = self.state_obs
        self.info["next_state"] = self.state_obs.copy()  # ...

        if self.task_config.blender_render:
            print("publish blender image")
            from env.env_util import blender_color_cam

            blender_context = blender_color_cam.GetMyContextFromRoot(self.context)
            blender_color_cam.Publish(blender_context)

        try:
            self.contact_results = self.contact_measurement_port.Eval(context)
        except:
            print("nan in contact results")
            pass

        # detect collisions and reset if needed
        for i in range(self.contact_results.num_point_pair_contacts()):
            contact_info_i = self.contact_results.point_pair_contact_info(i)
            body_ia_index = contact_info_i.bodyA_index()
            body_ib_index = contact_info_i.bodyB_index()

            body_a_name = self.plant.get_body(body_ia_index).name()
            body_b_name = self.plant.get_body(body_ib_index).name()
            if self.check_collision_name(body_a_name, body_b_name):
                color_print(
                    (f"{body_a_name} collide with {body_b_name} reset"),
                    "light_slate_blue",
                )
                self.need_termination = True
                self.info["collision"] = True  # emergent reset

            if self.task_config.env.debug_contacts:
                if "finger" not in body_a_name and "finger" not in body_b_name:
                    color_print(
                        "body A: {} body B: {} force: {}".format(
                            body_a_name,
                            body_b_name,
                            np.round(contact_info_i.contact_force(), 3),
                        ),
                        "light_slate_blue",
                    )

        # get jacobian
        self.joint_jacobian = self.plant.CalcJacobianSpatialVelocity(
            self.plant_context,
            JacobianWrtVariable.kQDot,
            self.plant.GetBodyByName("panda_hand").body_frame(),
            [0, 0, 0],
            self.plant.world_frame(),
            self.plant.world_frame(),
        )[:, :7]

        self._get_robot_freq_data()
        self.info["external_torque"] = self.external_torque  # for each joint
        # for the end effector, rotation first
        self.info["endeffector_wrench_diff"] = self.endeffector_wrench_diff
        self.info["endeffector_wrench"] = self.endeffector_wrench

        if len(self.history_wrench_obs) > 3:
            self.info["history_wrench_obs"] = np.concatenate(self.history_wrench_obs)
            self.at_contact = np.linalg.norm(self.info["history_wrench_obs"]) > 0.0
            if self.task_config.use_controller_portswitch and self.task_config.automatic_controller_switch:

                # testing purpose
                if self.at_contact:
                    self.controller_switch_port.FixValue(self.context, InputPortIndex(1))
                else:
                    self.controller_switch_port.FixValue(self.context, InputPortIndex(2))

        # compute keypoint information
        if hasattr(self, "tool_name"):
            if not hasattr(self, "tool_keypoints"):
                self.get_keypoint_information()
            self.update_keypoint_information()

    def render_two_views(self):
        # for recording video purpose
        base_cam_img = self._render("panda_link0_cam")
        wrist_cam_img = self._render("panda_hand_cam")
        return base_cam_img, wrist_cam_img

    def render_any_view(self, pose):
        body = self.plant.GetBodyByName("dummy_cam_body")
        self.plant.SetFreeBodyPose(self.plant_context, body, RigidTransform(pose))
        return self._render("dummy_cam_body_cam")

    def get_combined_mask(self, cam):
        render_port = self.system.GetOutputPort(cam + "_label_image")
        maskdata = render_port.Eval(self.context).data.copy()

        # simulate whatever is above the table and not the robot
        object_bodies = self.plant.GetBodyIndices(self.object_asset)
        tool_bodies = self.plant.GetBodyIndices(self.tool_asset)
        # masked_ids  = [int(x) for x in object_bodies] + [int(x) for x in tool_bodies]

        obj_masked_ids = [int(x) for x in object_bodies]
        tool_masked_ids = [int(x) for x in tool_bodies]
        obj_mask = sum([maskdata == x for x in obj_masked_ids]) > 0
        tool_mask = sum([maskdata == x for x in tool_masked_ids]) > 0

        filtered_maskdata = obj_mask | tool_mask

        # as auxiliary info
        full_mask = np.zeros(obj_mask.shape)
        full_mask[obj_mask] = 1.0
        full_mask[tool_mask] = 2.0
        self.full_mask_data.append(full_mask)

        return filtered_maskdata

    def _render(self, cam):
        """rendering the image of the scene based on the spec"""
        data = []
        if "RGB" in self.task_config.env.img_type:
            render_port = self.system.GetOutputPort(cam + "_rgb_image")
            rgbdata = render_port.Eval(self.context).data.copy()[..., :3] / 255.0
            data.append(rgbdata)

        if "D" in self.task_config.env.img_type:
            render_port = self.system.GetOutputPort(cam + "_depth_image")
            depthdata = render_port.Eval(self.context).data.copy()
            depthdata[depthdata == np.inf] = 3
            data.append(depthdata)

        if "M" in self.task_config.env.img_type:
            filtered_maskdata = self.get_combined_mask(cam)
            # filtered_maskdata = filtered_maskdata[...,None].astype(np.bool)
            data.append(filtered_maskdata)

        if "POINTS" in self.task_config.env.img_type:
            render_port = self.system.GetOutputPort(cam + "_point_cloud")
            pointdata = render_port.Eval(self.context).data.copy()
            data.append(pointdata)
        else:
            data = np.concatenate(data, axis=-1)

        if self.task_config.env.debug_obs_image and self.time < 3:
            visualize_rendered_image(data)

        return data

    def terminate_mechanism(self):
        pass

    def export_context(self):
        """export the current simulator context"""
        context = self.simulator.get_context()
        return context.Clone()

    def load_context(self, context):
        """load the context to the current simulation"""
        self.context.SetTimeStateAndParametersFrom(context)
        self.simulator.Initialize()

    def seed(self, seed=None):
        """Implements gym.Env.seed using Drake's RandomGenerator."""
        if seed:
            self.generator = RandomGenerator(seed)
        else:
            seed = self.generator()

        return [seed]

    def visualize_task_keypoints(self, point_sets, pointset_name=""):
        """visualize the keypoint locations"""
        if not self.task_config.env.vis_keypoints:
            return
        from env.env_util import meshcat, visualizer

        pcd = o3d.geometry.PointCloud()
        points = np.stack(point_sets, axis=0)
        point_colors = np.stack(([1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 0], [1, 0, 1], [0, 1, 1]), axis=0)[
            : len(point_sets)
        ]
        pcd.points = o3d.utility.Vector3dVector(points)

        pcd.colors = o3d.utility.Vector3dVector(point_colors[: len(points)])
        draw_open3d_point_cloud(meshcat, pointset_name, pcd, point_size=0.003, sphere_vis=True)

env/frankadrakeenv.py
liruiw/Fleet-Tools
import gym
import numpy as np
from env.env_util import *
from .frankadrakeenv import FrankaDrakeEnv
import IPython
from core.utils import *


class FrankaDrakeHammerEnv(FrankaDrakeEnv):
    def __init__(self, task_config):
        self.tool_class_name = task_config.tool_class_name
        super().__init__(task_config)

    def reset(self, done=False):
        self.init_context()
        from env.env_object_related import extra_force_system

        # extra force to make sure the pin does not go down due to gravity
        g = -self.plant.gravity_field().gravity_vector()[2]
        extra_force_system.wrench = np.concatenate(
            (
                [
                    0,
                    0,
                    self.plant.GetBodyByName(self.object_name).get_mass(self.context) * g - 1e-10,
                ],
                [0, 0, 0],
            )
        )

        for _ in range(self.task_config.sim.warmstart_iter):
            sim_done = self.advance_time(self.context)  #
            if sim_done:  # terminate because of simulation error
                return [], 0, True, self.info

        self.action = np.zeros(self.action_num)
        self._get_obs(self.context)
        obs = self.select_obs()
        reward = self._get_reward(self.context, self.action)
        return obs, reward, done, self.info

    def step(self, action):
        self.time_step += 1
        sim_done = self.advance_time(self.context)

        if sim_done:  # terminate because of simulation error
            return [], 0, True, self.info
        self._get_obs(self.context)

        obs = self.select_obs()
        reward = self._get_reward(self.context, action)
        done = sim_done or (reward > 0) or self.need_termination

        return obs, reward, done, self.info

    def check_collision_name(self, body_a_name, body_b_name):
        collision = super().check_collision_name(body_a_name, body_b_name)
        # anything hitting the table or the box would mean failure.
        table_collision = "link" == body_a_name or "link" == body_b_name
        return collision or table_collision

    def _get_obs(self, context, render=True):
        from env.env_object_related import extra_force_system

        super()._get_obs(context, render=render)
        contact_results = self.contact_measurement_port.Eval(self.context)
        g = -self.plant.gravity_field().gravity_vector()[2]
        z_force = self.plant.GetBodyByName(self.object_name).get_mass(self.context) * g - 1e-10
        if contact_results.num_point_pair_contacts() == 1:  # manually set the downward force
            extra_force_system.wrench = np.concatenate(([0, 0, z_force], [0, 0, 0]))
        else:
            extra_force_system.wrench = np.concatenate(([0, 0, z_force], [0, 0, 0]))

    def _get_reward(self, context, action):
        """hammer environment reward is based on how much the pin is pressed"""
        from env.env_object_related import bolt_joint

        hammer_state = bolt_joint.get_translation(self.plant_context)
        rew = float(hammer_state > -0.015 and self.time > 2)
        return rew

    def set_object_pose(self, tf):
        return

    def _get_progress(self, context):
        pass

env/frankadrakehammerenv.py
liruiw/Fleet-Tools
import gym
import numpy as np
from pydrake.math import RollPitchYaw, RigidTransform
from env.env_util import MakeManipulationStationTools
from env.env_util import *
from .frankadrakeenv import FrankaDrakeEnv
import IPython
from core.utils import *


class FrankaDrakeKnifeEnv(FrankaDrakeEnv):
    def __init__(self, task_config):
        self.tool_class_name = task_config.tool_class_name
        super().__init__(task_config)
        self.init_dist = 0

    def get_combined_mask(self, cam):
        """the policy needs to see both objects in this task"""
        render_port = self.system.GetOutputPort(cam + "_label_image")
        maskdata = render_port.Eval(self.context).data.copy()

        # simulate whatever is above the table and not the robot
        object_bodies = self.plant.GetBodyIndices(self.object_asset)
        object_bodies += [self.plant.GetBodyByName(self.object_name + "_2").index()]
        tool_bodies = self.plant.GetBodyIndices(self.tool_asset)
        obj_masked_ids = [int(x) for x in object_bodies]
        tool_masked_ids = [int(x) for x in tool_bodies]
        obj_mask = sum([maskdata == x for x in obj_masked_ids]) > 0
        tool_mask = sum([maskdata == x for x in tool_masked_ids]) > 0
        filtered_maskdata = obj_mask | tool_mask

        # as auxiliary info
        full_mask = np.zeros(obj_mask.shape)
        full_mask[obj_mask] = 1.0
        full_mask[tool_mask] = 2.0
        self.full_mask_data.append(full_mask)

        return filtered_maskdata

    def reset(self, done=False):
        self.init_context()
        # reset the other object's pose based on the first object
        body = self.plant.GetBodyByName(self.object_name + "_2")
        tf = self.plant.GetFrameByName(self.object_name).CalcPoseInWorld(self.plant_context)

        # sample a random delta pose
        x_delta = np.random.choice([-1, 1]) * np.random.uniform(0.04, 0.07)
        y_delta = np.random.choice([-1, 1]) * np.random.uniform(0.04, 0.07)
        delta_translation = np.array([x_delta, y_delta, 0.1]) + tf.translation()
        tf = RigidTransform(RollPitchYaw(np.random.uniform(-np.pi, np.pi), 0, 0), delta_translation)
        self.plant.SetFreeBodyPose(self.plant_context, body, tf)  # overlap and then split

        for _ in range(self.task_config.sim.warmstart_iter):
            sim_done = self.advance_time(self.context)  #
            if sim_done:  # terminate because of simulation error
                return [], 0, True, {}

        self.action = np.zeros(self.action_num)
        self.obj_pose1 = self.plant.GetFrameByName(self.object_name).CalcPoseInWorld(self.plant_context).GetAsMatrix4()
        self.obj_pose2 = (
            self.plant.GetFrameByName(self.object_name + "_2").CalcPoseInWorld(self.plant_context).GetAsMatrix4()
        )
        self._get_obs(self.context)
        self.init_dist = np.linalg.norm(self.obj_pose1[:3, 3] - self.obj_pose2[:3, 3])
        obs = self.select_obs()
        reward = self._get_reward(self.context, self.action)
        return obs, reward, done, self.info

    def step(self, action):
        self.time_step += 1
        sim_done = self.advance_time(self.context)
        if sim_done:  # terminate because of simulation error
            return [], 0, True, self.info

        self.obj_pose1 = self.plant.GetFrameByName(self.object_name).CalcPoseInWorld(self.plant_context).GetAsMatrix4()
        self.obj_pose2 = (
            self.plant.GetFrameByName(self.object_name + "_2").CalcPoseInWorld(self.plant_context).GetAsMatrix4()
        )
        self._get_obs(self.context)

        obs = self.select_obs()
        reward = self._get_reward(self.context, action)
        done = sim_done or (reward > 0) or self.need_termination

        return obs, reward, done, self.info

    def _get_reward(self, context, action):
        """knife (planar push) environment reward is based on how much two objects are pushed from one another"""
        translation_dist = np.linalg.norm(self.obj_pose1[:3, 3] - self.obj_pose2[:3, 3])
        reward = float((translation_dist - self.init_dist > 0.05) and (self.time > 5.0))

        return reward

    def update_keypoint_information(self):
        """get two keypoint positions and the direction between these two points"""
        super().update_keypoint_information()
        self.curr_tool_obj_keypoints = (self.obj_pose1[:3, 3] + self.obj_pose2[:3, 3]) / 2

    def _get_progress(self, context):
        pass

env/frankadrakeknifeenv.py
liruiw/Fleet-Tools
import numpy as np
from env.env_util import *
from .frankadrakeenv import FrankaDrakeEnv
from core.utils import *


class FrankaDrakeSpatulaEnv(FrankaDrakeEnv):
    def __init__(self, task_config):
        self.tool_class_name = task_config.tool_class_name
        super().__init__(task_config)

    def reset(self, done=False):
        self.init_context()
        if self.use_pose_controller() or self.task_config.use_controller_portswitch:
            self.robot_controller.set_joint_stiffness("normal")

        for _ in range(self.task_config.sim.warmstart_iter):
            sim_done = self.advance_time(self.context)  #
            if sim_done:  # terminate because of simulation error
                return [], 0, True, self.info

        self.action = np.zeros(self.action_num)
        self._get_obs(self.context)
        obs = self.select_obs()
        reward = self._get_reward(self.context, self.action)
        self.reset_env_time = self.time
        return obs, reward, done, self.info

    def transform_keypoints_env(self):
        self.curr_tool_keypoint_head = self.tool_pose[:3, :3].dot(self.tool_keypoint_head) + self.tool_pose[:3, 3]
        self.curr_tool_keypoint_tail = self.tool_pose[:3, :3].dot(self.tool_keypoint_tail) + self.tool_pose[:3, 3]
        self.curr_tool_keypoint_side = self.tool_pose[:3, :3].dot(self.tool_keypoint_side) + self.tool_pose[:3, 3]

        # transform the object keypoint to be the bottom point of the object
        self.curr_tool_obj_keypoints = self.object_pose[:3, 3]
        self.curr_tool_obj_keypoints[2] = self.table_height + 0.005  #

    def step(self, action):
        self.time_step += 1
        sim_done = self.advance_time(self.context)

        if sim_done:  # terminate because of simulation error
            return [], 0, True, self.info
        self._get_obs(self.context)

        obs = self.select_obs()
        reward = self._get_reward(self.context, action)
        done = sim_done or (reward > 0) or self.need_termination
        self.terminate_mechanism()
        return obs, reward, done, self.info

    def _get_obs(self, context, render=True):
        super()._get_obs(context, render=render)
        contact_results = self.contact_results

        for i in range(contact_results.num_hydroelastic_contacts()):
            contact_info_i = contact_results.hydroelastic_contact_info(i)
            contact_surface = contact_info_i.contact_surface()
            body_ia_index = contact_surface.id_M()
            body_ib_index = contact_surface.id_N()  #
            body_a_name = self.scene_inspector.GetName(body_ia_index)
            body_b_name = self.scene_inspector.GetName(body_ib_index)

    def _get_reward(self, context, action):
        """spatula environment reward is based on if the object is lifted with the spatula"""
        obj_pose = self._get_object_pose(context)
        obj_rot = mat2euler(obj_pose[:3, :3])
        reward = float((np.abs(obj_rot[0]) > np.pi / 2.0 - 0.2))
        return reward

    def _get_progress(self, context):
        pass

env/frankadrakespatulaenv.py
liruiw/Fleet-Tools
import numpy as np

from pydrake.math import RollPitchYaw, RigidTransform
from env.env_util import *
from .frankadrakeenv import FrankaDrakeEnv
import IPython
from core.utils import *


class FrankaDrakeWrenchEnv(FrankaDrakeEnv):
    def __init__(self, task_config):
        self.tool_class_name = task_config.tool_class_name
        super().__init__(task_config)

    def reset(self, done=False):
        self.init_context()
        if self.use_pose_controller() or self.task_config.use_controller_portswitch:
            self.robot_controller.set_joint_stiffness("compliant")
            self.robot_controller.set_pose_stiffness([400, 40, 100, 20])

        # smooth out contact force scaling a bit
        self.history_wrench_force = deque([0] * 10, maxlen=10)
        for _ in range(self.task_config.sim.warmstart_iter):
            sim_done = self.advance_time(self.context)  #
            if sim_done:  # terminate because of simulation error
                return [], 0, True, self.info

        self.action = np.zeros(self.action_num)
        self._get_obs(self.context)
        obs = self.select_obs()
        reward = self._get_reward(self.context, self.action)

        from env.env_object_related import screw_joint_

        screw_joint_.set_translation(self.plant_context, 0.00)
        screw_joint_.set_angular_velocity(self.plant_context, 0)
        screw_joint_.set_translational_velocity(self.plant_context, 0)

        return obs, reward, done, self.info

    def render_two_views(self):
        pose = OVERHEAD_CAM_EXTR.copy()
        pose[:3, 3] = [0.5, 0.0, 0.25]
        base_cam_img = self._render("panda_link0_cam")
        wrist_cam_img = self.render_any_view(pose)  # not much need to vis
        return base_cam_img, wrist_cam_img

    def _get_object_pose(self, context=None):
        if context is None:
            context = self.context
        plant_context = self.system.GetSubsystemContext(self.plant, self.context)
        # the nut is the actual object pose
        self.object_pose = (
            self.plant.GetFrameByName(self.object_name.replace("_body", "_nut_body"))
            .CalcPoseInWorld(plant_context)
            .GetAsMatrix4()
        )
        return self.object_pose

    def step(self, action):

        self.time_step += 1
        cur_pose = self.get_ee_pos().GetAsMatrix4()
        sim_done = self.advance_time(self.context)
        if sim_done:  # terminate because of simulation error
            return [], 0, True, self.info
        self._get_obs(self.context)

        obs = self.select_obs()
        reward = self._get_reward(self.context, action)
        done = sim_done or (reward > 0) or self.need_termination

        return obs, reward, done, self.info

    def _get_obs(self, context, render=True):
        super()._get_obs(context, render=render)
        from env.env_object_related import screw_joint_, extra_force_system
        from env.env_util import meshcat

        contact_results = self.contact_results
        extra_force_system.wrench = np.zeros(6)
        self.nut_contacts = []  # might be used later

        if self.task_config.sim.compliant_hydroelastic_model:
            if contact_results.num_hydroelastic_contacts() > 0:
                # scale the force
                for i in range(contact_results.num_hydroelastic_contacts()):
                    contact_info_i = contact_results.hydroelastic_contact_info(i)
                    contact_surface = contact_info_i.contact_surface()
                    body_ia_index = contact_surface.id_M()
                    body_ib_index = contact_surface.id_N()  #
                    body_a_name = self.scene_inspector.GetName(body_ia_index)
                    body_b_name = self.scene_inspector.GetName(body_ib_index)

                    if "nut_body_link" in body_a_name or "nut_body_link" in body_b_name:
                        contact_spatial_force = contact_info_i.F_Ac_W()  # on body B expressed in world, contact_force
                        contact_force = contact_spatial_force.translational()
                        contact_torque = contact_spatial_force.rotational()

                        self.nut_contacts.append(contact_info_i)
                        force_norm = min(np.linalg.norm(contact_force[:2]), 0.1)
                        self.history_wrench_force.append(force_norm)
                        force_norm = np.mean(self.history_wrench_force)
                        torque = np.array([0, 0, force_norm])

        else:
            if contact_results.num_point_pair_contacts() > 1:
                # scale the force
                for i in range(contact_results.num_point_pair_contacts()):
                    contact_info_i = contact_results.point_pair_contact_info(i)
                    body_ia_index = contact_info_i.bodyA_index()
                    body_ib_index = contact_info_i.bodyB_index()  #
                    if (
                        "nut_body_link" in self.plant.get_body(body_ib_index).name()
                        or "nut_body_link" in self.plant.get_body(body_ia_index).name()
                    ):
                        contact_force = contact_info_i.contact_force()  # on body B expressed in world, contact_force
                        self.nut_contacts.append(contact_info_i)
                        force_norm = min(np.linalg.norm(contact_force[:2]), 0.1)
                        self.history_wrench_force.append(force_norm)
                        force_norm = np.mean(self.history_wrench_force)
                        torque = np.array([0, 0, force_norm])

    def set_object_pose(self, tf):
        return

    def _get_reward(self, context, action):
        """compute reward based on how much the screw is turned"""
        from env.env_object_related import screw_joint_, extra_force_system

        rotation = screw_joint_.get_rotation(self.plant_context)
        translation = screw_joint_.get_translation(self.plant_context)
        rew = float(rotation < -0.4)
        return rew  #  rew

    def _get_progress(self, context):
        pass

env/frankadrakewrenchenv.py
liruiw/Fleet-Tools
# From Russ's code in manipulation repo
from functools import partial
import os
import sys
import time

from IPython.display import display, HTML, Javascript
import numpy as np


from pydrake.geometry import Cylinder, Rgba, Sphere
from pydrake.perception import PointCloud, Fields, BaseField
from pydrake.all import BoundingBoxConstraint

# imports for the pose sliders
from collections import namedtuple
from pydrake.all import AbstractValue
from pydrake.all import RollPitchYaw, RigidTransform, RotationMatrix
from pydrake.all import LeafSystem, PublishEvent

# imports for the joint sliders
from pydrake.all import JointIndex
from pydrake.all import PublishEvent, VectorSystem

from pydrake.all import *

# Some GUI code that will be moved into Drake.
import IPython


class DrawContactWrench(LeafSystem):
    def __init__(self, meshcat):
        LeafSystem.__init__(self)
        self.DeclareAbstractInputPort("contact_results", AbstractValue.Make(ContactResults()))
        self.DeclareForcedPublishEvent(self.Publish)
        self.meshcat = meshcat

    def Publish(self, context):
        results = self.get_input_port().Eval(context)

        for i in range(results.num_point_pair_contacts()):
            info = results.point_pair_contact_info(i)
            self.meshcat.SetObject(
                f"contact_{i}", MeshcatCone(height, height * mu, height * mu), rgba=Rgba(0.1, 0.9, 0.1, 1.0)
            )
            p_WC = info.contact_point()
            R_WC = RotationMatrix.MakeFromOneVector(info.point_pair().nhat_BA_W, 2)
            X_WC = RigidTransform(R_WC, p_WC)
            self.meshcat.SetTransform(f"contact_{i}", X_WC)

            X_WB = RigidTransform()  # box center is at world origin
            self.meshcat.SetObject(
                f"box_center/contact_{i}/translational",
                MeshcatCone(height, height * mu, height * mu),
                rgba=Rgba(0.7, 0.1, 0.1, 1.0),
            )
            self.meshcat.SetTransform("box_center", X_WB)
            R_BC = X_WB.inverse().rotation() @ R_WC
            self.meshcat.SetTransform(f"box_center/contact_{i}", RigidTransform(R_BC))
            p_CB_C = -(R_WC @ (X_WB.inverse() @ p_WC))
            A_CCtau = VectorToSkewSymmetric(p_CB_C)
            # Three.js does not fully support non-uniform scaling
            # https://github.com/mrdoob/three.js/issues/15079.
            # So I cannot simply draw the meshcatcone and transform it.
            # Instead I will create the vertices myself.
            cloud.mutable_xyzs()[:] = A_CCtau @ vertices
            self.meshcat.SetObject(
                f"box_center/contact_{i}/rotational_cloud", cloud, point_size=0.002, rgba=Rgba(0, 1, 0, 1)
            )
            # TODO(russt): Figure out colors.  The surface mesh example works ok in meshcat_manual_test...
            self.meshcat.SetTriangleMesh(
                f"box_center/contact_{i}/rotational", A_CCtau @ vertices, faces, rgba=Rgba(0, 0, 1, 1)
            )

        clear_output(wait=True)


class MeshcatSliders(LeafSystem):
    """
    A system that outputs the ``value``s from meshcat sliders.

    .. pydrake_system::

      name: MeshcatSliderSystem
      output_ports:
      - slider_group_0
      - ...
      - slider_group_{N-1}
    """

    def __init__(self, meshcat, slider_names):
        """
        An output port is created for each element in the list `slider_names`.
        Each element of `slider_names` must itself be an iterable collection
        (list, tuple, set, ...) of strings, with the names of sliders that have
        *already* been added to Meshcat via Meshcat.AddSlider().

        The same slider may be used in multiple ports.
        """
        LeafSystem.__init__(self)

        self._meshcat = meshcat
        self._sliders = slider_names
        for i, slider_iterable in enumerate(self._sliders):
            port = self.DeclareVectorOutputPort(
                f"slider_group_{i}",
                len(slider_iterable),
                partial(self.DoCalcOutput, port_index=i),
            )
            port.disable_caching_by_default()

    def DoCalcOutput(self, context, output, port_index):
        for i, slider in enumerate(self._sliders[port_index]):
            output[i] = self._meshcat.GetSliderValue(slider)


class MeshcatPoseSliders(LeafSystem):
    """
    Provides a set of ipywidget sliders (to be used in a Jupyter notebook) with
    one slider for each of roll, pitch, yaw, x, y, and z.  This can be used,
    for instance, as an interface to teleoperate the end-effector of a robot.

    .. pydrake_system::

        name: PoseSliders
        output_ports:
        - pose
    """

    # TODO(russt): Use namedtuple defaults parameter once we are Python >= 3.7.
    Visible = namedtuple("Visible", ("roll", "pitch", "yaw", "x", "y", "z"))
    Visible.__new__.__defaults__ = (True, True, True, True, True, True)
    MinRange = namedtuple("MinRange", ("roll", "pitch", "yaw", "x", "y", "z"))
    MinRange.__new__.__defaults__ = (-np.pi, -np.pi, -np.pi, -1.0, -1.0, -1.0)
    MaxRange = namedtuple("MaxRange", ("roll", "pitch", "yaw", "x", "y", "z"))
    MaxRange.__new__.__defaults__ = (np.pi, np.pi, np.pi, 1.0, 1.0, 1.0)
    Value = namedtuple("Value", ("roll", "pitch", "yaw", "x", "y", "z"))
    Value.__new__.__defaults__ = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0)

    def __init__(
        self,
        meshcat,
        visible=Visible(),
        min_range=MinRange(),
        max_range=MaxRange(),
        value=Value(),
    ):
        """
        Args:
            meshcat: A Meshcat instance.
            visible: An object with boolean elements for 'roll', 'pitch',
                     'yaw', 'x', 'y', 'z'; the intention is for this to be the
                     PoseSliders.Visible() namedtuple.  Defaults to all true.
            min_range, max_range, value: Objects with float values for 'roll',
                      'pitch', 'yaw', 'x', 'y', 'z'; the intention is for the
                      caller to use the PoseSliders.MinRange, MaxRange, and
                      Value namedtuples.  See those tuples for default values.
        """
        LeafSystem.__init__(self)
        port = self.DeclareAbstractOutputPort("pose", lambda: AbstractValue.Make(RigidTransform()), self.DoCalcOutput)

        # The widgets themselves have undeclared state.  For now, we accept it,
        # and simply disable caching on the output port.
        # TODO(russt): consider implementing the more elaborate methods seen
        # in, e.g., LcmMessageSubscriber.
        port.disable_caching_by_default()

        self._meshcat = meshcat
        self._visible = visible
        self._value = list(value)

        for i in range(6):
            if visible[i]:
                meshcat.AddSlider(
                    min=min_range[i],
                    max=max_range[i],
                    value=value[i],
                    step=0.01,
                    name=value._fields[i],
                )

    def __del__(self):
        for s in ["roll", "pitch", "yaw", "x", "y", "z"]:
            if visible[s]:
                self._meshcat.DeleteSlider(s)

    def SetPose(self, pose):
        """
        Sets the current value of the sliders.

        Args:
            pose: Any viable argument for the RigidTransform
                  constructor.
        """
        tf = RigidTransform(pose)
        self.SetRpy(RollPitchYaw(tf.rotation()))
        self.SetXyz(tf.translation())

    def SetRpy(self, rpy):
        """
        Sets the current value of the sliders for roll, pitch, and yaw.

        Args:
            rpy: An instance of drake.math.RollPitchYaw
        """
        self._value[0] = rpy.roll_angle()
        self._value[1] = rpy.pitch_angle()
        self._value[2] = rpy.yaw_angle()
        for i in range(3):
            if self._visible[i]:
                self._meshcat.SetSliderValue(self._visible._fields[i], self._value[i])

    def SetXyz(self, xyz):
        """
        Sets the current value of the sliders for x, y, and z.

        Args:
            xyz: A 3 element iterable object with x, y, z.
        """
        self._value[3:] = xyz
        for i in range(3, 6):
            if self._visible[i]:
                self._meshcat.SetSliderValue(self._visible._fields[i], self._value[i])

    def _update_values(self):
        changed = False
        for i in range(6):
            if self._visible[i]:
                old_value = self._value[i]
                self._value[i] = self._meshcat.GetSliderValue(self._visible._fields[i])
                changed = changed or self._value[i] != old_value
        return changed

    def _get_transform(self):
        return RigidTransform(
            RollPitchYaw(self._value[0], self._value[1], self._value[2]),
            self._value[3:],
        )

    def DoCalcOutput(self, context, output):
        """Constructs the output values from the sliders."""
        self._update_values()
        output.set_value(self._get_transform())

    def Run(self, publishing_system, root_context, callback):
        # Calls callback(root_context, pose), then publishing_system.Publish()
        # each time the sliders change value.
        if not running_as_notebook:
            return

        publishing_context = publishing_system.GetMyContextFromRoot(root_context)

        print("Press the 'Stop PoseSliders' button in Meshcat to continue.")
        self._meshcat.AddButton("Stop PoseSliders")
        while self._meshcat.GetButtonClicks("Stop PoseSliders") < 1:
            if self._update_values():
                callback(root_context, self._get_transform())
                publishing_system.Publish(publishing_context)
            time.sleep(0.1)

        self._meshcat.DeleteButton("Stop PoseSliders")


class WsgButton(LeafSystem):
    def __init__(self, meshcat):
        LeafSystem.__init__(self)
        port = self.DeclareVectorOutputPort("wsg_position", 1, self.DoCalcOutput)
        port.disable_caching_by_default()
        self._meshcat = meshcat
        self._button = "Open/Close Gripper"
        meshcat.AddButton(self._button)

    def __del__(self):
        self._meshcat.DeleteButton(self._button)

    def DoCalcOutput(self, context, output):
        position = 0.107  # open
        if (self._meshcat.GetButtonClicks(self._button) % 2) == 1:
            position = 0.002  # close
        output.SetAtIndex(0, position)


# TODO(russt): Add floating base support (as in manipulation.jupyter_widgets
# MakeJointSlidersThatPublish)
class MeshcatJointSliders(LeafSystem):
    """
    Adds one slider per joint of the MultibodyPlant.  Any positions that are
    not associated with joints (e.g. floating-base "mobilizers") are held
    constant at the default value obtained from robot.CreateDefaultContext().

    .. pydrake_system::

        name: JointSliders
        output_ports:
        - positions

    In addition to being used inside a Diagram that is being simulated with
    Simulator, this class also offers a `Run` method that runs its own simple
    event loop, querying the slider values and calling `Publish`.  It does not
    simulate any state dynamics.
    """

    def __init__(
        self,
        meshcat,
        plant,
        root_context=None,
        lower_limit=-10.0,
        upper_limit=10.0,
        resolution=0.01,
    ):
        """
        Creates an meshcat slider for each joint in the plant.

        Args:
            meshcat:      A Meshcat instance.
            plant:        A MultibodyPlant. publishing_system: The System whose
                          Publish method will be called.  Can be the entire
                          Diagram, but can also be a subsystem.
            root_context: A mutable root Context of the Diagram containing the
                          ``plant``; we will extract the subcontext's using
                          `GetMyContextFromRoot`.
            lower_limit:  A scalar or vector of length robot.num_positions().
                          The lower limit of the slider will be the maximum
                          value of this number and any limit specified in the
                          Joint.

            upper_limit:  A scalar or vector of length robot.num_positions().
                          The upper limit of the slider will be the minimum
                          value of this number and any limit specified in the
                          Joint.

            resolution:   A scalar or vector of length robot.num_positions()
                          that specifies the step argument of the FloatSlider.
        """
        LeafSystem.__init__(self)

        def _broadcast(x, num):
            x = np.array(x)
            assert len(x.shape) <= 1
            return np.array(x) * np.ones(num)

        lower_limit = _broadcast(lower_limit, plant.num_positions())
        upper_limit = _broadcast(upper_limit, plant.num_positions())
        resolution = _broadcast(resolution, plant.num_positions())

        self._meshcat = meshcat
        self._plant = plant
        plant_context = plant.GetMyContextFromRoot(root_context) if root_context else plant.CreateDefaultContext()

        self._sliders = {}
        self._positions = plant.GetPositions(plant_context)
        slider_num = 0
        for i in range(plant.num_joints()):
            joint = plant.get_joint(JointIndex(i))
            low = joint.position_lower_limits()
            upp = joint.position_upper_limits()
            for j in range(joint.num_positions()):
                index = joint.position_start() + j
                description = joint.name()
                if joint.num_positions() > 1:
                    description += "_" + joint.position_suffix(j)
                meshcat.AddSlider(
                    value=self._positions[index],
                    min=max(low[j], lower_limit[slider_num]),
                    max=min(upp[j], upper_limit[slider_num]),
                    step=resolution[slider_num],
                    name=description,
                )
                self._sliders[index] = description
                slider_num += 1

        port = self.DeclareVectorOutputPort("positions", plant.num_positions(), self.DoCalcOutput)
        port.disable_caching_by_default()

    def DoCalcOutput(self, context, output):
        output.SetFromVector(self._positions)
        for i, s in self._sliders.items():
            output[i] = self._meshcat.GetSliderValue(s)

    def Run(self, publishing_system, root_context, callback=None):
        """
        Args:
            publishing_system:  The system to call publish on.  Probably a
                          MeshcatVisualizer.
            root_context: A mutable root Context of the Diagram containing both
                          the ``plant`` and the ``publishing_system``; we will
                          extract the subcontext's using `GetMyContextFromRoot`.
            callback: callback(plant_context) will be called whenever the
                      slider values change.
        """
        if not running_as_notebook:
            return
        print("Press the 'Stop JointSliders' button in Meshcat to continue.")
        self._meshcat.AddButton("Stop JointSliders")

        plant_context = self._plant.GetMyContextFromRoot(root_context)
        publishing_context = publishing_system.GetMyContextFromRoot(root_context)

        publishing_system.Publish(publishing_context)
        while self._meshcat.GetButtonClicks("Stop JointSliders") < 1:
            old_positions = self._plant.GetPositions(plant_context)
            positions = self._positions
            for i, s in self._sliders.items():
                positions[i] = self._meshcat.GetSliderValue(s)
            if not np.array_equal(positions, old_positions):
                self._plant.SetPositions(plant_context, positions)
                if callback:
                    callback(plant_context)
                publishing_system.Publish(publishing_context)
            time.sleep(0.1)

        self._meshcat.DeleteButton("Stop JointSliders")


def AddMeshcatVector(meshcat, path, color, length=0.25, radius=0.01, opacity=1.0, X_PT=RigidTransform()):

    meshcat.SetTransform(path, X_PT)
    meshcat.SetObject(path, Cylinder(radius, length), color)


def AddMeshcatTriad(meshcat, path, length=0.25, radius=0.01, opacity=1.0, X_PT=RigidTransform()):
    meshcat.SetTransform(path, X_PT)
    # x-axis
    X_TG = RigidTransform(RotationMatrix.MakeYRotation(np.pi / 2), [length / 2.0, 0, 0])
    meshcat.SetTransform(path + "/x-axis", X_TG)
    meshcat.SetObject(path + "/x-axis", Cylinder(radius, length), Rgba(1, 0, 0, opacity))

    # y-axis
    X_TG = RigidTransform(RotationMatrix.MakeXRotation(np.pi / 2), [0, length / 2.0, 0])
    meshcat.SetTransform(path + "/y-axis", X_TG)
    meshcat.SetObject(path + "/y-axis", Cylinder(radius, length), Rgba(0, 1, 0, opacity))

    # z-axis
    X_TG = RigidTransform([0, 0, length / 2.0])
    meshcat.SetTransform(path + "/z-axis", X_TG)
    meshcat.SetObject(path + "/z-axis", Cylinder(radius, length), Rgba(0, 0, 1, opacity))


def draw_open3d_point_cloud(meshcat, path, pcd, normals_scale=0.0, point_size=0.001, sphere_vis=False):
    pts = np.asarray(pcd.points)
    colors = np.asarray(pcd.colors)
    # assert pcd.has_colors()  # TODO(russt): handle this case better

    if sphere_vis:
        for idx, pt in enumerate(pts):
            meshcat_geom = Sphere(point_size)
            color = colors[idx]
            point_path = path + f"_point_{idx}"
            meshcat.SetObject(point_path, Sphere(point_size), Rgba(color[0], color[1], color[2], 1.0))
            transform = np.eye(4)
            transform[:3, 3] = pt
            meshcat.SetTransform(point_path, RigidTransform(pt))

    else:
        cloud = PointCloud(pts.shape[0], Fields(BaseField.kXYZs | BaseField.kRGBs))
        cloud.mutable_xyzs()[:] = pts.T
        cloud.mutable_rgbs()[:] = 255 * np.asarray(pcd.colors).T
        meshcat.SetObject(path, cloud, point_size=point_size)

    if pcd.has_normals() and normals_scale > 0.0:
        assert "need to implement LineSegments in meshcat c++"
        normals = np.asarray(pcd.normals)
        vertices = np.hstack((pts, pts + normals_scale * normals)).reshape(-1, 3).T
        meshcat["normals"].set_object(g.Points(g.PointsGeometry(vertices), g.MeshBasicMaterial(color=0x000000)))
        # .set_object(meshcat_geometry.PointCloud(cloud.T, color.T))


# Writing System.
class MeshcatWriter(LeafSystem):
    def __init__(
        self,
        meshcat,
        drawing_body_index,
        canvas_body_index,
        rgba,
        line_width=0.003,
        draw_threshold=0.01,
    ):
        LeafSystem.__init__(self)
        self._meshcat = meshcat
        self._drawing_body_index = drawing_body_index
        self._canvas_body_index = canvas_body_index
        self._rgba = rgba
        self._line_width = line_width
        self._draw_threshold = draw_threshold

        self.DeclareAbstractInputPort("contact_results", AbstractValue.Make(ContactResults()))
        self.DeclarePeriodicDiscreteUpdateEvent(0.01, 0, self.MaybeDraw)

        self._p_WLastDraw_index = self.DeclareDiscreteState(3)
        self._was_in_contact_index = self.DeclareDiscreteState([0])
        self._num_drawn_index = self.DeclareDiscreteState([0])

    def MaybeDraw(self, context, discrete_state):
        results = self.get_input_port().Eval(context)

        # The point on the canvas that most deeply penetrates the chalk.
        p_WDraw = None
        for i in range(results.num_point_pair_contacts()):
            info = results.point_pair_contact_info(i)
            if info.bodyA_index() == self._drawing_body_index and info.bodyB_index() == self._canvas_body_index:
                p_WDraw = info.point_pair().p_WCb
                break
            elif info.bodyB_index() == self._drawing_body_index and info.bodyA_index() == self._canvas_body_index:
                p_WDraw = info.point_pair().p_WCa
                break

        if p_WDraw is not None:
            p_WLastDraw = context.get_discrete_state(self._p_WLastDraw_index).get_value()
            was_in_contact = context.get_discrete_state(self._was_in_contact_index)[0]
            num_drawn = context.get_discrete_state(self._num_drawn_index)[0]

            length = np.linalg.norm(p_WDraw - p_WLastDraw)
            if was_in_contact and length > self._draw_threshold:
                self._meshcat.SetObject(f"writer/{num_drawn}", Capsule(self._line_width, length), self._rgba)
                p_WMidpoint = (p_WDraw + p_WLastDraw) / 2
                X_WCapsule = RigidTransform(
                    RotationMatrix.MakeFromOneVector(p_WDraw - p_WLastDraw, 2),
                    p_WMidpoint,
                )
                self._meshcat.SetTransform(f"writer/{num_drawn}", X_WCapsule)
                discrete_state.set_value(self._num_drawn_index, [num_drawn + 1])
                discrete_state.set_value(self._p_WLastDraw_index, p_WDraw)
            elif not was_in_contact:
                discrete_state.set_value(self._p_WLastDraw_index, p_WDraw)

            discrete_state.set_value(self._was_in_contact_index, [1])
        else:
            discrete_state.set_value(self._was_in_contact_index, [0])


def plot_surface(
    meshcat,
    path,
    X,
    Y,
    Z,
    rgba=Rgba(0.87, 0.6, 0.6, 1.0),
    wireframe=False,
    wireframe_line_width=1.0,
):
    (rows, cols) = Z.shape
    assert np.array_equal(X.shape, Y.shape)
    assert np.array_equal(X.shape, Z.shape)

    vertices = np.empty((rows * cols, 3), dtype=np.float32)
    vertices[:, 0] = X.reshape((-1))
    vertices[:, 1] = Y.reshape((-1))
    vertices[:, 2] = Z.reshape((-1))

    # Vectorized faces code from https://stackoverflow.com/questions/44934631/making-grid-triangular-mesh-quickly-with-numpy  # noqa
    faces = np.empty((rows - 1, cols - 1, 2, 3), dtype=np.uint32)
    r = np.arange(rows * cols).reshape(rows, cols)
    faces[:, :, 0, 0] = r[:-1, :-1]
    faces[:, :, 1, 0] = r[:-1, 1:]
    faces[:, :, 0, 1] = r[:-1, 1:]
    faces[:, :, 1, 1] = r[1:, 1:]
    faces[:, :, :, 2] = r[1:, :-1, None]
    faces.shape = (-1, 3)

    # TODO(Russ): support per vertex / Colormap colors.
    meshcat.SetTriangleMesh(path, vertices.T, faces.T, rgba, wireframe, wireframe_line_width)


def plot_mathematical_program(meshcat, path, prog, X, Y, result=None, point_size=0.05):
    assert prog.num_vars() == 2
    assert X.size == Y.size

    N = X.size
    values = np.vstack((X.reshape(-1), Y.reshape(-1)))
    costs = prog.GetAllCosts()

    # Vectorized multiply for the quadratic form.
    # Z = (D*np.matmul(Q,D)).sum(0).reshape(nx, ny)

    if costs:
        Z = prog.EvalBindingVectorized(costs[0], values)
        for b in costs[1:]:
            Z = Z + prog.EvalBindingVectorized(b, values)

    cv = f"{path}/constraints"
    for binding in prog.GetAllConstraints():
        if isinstance(binding.evaluator(), BoundingBoxConstraint):
            c = binding.evaluator()
            var_indices = [int(prog.decision_variable_index()[v.get_id()]) for v in binding.variables()]
            satisfied = np.array(c.CheckSatisfiedVectorized(values[var_indices, :], 0.001)).reshape(1, -1)
            if costs:
                Z[~satisfied] = np.nan

            v = f"{cv}/{type(c).__name__}"
            Zc = np.zeros(Z.shape)
            Zc[satisfied] = np.nan
            plot_surface(
                meshcat,
                v,
                X,
                Y,
                Zc.reshape(X.shape),
                rgba=Rgba(1.0, 0.2, 0.2, 1.0),
                wireframe=True,
            )
        else:
            Zc = prog.EvalBindingVectorized(binding, values)
            evaluator = binding.evaluator()
            low = evaluator.lower_bound()
            up = evaluator.upper_bound()
            cvb = f"{cv}/{type(evaluator).__name__}"
            for index in range(Zc.shape[0]):
                # TODO(russt): Plot infeasible points in a different color.
                infeasible = np.logical_or(Zc[index, :] < low[index], Zc[index, :] > up[index])
                plot_surface(
                    meshcat,
                    f"{cvb}/{index}",
                    X,
                    Y,
                    Zc[index, :].reshape(X.shape),
                    rgba=Rgba(1.0, 0.3, 1.0, 1.0),
                    wireframe=True,
                )

    if costs:
        plot_surface(
            meshcat,
            f"{path}/objective",
            X,
            Y,
            Z.reshape(X.shape),
            rgba=Rgba(0.3, 1.0, 0.3, 1.0),
            wireframe=True,
        )

    if result:
        v = f"{path}/solution"
        meshcat.SetObject(v, Sphere(point_size), Rgba(0.3, 1.0, 0.3, 1.0))
        x_solution = result.get_x_val()
        meshcat.SetTransform(v, RigidTransform([x_solution[0], x_solution[1], result.get_optimal_cost()]))

env/meshcat_cpp_utils.py
liruiw/Fleet-Tools
import argparse

"""
Util to compile command line arguments for core script
"""


def get_parser():
    # Global Parameters
    parser = argparse.ArgumentParser(fromfile_prefix_chars="@", description="Arguments")
    parser.add_argument("--env_name", default="FrankaDrakeSpatulaEnv", help="Choice of environment.")
    parser.add_argument("--logdir", default="runs", help="exterior log directory")
    parser.add_argument("--logdir_suffix", default="", help="log directory suffix")
    parser.add_argument("--cuda", action="store_true", help="run on CUDA (default: False)")

    parser.add_argument("--teleop", action="store_true", help="use teleop control")
    parser.add_argument(
        "--eval",
        default="",
        help="if doing eval, this is the logdir from which to load model weights",
    )
    parser.add_argument("--seed", type=int, default=123456, help="random seed (default: 123456)")

    # rollout related
    parser.add_argument("--agent", type=str, default="BC", help="Type of parallel agent;")
    parser.add_argument(
        "--save_demo_suffix",
        type=str,
        default="",
        help="The suffix for which demo to save",
    )
    parser.add_argument("--expert", type=str, default="AnalyticExpert", help="Type of parallel expert;")
    parser.add_argument("--run_expert", action="store_true", help="use expert for action")
    parser.add_argument(
        "--teleop_type",
        type=str,
        default="keyboard",
        help="Type of teleop method, mouse, keyboard, vr",
    )

    parser.add_argument("--num_envs", type=int, default=10, help="number of robots")
    parser.add_argument("--log_freq", type=int, default=100, help="log frequency")
    parser.add_argument("--render", action="store_true", help="whether or not to render")
    parser.add_argument("--meshcat", action="store_true", help="use meshcat")
    parser.add_argument("--visdom", action="store_true", help="use visdom")
    parser.add_argument("--num_steps", type=int, default=100000, help="maximum number of timesteps ")
    parser.add_argument("--num_episode", type=int, default=1000, help="maximum number of episodes")

    parser.add_argument("--pretrained", type=str, default="", help="pretrained model to load")
    parser.add_argument(
        "--load_pretrained_model",
        type=str,
        default="both",
        help="pretrained policy or representation to load",
    )
    parser.add_argument("--demo_number", type=int, default=1000, help="number of demonstrations needed")
    parser.add_argument("--num_workers", type=int, default=2, help="number of workers for dataset")

    # train eval related
    parser.add_argument("--train_task", type=str, default="", help="the train / eval task")
    parser.add_argument(
        "--num_train_steps",
        type=int,
        default=10000,
        help="maximum number of training timesteps (default: 10000)",
    )
    parser.add_argument("--training", action="store_true", help="training from offline data")
    parser.add_argument(
        "--test_on_train_scenes",
        action="store_true",
        help="load the training scenes for testing",
    )
    parser.add_argument(
        "--save_demonstrations",
        action="store_true",
        help="saving the training scenes for testing",
    )
    parser.add_argument("--load_demonstrations", action="store_true", help="load demonstration data")
    parser.add_argument(
        "--record_video",
        action="store_true",
        help="record video fro either expert rollout or policy rollout",
    )
    parser.add_argument(
        "--config_suffix",
        type=str,
        default="",
        help="suffix for the training / env config",
    )
    parser.add_argument(
        "--demonstration_dir",
        type=str,
        default="",
        help="the directory where demonstrations are stored",
    )
    parser.add_argument("--fixed_encoder", action="store_true", help="fix the representation encoder")
    parser.add_argument(
        "--start_episode_position",
        type=int,
        default=0,
        help="loaded episode position for the replay buffer",
    )
    parser.add_argument("--save_dir", type=str, default="", help="the directory for the log experiment")
    parser.add_argument(
        "--eval_save_video",
        action="store_true",
        help="save the overhead view video",
        default=True,
    )
    parser.add_argument(
        "--reinit_num",
        type=int,
        default=40,
        help="number of rounds before reiniting ray",
    )

    return parser

experiments/args.py
liruiw/Fleet-Tools
import math
from torch.optim.lr_scheduler import MultiStepLR
import argparse
import os
import IPython
import numpy as np
import cv2
import copy
import json


def process_image(img):
    """process RGB image"""
    color_img = img[..., :3]
    img = cv2.resize(img, IMG_SIZE)
    return img


def load_image(path):
    wrist_img_color = cv2.imread(f"{path}_wrist_color.png", cv2.IMREAD_UNCHANGED)[:, :, [2, 1, 0]] / 255.0
    wrist_img_depth = cv2.imread(f"{path}_wrist_depth.png", cv2.IMREAD_UNCHANGED)[..., None] / 1000.0
    wrist_img_mask = cv2.imread(f"{path}_wrist_mask.png", cv2.IMREAD_UNCHANGED)[..., None]

    overhead_img_color = cv2.imread(f"{path}_overhead_color.png", cv2.IMREAD_UNCHANGED)[:, :, [2, 1, 0]] / 255.0
    overhead_img_depth = cv2.imread(f"{path}_overhead_depth.png", cv2.IMREAD_UNCHANGED)[..., None] / 1000.0
    overhead_img_mask = cv2.imread(f"{path}_overhead_mask.png", cv2.IMREAD_UNCHANGED)[..., None]

    wrist_img_obs = np.concatenate((wrist_img_color, wrist_img_depth, wrist_img_mask), axis=-1)
    overhead_img_obs = np.concatenate(
        (overhead_img_color, overhead_img_depth, overhead_img_mask),
        axis=-1,
    )
    overhead_img_obs = process_image(overhead_img_obs)
    wrist_img_obs = process_image(wrist_img_obs)

    new_data_obs = [overhead_img_obs, wrist_img_obs]
    return np.array(new_data_obs)


def mkdir_if_missing(dst_dir):
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)


parser = argparse.ArgumentParser()
parser.add_argument("--saved_path", "-s", type=str, default="assets/demonstrations")
parser.add_argument("--env_name", "-e", type=str, default="wrench")
parser.add_argument("--tool_idx", "-t", type=int, default=1)
parser.add_argument("--max_num", type=int, default=100000)
parser.add_argument("--output_path", "-o", type=str, default="assets/demonstrations/processed")
parser.add_argument("--processed_output_path", type=str, default="../metaworld_weights/data/demo_drake/")

args = parser.parse_args()

IMG_SIZE = (128, 128)
saved_path = os.path.join(args.saved_path, args.env_name + f"tool_{args.tool_idx}")
output_path = os.path.join(args.output_path, args.env_name + f"tool_{args.tool_idx}")
mkdir_if_missing(saved_path)
mkdir_if_missing(output_path)

TEST_EPS_NUM = 200

# read the offline folder and just save state / action
episode_dirs = sorted(os.listdir(saved_path), key=lambda s: int(s.split("_")[1]))[:-TEST_EPS_NUM]
# note! the last ten is used for testing
MAX_NUM = args.max_num
USE_IMAGE = False
images = []
traj_datas = {}
traj_datas["tool_name"] = []
traj_datas["object_name"] = []
traj_datas["task_name"] = []

for path in episode_dirs:
    curr_num = len(images)
    episode_path = os.path.join(saved_path, path)
    if "traj_data.npz" in os.listdir(episode_path):
        traj_data = np.load(os.path.join(episode_path, "traj_data.npz"))
        traj_data = dict(traj_data)
        scene_descriptions = json.load(open(os.path.join(episode_path, "scene_descriptions.json")))
        # traj_data["global_scene_index"] =
        traj_data["tool_name"] = [scene_descriptions["tool_name"]] * len(traj_data["action"])
        traj_data["object_name"] = [scene_descriptions["object_name"]] * len(traj_data["action"])
        traj_data["task_name"] = [saved_path[saved_path.find("for") + len("for") : saved_path.rfind("tool")]] * len(
            traj_data["action"]
        )

        try:
            traj_datas = {
                k: np.concatenate((traj_datas[k], v), axis=0) if k in traj_datas else v for k, v in traj_data.items()
            }
        except:
            print("load trajectory folder: {} failed".format(episode_path))
            traj_datas = {k: v[:curr_num] for k, v in traj_data.items()}
            images = images[:curr_num]

    if len(traj_datas["action"]) > MAX_NUM:
        break

if USE_IMAGE:
    images = np.array(images)
    wrist_view_only = False
    if wrist_view_only:
        images = images[:, 1]  # , ..., :3]
    else:
        images = np.concatenate((images[:, 0], images[:, 1]), axis=-1)

# read images
mkdir_if_missing(f"{ output_path}/")
print(f"save to {output_path}")
print("number of saved trajectories:", len(episode_dirs), traj_datas["action"].shape, path)

key_list = [
    "state",
    "action",
    "reward",
    "next_state",
    "done",
    "obs",
    "timestep",
    "joint_pos",
    "hand_camera_pose",
    "base_camera_pose",
    "ee_pose",
    "cam_intr",
    "curr_keypoint",
    "goal_keypoint",
    "goal_pose",
    "keypoint_wrench",
    "external_torque",
    "endeffector_wrench",
    "history_endeffector_wrench",
    "point_cloud",
    "tool_rel_pose",
    "img_feature",
]

# for idx, combined_pc in enumerate(traj_datas["point_cloud"][::100]):
#     pcd = o3d.geometry.PointCloud()
#     pcd.points = o3d.utility.Vector3dVector(combined_pc[:3].T)
#     pc_colors = np.zeros_like(pcd.points)
#     pc_colors[combined_pc.T[:, -1] == 1] = (1,0,0)
#     pc_colors[combined_pc.T[:, -1] == 0] = (0,1,0)
#     pcd.colors = o3d.utility.Vector3dVector(pc_colors)

#     o3d.visualization.draw_geometries([pcd]) #

data_dict = {
    "tool_name": traj_datas["tool_name"],
    "task_name": traj_datas["task_name"],
    "object_name": traj_datas["object_name"],
}

for key in key_list:
    data_dict[key] = traj_datas[key]

np.savez(f"{ output_path}/demo_state.npz", **data_dict)
print("saving offline dataset on:", f"{output_path}/demo_state.npz")
output_path2 = f"{args.processed_output_path}/FrankaDrake{args.env_name.capitalize()}Env-Tool{args.tool_idx}"
mkdir_if_missing(f"{output_path2}/")

# save in two places
np.savez(f"{output_path2}/demo_state.npz", **data_dict)
print("saving offline dataset on:", f"{ output_path2}/demo_state.npz")

scripts/collapse_dataset.py
liruiw/Fleet-Tools
import argparse
import logging
from lxml import etree as ET
import numpy as np
import os
import trimesh
import sys

import typing
import numpy as np
from scipy import linalg

import trimesh
import trimesh.remesh
from trimesh.visual.material import SimpleMaterial
import open3d
import IPython
import open3d as o3d
import json
from PIL import Image
from generate_tool_sdf_from_mesh import *

EPS = 10e-10


 

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate SDF from a mesh.")
    parser.add_argument("mesh_file", type=str, help="Path to mesh file.")
    parser.add_argument("--preview", default=False, action="store_true", help="Preview decomp with a trimesh window?")
    parser.add_argument(
        "--scale", type=float, default=1, help="Scale factor to convert the specified mesh's coordinates to meters."
    )
    parser.add_argument("--size", type=float, default=0.15, help="Diameter for the output mesh.")
    parser.add_argument(
        "--density", type=float, default=2, help="Assumed density in kg/m^3 of object, for inertia calculation."
    )
    parser.add_argument(
        "--do_visual_simplification",
        default=False,
        action="store_true",
        help="Do additional visual simplification of mesh. Requires open3d. Probably won't preserve materials.",
    )
    parser.add_argument(
        "--target_tris",
        type=int,
        default=1000,
        help="If we do visual simplification, we decimate to this # of triangles.",
    )
    parser.add_argument("--resolution", type=int, default=100000, help="VHACD voxel resolution.")
    parser.add_argument("--maxhulls", type=int, default=50, help="VHACD max # of convex hulls.")
    parser.add_argument("--minVolumePerCH", type=float, default=0.001, help="VHACD min convex hull volume.")
    parser.add_argument("--maxNumVerticesPerCH", type=int, default=50, help="VHACD voxel resolution.")
    parser.add_argument(
        "--loglevel",
        type=str,
        default="INFO",
        choices=["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG"],
        help="Log level.",
    )
    parser.add_argument("--prefix", type=str, default="")
    # manually tuned
    parser.add_argument(
        "--up", type=int, default=-10, help="Scale factor to convert the specified mesh's coordinates to meters."
    )
    parser.add_argument("--forward", type=int, default=-10, help="Diameter for the output mesh.")
    parser.add_argument("--flip_z", default=False, action="store_true", help="Preview decomp with a trimesh window?")
    parser.add_argument(
        "--select_keypoint", default=False, action="store_true", help="Preview decomp with a trimesh window?"
    )
    parser.add_argument("--additional_transform", type=str, default="")
    parser.add_argument("--rename_idx", type=str, default="")
    parser.add_argument("--no_convex_decomp", action="store_true", help="Preview decomp with a trimesh window?")
    # parser.add_argument('--resize', action="store_true", help="Preview decomp with a trimesh window?")

    args = parser.parse_args()
    logging.basicConfig(level=args.loglevel)
    input_mesh_path = os.path.abspath(args.mesh_file)

    args.prefix = args.mesh_file.split("/")[2] + "_"  # add prefix
    dir_path, mesh_filename = os.path.split(input_mesh_path)
    mesh_minus_ext, _ = os.path.splitext(mesh_filename)
    sdf_path = os.path.join(dir_path, mesh_minus_ext + ".sdf")

    # Load in and prescale mesh.
    obj = Object3D(input_mesh_path)
    obj.recompute_normals()
    open3d_mesh = obj.meshes[0].as_open3d
    open3d_mesh.compute_vertex_normals()

    o3d.io.write_triangle_mesh(input_mesh_path.replace(".obj", "_open3d.obj"), open3d_mesh, write_vertex_colors=True)
    obj.meshes[0].export(input_mesh_path, include_texture=True, include_normals=True)
    print("save mesh:", input_mesh_path)

scripts/generate_normals_for_mesh.py
liruiw/Fleet-Tools
import argparse
import logging
from lxml import etree as ET
import numpy as np
import os
import trimesh
import sys


import typing

import numpy as np
from scipy import linalg

import trimesh
import trimesh.remesh
from trimesh.visual.material import SimpleMaterial
import open3d
from generate_tool_sdf_from_mesh import *


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate SDF from a mesh.")
    parser.add_argument("mesh_file", type=str, help="Path to mesh file.")
    parser.add_argument("--preview", default=False, action="store_true", help="Preview decomp with a trimesh window?")
    parser.add_argument(
        "--scale", type=float, default=1, help="Scale factor to convert the specified mesh's coordinates to meters."
    )
    parser.add_argument("--size", type=float, default=0.1, help="Diameter for the output mesh.")
    parser.add_argument(
        "--density", type=float, default=200, help="Assumed density in kg/m^3 of object, for inertia calculation."
    )
    parser.add_argument(
        "--do_visual_simplification",
        default=False,
        action="store_true",
        help="Do additional visual simplification of mesh. Requires open3d. Probably won't preserve materials.",
    )
    parser.add_argument(
        "--target_tris",
        type=int,
        default=1000,
        help="If we do visual simplification, we decimate to this # of triangles.",
    )
    parser.add_argument("--resolution", type=int, default=100000, help="VHACD voxel resolution.")
    parser.add_argument("--maxhulls", type=int, default=50, help="VHACD max # of convex hulls.")
    parser.add_argument("--minVolumePerCH", type=float, default=0.001, help="VHACD min convex hull volume.")
    parser.add_argument("--maxNumVerticesPerCH", type=int, default=50, help="VHACD voxel resolution.")
    parser.add_argument(
        "--loglevel",
        type=str,
        default="INFO",
        choices=["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG"],
        help="Log level.",
    )
    parser.add_argument("--prefix", type=str, default="")
    # manually tuned
    parser.add_argument(
        "--up", type=int, default=-10, help="Scale factor to convert the specified mesh's coordinates to meters."
    )
    parser.add_argument("--forward", type=int, default=-10, help="Diameter for the output mesh.")
    parser.add_argument("--flip_z", default=False, action="store_true", help="Preview decomp with a trimesh window?")
    parser.add_argument(
        "--select_keypoint", default=False, action="store_true", help="Preview decomp with a trimesh window?"
    )
    parser.add_argument("--additional_transform", type=str, default="")
    parser.add_argument("--rename_idx", type=str, default="")
    parser.add_argument("--no_convex_decomp", action="store_true", help="Preview decomp with a trimesh window?")

    parser.add_argument(
        "--no_resize",
        default=False,
        action="store_true",
        help="Avoid resizing?",
    )
    parser.add_argument(
        "--no_recenter",
        default=False,
        action="store_true",
        help="Avoid recentering?",
    )

    args = parser.parse_args()
    logging.basicConfig(level=args.loglevel)
    mesh_path = os.path.abspath(args.mesh_file)

    if len(args.rename_idx) > 0:
        mesh_path = rename_model_normalized_folder(args.mesh_file, args.mesh_file.split("/")[-1], int(args.rename_idx))

    if not os.path.exists(mesh_path):
        logging.error("No mesh found at %s" % mesh_path)
        sys.exit(-1)

    gen_sdf_func = create_sdf_with_mesh if args.no_convex_decomp else create_sdf_with_convex_decomp
    obj = gen_sdf_func(
        mesh_path,
        scale=args.scale,
        size=args.size,
        do_visual_simplification=args.do_visual_simplification,
        target_tris=args.target_tris,
        preview_with_trimesh=args.preview,
        density=args.density,
        resolution=args.resolution,
        maxhulls=args.maxhulls,
        maxNumVerticesPerCH=args.maxNumVerticesPerCH,
        minVolumePerCH=args.minVolumePerCH,
        pca=1,
        prefix=args.prefix,
        flip_z=args.flip_z,
        additional_transform=args.additional_transform,
        resize=not args.no_resize,
        recenter=not args.no_recenter,
    )
    input_mesh_path = mesh_path.replace(".obj", "_corr.obj")
    obj = Object3D(input_mesh_path)
    obj.recompute_normals()
    open3d_mesh = obj.meshes[0].as_open3d
    open3d_mesh.compute_vertex_normals()

    o3d.io.write_triangle_mesh(input_mesh_path.replace(".obj", "_open3d.obj"), open3d_mesh, write_vertex_colors=True)
    obj.meshes[0].export(input_mesh_path, include_texture=True, include_normals=True)

    if args.select_keypoint:
        mesh = o3d.io.read_triangle_mesh(input_mesh_path)
        pcd = mesh.sample_points_uniformly(number_of_points=10000)

        vis = o3d.visualization.VisualizerWithEditing()
        vis.create_window()
        vis.add_geometry(pcd)
        vis.run()  # user picks points
        vis.destroy_window()
        print("")

        # vis = o3d.visualization.draw_geometries_with_editing([pcd])
        picked_points = vis.get_picked_points()
        picked_points = np.unique(picked_points)
        xyz = np.asarray(pcd.points)
        picked_points = xyz[picked_points]

        pointcloudfile = mesh_path.replace(".obj", "_points.npy")
        np.save(pointcloudfile, xyz)
        keypoint_description_file = mesh_path.replace(".obj", "_info.json")
        keypoint_info = {"keypoints": picked_points.tolist(), "extents": obj.extents.tolist()}
        with open(keypoint_description_file, "w") as f:
            json.dump(keypoint_info, f, indent=4, sort_keys=True)

scripts/generate_object_sdf_from_mesh.py
liruiw/Fleet-Tools
import argparse
import logging
from lxml import etree as ET
import numpy as np
import os
import trimesh
import sys

import typing
import numpy as np
from scipy import linalg

import trimesh
import trimesh.remesh
from trimesh.visual.material import SimpleMaterial
import open3d
import IPython
import open3d as o3d
import json
from PIL import Image

EPS = 10e-10


def compute_vertex_normals(vertices, faces):
    normals = np.ones_like(vertices)
    triangles = vertices[faces]
    triangle_normals = np.cross(triangles[:, 1] - triangles[:, 0], triangles[:, 2] - triangles[:, 0])
    triangle_normals /= linalg.norm(triangle_normals, axis=1)[:, None] + EPS
    normals[faces[:, 0]] += triangle_normals
    normals[faces[:, 1]] += triangle_normals
    normals[faces[:, 2]] += triangle_normals
    normals /= linalg.norm(normals, axis=1)[:, None] + 0
    return normals


def are_trimesh_normals_corrupt(trimesh):
    corrupt_normals = linalg.norm(trimesh.vertex_normals, axis=1) == 0.0
    return corrupt_normals.sum() > 0


def subdivide_mesh(mesh):
    attributes = {}
    if hasattr(mesh.visual, "uv"):
        attributes = {"uv": mesh.visual.uv}
    vertices, faces, attributes = trimesh.remesh.subdivide(mesh.vertices, mesh.faces, attributes=attributes)
    mesh.vertices = vertices
    mesh.faces = faces
    if "uv" in attributes:
        mesh.visual.uv = attributes["uv"]
    return mesh


class Object3D(object):
    """Represents a graspable object."""

    def __init__(self, path, load_materials=False):
        scene = trimesh.load(
            str(path),
            ignore_mtl=not load_materials,
            skip_materials=True,
            force="mesh",
            default_material=SimpleMaterial(color=(1.0, 0.0, 1.0)),
        )
        # if isinstance(scene, trimesh.Trimesh):
        #     scene = trimesh.Scene(scene)

        self.meshes: typing.List[trimesh.Trimesh] = [scene]  # list(scene.dump())
        self.recompute_normals()
        self.path = path
        self.scale = 1.0
        # self.fix()

    def fix(self):
        for idx, mesh in enumerate(self.meshes):
            trimesh.repair.fill_holes(self.meshes[idx])
            trimesh.repair.fix_inversion(self.meshes[idx])
            trimesh.repair.fix_normals(self.meshes[idx])
            self.meshes[idx].remove_unreferenced_vertices()
            self.meshes[idx].remove_degenerate_faces()
            self.meshes[idx].remove_duplicate_faces()

    def to_scene(self):
        return trimesh.Scene(self.meshes)

    def are_normals_corrupt(self):
        for mesh in self.meshes:
            if are_trimesh_normals_corrupt(mesh):
                return True

        return False

    def recompute_normals(self):
        for mesh in self.meshes:
            mesh.vertex_normals = compute_vertex_normals(mesh.vertices, mesh.faces)

        return self

    def rescale(self, scale=1.0):
        """Set scale of object mesh.
        :param scale
        """
        self.scale = scale
        for mesh in self.meshes:
            mesh.apply_scale(self.scale)

        return self

    def resize(self, size, ref="diameter"):
        """Set longest of all three lengths in Cartesian space.
        :param size
        """
        if ref == "diameter":
            ref_scale = self.bounding_diameter
        else:
            ref_scale = self.bounding_size

        self.scale = size / ref_scale
        for mesh in self.meshes:
            mesh.apply_scale(self.scale)

        return self

    @property
    def bounding_size(self):
        return max(self.extents)

    @property
    def bounding_diameter(self):
        centroid = self.bounds.mean(axis=0)
        max_radius = linalg.norm(self.vertices - centroid, axis=1).max()
        return max_radius * 2

    @property
    def bounding_radius(self):
        return self.bounding_diameter / 2.0

    @property
    def extents(self):
        min_dim = np.min(self.vertices, axis=0)
        max_dim = np.max(self.vertices, axis=0)
        return max_dim - min_dim

    @property
    def bounds(self):
        min_dim = np.min(self.vertices, axis=0)
        max_dim = np.max(self.vertices, axis=0)
        return np.stack((min_dim, max_dim), axis=0)

    def recenter(self, method="bounds"):
        if method == "mean":
            # Center the mesh.
            vertex_mean = np.mean(self.vertices, 0)
            translation = -vertex_mean
        elif method == "bounds":
            center = self.bounds.mean(axis=0)
            translation = -center
        else:
            raise ValueError(f"Unknown method {method!r}")

        for mesh in self.meshes:
            mesh.apply_translation(translation)

        return self

    def normalize_direction(self, flip_z=False):
        # make sure the longest dimension facing upward (+z) and the second longest dimension points at positive x

        up = np.zeros(3)
        forward = np.zeros(3)

        sorted_idx = np.argsort(self.extents)
        up[sorted_idx[-1]] = 1
        if flip_z:
            up[sorted_idx[-1]] = -1

        forward[sorted_idx[-2]] = 1

        side = np.cross(forward, up)
        # R = np.stack([side, up, -forward], axis=-1)
        R = np.stack([forward, side, up], axis=-1).T  # transpose
        transform = np.eye(4)
        transform[:3, :3] = R
        print("apply rotation", R)
        self.meshes[0] = self.meshes[0].apply_transform(transform)
        return self

    def apply_transform(self, trans):
        # make sure the longest dimension facing upward (+z) and the second longest dimension points at positive x

        self.meshes[0] = self.meshes[0].apply_transform(trans)
        return self

    @property
    def vertices(self):
        return np.concatenate([mesh.vertices for mesh in self.meshes])


def do_visual_mesh_simplification(input_mesh_path, target_tris=1000):
    """
    Given an obj path, simplifies the geometry to make it easier to render
    by creating a mesh alongside it with a "_simple_vis.<ext>" postfix.
    - Looks for a texture file at a fixed relative path

    Args:
    - input_mesh_path: String path to mesh file. Only 'obj' format is tested,
        but others might work.
    - target_tris: Currently unusued, but would be target for mesh decimation.
    Returns:
    - output obj file path
    """
    import open3d
    import cv2
    import imutils

    # TODO(gizatt) What gives, open3d is trashing my models...
    logging.warning("As of writing, this is sometimes creating terrible models.")

    mesh_minus_ext, mesh_ext = os.path.splitext(input_mesh_path)
    output_mesh_path = mesh_minus_ext + "_simple_vis" + mesh_ext

    mesh = open3d.io.read_triangle_mesh(input_mesh_path)
    mesh.compute_vertex_normals()
    simplified_mesh = mesh.simplify_quadric_decimation(target_tris)
    simplified_mesh.compute_vertex_normals()
    open3d.io.write_triangle_mesh(output_mesh_path, simplified_mesh)
    return output_mesh_path


def do_collision_mesh_simplification_coacd(mesh, mesh_name, mesh_dir, preview_with_trimesh=False, **kwargs):
    """
    Given a mesh, performs a convex decomposition of it with
    coacd, saving all the parts in a subfolder named
    `<mesh_filename>_parts`.

    Args:
    - input_mesh_path: String path to mesh file to decompose. Only
        'obj' format is currently tested, but other formats supported
        by trimesh might work.
    - preview_with_trimesh: Whether to open (and block on) a window to preview
    the decomposition.
    - A set of control kwargs, plus any additional kwargs, are passed to the convex
      decomposition routine 'vhacd'; you can run `testVHACD --help` to see options.

    Returns:
    - List of generated mesh file parts, in obj format.
    """
    # IPython.embed()

    # Create a subdir for the convex decomp parts.

    mesh_parts_folder = mesh_name + "_parts"
    out_dir = os.path.join(mesh_dir, mesh_parts_folder)
    mesh_full_path = os.path.join(mesh_dir, mesh_name)
    os.makedirs(out_dir, exist_ok=True)

    if len(os.listdir(out_dir)) > 0:
        print("WARNING: {} is not empty".format(out_dir))
        d = input("** DELETE? [y/n]")
        if d == 'y':
            for fn in os.listdir(out_dir):
                os.remove(os.path.join(out_dir, fn))
                print("\tRemoved {}...".format(fn))
        else:
            raise ValueError("{} should be empty".format(out_dir))

    if preview_with_trimesh:
        logging.info("Showing mesh before decomp. Close window to proceed.")
        # mesh.set_camera(angles=(1, 0, 0), distance=1, center=(0,0,0))
        # mesh.show(angles=(0, 1, 0), distance=1, center=(0,0,0))
        scene = trimesh.scene.scene.Scene()
        scene.add_geometry(mesh)
        scene.set_camera(angles=(1, 0, 0), distance=0.3, center=(0, 0, 0))
        # scene.viewer.toggle_axis()
        scene.show()
    try:
        convex_pieces = []
        os.system(f"coacd -i {mesh_full_path}_corr.obj -o {out_dir}")
        logging.info("Performing convex decomposition with CoACD.")

        for file in sorted(os.listdir(out_dir)):
            part = trimesh.load(os.path.join(out_dir, file))
            convex_pieces += [part]

    except Exception as e:
        logging.error("Problem performing decomposition: %s", e)

    if preview_with_trimesh:
        # Display the convex decomp, giving each a random colors
        # to make them easier to distinguish.
        for part in convex_pieces:
            this_color = trimesh.visual.random_color()
            part.visual.face_colors[:] = this_color
        scene = trimesh.scene.scene.Scene()
        for part in convex_pieces:
            scene.add_geometry(part)

        logging.info("Showing mesh convex decomp into %d parts. Close window to proceed." % (len(convex_pieces)))
        scene.set_camera(angles=(1, 0, 0), distance=0.3, center=(0, 0, 0))

        # viewer = trimesh.viewer.SceneViewer(scene)
        # viewer.init_gl()
        # viewer.on_draw()
        scene.show()

    # rewrite the mesh with new names
    os.system(f"rm {out_dir}/*")
    out_paths = []
    for k, part in enumerate(convex_pieces):
        piece_name = "%s_convex_piece_%03d.obj" % (mesh_name, k)
        full_path = os.path.join(out_dir, piece_name)

        trimesh.exchange.export.export_mesh(part, full_path)
        out_paths.append(full_path)
    return out_paths


def create_sdf_with_mesh(
    input_mesh_path,
    scale=1.0,
    do_visual_simplification=False,
    size=0.1,
    target_tris=1000,
    preview_with_trimesh=False,
    density=2000,
    recenter=True,
    resize=True,
    bound_type="extents",
    prefix="",
    flip_z=False,
    additional_transform="",
    **kwargs,
):
    """
    Given an input mesh file, produces an SDF if the same directory that:
    - Uses the mesh as its sole visual geometry.
    - Performs a convex decomposition of the mesh, and uses those pieces
    as the collision geometry.
    - Inserts inertia for the object, calculated from the original
    mesh assuming a constant density.

    The SDF is saved as `<mesh_file>.sdf` next to the mesh file.

    Args:
    - input_mesh_path: Path to the mesh file
    - preview_with_trimesh: Whether to show 3D previews of pre/post decomposition.
    - density: Assumed density of the object, in kg/m^3.
    - kwargs: Passed through to do_collision_mesh_simplification as convex decomp args.
    """
    # Get mesh name.
    dir_path, mesh_filename = os.path.split(input_mesh_path)
    mesh_minus_ext, _ = os.path.splitext(mesh_filename)
    sdf_path = os.path.join(dir_path, mesh_minus_ext + ".sdf")

    # Load in and prescale mesh.
    obj = Object3D(input_mesh_path)  # trimesh.load(input_mesh_path, skip_materials=True, force="mesh")

    if resize and size > 0:
        if bound_type == "diameter":
            object_scale = size / obj.bounding_diameter
        elif bound_type == "extents":
            object_scale = size / obj.bounding_size
        else:
            raise ValueError(f"Unkown size_type {bound_type!r}")

        obj.rescale(object_scale)

    if recenter:
        obj.recenter("bounds")
    # scale = object_scale
    mesh = obj.meshes[0]
    input_mesh_path = input_mesh_path.replace(".obj", "_corr.obj")
    mesh = mesh.process(validate=True)
    # add texture

    if not hasattr(mesh.visual, "uv"):
        # load template mesh and assign materials
        # template_mesh = trimesh.load("assets/template_mesh/textured_simple.obj")
        from trimesh.visual import TextureVisuals  # template_mesh.visual #

        # assign random texture coordinates
        open3d_mesh = mesh.as_open3d
        open3d_mesh.paint_uniform_color((0.8, 0.8, 0.8))
        v_uv = np.random.rand(len(open3d_mesh.triangles) * 3, 2)
        open3d_mesh.triangle_uvs = o3d.utility.Vector2dVector(v_uv)
        o3d.io.write_triangle_mesh(input_mesh_path, open3d_mesh)
    else:
        mesh.export(input_mesh_path, include_texture=True)

    # old version that writes with trimesh
    # Generate SDF file and the robot and link elements.
    robot_name = prefix + mesh_minus_ext
    root_item = ET.Element("sdf", version="1.5", nsmap={"drake": "drake.mit.edu"})
    model_item = ET.SubElement(root_item, "model", name=robot_name)
    link_name = "{}_body_link".format(robot_name)
    link_item = ET.SubElement(model_item, "link", name=link_name)
    pose_item = ET.SubElement(link_item, "pose")
    pose_item.text = "0 0 0 0 0 0"

    # Set up object inertia.
    mass, I = calc_mesh_inertia(mesh, density=density)
    inertial_item = ET.SubElement(link_item, "inertial")
    mass_item = ET.SubElement(inertial_item, "mass")
    mass_item.text = "{:.4E}".format(mass)
    inertia_item = ET.SubElement(inertial_item, "inertia")
    for i in range(3):
        for j in range(i, 3):
            item = ET.SubElement(inertia_item, "i" + "xyz"[i] + "xyz"[j])
            item.text = "{:.4E}".format(I[i, j])

    # Set up object visual geometry.
    visual_mesh_filename = input_mesh_path
    logging.warn("%s -> %s", visual_mesh_filename, sdf_path)
    visual_mesh_filename = os.path.relpath(visual_mesh_filename, dir_path)
    visual_item = ET.SubElement(link_item, "visual", name="visual")
    geometry_item = ET.SubElement(visual_item, "geometry")
    mesh_item = ET.SubElement(geometry_item, "mesh")
    uri_item = ET.SubElement(mesh_item, "uri")
    uri_item.text = visual_mesh_filename
    scale_item = ET.SubElement(mesh_item, "scale")
    scale_item.text = "{:.4E} {:.4E} {:.4E}".format(scale, scale, scale)

    # Set up object collision geometry.
    collision_paths = [input_mesh_path]
    for k, new_model_path in enumerate(collision_paths):
        new_model_path = os.path.relpath(new_model_path, dir_path)
        # Create a new XML subtree for the collection of meshes
        # we just created. I *think* each convex piece needs
        # to be in its own collision tag, otherwise Drake
        # seems to be ignoring them...
        collision_item = ET.SubElement(link_item, "collision", name="collision_%04d" % k)
        geometry_item = ET.SubElement(collision_item, "geometry")
        mesh_item = ET.SubElement(geometry_item, "mesh")
        uri_item = ET.SubElement(mesh_item, "uri")
        uri_item.text = new_model_path

    logging.info("Writing SDF to %s" % sdf_path)
    ET.ElementTree(root_item).write(sdf_path, pretty_print=True)
    return obj


def do_collision_mesh_simplification(mesh, mesh_name, mesh_dir, preview_with_trimesh=False, **kwargs):
    """
    Given a mesh, performs a convex decomposition of it with
    trimesh _ vhacd, saving all the parts in a subfolder named
    `<mesh_filename>_parts`.

    Args:
    - input_mesh_path: String path to mesh file to decompose. Only
        'obj' format is currently tested, but other formats supported
        by trimesh might work.
    - preview_with_trimesh: Whether to open (and block on) a window to preview
    the decomposition.
    - A set of control kwargs, plus any additional kwargs, are passed to the convex
      decomposition routine 'vhacd'; you can run `testVHACD --help` to see options.

    Returns:
    - List of generated mesh file parts, in obj format.
    """

    # Create a subdir for the convex decomp parts.
    mesh_parts_folder = mesh_name + "_parts"
    out_dir = os.path.join(mesh_dir, mesh_parts_folder)
    os.makedirs(out_dir, exist_ok=True)

    if preview_with_trimesh:
        logging.info("Showing mesh before decomp. Close window to proceed.")
        # mesh.set_camera(angles=(1, 0, 0), distance=1, center=(0,0,0))
        # mesh.show(angles=(0, 1, 0), distance=1, center=(0,0,0))
        scene = trimesh.scene.scene.Scene()
        scene.add_geometry(mesh)
        scene.set_camera(angles=(1, 0, 0), distance=0.3, center=(0, 0, 0))
        # scene.viewer.toggle_axis()
        scene.show()
    try:
        convex_pieces = []
        logging.info("Performing convex decomposition. If this runs too long, try decreasing --resolution.")
        convex_pieces_new = trimesh.decomposition.convex_decomposition(mesh, **kwargs)
        if not isinstance(convex_pieces_new, list):
            convex_pieces_new = [convex_pieces_new]
        convex_pieces += convex_pieces_new
    except Exception as e:
        logging.error("Problem performing decomposition: %s", e)

    if preview_with_trimesh:
        # Display the convex decomp, giving each a random colors
        # to make them easier to distinguish.
        for part in convex_pieces:
            this_color = trimesh.visual.random_color()
            part.visual.face_colors[:] = this_color
        scene = trimesh.scene.scene.Scene()
        for part in convex_pieces:
            scene.add_geometry(part)

        logging.info("Showing mesh convex decomp into %d parts. Close window to proceed." % (len(convex_pieces)))
        scene.set_camera(angles=(1, 0, 0), distance=0.3, center=(0, 0, 0))

        # viewer = trimesh.viewer.SceneViewer(scene)
        # viewer.init_gl()
        # viewer.on_draw()
        scene.show()

    out_paths = []
    for k, part in enumerate(convex_pieces):
        piece_name = "%s_convex_piece_%03d.obj" % (mesh_name, k)
        full_path = os.path.join(out_dir, piece_name)
        trimesh.exchange.export.export_mesh(part, full_path)
        out_paths.append(full_path)
    return out_paths


def calc_mesh_inertia(mesh, density=2000):
    """
    Given a mesh, calculates its total mass and inertia assuming
    a fixed density.
    Args:
    - mesh: A trimesh mesh.
    - density: Density of object in kg/m^3, used for inertia calculation.
    Returns: (mass, inertia)
    - out_paths: List of generated mesh file parts, in obj format.
    - inertia: total inertia of the input mesh.
    """
    mesh.density = density
    I = mesh.moment_inertia
    return mesh.mass, mesh.moment_inertia


def rotX(rotx):
    RotX = np.array(
        [
            [1, 0, 0, 0],
            [0, np.cos(rotx), -np.sin(rotx), 0],
            [0, np.sin(rotx), np.cos(rotx), 0],
            [0, 0, 0, 1],
        ]
    )
    return RotX


def create_sdf_with_convex_decomp(
    input_mesh_path,
    scale=1.0,
    do_visual_simplification=False,
    size=0.1,
    target_tris=1000,
    preview_with_trimesh=False,
    density=2000,
    recenter=True,
    resize=True,
    bound_type="extents",
    prefix="",
    flip_z=False,
    additional_transform="",
    **kwargs,
):
    """
    Given an input mesh file, produces an SDF if the same directory that:
    - Uses the mesh as its sole visual geometry.
    - Performs a convex decomposition of the mesh, and uses those pieces
    as the collision geometry.
    - Inserts inertia for the object, calculated from the original
    mesh assuming a constant density.

    The SDF is saved as `<mesh_file>.sdf` next to the mesh file.

    Args:
    - input_mesh_path: Path to the mesh file
    - preview_with_trimesh: Whether to show 3D previews of pre/post decomposition.
    - density: Assumed density of the object, in kg/m^3.
    - kwargs: Passed through to do_collision_mesh_simplification as convex decomp args.
    """
    # Get mesh name.
    dir_path, mesh_filename = os.path.split(input_mesh_path)
    mesh_minus_ext, _ = os.path.splitext(mesh_filename)
    sdf_path = os.path.join(dir_path, mesh_minus_ext + ".sdf")

    # Load in and prescale mesh.
    obj = Object3D(input_mesh_path)  # trimesh.load(input_mesh_path, skip_materials=True, force="mesh")

    if resize and size > 0:
        if bound_type == "diameter":
            object_scale = size / obj.bounding_diameter
        elif bound_type == "extents":
            object_scale = size / obj.bounding_size
        else:
            raise ValueError(f"Unkown size_type {bound_type!r}")
        if scale != 1:
            object_scale = scale
        obj.rescale(object_scale)
    # obj.normalize_direction(flip_z)
    if recenter:
        obj.recenter("bounds")

    # one off rotate
    if len(additional_transform) > 0:
        obj.apply_transform(eval(additional_transform)(np.pi / 2))

    if flip_z:
        obj.apply_transform(rotX(np.pi))

    #
    mesh = obj.meshes[0]
    input_mesh_path = input_mesh_path.replace(".obj", "_corr.obj")
    mesh = mesh.process(validate=True)
    # trimesh.exchange.export.export_obj(mesh, input_mesh_path, include_texture=True)
    # add texture
    # IPython.embed()
    # if not hasattr(mesh.visual, 'uv'):
    # load template mesh and assign materials
    # template_mesh = trimesh.load("assets/template_mesh/textured_simple.obj")
    from trimesh.visual import TextureVisuals  # template_mesh.visual #

    # assign random texture coordinates
    open3d_mesh = mesh.as_open3d
    open3d_mesh.paint_uniform_color((0.8, 0.8, 0.8))
    v_uv = np.random.rand(len(open3d_mesh.triangles) * 3, 2)
    open3d_mesh.triangle_uvs = o3d.utility.Vector2dVector(v_uv)
    o3d.io.write_triangle_mesh(input_mesh_path, open3d_mesh)
    # else:
    #     mesh.export(input_mesh_path, include_texture=True)

    # old version that writes with trimesh
    # Generate SDF file and the robot and link elements.
    robot_name = prefix + mesh_minus_ext
    root_item = ET.Element("sdf", version="1.5", nsmap={"drake": "drake.mit.edu"})
    model_item = ET.SubElement(root_item, "model", name=robot_name)
    link_name = "{}_body_link".format(robot_name)
    link_item = ET.SubElement(model_item, "link", name=link_name)
    pose_item = ET.SubElement(link_item, "pose")
    pose_item.text = "0 0 0 0 0 0"

    # Set up object inertia.
    mass, I = calc_mesh_inertia(mesh, density=density)
    # all 0s for off diagonal
    I[0, 1] = I[0, 2] = I[1, 0] = I[1, 2] = I[2, 0] = I[2, 1] = 0
    I = np.abs(I)  # all positive diagonal
    inertial_item = ET.SubElement(link_item, "inertial")
    mass_item = ET.SubElement(inertial_item, "mass")
    mass_item.text = "{:.4E}".format(mass)
    inertia_item = ET.SubElement(inertial_item, "inertia")
    for i in range(3):
        for j in range(i, 3):
            item = ET.SubElement(inertia_item, "i" + "xyz"[i] + "xyz"[j])
            item.text = "{:.4E}".format(I[i, j])

    # Set up object visual geometry.
    if do_visual_simplification:
        visual_mesh_filename = do_visual_mesh_simplification(input_mesh_path, target_tris=target_tris)
    else:
        visual_mesh_filename = input_mesh_path

    logging.warn("%s -> %s", visual_mesh_filename, sdf_path)
    visual_mesh_filename = os.path.relpath(visual_mesh_filename, dir_path)
    visual_item = ET.SubElement(link_item, "visual", name="visual")
    geometry_item = ET.SubElement(visual_item, "geometry")
    mesh_item = ET.SubElement(geometry_item, "mesh")
    uri_item = ET.SubElement(mesh_item, "uri")
    uri_item.text = visual_mesh_filename
    # scale_item = ET.SubElement(mesh_item, "scale") # don't need the scale, already applied
    # scale_item.text = "{:.4E} {:.4E} {:.4E}".format(scale, scale, scale)

    # Set up object collision geometry.
    # do_collision_mesh_simplification
    collision_paths = do_collision_mesh_simplification_coacd(
        mesh,
        mesh_name=mesh_minus_ext,
        mesh_dir=dir_path,
        preview_with_trimesh=preview_with_trimesh,
        **kwargs,
    )
    for k, new_model_path in enumerate(collision_paths):
        new_model_path = os.path.relpath(new_model_path, dir_path)
        collision_item = ET.SubElement(link_item, "collision", name="collision_%04d" % k)
        geometry_item = ET.SubElement(collision_item, "geometry")
        mesh_item = ET.SubElement(geometry_item, "mesh")
        uri_item = ET.SubElement(mesh_item, "uri")
        uri_item.text = new_model_path
        ET.SubElement(mesh_item, "{drake.mit.edu}declare_convex")

    logging.info("Writing SDF to %s" % sdf_path)
    ET.ElementTree(root_item).write(sdf_path, pretty_print=True)
    return obj


def rename_model_normalized_folder(path, obj_name, idx):
    print(
        "rename object folder:",
        f"{path}/model_normalized.obj",
        os.path.exists(f"{path}/model_normalized.obj"),
    )
    if os.path.exists(f"{path}/model_normalized.obj"):
        new_model_path = path.replace(obj_name, f"model{idx}")  #
        os.system(f"mkdir {new_model_path}  ")
        os.system(f"cp {path}/model_normalized.obj {new_model_path}/model{idx}.obj ")  # {path}{obj_name}.obj
        # better rename it as modelx
        # os.system(f'rm {path}/model_normalized* ')
        # print(new_model_path)
        return f"{new_model_path}/model{idx}.obj"
    return f"{path}/{obj_name}.obj"


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate SDF from a mesh.")
    parser.add_argument("mesh_file", type=str, help="Path to mesh file.")
    parser.add_argument(
        "--preview",
        default=False,
        action="store_true",
        help="Preview decomp with a trimesh window?",
    )
    parser.add_argument(
        "--scale",
        type=float,
        default=1,
        help="Scale factor to convert the specified mesh's coordinates to meters.",
    )
    parser.add_argument("--size", type=float, default=0.15, help="Diameter for the output mesh.")
    parser.add_argument(
        "--density",
        type=float,
        default=2,
        help="Assumed density in kg/m^3 of object, for inertia calculation.",
    )
    parser.add_argument(
        "--do_visual_simplification",
        default=False,
        action="store_true",
        help="Do additional visual simplification of mesh. Requires open3d. Probably won't preserve materials.",
    )
    parser.add_argument(
        "--target_tris",
        type=int,
        default=1000,
        help="If we do visual simplification, we decimate to this # of triangles.",
    )
    parser.add_argument("--resolution", type=int, default=100000, help="VHACD voxel resolution.")
    parser.add_argument("--maxhulls", type=int, default=500, help="VHACD max # of convex hulls.")
    parser.add_argument(
        "--minVolumePerCH",
        type=float,
        default=0.001,
        help="VHACD min convex hull volume.",
    )
    parser.add_argument("--maxNumVerticesPerCH", type=int, default=500, help="VHACD voxel resolution.")
    parser.add_argument(
        "--loglevel",
        type=str,
        default="INFO",
        choices=["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG"],
        help="Log level.",
    )
    parser.add_argument("--prefix", type=str, default="")
    # manually tuned
    parser.add_argument(
        "--up",
        type=int,
        default=-10,
        help="Scale factor to convert the specified mesh's coordinates to meters.",
    )
    parser.add_argument("--forward", type=int, default=-10, help="Diameter for the output mesh.")
    parser.add_argument(
        "--flip_z",
        default=False,
        action="store_true",
        help="Preview decomp with a trimesh window?",
    )
    parser.add_argument(
        "--select_keypoint",
        default=False,
        action="store_true",
        help="Preview decomp with a trimesh window?",
    )
    parser.add_argument("--additional_transform", type=str, default="")
    parser.add_argument("--rename_idx", type=str, default="")
    parser.add_argument(
        "--no_convex_decomp",
        action="store_true",
        help="Preview decomp with a trimesh window?",
    )
    parser.add_argument(
        "--no_resize",
        default=False,
        action="store_true",
        help="Avoid resizing?",
    )
    parser.add_argument(
        "--no_recenter",
        default=False,
        action="store_true",
        help="Avoid recentering?",
    )
    # parser.add_argument('--resize', action="store_true", help="Preview decomp with a trimesh window?")

    args = parser.parse_args()
    logging.basicConfig(level=args.loglevel)
    mesh_path = os.path.abspath(args.mesh_file)

    if len(args.rename_idx) > 0:
        mesh_path = rename_model_normalized_folder(args.mesh_file, args.mesh_file.split("/")[-1], int(args.rename_idx))

    args.prefix = args.mesh_file.split("/")[2] + "_"  # add prefix
    gen_sdf_func = create_sdf_with_mesh if args.no_convex_decomp else create_sdf_with_convex_decomp
    obj = gen_sdf_func(
        mesh_path,
        scale=args.scale,
        size=args.size,
        do_visual_simplification=args.do_visual_simplification,
        target_tris=args.target_tris,
        preview_with_trimesh=args.preview,
        density=args.density,
        resolution=args.resolution,
        maxhulls=args.maxhulls,
        maxNumVerticesPerCH=args.maxNumVerticesPerCH,
        minVolumePerCH=args.minVolumePerCH,
        pca=1,
        prefix=args.prefix,
        flip_z=args.flip_z,
        additional_transform=args.additional_transform,
        resize=not args.no_resize,
        recenter=not args.no_recenter,
    )

    input_mesh_path = mesh_path.replace(".obj", "_corr.obj")
    obj = Object3D(input_mesh_path)
    obj.recompute_normals()
    open3d_mesh = obj.meshes[0].as_open3d
    open3d_mesh.compute_vertex_normals()

    o3d.io.write_triangle_mesh(input_mesh_path.replace(".obj", "_open3d.obj"), open3d_mesh, write_vertex_colors=True)
    obj.meshes[0].export(input_mesh_path, include_texture=True, include_normals=True)


    # if args.preview and not args.no_convex_decomp:
    #     # remaining one degree of freedom for the heuristicsq
    #     import cv2

    #     cv2.imshow("flip?", np.zeros((64, 64, 3)))
    #     k = cv2.waitKey(0)

    #     if k == ord("f"):
    #         args.flip_z = True
    #         print("flip")
    #         obj = create_sdf_with_convex_decomp(
    #             mesh_path,
    #             scale=args.scale,
    #             size=args.size,
    #             do_visual_simplification=args.do_visual_simplification,
    #             target_tris=args.target_tris,
    #             preview_with_trimesh=False,
    #             density=args.density,
    #             resolution=args.resolution,
    #             maxhulls=args.maxhulls,
    #             maxNumVerticesPerCH=args.maxNumVerticesPerCH,
    #             minVolumePerCH=args.minVolumePerCH,
    #             pca=1,
    #             prefix=args.prefix,
    #             flip_z=args.flip_z,
    #             additional_transform=args.additional_transform,
    #         )

    if args.select_keypoint:
        mesh = o3d.io.read_triangle_mesh(mesh_path.replace(".obj", "_corr.obj"))
        pcd = mesh.sample_points_uniformly(number_of_points=10000)

        vis = o3d.visualization.VisualizerWithEditing()
        vis.create_window()
        vis.add_geometry(pcd)
        vis.run()  # user picks points
        vis.destroy_window()
        print("")

        # vis = o3d.visualization.draw_geometries_with_editing([pcd])
        picked_points = vis.get_picked_points()
        xyz = np.asarray(pcd.points)

        picked_points = xyz[picked_points]
        # save as json  try two-point method
        pointcloudfile = mesh_path.replace(".obj", "_points.npy")
        np.save(pointcloudfile, xyz)
        keypoint_description_file = mesh_path.replace(".obj", "_info.json")
        keypoint_info = {
            "keypoints": picked_points.tolist(),
            "extents": obj.extents.tolist(),
        }

        with open(keypoint_description_file, "w") as f:
            json.dump(keypoint_info, f, indent=4, sort_keys=True)

        # also save sampled points

scripts/generate_tool_sdf_from_mesh.py
liruiw/Fleet-Tools
import argparse
import logging
import numpy as np
import os
import sys
import webbrowser

from pydrake.all import (
    AddMultibodyPlantSceneGraph,
    DiagramBuilder,
    FindResourceOrThrow,
    Meshcat,
    MeshcatVisualizer,
    MeshcatVisualizerParams,
    Parser,
    RigidTransform,
    Rgba,
    Role,
    Simulator,
)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Visualize SDF and its collision geometry in meshcat.")
    parser.add_argument("sdf", type=str, help="Path to SDF.")
    parser.add_argument("--noopen", action="store_true", help="Open webbrowser automatically.")

    args = parser.parse_args()

    sdf = os.path.abspath(args.sdf)
    if not os.path.exists(sdf):
        logging.error("No SDF found at %s" % sdf)
        sys.exit(-1)

    # Start the visualizer.
    meshcat = Meshcat()
    web_url = meshcat.web_url()
    print("Meshcat available at %s" % web_url)
    if not args.noopen:
        webbrowser.open_new(web_url)

    # Build a diagram to do visualization.
    builder = DiagramBuilder()
    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.01)
    # Add a single object into it.
    parser = Parser(plant)
    model_id = parser.AddModelFromFile(sdf)
    body_ids = plant.GetBodyIndices(model_id)
    for body_id in body_ids:
        body = plant.get_body(body_id)
        plant.WeldFrames(plant.world_frame(), body.body_frame())
    plant.Finalize()

    # Add two visualizers with the same meshcat instance:
    #  one for visual, and one for collision.
    visual_params = MeshcatVisualizerParams(
        role=Role.kIllustration, prefix="visual", default_color=Rgba(0.1, 0.6, 1.0, 0.3)
    )
    # visual_viz = MeshcatVisualizerCpp.AddToBuilder(builder, scene_graph, meshcat, visual_params)
    visual_viz = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat, visual_params)
    collision_params = MeshcatVisualizerParams(
        role=Role.kProximity, prefix="collision", default_color=Rgba(0.8, 0.5, 0.2, 0.3)
    )
    # collision_viz = MeshcatVisualizerCpp.AddToBuilder(builder, scene_graph, meshcat, collision_params)
    collision_viz = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat, collision_params)

    meshcat.ResetRenderMode()

    diagram = builder.Build()
    simulator = Simulator(diagram)
    context = simulator.get_mutable_context()

    simulator.set_target_realtime_rate(1.0)

    meshcat.AddButton("Stop Simulation")
    while meshcat.GetButtonClicks("Stop Simulation") < 1:
        simulator.AdvanceTo(simulator.get_context().get_time() + 2.0)
    meshcat.DeleteButton("Stop Simulation")

scripts/inspect_sdf_in_meshcat.py
liruiw/Fleet-Tools
import os
import glob
import argparse

# from colmap_depth import extract_all_depths
import numpy as np
import cv2

import open3d as o3d
import os
import IPython
import json


def mkdir_if_missing(dst_dir):
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)


parser = argparse.ArgumentParser()
parser.add_argument("--dir", type=str)
parser.add_argument("--mesh_file", type=str)

# parser.add_argument('--write', required=True)
args = parser.parse_args()
datapath = args.dir

o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Debug)
mesh = o3d.io.read_triangle_mesh(args.mesh_file)
pcd = mesh.sample_points_uniformly(number_of_points=50000)


viewer = o3d.visualization.VisualizerWithEditing()
viewer.create_window()
viewer.add_geometry(pcd)
opt = viewer.get_render_option()
opt.show_coordinate_frame = True
# opt.background_color = np.asarray([0.5, 0.5, 0.5])
viewer.run()
viewer.destroy_window()

scripts/select_keypoint.py
liruiw/Fleet-Tools
import os
import glob
import argparse

import numpy as np
import cv2

import open3d as o3d
import os
import IPython
import json


def mkdir_if_missing(dst_dir):
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)


parser = argparse.ArgumentParser()
parser.add_argument("--mesh_file", "-m", type=str)

# parser.add_argument('--write', required=True)
args = parser.parse_args()

o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Debug)
mesh = o3d.io.read_triangle_mesh(args.mesh_file)
# o3d.visualization.draw_geometries([mesh])

pcd = mesh.sample_points_uniformly(number_of_points=2000)
pcd.paint_uniform_color([0.5, 0.5, 0.5])
# o3d.visualization.draw_geometries([pcd])

meshes = [pcd]
keypoint_info = json.load(open(args.mesh_file.replace("_corr.obj", "_info.json")))
tool_keypoints = keypoint_info["keypoints"]
colors = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
for idx in range(len(tool_keypoints)):
    mesh_sphere = o3d.geometry.TriangleMesh.create_sphere(radius=0.001)
    mesh_sphere.translate(tool_keypoints[idx])
    mesh_sphere.paint_uniform_color(colors[idx])
    meshes.append(mesh_sphere)


viewer = o3d.visualization.Visualizer()
viewer.create_window()
print("meshes:", len(meshes))
for m in meshes:
    viewer.add_geometry(m)
opt = viewer.get_render_option()
opt.show_coordinate_frame = True
viewer.run()
viewer.destroy_window()

scripts/vis_keypoint.py
liruiw/Fleet-Tools
import os
import sys
from setuptools import setup, find_packages

print(
    "Installing. \n Package intended for use with provided conda env. See setup instructions here: https://github.com/liruiw/FleetDrake"
)

if sys.version_info.major != 3:
    print(
        "This Python is only compatible with Python 3, but you are running "
        "Python {}. The installation will likely fail.".format(sys.version_info.major)
    )


def read(fname):
    return open(os.path.join(os.path.dirname(__file__), fname)).read()


setup(
    name="fleetdrake",
    version="1.0.0",
    packages=find_packages(),
    description="Training and experiment with algorithms for environments in Drake",
    long_description=read("README.md"),
    url="https://github.com/liruiw/FleetDrake.git",
    author="Lirui Wang",
)

setup.py
liruiw/Fleet-Tools