diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bbbced6 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +out +datasets +results \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..693cf9f --- /dev/null +++ b/LICENSE @@ -0,0 +1,50 @@ +EGOPOSE +SOFTWARE LICENSE AGREEMENT +ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY + +BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS LICENSE AGREEMENT. IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR DOWNLOAD THE SOFTWARE. + +This is a license agreement ("Agreement") between your academic institution or non profit organization or self (called "Licensee" or "You" in this Agreement) and Carnegie Mellon University (called "Licensor" in this Agreement). All rights not specifically granted to you in this Agreement are reserved for Licensor. + +RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: +Licensor retains exclusive ownership of any copy of the Software (as defined below) licensed under this Agreement and hereby grants to Licensee a personal, non-exclusive, +non-transferable license to use the Software for noncommercial research purposes, without the right to sublicense, pursuant to the terms and conditions of this Agreement. As used in this Agreement, the term "Software" means (i) the actual copy of all or any portion of code for program routines made accessible to Licensee by Licensor pursuant to this Agreement, inclusive of backups, updates, and/or merged copies permitted hereunder or subsequently supplied by Licensor, including all or any file structures, programming instructions, user interfaces and screen formats and sequences as well as any and all documentation and instructions related to it, and (ii) all or any derivatives and/or modifications created or made by You to any of the items specified in (i). + +CONFIDENTIALITY: Licensee acknowledges that the Software is proprietary to Licensor, and as such, Licensee agrees to receive all such materials in confidence and use the Software only in accordance with the terms of this Agreement. Licensee agrees to use reasonable effort to protect the Software from unauthorized use, reproduction, distribution, or publication. + +COPYRIGHT: The Software is owned by Licensor and is protected by United States copyright laws and applicable international treaties and/or conventions. + +PERMITTED USES: The Software may be used for your own noncommercial internal research purposes. You understand and agree that Licensor is not obligated to implement any suggestions and/or feedback you might provide regarding the Software, but to the extent Licensor does so, you are not entitled to any compensation related thereto. + +DERIVATIVES: You may create derivatives of or make modifications to the Software, however, You agree that all and any such derivatives and modifications will be owned by Licensor and become a part of the Software licensed to You under this Agreement. You may only use such derivatives and modifications for your own noncommercial internal research purposes, and you may not otherwise use, distribute or copy such derivatives and modifications in violation of this Agreement. + +BACKUPS: If Licensee is an organization, it may make that number of copies of the Software necessary for internal noncommercial use at a single site within its organization provided that all information appearing in or on the original labels, including the copyright and trademark notices are copied onto the labels of the copies. + +USES NOT PERMITTED: You may not distribute, copy or use the Software except as explicitly permitted herein. Licensee has not been granted any trademark license as part of this Agreement and may not use the name or mark "AB3DMOT", "Carnegie Mellon", or any renditions thereof without the prior written permission of Licensor. + +You may not sell, rent, lease, sublicense, lend, time-share or transfer, in whole or in part, or provide third parties access to prior or present versions (or any parts thereof) of the Software. + +ASSIGNMENT: You may not assign this Agreement or your rights hereunder without the prior written consent of Licensor. Any attempted assignment without such consent shall be null and void. + +TERM: The term of the license granted by this Agreement is from Licensee's acceptance of this Agreement by clicking "I Agree" below or by using the Software until terminated as provided below. + +The Agreement automatically terminates without notice if you fail to comply with any provision of this Agreement. Licensee may terminate this Agreement by ceasing using the Software. Upon any termination of this Agreement, Licensee will delete any and all copies of the Software. You agree that all provisions which operate to protect the proprietary rights of Licensor shall remain in force should breach occur and that the obligation of confidentiality described in this Agreement is binding in perpetuity and, as such, survives the term of the Agreement. + +FEE: Provided Licensee abides completely by the terms and conditions of this Agreement, there is no fee due to Licensor for Licensee's use of the Software in accordance with this Agreement. + +DISCLAIMER OF WARRANTIES: THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON-INFRINGEMENT. LICENSEE BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND RELATED MATERIALS. + +SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is provided as part of this Agreement. + +EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent permitted under applicable law, Licensor shall not be liable for direct, indirect, special, incidental, or consequential damages or lost profits related to Licensee's use of and/or inability to use the Software, even if Licensor is advised of the possibility of such damage. + +EXPORT REGULATION: Licensee agrees to comply with any and all applicable +U.S. export control laws, regulations, and/or other laws related to embargoes and sanction programs administered by the Office of Foreign Assets Control. + +SEVERABILITY: If any provision(s) of this Agreement shall be held to be invalid, illegal, or unenforceable by a court or other tribunal of competent jurisdiction, the validity, legality and enforceability of the remaining provisions shall not in any way be affected or impaired thereby. + +NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right or remedy under this Agreement shall be construed as a waiver of any future or other exercise of such right or remedy by Licensor. + +GOVERNING LAW: This Agreement shall be construed and enforced in accordance with the laws of the Commonwealth of Pennsylvania without reference to conflict of laws principles. You consent to the personal jurisdiction of the courts of this County and waive their rights to venue outside of Allegheny County, Pennsylvania. + +ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and entire agreement between Licensee and Licensor as to the matter set forth herein and supersedes any previous agreements, understandings, and arrangements between the parties relating hereto. diff --git a/README.md b/README.md new file mode 100644 index 0000000..151636e --- /dev/null +++ b/README.md @@ -0,0 +1,66 @@ +# EgoPose +![Loading EgoPose demo gif](https://www.ye-yuan.com/wp-content/uploads/2019/07/ego_pose_rs.gif "EgoPose demo gif") +![Loading EgoPose demo gif](https://www.ye-yuan.com/wp-content/uploads/2019/07/ego_forecast_rs.gif "EgoPose demo gif") +--- +This repo contains the official implementation of our paper: + +Ego-Pose Estimation and Forecasting as Real-Time PD Control +Ye Yuan, Kris Kitani. ICCV 2019. +[[project page](https://www.ye-yuan.com/ego-pose)] [[paper](https://arxiv.org/pdf/1906.03173.pdf)] [[video](https://youtu.be/968IIDZeWE0)] + +# Installation +### Dataset +* Download the dataset from this [link](https://drive.google.com/file/d/1vzxVHAtfvfIEDreqYvHulhtNwHcomotV/view?usp=sharing) and place the unzipped dataset folder inside the repo as "EgoPose/datasets". Please see the README.txt inside the folder for details about the dataset. +### Environment +* **Supported OS:** MacOS, Linux +* **Packages:** + * Python >= 3.6 + * PyTorch >= 0.4 ([https://pytorch.org/](https://pytorch.org/)) + * gym ([https://github.com/openai/gym](https://github.com/openai/gym)) + * mujoco-py ([https://github.com/openai/mujoco-py](https://github.com/openai/mujoco-py)) + * OpenCV: ```conda install -c menpo opencv``` + * Tensorflow, OpenGL, yaml: + ```conda install tensorflow pyopengl pyyaml``` +* **Additional setup:** + * For linux, set the following environment variable to greatly improve multi-threaded sampling performance: + ```export OMP_NUM_THREADS=1``` +* **Note**: All scripts should be run from the root of this repo. +# Quick Demo +### Pretrained Models +* Please first download our pretrained models from this [link](https://drive.google.com/file/d/1DE-uSUk4JMDtL9aQY2R5rAd3_yPRUIIH/view?usp=sharing) and place the unzipped results folder inside the repo as "EgoPose/results". + +### Ego-Pose Estimation +* To visualize the results for MoCap data: + ```python ego_pose/eval_pose.py --egomimic-cfg subject_03 --statereg-cfg subject_03 --mode vis``` + Here we use the config file for subject_03. Note that in the visualization, the red humanoid represents the GT. + +* To visualize the results for in-the-wild data: + ```python ego_pose/eval_pose_wild.py --egomimic-cfg cross_01 --statereg-cfg cross_01 --data wild_01 --mode vis``` + Here we use the config file for cross-subject model (cross_01) and test it on in-the-wild data (wild_01). + +* Keyboard shortcuts for the visualizer: [keymap.md](https://github.com/Khrylx/EgoPose/blob/master/docs/keymap.md) +### Ego-Pose Forecasting +* To visualize the results for MoCap data: + ```python ego_pose/eval_forecast.py --egoforecast-cfg subject_03 --mode vis``` + +* To visualize the results for in-the-wild data: + ```python ego_pose/eval_forecast_wild.py --egoforecast-cfg cross_01 --data wild_01 --mode vis``` + + +# Training and Testing +* If you are interested in training and testing with our code, please see [train_and_test.md](https://github.com/Khrylx/EgoPose/blob/master/docs/train_and_test.md). + +# Citation +If you find our work useful in your research, please consider citing our paper [Ego-Pose Estimation and Forecasting as Real-Time PD Control](https://www.ye-yuan.com/ego-pose): +``` +@inproceedings{yuan2019ego, + title={Ego-Pose Estimation and Forecasting as Real-Time PD Control}, + author={Yuan, Ye and Kitani, Kris}, + booktitle={Proceedings of the IEEE International Conference on Computer Vision}, + year={2019} +} +``` + +# License + +The software in this repo is freely available for free non-commercial use. Please see the [license](https://github.com/Khrylx/EgoPose/blob/master/LICENSE) for further details. diff --git a/agents/__init__.py b/agents/__init__.py new file mode 100644 index 0000000..0c40fc9 --- /dev/null +++ b/agents/__init__.py @@ -0,0 +1,3 @@ +from agents.agent_pg import AgentPG +from agents.agent_ppo import AgentPPO +from agents.agent_trpo import AgentTRPO diff --git a/agents/agent.py b/agents/agent.py new file mode 100644 index 0000000..1e8e15c --- /dev/null +++ b/agents/agent.py @@ -0,0 +1,122 @@ +import multiprocessing +from core import LoggerRL, TrajBatch +from utils.memory import Memory +from utils.torch import * +import math +import time + + +class Agent: + + def __init__(self, env, policy_net, value_net, dtype, device, custom_reward=None, + mean_action=False, render=False, running_state=None, num_threads=1): + self.env = env + self.policy_net = policy_net + self.value_net = value_net + self.dtype = dtype + self.device = device + self.custom_reward = custom_reward + self.mean_action = mean_action + self.running_state = running_state + self.render = render + self.num_threads = num_threads + self.noise_rate = 1.0 + self.traj_cls = TrajBatch + self.logger_cls = LoggerRL + self.sample_modules = [policy_net] + self.update_modules = [policy_net, value_net] + + def sample_worker(self, pid, queue, min_batch_size): + torch.randn(pid) + if hasattr(self.env, 'np_random'): + self.env.np_random.rand(pid) + memory = Memory() + logger = LoggerRL() + + while logger.num_steps < min_batch_size: + state = self.env.reset() + if self.running_state is not None: + state = self.running_state(state) + logger.start_episode(self.env) + self.pre_episode() + + for t in range(10000): + state_var = tensor(state).unsqueeze(0) + vs_out = self.trans_policy(state_var) + mean_action = self.mean_action or np.random.binomial(1, 1 - self.noise_rate) + action = self.policy_net.select_action(vs_out, mean_action)[0].numpy() + action = int(action) if self.policy_net.type == 'discrete' else action.astype(np.float64) + next_state, env_reward, done, info = self.env.step(action) + if self.running_state is not None: + next_state = self.running_state(next_state) + if self.custom_reward is not None: + c_reward, c_info = self.custom_reward(self.env, state, action, info) + reward = c_reward + else: + c_reward, c_info = 0.0, np.array([0.0]) + reward = env_reward + logger.step(self.env, env_reward, c_reward, c_info) + + mask = 0 if done else 1 + exp = 1 - mean_action + self.push_memory(memory, state, action, mask, next_state, reward, exp) + + if pid == 0 and self.render: + self.env.render() + if done: + break + state = next_state + + logger.end_episode(self.env) + logger.end_sampling() + + if queue is not None: + queue.put([pid, memory, logger]) + else: + return memory, logger + + def pre_episode(self): + return + + def push_memory(self, memory, state, action, mask, next_state, reward, exp): + memory.push(state, action, mask, next_state, reward, exp) + + def pre_sample(self): + return + + def sample(self, min_batch_size): + t_start = time.time() + self.pre_sample() + to_test(*self.sample_modules) + with to_cpu(*self.sample_modules): + with torch.no_grad(): + thread_batch_size = int(math.floor(min_batch_size / self.num_threads)) + queue = multiprocessing.Queue() + memories = [None] * self.num_threads + loggers = [None] * self.num_threads + for i in range(self.num_threads-1): + worker_args = (i+1, queue, thread_batch_size) + worker = multiprocessing.Process(target=self.sample_worker, args=worker_args) + worker.start() + memories[0], loggers[0] = self.sample_worker(0, None, thread_batch_size) + + for i in range(self.num_threads - 1): + pid, worker_memory, worker_logger = queue.get() + memories[pid] = worker_memory + loggers[pid] = worker_logger + traj_batch = self.traj_cls(memories) + logger = self.logger_cls.merge(loggers) + + logger.sample_time = time.time() - t_start + return traj_batch, logger + + def trans_policy(self, states): + """transform states before going into policy net""" + return states + + def trans_value(self, states): + """transform states before going into value net""" + return states + + def set_noise_rate(self, noise_rate): + self.noise_rate = noise_rate diff --git a/agents/agent_pg.py b/agents/agent_pg.py new file mode 100644 index 0000000..880ba72 --- /dev/null +++ b/agents/agent_pg.py @@ -0,0 +1,57 @@ +from core import estimate_advantages +from agents.agent import Agent +from utils.torch import * +import time + + +class AgentPG(Agent): + + def __init__(self, gamma=0.99, tau=0.95, optimizer_policy=None, optimizer_value=None, + opt_num_epochs=1, value_opt_niter=1, **kwargs): + super().__init__(**kwargs) + self.gamma = gamma + self.tau = tau + self.optimizer_policy = optimizer_policy + self.optimizer_value = optimizer_value + self.opt_num_epochs = opt_num_epochs + self.value_opt_niter = value_opt_niter + + def update_value(self, states, returns): + """update critic""" + for _ in range(self.value_opt_niter): + values_pred = self.value_net(self.trans_value(states)) + value_loss = (values_pred - returns).pow(2).mean() + self.optimizer_value.zero_grad() + value_loss.backward() + self.optimizer_value.step() + + def update_policy(self, states, actions, returns, advantages, exps): + """update policy""" + # use a2c by default + ind = exps.nonzero().squeeze(1) + for _ in range(self.opt_num_epochs): + self.update_value(states, returns) + log_probs = self.policy_net.get_log_prob(self.trans_policy(states)[ind], actions[ind]) + policy_loss = -(log_probs * advantages[ind]).mean() + self.optimizer_policy.zero_grad() + policy_loss.backward() + self.optimizer_policy.step() + + def update_params(self, batch): + t0 = time.time() + to_train(*self.update_modules) + states = torch.from_numpy(batch.states).to(self.dtype).to(self.device) + actions = torch.from_numpy(batch.actions).to(self.dtype).to(self.device) + rewards = torch.from_numpy(batch.rewards).to(self.dtype).to(self.device) + masks = torch.from_numpy(batch.masks).to(self.dtype).to(self.device) + exps = torch.from_numpy(batch.exps).to(self.dtype).to(self.device) + with to_test(*self.update_modules): + with torch.no_grad(): + values = self.value_net(self.trans_value(states)) + + """get advantage estimation from the trajectories""" + advantages, returns = estimate_advantages(rewards, masks, values, self.gamma, self.tau) + + self.update_policy(states, actions, returns, advantages, exps) + + return time.time() - t0 diff --git a/agents/agent_ppo.py b/agents/agent_ppo.py new file mode 100644 index 0000000..a0e0295 --- /dev/null +++ b/agents/agent_ppo.py @@ -0,0 +1,65 @@ +import math +from utils.torch import * +from agents.agent_pg import AgentPG + + +class AgentPPO(AgentPG): + + def __init__(self, clip_epsilon=0.2, opt_batch_size=64, use_mini_batch=False, + policy_grad_clip=None, **kwargs): + super().__init__(**kwargs) + self.clip_epsilon = clip_epsilon + self.opt_batch_size = opt_batch_size + self.use_mini_batch = use_mini_batch + self.policy_grad_clip = policy_grad_clip + + def update_policy(self, states, actions, returns, advantages, exps): + """update policy""" + with to_test(*self.update_modules): + with torch.no_grad(): + fixed_log_probs = self.policy_net.get_log_prob(self.trans_policy(states), actions) + + optim_iter_num = int(math.ceil(states.shape[0] / self.opt_batch_size)) + for _ in range(self.opt_num_epochs): + if self.use_mini_batch: + perm = np.arange(states.shape[0]) + np.random.shuffle(perm) + perm = LongTensor(perm).to(self.device) + + states, actions, returns, advantages, fixed_log_probs, exps = \ + states[perm].clone(), actions[perm].clone(), returns[perm].clone(), advantages[perm].clone(), \ + fixed_log_probs[perm].clone(), exps[perm].clone() + + for i in range(optim_iter_num): + ind = slice(i * self.opt_batch_size, min((i + 1) * self.opt_batch_size, states.shape[0])) + states_b, actions_b, advantages_b, returns_b, fixed_log_probs_b, exps_b = \ + states[ind], actions[ind], advantages[ind], returns[ind], fixed_log_probs[ind], exps[ind] + ind = exps_b.nonzero().squeeze(1) + self.update_value(states_b, returns_b) + surr_loss = self.ppo_loss(states_b, actions_b, advantages_b, fixed_log_probs_b, ind) + self.optimizer_policy.zero_grad() + surr_loss.backward() + self.clip_policy_grad() + self.optimizer_policy.step() + else: + ind = exps.nonzero().squeeze(1) + self.update_value(states, returns) + surr_loss = self.ppo_loss(states, actions, advantages, fixed_log_probs, ind) + self.optimizer_policy.zero_grad() + surr_loss.backward() + self.clip_policy_grad() + self.optimizer_policy.step() + + def clip_policy_grad(self): + if self.policy_grad_clip is not None: + for params, max_norm in self.policy_grad_clip: + torch.nn.utils.clip_grad_norm_(params, max_norm) + + def ppo_loss(self, states, actions, advantages, fixed_log_probs, ind): + log_probs = self.policy_net.get_log_prob(self.trans_policy(states)[ind], actions[ind]) + ratio = torch.exp(log_probs - fixed_log_probs[ind]) + advantages = advantages[ind] + surr1 = ratio * advantages + surr2 = torch.clamp(ratio, 1.0 - self.clip_epsilon, 1.0 + self.clip_epsilon) * advantages + surr_loss = -torch.min(surr1, surr2).mean() + return surr_loss diff --git a/agents/agent_trpo.py b/agents/agent_trpo.py new file mode 100644 index 0000000..ffe1641 --- /dev/null +++ b/agents/agent_trpo.py @@ -0,0 +1,137 @@ +import scipy.optimize +from agents.agent_pg import AgentPG +from utils import * + + +def conjugate_gradients(Avp_f, b, nsteps, rdotr_tol=1e-10): + x = zeros(b.size()) + if b.is_cuda: + x.to(b.get_device()) + r = b.clone() + p = b.clone() + rdotr = torch.dot(r, r) + for i in range(nsteps): + Avp = Avp_f(p) + alpha = rdotr / torch.dot(p, Avp) + x += alpha * p + r -= alpha * Avp + new_rdotr = torch.dot(r, r) + betta = new_rdotr / rdotr + p = r + betta * p + rdotr = new_rdotr + if rdotr < rdotr_tol: + break + return x + + +def line_search(model, f, x, fullstep, expected_improve_full, max_backtracks=10, accept_ratio=0.1): + fval = f(True).item() + + for stepfrac in [.5**x for x in range(max_backtracks)]: + x_new = x + stepfrac * fullstep + set_flat_params_to(model, x_new) + fval_new = f(True).item() + actual_improve = fval - fval_new + expected_improve = expected_improve_full * stepfrac + ratio = actual_improve / expected_improve + + if ratio > accept_ratio: + return True, x_new + return False, x + + +class AgentTRPO(AgentPG): + + def __init__(self, max_kl=1e-2, damping=1e-2, use_fim=True, **kwargs): + super().__init__(**kwargs) + self.max_kl = max_kl + self.damping = damping + self.use_fim = use_fim + + def update_value(self, states, returns): + + def get_value_loss(flat_params): + set_flat_params_to(self.value_net, tensor(flat_params)) + for param in self.value_net.parameters(): + if param.grad is not None: + param.grad.data.fill_(0) + values_pred = self.value_net(self.trans_value(states)) + value_loss = (values_pred - returns).pow(2).mean() + + # weight decay + for param in self.value_net.parameters(): + value_loss += param.pow(2).sum() * 1e-3 + value_loss.backward() + return value_loss.item(), get_flat_grad_from(self.value_net.parameters()).cpu().numpy() + + flat_params, _, opt_info = scipy.optimize.fmin_l_bfgs_b(get_value_loss, + get_flat_params_from(self.value_net).detach().cpu().numpy(), + maxiter=25) + set_flat_params_to(self.value_net, tensor(flat_params)) + + def update_policy(self, states, actions, returns, advantages, exps): + """update policy""" + ind = exps.nonzero().squeeze(1) + self.update_value(states, returns) + + with torch.no_grad(): + fixed_log_probs = self.policy_net.get_log_prob(self.trans_policy(states)[ind], actions[ind]) + """define the loss function for TRPO""" + + def get_loss(volatile=False): + with torch.set_grad_enabled(not volatile): + log_probs = self.policy_net.get_log_prob(self.trans_policy(states[ind]), actions[ind]) + action_loss = -advantages[ind] * torch.exp(log_probs - fixed_log_probs) + return action_loss.mean() + + """use fisher information matrix for Hessian*vector""" + + def Fvp_fim(v): + M, mu, info = self.policy_net.get_fim(self.trans_policy(states)[ind]) + mu = mu.view(-1) + filter_input_ids = set([info['std_id']]) if self.policy_net.type == 'gaussian' else set() + + t = ones(mu.size(), requires_grad=True) + mu_t = (mu * t).sum() + Jt = compute_flat_grad(mu_t, self.policy_net.parameters(), filter_input_ids=filter_input_ids, create_graph=True) + Jtv = (Jt * v).sum() + Jv = torch.autograd.grad(Jtv, t)[0] + MJv = M * Jv.detach() + mu_MJv = (MJv * mu).sum() + JTMJv = compute_flat_grad(mu_MJv, self.policy_net.parameters(), filter_input_ids=filter_input_ids).detach() + JTMJv /= states.shape[0] + if self.policy_net.type == 'gaussian': + std_index = info['std_index'] + JTMJv[std_index: std_index + M.shape[0]] += 2 * v[std_index: std_index + M.shape[0]] + return JTMJv + v * self.damping + + """directly compute Hessian*vector from KL""" + + def Fvp_direct(v): + kl = self.policy_net.get_kl(self.trans_policy(states)[ind]) + kl = kl.mean() + + grads = torch.autograd.grad(kl, self.policy_net.parameters(), create_graph=True) + flat_grad_kl = torch.cat([grad.view(-1) for grad in grads]) + + kl_v = (flat_grad_kl * v).sum() + grads = torch.autograd.grad(kl_v, self.policy_net.parameters()) + flat_grad_grad_kl = torch.cat([grad.contiguous().view(-1) for grad in grads]).detach() + + return flat_grad_grad_kl + v * self.damping + + Fvp = Fvp_fim if self.use_fim else Fvp_direct + + loss = get_loss() + grads = torch.autograd.grad(loss, self.policy_net.parameters()) + loss_grad = torch.cat([grad.view(-1) for grad in grads]).detach() + stepdir = conjugate_gradients(Fvp, -loss_grad, 10) + + shs = 0.5 * (stepdir.dot(Fvp(stepdir))) + lm = math.sqrt(self.max_kl / shs) + fullstep = stepdir * lm + expected_improve = -loss_grad.dot(fullstep) + + prev_params = get_flat_params_from(self.policy_net) + success, new_params = line_search(self.policy_net, get_loss, prev_params, fullstep, expected_improve) + set_flat_params_to(self.policy_net, new_params) diff --git a/assets/mujoco_models/common/grass.png b/assets/mujoco_models/common/grass.png new file mode 100644 index 0000000..8844360 Binary files /dev/null and b/assets/mujoco_models/common/grass.png differ diff --git a/assets/mujoco_models/common/materials.xml b/assets/mujoco_models/common/materials.xml new file mode 100644 index 0000000..a20078b --- /dev/null +++ b/assets/mujoco_models/common/materials.xml @@ -0,0 +1,27 @@ + + + + + + + + + + + + + + + + + + + + + + + diff --git a/assets/mujoco_models/common/sky1.png b/assets/mujoco_models/common/sky1.png new file mode 100644 index 0000000..f8a4cc2 Binary files /dev/null and b/assets/mujoco_models/common/sky1.png differ diff --git a/assets/mujoco_models/common/skybox.xml b/assets/mujoco_models/common/skybox.xml new file mode 100644 index 0000000..d2a187c --- /dev/null +++ b/assets/mujoco_models/common/skybox.xml @@ -0,0 +1,8 @@ + + + + + + + diff --git a/assets/mujoco_models/common/visual.xml b/assets/mujoco_models/common/visual.xml new file mode 100644 index 0000000..ede15ad --- /dev/null +++ b/assets/mujoco_models/common/visual.xml @@ -0,0 +1,7 @@ + + + + + + + diff --git a/assets/mujoco_models/humanoid_1205_orig.xml b/assets/mujoco_models/humanoid_1205_orig.xml new file mode 100644 index 0000000..2411a36 --- /dev/null +++ b/assets/mujoco_models/humanoid_1205_orig.xml @@ -0,0 +1,193 @@ + + + + + + + + + + + + + + diff --git a/assets/mujoco_models/humanoid_1205_v1.xml b/assets/mujoco_models/humanoid_1205_v1.xml new file mode 100644 index 0000000..559e167 --- /dev/null +++ b/assets/mujoco_models/humanoid_1205_v1.xml @@ -0,0 +1,193 @@ + + + + + + + + + + + + + + diff --git a/assets/mujoco_models/humanoid_1205_vis.xml b/assets/mujoco_models/humanoid_1205_vis.xml new file mode 100644 index 0000000..084e15b --- /dev/null +++ b/assets/mujoco_models/humanoid_1205_vis.xml @@ -0,0 +1,261 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/assets/mujoco_models/humanoid_1205_vis_double_v1.xml b/assets/mujoco_models/humanoid_1205_vis_double_v1.xml new file mode 100644 index 0000000..2d7ea4b --- /dev/null +++ b/assets/mujoco_models/humanoid_1205_vis_double_v1.xml @@ -0,0 +1,265 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/assets/mujoco_models/humanoid_1205_vis_estimate_v1.xml b/assets/mujoco_models/humanoid_1205_vis_estimate_v1.xml new file mode 100644 index 0000000..7e7dcde --- /dev/null +++ b/assets/mujoco_models/humanoid_1205_vis_estimate_v1.xml @@ -0,0 +1,1562 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/assets/mujoco_models/humanoid_1205_vis_forecast_v1.xml b/assets/mujoco_models/humanoid_1205_vis_forecast_v1.xml new file mode 100644 index 0000000..5b70280 --- /dev/null +++ b/assets/mujoco_models/humanoid_1205_vis_forecast_v1.xml @@ -0,0 +1,1562 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/assets/mujoco_models/humanoid_1205_vis_ghost_v1.xml b/assets/mujoco_models/humanoid_1205_vis_ghost_v1.xml new file mode 100644 index 0000000..bb05149 --- /dev/null +++ b/assets/mujoco_models/humanoid_1205_vis_ghost_v1.xml @@ -0,0 +1,264 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/assets/mujoco_models/humanoid_1205_vis_multi_v1.xml b/assets/mujoco_models/humanoid_1205_vis_multi_v1.xml new file mode 100644 index 0000000..c7049eb --- /dev/null +++ b/assets/mujoco_models/humanoid_1205_vis_multi_v1.xml @@ -0,0 +1,2388 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/assets/mujoco_models/humanoid_1205_vis_single_v1.xml b/assets/mujoco_models/humanoid_1205_vis_single_v1.xml new file mode 100644 index 0000000..c4027f0 --- /dev/null +++ b/assets/mujoco_models/humanoid_1205_vis_single_v1.xml @@ -0,0 +1,146 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/assets/mujoco_models/template/humanoid_template.xml b/assets/mujoco_models/template/humanoid_template.xml new file mode 100644 index 0000000..6822e46 --- /dev/null +++ b/assets/mujoco_models/template/humanoid_template.xml @@ -0,0 +1,26 @@ + + + + + + + + + + + + + + diff --git a/config/egoforecast/cross_01.yml b/config/egoforecast/cross_01.yml new file mode 100644 index 0000000..bd69eb7 --- /dev/null +++ b/config/egoforecast/cross_01.yml @@ -0,0 +1,142 @@ +# data +meta_id: meta_cross_01 +cnn_feat: 'cross_01' +expert_feat: 'cross_01' +fr_margin: 30 + +# ego mimic net +ego_mimic_cfg: cross_01 +ego_mimic_iter: 6000 + +# training parameters +gamma: 0.95 +tau: 0.95 +policy_htype: relu +policy_hsize: [300, 200] +policy_v_hdim: 512 +policy_s_net: 'lstm' +policy_s_hdim: 128 +policy_optimizer: 'Adam' +policy_lr: 5.e-5 +policy_momentum: 0.0 +policy_weightdecay: 0.0 +value_htype: relu +value_hsize: [300, 200] +value_v_hdim: 512 +value_s_net: 'lstm' +value_s_hdim: 128 +value_optimizer: 'Adam' +value_lr: 3.e-4 +value_momentum: 0.0 +value_weightdecay: 0.0 +clip_epsilon: 0.2 +min_batch_size: 50000 +num_optim_epoch: 10 +log_std: -2.3 +fix_std: false +max_iter_num: 6000 +seed: 1 +save_model_interval: 100 +reward_id: 'quat_v3' +reward_weights: + w_p: 0.3 + w_v: 0.0 + w_e: 0.5 + w_rp: 0.1 + w_rv: 0.1 + k_p: 2 + k_v: 0.005 + k_e: 20 + k_rh: 300 + k_rq: 300 + k_rl: 1.0 + k_ra: 0.1 + decay: true +end_reward: false + +# expert and environment +mujoco_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis +env_episode_len: 90 +obs_coord: 'heading' +root_deheading: true + + +jkp_multiplier: 0.5 +joint_params: + # ["name", "k_p", "k_d", "a_ref", "a_scale", "torque_limit"] + - ["Spine_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Neck_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_y" , 400.0, 40.0, -80.0, 1.0, 100.0] + - ["RightArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightForeArm_z" , 300.0, 30.0, 45.0, 1.0, 60.0] + - ["RightHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_y" , 400.0, 40.0, 80.0, 1.0, 100.0] + - ["LeftArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftForeArm_z" , 300.0, 30.0, -45.0, 1.0, 60.0] + - ["LeftHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["RightFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["LeftFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + +body_params: + # ["name","diff weight"] + - ["Spine" , 1.0] + - ["Spine1" , 1.0] + - ["Spine2" , 1.0] + - ["Spine3" , 1.0] + - ["Neck" , 1.0] + - ["Head" , 1.0] + - ["RightShoulder" , 1.0] + - ["RightArm" , 1.0] + - ["RightForeArm" , 1.0] + - ["RightHand" , 1.0] + - ["LeftShoulder" , 1.0] + - ["LeftArm" , 1.0] + - ["LeftForeArm" , 1.0] + - ["LeftHand" , 1.0] + - ["RightUpLeg" , 1.0] + - ["RightLeg" , 1.0] + - ["RightFoot" , 1.0] + - ["LeftUpLeg" , 1.0] + - ["LeftLeg" , 1.0] + - ["LeftFoot" , 1.0] \ No newline at end of file diff --git a/config/egoforecast/subject_01.yml b/config/egoforecast/subject_01.yml new file mode 100644 index 0000000..34f74e8 --- /dev/null +++ b/config/egoforecast/subject_01.yml @@ -0,0 +1,142 @@ +# data +meta_id: meta_subject_01 +cnn_feat: 'subject_01' +expert_feat: 'subject_01' +fr_margin: 30 + +# ego mimic net +ego_mimic_cfg: subject_01 +ego_mimic_iter: 3000 + +# training parameters +gamma: 0.95 +tau: 0.95 +policy_htype: relu +policy_hsize: [300, 200] +policy_v_hdim: 128 +policy_s_net: 'lstm' +policy_s_hdim: 128 +policy_optimizer: 'Adam' +policy_lr: 5.e-5 +policy_momentum: 0.0 +policy_weightdecay: 0.0 +value_htype: relu +value_hsize: [300, 200] +value_v_hdim: 128 +value_s_net: 'lstm' +value_s_hdim: 128 +value_optimizer: 'Adam' +value_lr: 3.e-4 +value_momentum: 0.0 +value_weightdecay: 0.0 +clip_epsilon: 0.2 +min_batch_size: 50000 +num_optim_epoch: 10 +log_std: -2.3 +fix_std: true +max_iter_num: 3000 +seed: 1 +save_model_interval: 100 +reward_id: 'quat_v3' +reward_weights: + w_p: 0.3 + w_v: 0.0 + w_e: 0.5 + w_rp: 0.1 + w_rv: 0.1 + k_p: 2 + k_v: 0.005 + k_e: 20 + k_rh: 300 + k_rq: 300 + k_rl: 1.0 + k_ra: 0.1 + decay: true +end_reward: false + +# expert and environment +mujoco_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis +env_episode_len: 90 +obs_coord: 'heading' +root_deheading: true + + +jkp_multiplier: 0.5 +joint_params: + # ["name", "k_p", "k_d", "a_ref", "a_scale", "torque_limit"] + - ["Spine_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Neck_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_y" , 400.0, 40.0, -80.0, 1.0, 100.0] + - ["RightArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightForeArm_z" , 300.0, 30.0, 45.0, 1.0, 60.0] + - ["RightHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_y" , 400.0, 40.0, 80.0, 1.0, 100.0] + - ["LeftArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftForeArm_z" , 300.0, 30.0, -45.0, 1.0, 60.0] + - ["LeftHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["RightFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["LeftFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + +body_params: + # ["name","diff weight"] + - ["Spine" , 1.0] + - ["Spine1" , 1.0] + - ["Spine2" , 1.0] + - ["Spine3" , 1.0] + - ["Neck" , 1.0] + - ["Head" , 1.0] + - ["RightShoulder" , 1.0] + - ["RightArm" , 1.0] + - ["RightForeArm" , 1.0] + - ["RightHand" , 1.0] + - ["LeftShoulder" , 1.0] + - ["LeftArm" , 1.0] + - ["LeftForeArm" , 1.0] + - ["LeftHand" , 1.0] + - ["RightUpLeg" , 1.0] + - ["RightLeg" , 1.0] + - ["RightFoot" , 1.0] + - ["LeftUpLeg" , 1.0] + - ["LeftLeg" , 1.0] + - ["LeftFoot" , 1.0] \ No newline at end of file diff --git a/config/egoforecast/subject_02.yml b/config/egoforecast/subject_02.yml new file mode 100644 index 0000000..e00fc9d --- /dev/null +++ b/config/egoforecast/subject_02.yml @@ -0,0 +1,142 @@ +# data +meta_id: meta_subject_02 +cnn_feat: 'subject_02' +expert_feat: 'subject_02' +fr_margin: 30 + +# ego mimic net +ego_mimic_cfg: subject_02 +ego_mimic_iter: 3000 + +# training parameters +gamma: 0.95 +tau: 0.95 +policy_htype: relu +policy_hsize: [300, 200] +policy_v_hdim: 128 +policy_s_net: 'lstm' +policy_s_hdim: 128 +policy_optimizer: 'Adam' +policy_lr: 5.e-5 +policy_momentum: 0.0 +policy_weightdecay: 0.0 +value_htype: relu +value_hsize: [300, 200] +value_v_hdim: 128 +value_s_net: 'lstm' +value_s_hdim: 128 +value_optimizer: 'Adam' +value_lr: 3.e-4 +value_momentum: 0.0 +value_weightdecay: 0.0 +clip_epsilon: 0.2 +min_batch_size: 50000 +num_optim_epoch: 10 +log_std: -2.3 +fix_std: true +max_iter_num: 3000 +seed: 1 +save_model_interval: 100 +reward_id: 'quat_v3' +reward_weights: + w_p: 0.3 + w_v: 0.0 + w_e: 0.5 + w_rp: 0.1 + w_rv: 0.1 + k_p: 2 + k_v: 0.005 + k_e: 20 + k_rh: 300 + k_rq: 300 + k_rl: 1.0 + k_ra: 0.1 + decay: true +end_reward: false + +# expert and environment +mujoco_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis +env_episode_len: 90 +obs_coord: 'heading' +root_deheading: true + + +jkp_multiplier: 0.5 +joint_params: + # ["name", "k_p", "k_d", "a_ref", "a_scale", "torque_limit"] + - ["Spine_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Neck_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_y" , 400.0, 40.0, -80.0, 1.0, 100.0] + - ["RightArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightForeArm_z" , 300.0, 30.0, 45.0, 1.0, 60.0] + - ["RightHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_y" , 400.0, 40.0, 80.0, 1.0, 100.0] + - ["LeftArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftForeArm_z" , 300.0, 30.0, -45.0, 1.0, 60.0] + - ["LeftHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["RightFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["LeftFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + +body_params: + # ["name","diff weight"] + - ["Spine" , 1.0] + - ["Spine1" , 1.0] + - ["Spine2" , 1.0] + - ["Spine3" , 1.0] + - ["Neck" , 1.0] + - ["Head" , 1.0] + - ["RightShoulder" , 1.0] + - ["RightArm" , 1.0] + - ["RightForeArm" , 1.0] + - ["RightHand" , 1.0] + - ["LeftShoulder" , 1.0] + - ["LeftArm" , 1.0] + - ["LeftForeArm" , 1.0] + - ["LeftHand" , 1.0] + - ["RightUpLeg" , 1.0] + - ["RightLeg" , 1.0] + - ["RightFoot" , 1.0] + - ["LeftUpLeg" , 1.0] + - ["LeftLeg" , 1.0] + - ["LeftFoot" , 1.0] \ No newline at end of file diff --git a/config/egoforecast/subject_03.yml b/config/egoforecast/subject_03.yml new file mode 100644 index 0000000..255ee8d --- /dev/null +++ b/config/egoforecast/subject_03.yml @@ -0,0 +1,142 @@ +# data +meta_id: meta_subject_03 +cnn_feat: 'subject_03' +expert_feat: 'subject_03' +fr_margin: 30 + +# ego mimic net +ego_mimic_cfg: subject_03 +ego_mimic_iter: 3000 + +# training parameters +gamma: 0.95 +tau: 0.95 +policy_htype: relu +policy_hsize: [300, 200] +policy_v_hdim: 128 +policy_s_net: 'lstm' +policy_s_hdim: 128 +policy_optimizer: 'Adam' +policy_lr: 5.e-5 +policy_momentum: 0.0 +policy_weightdecay: 0.0 +value_htype: relu +value_hsize: [300, 200] +value_v_hdim: 128 +value_s_net: 'lstm' +value_s_hdim: 128 +value_optimizer: 'Adam' +value_lr: 3.e-4 +value_momentum: 0.0 +value_weightdecay: 0.0 +clip_epsilon: 0.2 +min_batch_size: 50000 +num_optim_epoch: 10 +log_std: -2.3 +fix_std: true +max_iter_num: 3000 +seed: 1 +save_model_interval: 100 +reward_id: 'quat_v3' +reward_weights: + w_p: 0.3 + w_v: 0.0 + w_e: 0.5 + w_rp: 0.1 + w_rv: 0.1 + k_p: 2 + k_v: 0.005 + k_e: 20 + k_rh: 300 + k_rq: 300 + k_rl: 1.0 + k_ra: 0.1 + decay: true +end_reward: false + +# expert and environment +mujoco_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis +env_episode_len: 90 +obs_coord: 'heading' +root_deheading: true + + +jkp_multiplier: 0.5 +joint_params: + # ["name", "k_p", "k_d", "a_ref", "a_scale", "torque_limit"] + - ["Spine_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Neck_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_y" , 400.0, 40.0, -80.0, 1.0, 100.0] + - ["RightArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightForeArm_z" , 300.0, 30.0, 45.0, 1.0, 60.0] + - ["RightHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_y" , 400.0, 40.0, 80.0, 1.0, 100.0] + - ["LeftArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftForeArm_z" , 300.0, 30.0, -45.0, 1.0, 60.0] + - ["LeftHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["RightFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["LeftFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + +body_params: + # ["name","diff weight"] + - ["Spine" , 1.0] + - ["Spine1" , 1.0] + - ["Spine2" , 1.0] + - ["Spine3" , 1.0] + - ["Neck" , 1.0] + - ["Head" , 1.0] + - ["RightShoulder" , 1.0] + - ["RightArm" , 1.0] + - ["RightForeArm" , 1.0] + - ["RightHand" , 1.0] + - ["LeftShoulder" , 1.0] + - ["LeftArm" , 1.0] + - ["LeftForeArm" , 1.0] + - ["LeftHand" , 1.0] + - ["RightUpLeg" , 1.0] + - ["RightLeg" , 1.0] + - ["RightFoot" , 1.0] + - ["LeftUpLeg" , 1.0] + - ["LeftLeg" , 1.0] + - ["LeftFoot" , 1.0] \ No newline at end of file diff --git a/config/egoforecast/subject_04.yml b/config/egoforecast/subject_04.yml new file mode 100644 index 0000000..4bbd375 --- /dev/null +++ b/config/egoforecast/subject_04.yml @@ -0,0 +1,142 @@ +# data +meta_id: meta_subject_04 +cnn_feat: 'subject_04' +expert_feat: 'subject_04' +fr_margin: 30 + +# ego mimic net +ego_mimic_cfg: subject_04 +ego_mimic_iter: 3000 + +# training parameters +gamma: 0.95 +tau: 0.95 +policy_htype: relu +policy_hsize: [300, 200] +policy_v_hdim: 128 +policy_s_net: 'lstm' +policy_s_hdim: 128 +policy_optimizer: 'Adam' +policy_lr: 5.e-5 +policy_momentum: 0.0 +policy_weightdecay: 0.0 +value_htype: relu +value_hsize: [300, 200] +value_v_hdim: 128 +value_s_net: 'lstm' +value_s_hdim: 128 +value_optimizer: 'Adam' +value_lr: 3.e-4 +value_momentum: 0.0 +value_weightdecay: 0.0 +clip_epsilon: 0.2 +min_batch_size: 50000 +num_optim_epoch: 10 +log_std: -2.3 +fix_std: true +max_iter_num: 3000 +seed: 1 +save_model_interval: 100 +reward_id: 'quat_v3' +reward_weights: + w_p: 0.3 + w_v: 0.0 + w_e: 0.5 + w_rp: 0.1 + w_rv: 0.1 + k_p: 2 + k_v: 0.005 + k_e: 20 + k_rh: 300 + k_rq: 300 + k_rl: 1.0 + k_ra: 0.1 + decay: true +end_reward: false + +# expert and environment +mujoco_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis +env_episode_len: 90 +obs_coord: 'heading' +root_deheading: true + + +jkp_multiplier: 0.5 +joint_params: + # ["name", "k_p", "k_d", "a_ref", "a_scale", "torque_limit"] + - ["Spine_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Neck_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_y" , 400.0, 40.0, -80.0, 1.0, 100.0] + - ["RightArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightForeArm_z" , 300.0, 30.0, 45.0, 1.0, 60.0] + - ["RightHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_y" , 400.0, 40.0, 80.0, 1.0, 100.0] + - ["LeftArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftForeArm_z" , 300.0, 30.0, -45.0, 1.0, 60.0] + - ["LeftHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["RightFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["LeftFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + +body_params: + # ["name","diff weight"] + - ["Spine" , 1.0] + - ["Spine1" , 1.0] + - ["Spine2" , 1.0] + - ["Spine3" , 1.0] + - ["Neck" , 1.0] + - ["Head" , 1.0] + - ["RightShoulder" , 1.0] + - ["RightArm" , 1.0] + - ["RightForeArm" , 1.0] + - ["RightHand" , 1.0] + - ["LeftShoulder" , 1.0] + - ["LeftArm" , 1.0] + - ["LeftForeArm" , 1.0] + - ["LeftHand" , 1.0] + - ["RightUpLeg" , 1.0] + - ["RightLeg" , 1.0] + - ["RightFoot" , 1.0] + - ["LeftUpLeg" , 1.0] + - ["LeftLeg" , 1.0] + - ["LeftFoot" , 1.0] \ No newline at end of file diff --git a/config/egoforecast/subject_05.yml b/config/egoforecast/subject_05.yml new file mode 100644 index 0000000..4883494 --- /dev/null +++ b/config/egoforecast/subject_05.yml @@ -0,0 +1,142 @@ +# data +meta_id: meta_subject_05 +cnn_feat: 'subject_05' +expert_feat: 'subject_05' +fr_margin: 30 + +# ego mimic net +ego_mimic_cfg: subject_05 +ego_mimic_iter: 3000 + +# training parameters +gamma: 0.95 +tau: 0.95 +policy_htype: relu +policy_hsize: [300, 200] +policy_v_hdim: 128 +policy_s_net: 'lstm' +policy_s_hdim: 128 +policy_optimizer: 'Adam' +policy_lr: 5.e-5 +policy_momentum: 0.0 +policy_weightdecay: 0.0 +value_htype: relu +value_hsize: [300, 200] +value_v_hdim: 128 +value_s_net: 'lstm' +value_s_hdim: 128 +value_optimizer: 'Adam' +value_lr: 3.e-4 +value_momentum: 0.0 +value_weightdecay: 0.0 +clip_epsilon: 0.2 +min_batch_size: 50000 +num_optim_epoch: 10 +log_std: -2.3 +fix_std: true +max_iter_num: 3000 +seed: 1 +save_model_interval: 100 +reward_id: 'quat_v3' +reward_weights: + w_p: 0.3 + w_v: 0.0 + w_e: 0.5 + w_rp: 0.1 + w_rv: 0.1 + k_p: 2 + k_v: 0.005 + k_e: 20 + k_rh: 300 + k_rq: 300 + k_rl: 1.0 + k_ra: 0.1 + decay: true +end_reward: false + +# expert and environment +mujoco_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis +env_episode_len: 90 +obs_coord: 'heading' +root_deheading: true + + +jkp_multiplier: 0.5 +joint_params: + # ["name", "k_p", "k_d", "a_ref", "a_scale", "torque_limit"] + - ["Spine_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Neck_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_y" , 400.0, 40.0, -80.0, 1.0, 100.0] + - ["RightArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightForeArm_z" , 300.0, 30.0, 45.0, 1.0, 60.0] + - ["RightHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_y" , 400.0, 40.0, 80.0, 1.0, 100.0] + - ["LeftArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftForeArm_z" , 300.0, 30.0, -45.0, 1.0, 60.0] + - ["LeftHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["RightFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["LeftFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + +body_params: + # ["name","diff weight"] + - ["Spine" , 1.0] + - ["Spine1" , 1.0] + - ["Spine2" , 1.0] + - ["Spine3" , 1.0] + - ["Neck" , 1.0] + - ["Head" , 1.0] + - ["RightShoulder" , 1.0] + - ["RightArm" , 1.0] + - ["RightForeArm" , 1.0] + - ["RightHand" , 1.0] + - ["LeftShoulder" , 1.0] + - ["LeftArm" , 1.0] + - ["LeftForeArm" , 1.0] + - ["LeftHand" , 1.0] + - ["RightUpLeg" , 1.0] + - ["RightLeg" , 1.0] + - ["RightFoot" , 1.0] + - ["LeftUpLeg" , 1.0] + - ["LeftLeg" , 1.0] + - ["LeftFoot" , 1.0] \ No newline at end of file diff --git a/config/egomimic/cross_01.yml b/config/egomimic/cross_01.yml new file mode 100644 index 0000000..f7e9f25 --- /dev/null +++ b/config/egomimic/cross_01.yml @@ -0,0 +1,136 @@ +# data +meta_id: meta_cross_01 +cnn_feat: 'cross_01' +expert_feat: 'cross_01' +fr_margin: 10 + +# state net +state_net_cfg: cross_01 +state_net_iter: 100 + +# training parameters +gamma: 0.95 +tau: 0.95 +policy_htype: relu +policy_hsize: [300, 200] +policy_v_hdim: 128 +policy_optimizer: 'Adam' +policy_lr: 5.e-5 +policy_momentum: 0.0 +policy_weightdecay: 0.0 +value_htype: relu +value_hsize: [300, 200] +value_v_hdim: 128 +value_optimizer: 'Adam' +value_lr: 3.e-4 +value_momentum: 0.0 +value_weightdecay: 0.0 +clip_epsilon: 0.2 +min_batch_size: 50000 +num_optim_epoch: 10 +log_std: -2.3 +fix_std: true +max_iter_num: 6000 +seed: 1 +save_model_interval: 100 +reward_id: 'quat_v3' +reward_weights: + w_p: 0.5 + w_v: 0.0 + w_e: 0.3 + w_rp: 0.1 + w_rv: 0.1 + k_p: 2 + k_v: 0.005 + k_e: 20 + k_rh: 300 + k_rq: 300 + k_rl: 1.0 + k_ra: 0.1 + +# expert and environment +mujoco_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis +env_episode_len: 200 +obs_coord: 'heading' +root_deheading: true + + +jkp_multiplier: 0.5 +joint_params: + # ["name", "k_p", "k_d", "a_ref", "a_scale", "torque_limit"] + - ["Spine_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Neck_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_y" , 400.0, 40.0, -80.0, 1.0, 100.0] + - ["RightArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightForeArm_z" , 300.0, 30.0, 45.0, 1.0, 60.0] + - ["RightHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_y" , 400.0, 40.0, 80.0, 1.0, 100.0] + - ["LeftArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftForeArm_z" , 300.0, 30.0, -45.0, 1.0, 60.0] + - ["LeftHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["RightFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["LeftFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + +body_params: + # ["name","diff weight"] + - ["Spine" , 1.0] + - ["Spine1" , 1.0] + - ["Spine2" , 1.0] + - ["Spine3" , 1.0] + - ["Neck" , 1.0] + - ["Head" , 1.0] + - ["RightShoulder" , 1.0] + - ["RightArm" , 1.0] + - ["RightForeArm" , 1.0] + - ["RightHand" , 1.0] + - ["LeftShoulder" , 1.0] + - ["LeftArm" , 1.0] + - ["LeftForeArm" , 1.0] + - ["LeftHand" , 1.0] + - ["RightUpLeg" , 1.0] + - ["RightLeg" , 1.0] + - ["RightFoot" , 1.0] + - ["LeftUpLeg" , 1.0] + - ["LeftLeg" , 1.0] + - ["LeftFoot" , 1.0] \ No newline at end of file diff --git a/config/egomimic/subject_01.yml b/config/egomimic/subject_01.yml new file mode 100644 index 0000000..2684e1a --- /dev/null +++ b/config/egomimic/subject_01.yml @@ -0,0 +1,136 @@ +# data +meta_id: meta_subject_01 +cnn_feat: 'subject_01' +expert_feat: 'subject_01' +fr_margin: 10 + +# state net +state_net_cfg: subject_01 +state_net_iter: 100 + +# training parameters +gamma: 0.95 +tau: 0.95 +policy_htype: relu +policy_hsize: [300, 200] +policy_v_hdim: 128 +policy_optimizer: 'Adam' +policy_lr: 5.e-5 +policy_momentum: 0.0 +policy_weightdecay: 0.0 +value_htype: relu +value_hsize: [300, 200] +value_v_hdim: 128 +value_optimizer: 'Adam' +value_lr: 3.e-4 +value_momentum: 0.0 +value_weightdecay: 0.0 +clip_epsilon: 0.2 +min_batch_size: 50000 +num_optim_epoch: 10 +log_std: -2.3 +fix_std: true +max_iter_num: 3000 +seed: 1 +save_model_interval: 100 +reward_id: 'quat_v3' +reward_weights: + w_p: 0.5 + w_v: 0.0 + w_e: 0.3 + w_rp: 0.1 + w_rv: 0.1 + k_p: 2 + k_v: 0.005 + k_e: 20 + k_rh: 300 + k_rq: 300 + k_rl: 1.0 + k_ra: 0.1 + +# expert and environment +mujoco_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis +env_episode_len: 200 +obs_coord: 'heading' +root_deheading: true + + +jkp_multiplier: 0.5 +joint_params: + # ["name", "k_p", "k_d", "a_ref", "a_scale", "torque_limit"] + - ["Spine_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Neck_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_y" , 400.0, 40.0, -80.0, 1.0, 100.0] + - ["RightArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightForeArm_z" , 300.0, 30.0, 45.0, 1.0, 60.0] + - ["RightHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_y" , 400.0, 40.0, 80.0, 1.0, 100.0] + - ["LeftArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftForeArm_z" , 300.0, 30.0, -45.0, 1.0, 60.0] + - ["LeftHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["RightFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["LeftFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + +body_params: + # ["name","diff weight"] + - ["Spine" , 1.0] + - ["Spine1" , 1.0] + - ["Spine2" , 1.0] + - ["Spine3" , 1.0] + - ["Neck" , 1.0] + - ["Head" , 1.0] + - ["RightShoulder" , 1.0] + - ["RightArm" , 1.0] + - ["RightForeArm" , 1.0] + - ["RightHand" , 1.0] + - ["LeftShoulder" , 1.0] + - ["LeftArm" , 1.0] + - ["LeftForeArm" , 1.0] + - ["LeftHand" , 1.0] + - ["RightUpLeg" , 1.0] + - ["RightLeg" , 1.0] + - ["RightFoot" , 1.0] + - ["LeftUpLeg" , 1.0] + - ["LeftLeg" , 1.0] + - ["LeftFoot" , 1.0] \ No newline at end of file diff --git a/config/egomimic/subject_02.yml b/config/egomimic/subject_02.yml new file mode 100644 index 0000000..e472b2e --- /dev/null +++ b/config/egomimic/subject_02.yml @@ -0,0 +1,136 @@ +# data +meta_id: meta_subject_02 +cnn_feat: 'subject_02' +expert_feat: 'subject_02' +fr_margin: 10 + +# state net +state_net_cfg: 0301-3 +state_net_iter: 100 + +# training parameters +gamma: 0.95 +tau: 0.95 +policy_htype: relu +policy_hsize: [300, 200] +policy_v_hdim: 128 +policy_optimizer: 'Adam' +policy_lr: 5.e-5 +policy_momentum: 0.0 +policy_weightdecay: 0.0 +value_htype: relu +value_hsize: [300, 200] +value_v_hdim: 128 +value_optimizer: 'Adam' +value_lr: 3.e-4 +value_momentum: 0.0 +value_weightdecay: 0.0 +clip_epsilon: 0.2 +min_batch_size: 50000 +num_optim_epoch: 10 +log_std: -2.3 +fix_std: true +max_iter_num: 3000 +seed: 1 +save_model_interval: 100 +reward_id: 'quat_v3' +reward_weights: + w_p: 0.5 + w_v: 0.0 + w_e: 0.3 + w_rp: 0.1 + w_rv: 0.1 + k_p: 2 + k_v: 0.005 + k_e: 20 + k_rh: 300 + k_rq: 300 + k_rl: 1.0 + k_ra: 0.1 + +# expert and environment +mujoco_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis +env_episode_len: 200 +obs_coord: 'heading' +root_deheading: true + + +jkp_multiplier: 0.5 +joint_params: + # ["name", "k_p", "k_d", "a_ref", "a_scale", "torque_limit"] + - ["Spine_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Neck_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_y" , 400.0, 40.0, -80.0, 1.0, 100.0] + - ["RightArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightForeArm_z" , 300.0, 30.0, 45.0, 1.0, 60.0] + - ["RightHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_y" , 400.0, 40.0, 80.0, 1.0, 100.0] + - ["LeftArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftForeArm_z" , 300.0, 30.0, -45.0, 1.0, 60.0] + - ["LeftHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["RightFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["LeftFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + +body_params: + # ["name","diff weight"] + - ["Spine" , 1.0] + - ["Spine1" , 1.0] + - ["Spine2" , 1.0] + - ["Spine3" , 1.0] + - ["Neck" , 1.0] + - ["Head" , 1.0] + - ["RightShoulder" , 1.0] + - ["RightArm" , 1.0] + - ["RightForeArm" , 1.0] + - ["RightHand" , 1.0] + - ["LeftShoulder" , 1.0] + - ["LeftArm" , 1.0] + - ["LeftForeArm" , 1.0] + - ["LeftHand" , 1.0] + - ["RightUpLeg" , 1.0] + - ["RightLeg" , 1.0] + - ["RightFoot" , 1.0] + - ["LeftUpLeg" , 1.0] + - ["LeftLeg" , 1.0] + - ["LeftFoot" , 1.0] \ No newline at end of file diff --git a/config/egomimic/subject_03.yml b/config/egomimic/subject_03.yml new file mode 100644 index 0000000..ef936d8 --- /dev/null +++ b/config/egomimic/subject_03.yml @@ -0,0 +1,136 @@ +# data +meta_id: meta_subject_03 +cnn_feat: 'subject_03' +expert_feat: 'subject_03' +fr_margin: 10 + +# state net +state_net_cfg: subject_03 +state_net_iter: 100 + +# training parameters +gamma: 0.95 +tau: 0.95 +policy_htype: relu +policy_hsize: [300, 200] +policy_v_hdim: 128 +policy_optimizer: 'Adam' +policy_lr: 5.e-5 +policy_momentum: 0.0 +policy_weightdecay: 0.0 +value_htype: relu +value_hsize: [300, 200] +value_v_hdim: 128 +value_optimizer: 'Adam' +value_lr: 3.e-4 +value_momentum: 0.0 +value_weightdecay: 0.0 +clip_epsilon: 0.2 +min_batch_size: 50000 +num_optim_epoch: 10 +log_std: -2.3 +fix_std: true +max_iter_num: 3000 +seed: 1 +save_model_interval: 100 +reward_id: 'quat_v3' +reward_weights: + w_p: 0.5 + w_v: 0.0 + w_e: 0.3 + w_rp: 0.1 + w_rv: 0.1 + k_p: 2 + k_v: 0.005 + k_e: 20 + k_rh: 300 + k_rq: 300 + k_rl: 1.0 + k_ra: 0.1 + +# expert and environment +mujoco_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis +env_episode_len: 200 +obs_coord: 'heading' +root_deheading: true + + +jkp_multiplier: 0.5 +joint_params: + # ["name", "k_p", "k_d", "a_ref", "a_scale", "torque_limit"] + - ["Spine_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Neck_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_y" , 400.0, 40.0, -80.0, 1.0, 100.0] + - ["RightArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightForeArm_z" , 300.0, 30.0, 45.0, 1.0, 60.0] + - ["RightHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_y" , 400.0, 40.0, 80.0, 1.0, 100.0] + - ["LeftArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftForeArm_z" , 300.0, 30.0, -45.0, 1.0, 60.0] + - ["LeftHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["RightFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["LeftFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + +body_params: + # ["name","diff weight"] + - ["Spine" , 1.0] + - ["Spine1" , 1.0] + - ["Spine2" , 1.0] + - ["Spine3" , 1.0] + - ["Neck" , 1.0] + - ["Head" , 1.0] + - ["RightShoulder" , 1.0] + - ["RightArm" , 1.0] + - ["RightForeArm" , 1.0] + - ["RightHand" , 1.0] + - ["LeftShoulder" , 1.0] + - ["LeftArm" , 1.0] + - ["LeftForeArm" , 1.0] + - ["LeftHand" , 1.0] + - ["RightUpLeg" , 1.0] + - ["RightLeg" , 1.0] + - ["RightFoot" , 1.0] + - ["LeftUpLeg" , 1.0] + - ["LeftLeg" , 1.0] + - ["LeftFoot" , 1.0] \ No newline at end of file diff --git a/config/egomimic/subject_04.yml b/config/egomimic/subject_04.yml new file mode 100644 index 0000000..8169308 --- /dev/null +++ b/config/egomimic/subject_04.yml @@ -0,0 +1,136 @@ +# data +meta_id: meta_subject_04 +cnn_feat: 'subject_04' +expert_feat: 'subject_04' +fr_margin: 10 + +# state net +state_net_cfg: subject_04 +state_net_iter: 100 + +# training parameters +gamma: 0.95 +tau: 0.95 +policy_htype: relu +policy_hsize: [300, 200] +policy_v_hdim: 128 +policy_optimizer: 'Adam' +policy_lr: 5.e-5 +policy_momentum: 0.0 +policy_weightdecay: 0.0 +value_htype: relu +value_hsize: [300, 200] +value_v_hdim: 128 +value_optimizer: 'Adam' +value_lr: 3.e-4 +value_momentum: 0.0 +value_weightdecay: 0.0 +clip_epsilon: 0.2 +min_batch_size: 50000 +num_optim_epoch: 10 +log_std: -2.3 +fix_std: true +max_iter_num: 3000 +seed: 1 +save_model_interval: 100 +reward_id: 'quat_v3' +reward_weights: + w_p: 0.5 + w_v: 0.0 + w_e: 0.3 + w_rp: 0.1 + w_rv: 0.1 + k_p: 2 + k_v: 0.005 + k_e: 20 + k_rh: 300 + k_rq: 300 + k_rl: 1.0 + k_ra: 0.1 + +# expert and environment +mujoco_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis +env_episode_len: 200 +obs_coord: 'heading' +root_deheading: true + + +jkp_multiplier: 0.5 +joint_params: + # ["name", "k_p", "k_d", "a_ref", "a_scale", "torque_limit"] + - ["Spine_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Neck_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_y" , 400.0, 40.0, -80.0, 1.0, 100.0] + - ["RightArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightForeArm_z" , 300.0, 30.0, 45.0, 1.0, 60.0] + - ["RightHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_y" , 400.0, 40.0, 80.0, 1.0, 100.0] + - ["LeftArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftForeArm_z" , 300.0, 30.0, -45.0, 1.0, 60.0] + - ["LeftHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["RightFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["LeftFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + +body_params: + # ["name","diff weight"] + - ["Spine" , 1.0] + - ["Spine1" , 1.0] + - ["Spine2" , 1.0] + - ["Spine3" , 1.0] + - ["Neck" , 1.0] + - ["Head" , 1.0] + - ["RightShoulder" , 1.0] + - ["RightArm" , 1.0] + - ["RightForeArm" , 1.0] + - ["RightHand" , 1.0] + - ["LeftShoulder" , 1.0] + - ["LeftArm" , 1.0] + - ["LeftForeArm" , 1.0] + - ["LeftHand" , 1.0] + - ["RightUpLeg" , 1.0] + - ["RightLeg" , 1.0] + - ["RightFoot" , 1.0] + - ["LeftUpLeg" , 1.0] + - ["LeftLeg" , 1.0] + - ["LeftFoot" , 1.0] \ No newline at end of file diff --git a/config/egomimic/subject_05.yml b/config/egomimic/subject_05.yml new file mode 100644 index 0000000..8ed37e1 --- /dev/null +++ b/config/egomimic/subject_05.yml @@ -0,0 +1,136 @@ +# data +meta_id: meta_subject_05 +cnn_feat: 'subject_05' +expert_feat: 'subject_05' +fr_margin: 10 + +# state net +state_net_cfg: subject_05 +state_net_iter: 100 + +# training parameters +gamma: 0.95 +tau: 0.95 +policy_htype: relu +policy_hsize: [300, 200] +policy_v_hdim: 128 +policy_optimizer: 'Adam' +policy_lr: 5.e-5 +policy_momentum: 0.0 +policy_weightdecay: 0.0 +value_htype: relu +value_hsize: [300, 200] +value_v_hdim: 128 +value_optimizer: 'Adam' +value_lr: 3.e-4 +value_momentum: 0.0 +value_weightdecay: 0.0 +clip_epsilon: 0.2 +min_batch_size: 50000 +num_optim_epoch: 10 +log_std: -2.3 +fix_std: true +max_iter_num: 3000 +seed: 1 +save_model_interval: 100 +reward_id: 'quat_v3' +reward_weights: + w_p: 0.5 + w_v: 0.1 + w_e: 0.2 + w_rp: 0.1 + w_rv: 0.1 + k_p: 2 + k_v: 0.005 + k_e: 20 + k_rh: 300 + k_rq: 300 + k_rl: 1.0 + k_ra: 0.1 + +# expert and environment +mujoco_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis +env_episode_len: 200 +obs_coord: 'heading' +root_deheading: true + + +jkp_multiplier: 0.5 +joint_params: + # ["name", "k_p", "k_d", "a_ref", "a_scale", "torque_limit"] + - ["Spine_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine1_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine2_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_x" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_y" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Spine3_z" , 1000.0, 100.0, 0.0, 1.0, 200.0] + - ["Neck_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Neck_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["Head_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightArm_y" , 400.0, 40.0, -80.0, 1.0, 100.0] + - ["RightArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["RightForeArm_z" , 300.0, 30.0, 45.0, 1.0, 60.0] + - ["RightHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftShoulder_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_y" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftShoulder_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_x" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftArm_y" , 400.0, 40.0, 80.0, 1.0, 100.0] + - ["LeftArm_z" , 400.0, 40.0, 0.0, 1.0, 100.0] + - ["LeftForeArm_z" , 300.0, 30.0, -45.0, 1.0, 60.0] + - ["LeftHand_x" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_y" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["LeftHand_z" , 100.0, 10.0, 0.0, 1.0, 50.0] + - ["RightUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["RightLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["RightFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["RightFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftUpLeg_x" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_y" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftUpLeg_z" , 500.0, 50.0, 0.0, 1.0, 200.0] + - ["LeftLeg_x" , 500.0, 50.0, 45.0, 1.0, 150.0] + - ["LeftFoot_x" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_y" , 400.0, 40.0, 0.0, 1.0, 90.0] + - ["LeftFoot_z" , 400.0, 40.0, 0.0, 1.0, 90.0] + +body_params: + # ["name","diff weight"] + - ["Spine" , 1.0] + - ["Spine1" , 1.0] + - ["Spine2" , 1.0] + - ["Spine3" , 1.0] + - ["Neck" , 1.0] + - ["Head" , 1.0] + - ["RightShoulder" , 1.0] + - ["RightArm" , 1.0] + - ["RightForeArm" , 1.0] + - ["RightHand" , 1.0] + - ["LeftShoulder" , 1.0] + - ["LeftArm" , 1.0] + - ["LeftForeArm" , 1.0] + - ["LeftHand" , 1.0] + - ["RightUpLeg" , 1.0] + - ["RightLeg" , 1.0] + - ["RightFoot" , 1.0] + - ["LeftUpLeg" , 1.0] + - ["LeftLeg" , 1.0] + - ["LeftFoot" , 1.0] \ No newline at end of file diff --git a/config/statereg/cross_01.yml b/config/statereg/cross_01.yml new file mode 100644 index 0000000..59e063c --- /dev/null +++ b/config/statereg/cross_01.yml @@ -0,0 +1,16 @@ +meta_id: meta_cross_01 +seed: 1 +fr_num: 120 +v_hdim: 128 +mlp_dim: [300, 200] +cnn_fdim: 128 +lr: 1.e-4 +num_epoch: 100 +iter_method: sample +num_sample: 20000 +save_model_interval: 20 +fr_margin: 10 + +# misc +humanoid_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis \ No newline at end of file diff --git a/config/statereg/subject_01.yml b/config/statereg/subject_01.yml new file mode 100644 index 0000000..7d6d181 --- /dev/null +++ b/config/statereg/subject_01.yml @@ -0,0 +1,16 @@ +meta_id: meta_subject_01 +seed: 1 +fr_num: 120 +v_hdim: 128 +mlp_dim: [300, 200] +cnn_fdim: 128 +lr: 1.e-4 +num_epoch: 100 +iter_method: iter +shuffle: true +save_model_interval: 20 +fr_margin: 10 + +# misc +humanoid_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis \ No newline at end of file diff --git a/config/statereg/subject_02.yml b/config/statereg/subject_02.yml new file mode 100644 index 0000000..8732660 --- /dev/null +++ b/config/statereg/subject_02.yml @@ -0,0 +1,16 @@ +meta_id: meta_subject_02 +seed: 1 +fr_num: 120 +v_hdim: 128 +mlp_dim: [300, 200] +cnn_fdim: 128 +lr: 1.e-4 +num_epoch: 100 +iter_method: iter +shuffle: true +save_model_interval: 20 +fr_margin: 10 + +# misc +humanoid_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis \ No newline at end of file diff --git a/config/statereg/subject_03.yml b/config/statereg/subject_03.yml new file mode 100644 index 0000000..301e2c4 --- /dev/null +++ b/config/statereg/subject_03.yml @@ -0,0 +1,16 @@ +meta_id: meta_subject_03 +seed: 1 +fr_num: 120 +v_hdim: 128 +mlp_dim: [300, 200] +cnn_fdim: 128 +lr: 1.e-4 +num_epoch: 100 +iter_method: iter +shuffle: true +save_model_interval: 20 +fr_margin: 10 + +# misc +humanoid_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis \ No newline at end of file diff --git a/config/statereg/subject_04.yml b/config/statereg/subject_04.yml new file mode 100644 index 0000000..e4a069a --- /dev/null +++ b/config/statereg/subject_04.yml @@ -0,0 +1,16 @@ +meta_id: meta_subject_04 +seed: 1 +fr_num: 120 +v_hdim: 128 +mlp_dim: [300, 200] +cnn_fdim: 128 +lr: 1.e-4 +num_epoch: 100 +iter_method: iter +shuffle: true +save_model_interval: 20 +fr_margin: 10 + +# misc +humanoid_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis \ No newline at end of file diff --git a/config/statereg/subject_05.yml b/config/statereg/subject_05.yml new file mode 100644 index 0000000..678dfca --- /dev/null +++ b/config/statereg/subject_05.yml @@ -0,0 +1,16 @@ +meta_id: meta_subject_05 +seed: 1 +fr_num: 120 +v_hdim: 128 +mlp_dim: [300, 200] +cnn_fdim: 128 +lr: 1.e-4 +num_epoch: 100 +iter_method: iter +shuffle: true +save_model_interval: 20 +fr_margin: 10 + +# misc +humanoid_model: humanoid_1205_v1 +vis_model: humanoid_1205_vis \ No newline at end of file diff --git a/core/__init__.py b/core/__init__.py new file mode 100644 index 0000000..789871f --- /dev/null +++ b/core/__init__.py @@ -0,0 +1,8 @@ +from core.common import * +from core.critic import Value +from core.distributions import DiagGaussian, Categorical +from core.logger_rl import LoggerRL +from core.policy import Policy +from core.policy_disc import PolicyDiscrete +from core.policy_gaussian import PolicyGaussian +from core.trajbatch import TrajBatch diff --git a/core/common.py b/core/common.py new file mode 100644 index 0000000..e683a0e --- /dev/null +++ b/core/common.py @@ -0,0 +1,25 @@ +import torch +from utils import batch_to + + +def estimate_advantages(rewards, masks, values, gamma, tau): + device = rewards.device + rewards, masks, values = batch_to(torch.device('cpu'), rewards, masks, values) + tensor_type = type(rewards) + deltas = tensor_type(rewards.size(0), 1) + advantages = tensor_type(rewards.size(0), 1) + + prev_value = 0 + prev_advantage = 0 + for i in reversed(range(rewards.size(0))): + deltas[i] = rewards[i] + gamma * prev_value * masks[i] - values[i] + advantages[i] = deltas[i] + gamma * tau * prev_advantage * masks[i] + + prev_value = values[i, 0] + prev_advantage = advantages[i, 0] + + returns = values + advantages + advantages = (advantages - advantages.mean()) / advantages.std() + + advantages, returns = batch_to(device, advantages, returns) + return advantages, returns diff --git a/core/critic.py b/core/critic.py new file mode 100644 index 0000000..a44c557 --- /dev/null +++ b/core/critic.py @@ -0,0 +1,18 @@ +import torch.nn as nn +import torch + + +class Value(nn.Module): + def __init__(self, net, net_out_dim=None): + super().__init__() + self.net = net + if net_out_dim is None: + net_out_dim = net.out_dim + self.value_head = nn.Linear(net_out_dim, 1) + self.value_head.weight.data.mul_(0.1) + self.value_head.bias.data.mul_(0.0) + + def forward(self, x): + x = self.net(x) + value = self.value_head(x) + return value diff --git a/core/distributions.py b/core/distributions.py new file mode 100644 index 0000000..2265e51 --- /dev/null +++ b/core/distributions.py @@ -0,0 +1,47 @@ +import torch +from torch.distributions import Normal +from torch.distributions import Categorical as TorchCategorical + + +class DiagGaussian(Normal): + + def __init__(self, loc, scale): + super().__init__(loc, scale) + + def kl(self): + loc1 = self.loc + scale1 = self.scale + log_scale1 = self.scale.log() + loc0 = self.loc.detach() + scale0 = self.scale.detach() + log_scale0 = log_scale1.detach() + kl = log_scale1 - log_scale0 + (scale0.pow(2) + (loc0 - loc1).pow(2)) / (2.0 * scale1.pow(2)) - 0.5 + return kl.sum(1, keepdim=True) + + def log_prob(self, value): + return super().log_prob(value).sum(1, keepdim=True) + + def mean_sample(self): + return self.loc + + +class Categorical(TorchCategorical): + + def __init__(self, probs=None, logits=None): + super().__init__(probs, logits) + + def kl(self): + loc1 = self.loc + scale1 = self.scale + log_scale1 = self.scale.log() + loc0 = self.loc.detach() + scale0 = self.scale.detach() + log_scale0 = log_scale1.detach() + kl = log_scale1 - log_scale0 + (scale0.pow(2) + (loc0 - loc1).pow(2)) / (2.0 * scale1.pow(2)) - 0.5 + return kl.sum(1, keepdim=True) + + def log_prob(self, value): + return super().log_prob(value).unsqueeze(1) + + def mean_sample(self): + return self.probs.argmax(dim=1) diff --git a/core/logger_rl.py b/core/logger_rl.py new file mode 100644 index 0000000..6b19bb1 --- /dev/null +++ b/core/logger_rl.py @@ -0,0 +1,59 @@ +import math + + +class LoggerRL: + + def __init__(self): + self.num_steps = 0 + self.num_episodes = 0 + self.total_reward = 0 + self.min_episode_reward = math.inf + self.max_episode_reward = -math.inf + self.total_c_reward = 0 + self.min_c_reward = math.inf + self.max_c_reward = -math.inf + self.episode_reward = 0 + self.avg_episode_reward = 0 + self.avg_c_reward = 0 + self.total_c_info = 0 + self.avg_c_info = 0 + self.sample_time = 0 + + def start_episode(self, env): + self.episode_reward = 0 + + def step(self, env, reward, c_reward, c_info): + self.episode_reward += reward + self.total_c_reward += c_reward + self.total_c_info += c_info + self.min_c_reward = min(self.min_c_reward, c_reward) + self.max_c_reward = max(self.max_c_reward, c_reward) + self.num_steps += 1 + + def end_episode(self, env): + self.num_episodes += 1 + self.total_reward += self.episode_reward + self.min_episode_reward = min(self.min_episode_reward, self.episode_reward) + self.max_episode_reward = max(self.max_episode_reward, self.episode_reward) + + def end_sampling(self): + self.avg_episode_reward = self.total_reward / self.num_episodes + self.avg_c_reward = self.total_c_reward / self.num_steps + self.avg_c_info = self.total_c_info / self.num_steps + + @classmethod + def merge(cls, logger_list): + logger = cls() + logger.total_reward = sum([x.total_reward for x in logger_list]) + logger.num_episodes = sum([x.num_episodes for x in logger_list]) + logger.num_steps = sum([x.num_steps for x in logger_list]) + logger.avg_episode_reward = logger.total_reward / logger.num_episodes + logger.max_episode_reward = max([x.max_episode_reward for x in logger_list]) + logger.min_episode_reward = max([x.min_episode_reward for x in logger_list]) + logger.total_c_reward = sum([x.total_c_reward for x in logger_list]) + logger.avg_c_reward = logger.total_c_reward / logger.num_steps + logger.max_c_reward = max([x.max_c_reward for x in logger_list]) + logger.min_c_reward = min([x.min_c_reward for x in logger_list]) + logger.total_c_info = sum([x.total_c_info for x in logger_list]) + logger.avg_c_info = logger.total_c_info / logger.num_steps + return logger diff --git a/core/policy.py b/core/policy.py new file mode 100644 index 0000000..8ff40de --- /dev/null +++ b/core/policy.py @@ -0,0 +1,23 @@ +import torch.nn as nn + + +class Policy(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + """This function should return a distribution to sample action from""" + raise NotImplementedError + + def select_action(self, x, mean_action=False): + dist = self.forward(x) + action = dist.mean_sample() if mean_action else dist.sample() + return action + + def get_kl(self, x): + dist = self.forward(x) + return dist.kl() + + def get_log_prob(self, x, action): + dist = self.forward(x) + return dist.log_prob(action) diff --git a/core/policy_disc.py b/core/policy_disc.py new file mode 100644 index 0000000..9ca20b9 --- /dev/null +++ b/core/policy_disc.py @@ -0,0 +1,27 @@ +import torch.nn as nn +from utils.math import * +from core.distributions import Categorical +from core.policy import Policy + + +class PolicyDiscrete(Policy): + def __init__(self, net, action_num, net_out_dim=None): + super().__init__() + self.type = 'discrete' + if net_out_dim is None: + net_out_dim = net.out_dim + self.net = net + self.action_head = nn.Linear(net_out_dim, action_num) + self.action_head.weight.data.mul_(0.1) + self.action_head.bias.data.mul_(0.0) + + def forward(self, x): + x = self.net(x) + action_prob = torch.softmax(self.action_head(x), dim=1) + return Categorical(probs=action_prob) + + def get_fim(self, x): + action_prob = self.forward(x) + M = action_prob.pow(-1).view(-1).detach() + return M, action_prob, {} + diff --git a/core/policy_gaussian.py b/core/policy_gaussian.py new file mode 100644 index 0000000..cb6c36d --- /dev/null +++ b/core/policy_gaussian.py @@ -0,0 +1,40 @@ +import torch.nn as nn +from core.distributions import DiagGaussian +from core.policy import Policy +from utils.math import * + + +class PolicyGaussian(Policy): + def __init__(self, net, action_dim, net_out_dim=None, log_std=0, fix_std=False): + super().__init__() + self.type = 'gaussian' + self.net = net + if net_out_dim is None: + net_out_dim = net.out_dim + self.action_mean = nn.Linear(net_out_dim, action_dim) + self.action_mean.weight.data.mul_(0.1) + self.action_mean.bias.data.mul_(0.0) + self.action_log_std = nn.Parameter(torch.ones(1, action_dim) * log_std, requires_grad=not fix_std) + + def forward(self, x): + x = self.net(x) + action_mean = self.action_mean(x) + action_log_std = self.action_log_std.expand_as(action_mean) + action_std = torch.exp(action_log_std) + return DiagGaussian(action_mean, action_std) + + def get_fim(self, x): + dist = self.forward(x) + cov_inv = self.action_log_std.exp().pow(-2).squeeze(0).repeat(x.size(0)) + param_count = 0 + std_index = 0 + id = 0 + for name, param in self.named_parameters(): + if name == "action_log_std": + std_id = id + std_index = param_count + param_count += param.view(-1).shape[0] + id += 1 + return cov_inv.detach(), dist.loc, {'std_id': std_id, 'std_index': std_index} + + diff --git a/core/trajbatch.py b/core/trajbatch.py new file mode 100644 index 0000000..98d6ad6 --- /dev/null +++ b/core/trajbatch.py @@ -0,0 +1,16 @@ +import numpy as np + + +class TrajBatch: + + def __init__(self, memory_list): + memory = memory_list[0] + for x in memory_list[1:]: + memory.append(x) + self.batch = zip(*memory.sample()) + self.states = np.stack(next(self.batch)) + self.actions = np.stack(next(self.batch)) + self.masks = np.stack(next(self.batch)) + self.next_states = np.stack(next(self.batch)) + self.rewards = np.stack(next(self.batch)) + self.exps = np.stack(next(self.batch)) diff --git a/docs/keymap.md b/docs/keymap.md new file mode 100644 index 0000000..8616481 --- /dev/null +++ b/docs/keymap.md @@ -0,0 +1,18 @@ +## Keyboard Shortcuts for Visualizer +When visualizing results with evaluation scripts (eval_pose.py, eval_pose_wild.py, eval_forecast.py, eval_forecast_wild.py), a GUI visualizer is opened, which has the following shortcuts: + +| Key | Functionality | +| ------------- | ------------- | +| space | toggle motion autoplay | +| left arrow | previous frame (when not autoplay) | +| right arrow | next frame (when not autoplay) | +| w | restart the take | +| s | toggle reverse play | +| d | slow down motion | +| f | speed up motion | +| z | previous take | +| c | next take | +| q | previous prediction instance (forecast only)| +| e | next prediction instance (forecast only)| +| 1 | show ego mimic results| +| 2 | show state regression results (estimate only)| \ No newline at end of file diff --git a/docs/train_and_test.md b/docs/train_and_test.md new file mode 100644 index 0000000..c58f602 --- /dev/null +++ b/docs/train_and_test.md @@ -0,0 +1,86 @@ +# Training +**Note**: All scripts should be run from the root of this repo. +### Config Files +* Our code uses config files for training state regression ("config/statereg"), ego-pose estimation ("config/egomimic"), and ego-pose forecasting ("config/"). +Please see .yml files in these folders for parameter settings. + * The parameter **meta_id** specifies the meta file to be used in training, which defines training and testing data. + * For ego-pose estimation and state regression, The parameter **fr_margin** specifies the number of padding frames as mentioned in the implementation details of [the paper](https://arxiv.org/pdf/1906.03173.pdf). + For ego-pose forecasting, **fr_margin** corresponds to the number of past frames. +* When training, each config will generate a folder with the config name under "results/" including the following sub-folders: + * **models:** saves model checkpoints. + * **results:** contains saved test results. + * **log:** contains logs recoreded by the logger. + * **tb:** contains Tensorboard logs. + +### Training State Regression +The state regression network provides initial state estimation for our ego-pose estimation, and it is also used in the fail-safe mechanism to reset the humanoid state. +Additionally, it pretrains the ResNet-18 and is used to precompute the CNN features for ego-pose estimation and forecasting. + +* If you have setup your config files "config/statereg/subject_02.yml", run the following command to train the state regression network: + ```python ego_pose/state_reg.py --cfg subject_02``` + And it will save model checkpoints under "results/statereg/subject_02/models/". + +### Training Ego-Pose Estimation +1. **Generate CNN features.** we need to use trained state regression network to precompute CNN features. Run the following script: + ```python ego_pose/data_process/gen_cnn_feature.py --cfg subject_02 --iter 100 --out-id subject_02``` + It means using the state regression net from config "subject_02" and iter 100 to generate pickled cnn features "cnn_feat_subject_02.p", which will be saved under "datasets/features/". + Note that you may find some features already exist, but they are only compatible with predefined configs and pretrained models. So If you change the config or retrain the state regression network, you need to regenerate CNN features. + +2. **Generate expert features.** Expert means expert demonstrations, a term often used in imitation learning. + In our case, expert demonstrations include trajectory features such as joint orientations and velocities. And we need to precompute these expert features from GT MoCap data using this script: + ```python ego_pose/data_process/gen_expert.py --meta-id meta_subject_02 --out-id subject_02``` + It will generate a pickle file "expert_subject_02.p" under "datasets/features/". + +3. **Config file**. In your config "config/egomimic/subject_02.yml", make sure to set **meta_id** to "meta_subject_02", and **cnn_feat** and **expert** to "subject_02". + Also, set **state_net_cfg** to "subject_02", which is the statereg config you used to generate CNN features. + +4. Start training with the following script: + ```python ego_pose/ego_mimic.py --cfg subject_02``` + You can additionally use "--num-threads" and "--gpu-index" to choose the number of threads for sampling and the GPU you use. + +### Training Ego-Pose Forecasting +1. If you haven't generated CNN and expert features for the meta file you want to use, generate them following the same steps for training ego-pose estimation. + +2. Setup the config file "config/egoforecast/subject_02.yml". + +3. Start training with the following script: + ```python ego_pose/ego_forecast.py --cfg subject_02``` + + +# Testing +### Testing State Regression +* Run the following script to generate test results: + ```python ego_pose/state_reg.py --cfg subject_02 --iter 100 --mode test``` + It means using the model from config "subject_02" and iter 100 to generate predicted motion results, which will be saved as "results/statereg/subject_02/results/iter_100_test.p". + +### Testing Ego-Pose Estimation +* As ego-pose estimation relies on the state regression network and we are testing with precomputed CNN features, we first generate a lightweight version of the state regression net without CNN: + ```python ego_pose/state_reg.py --cfg subject_02 --iter 100 --mode save_inf``` + +* For testing ego-pose estimation with MoCap data, run the following script: + ```python ego_pose/ego_mimic_eval.py --cfg subject_02 --iter 3000``` + You can use "--render" if you want to visualize the test results generation. The results will be saved as "results/egomimic/subject_02/results/iter_3000_test.p". + +* For testing with in-the-wild data, we need to use the config "cross_01" which corresponds to the cross-subject model, so it will generalize better. (Of course, you need to train state regression and ego-pose estimation with config "cross_01" first, or use our pretrained models.) + * First make sure that you have generated CNN features using the meta file "meta_wild_01.yml": + ```python ego_pose/data_process/gen_cnn_feature.py --cfg cross_01 --iter 100 --meta-id meta_wild_01 --out-id wild_01``` + * Then you can run the testing script: + ```python ego_pose/ego_mimic_eval_wild.py --cfg cross_01 --iter 6000 --test-feat wild_01``` + The results will be saved as "results/egomimic/cross_01/results/iter_6000_wild_01.p". + +### Testing Ego-Pose Forecasting +* First, make sure you have generated testing results for corresponding ego-pose estimation config (specified in the forecasting config), as forecasting relies on the estimation to provide the initial state. + +* For testing ego-pose estimation with MoCap data, run the following script: + ```python ego_pose/ego_forecast_eval.py --cfg subject_02 --iter 3000 --mode save``` + The results will be saved as "results/egoforecast/subject_02/results/iter_3000_test.p". + +* For testing with in-the-wild data, we need to use the config "cross_01" which corresponds to the cross-subject model, so it will generalize better. (Of course, you need to train state regression and ego-pose forecasting with config "cross_01" first, or use our pretrained models.) + * Same as ego-pose estimation, first make sure that you have generated CNN features using the meta file "meta_wild_01.yml". + * Then you can run the testing script: + ```python ego_pose/ego_forecast_eval_wild.py --cfg cross_01 --iter 6000 --test-feat wild_01``` + The results will be saved as "results/egoforecast/cross_01/results/iter_6000_wild_01.p". + +# Visualizing Results +Please refer to the [quick demo](https://github.com/Khrylx/EgoPose#quick-demo) on how to visualize saved results. + diff --git a/ego_pose/core/agent_ego.py b/ego_pose/core/agent_ego.py new file mode 100644 index 0000000..2e99763 --- /dev/null +++ b/ego_pose/core/agent_ego.py @@ -0,0 +1,57 @@ +import time +from utils.torch import * +from agents import AgentPPO +from core.common import * +from ego_pose.core.trajbatch_ego import TrajBatchEgo + + +class AgentEgo(AgentPPO): + + def __init__(self, policy_vs_net=None, value_vs_net=None, **kwargs): + super().__init__(use_mini_batch=False, **kwargs) + self.traj_cls = TrajBatchEgo + self.policy_vs_net = policy_vs_net + self.value_vs_net = value_vs_net + self.sample_modules.append(policy_vs_net) + self.update_modules += [policy_vs_net, value_vs_net] + + def pre_sample(self): + self.policy_vs_net.set_mode('test') + + def pre_episode(self): + self.policy_vs_net.initialize(tensor(self.env.get_episode_cnn_feat())) + + def push_memory(self, memory, state, action, mask, next_state, reward, exp): + v_meta = np.array([self.env.expert_ind, self.env.start_ind]) + memory.push(state, action, mask, next_state, reward, exp, v_meta) + + def trans_policy(self, states): + return self.policy_vs_net(states) + + def trans_value(self, states): + return self.value_vs_net(states) + + def update_params(self, batch): + t0 = time.time() + to_train(*self.update_modules) + states = torch.from_numpy(batch.states).to(self.dtype).to(self.device) + actions = torch.from_numpy(batch.actions).to(self.dtype).to(self.device) + rewards = torch.from_numpy(batch.rewards).to(self.dtype).to(self.device) + masks = torch.from_numpy(batch.masks).to(self.dtype).to(self.device) + exps = torch.from_numpy(batch.exps).to(self.dtype).to(self.device) + v_metas = batch.v_metas + + self.policy_vs_net.set_mode('train') + self.value_vs_net.set_mode('train') + self.policy_vs_net.initialize((masks, self.env.cnn_feat, v_metas)) + self.value_vs_net.initialize((masks, self.env.cnn_feat, v_metas)) + with to_test(*self.update_modules): + with torch.no_grad(): + values = self.value_net(self.trans_value(states)) + + """get advantage estimation from the trajectories""" + advantages, returns = estimate_advantages(rewards, masks, values, self.gamma, self.tau) + + self.update_policy(states, actions, returns, advantages, exps) + + return time.time() - t0 diff --git a/ego_pose/core/agent_vgail.py b/ego_pose/core/agent_vgail.py new file mode 100644 index 0000000..0cb3126 --- /dev/null +++ b/ego_pose/core/agent_vgail.py @@ -0,0 +1,88 @@ +import time +from utils.torch import * +from core.common import * +from ego_pose.core.agent_ego import AgentEgo + + +class AgentVGAIL(AgentEgo): + + def __init__(self, discrim_net=None, discrim_vs_net=None, discrim_criterion=None, + optimizer_discrim=None, discrim_num_update=10, **kwargs): + super().__init__(**kwargs) + self.discrim_net = discrim_net + self.discrim_vs_net = discrim_vs_net + self.discrim_criterion = discrim_criterion + self.optimizer_discrim = optimizer_discrim + self.discrim_num_update = discrim_num_update + self.sample_modules += [discrim_net, discrim_vs_net] + self.update_modules += [discrim_net, discrim_vs_net] + + def pre_sample(self): + super().pre_sample() + self.discrim_vs_net.set_mode('test') + + def pre_episode(self): + super().pre_episode() + self.discrim_vs_net.initialize(tensor(self.env.get_episode_cnn_feat())) + + def update_params(self, batch): + t0 = time.time() + to_train(*self.update_modules) + states = torch.from_numpy(batch.states).to(self.dtype).to(self.device) + actions = torch.from_numpy(batch.actions).to(self.dtype).to(self.device) + rewards = torch.from_numpy(batch.rewards).to(self.dtype).to(self.device) + masks = torch.from_numpy(batch.masks).to(self.dtype).to(self.device) + exps = torch.from_numpy(batch.exps).to(self.dtype).to(self.device) + v_metas = batch.v_metas + + self.policy_vs_net.set_mode('train') + self.value_vs_net.set_mode('train') + self.policy_vs_net.initialize((masks, self.env.cnn_feat, v_metas)) + self.value_vs_net.initialize((masks, self.env.cnn_feat, v_metas)) + with to_test(*self.update_modules): + with torch.no_grad(): + values = self.value_net(self.trans_value(states)) + + """get advantage estimation from the trajectories""" + advantages, returns = estimate_advantages(rewards, masks, values, self.gamma, self.tau) + + self.update_policy(states, actions, returns, advantages, exps) + self.update_discriminator(states, masks, v_metas) + + return time.time() - t0 + + def update_discriminator(self, states, masks, v_metas): + """perform discriminator update""" + self.discrim_vs_net.set_mode('train') + self.discrim_vs_net.initialize((masks, self.env.cnn_feat, v_metas)) + expert_states = self.get_expert_states(v_metas, masks) + for i in range(self.discrim_num_update): + g_vs_out = self.discrim_vs_net(states) + e_vs_out = self.discrim_vs_net(expert_states) + g_o = self.discrim_net(g_vs_out) + e_o = self.discrim_net(e_vs_out) + + self.optimizer_discrim.zero_grad() + l_g = self.discrim_criterion(g_o, ones((g_o.size(0), 1), device=self.device)) + l_e = self.discrim_criterion(e_o, zeros((e_o.size(0), 1), device=self.device)) + discrim_loss = l_g + l_e + discrim_loss.backward() + torch.nn.utils.clip_grad_norm_(self.discrim_net.parameters(), 40) + self.optimizer_discrim.step() + e_loss = l_e.detach().item() + return e_loss + + def get_expert_states(self, v_metas, masks): + expert_states = [] + end_indice = np.where(masks.cpu().numpy() == 0)[0] + v_metas = v_metas[end_indice, :] + end_indice = np.insert(end_indice, 0, -1) + episode_lens = np.diff(end_indice) + for v_meta, len in zip(v_metas, episode_lens): + exp_ind, start_ind = v_meta + e_states = self.env.expert_arr[exp_ind]['obs'][start_ind: start_ind + len, :] + expert_states.append(e_states) + expert_states = np.vstack(expert_states) + if self.running_state is not None: + expert_states = (expert_states - self.running_state.rs.mean[None, :]) / self.running_state.rs.std[None, :] + return torch.from_numpy(expert_states).to(self.dtype).to(self.device) diff --git a/ego_pose/core/reward_function.py b/ego_pose/core/reward_function.py new file mode 100644 index 0000000..bb9d7e9 --- /dev/null +++ b/ego_pose/core/reward_function.py @@ -0,0 +1,81 @@ +from utils import * + + +def quat_space_reward_v3(env, state, action, info): + # reward coefficients + cfg = env.cfg + ws = cfg.reward_weights + w_p, w_v, w_e, w_rp, w_rv = ws.get('w_p', 0.5), ws.get('w_v', 0.1), ws.get('w_e', 0.2), ws.get('w_rp', 0.1), ws.get('w_rv', 0.1) + k_p, k_v, k_e = ws.get('k_p', 2), ws.get('k_v', 0.005), ws.get('k_e', 20) + k_rh, k_rq, k_rl, k_ra = ws.get('k_rh', 300), ws.get('k_rq', 300), ws.get('k_rl', 5.0), ws.get('k_ra', 0.5) + v_ord = ws.get('v_ord', 2) + # data from env + t = env.cur_t + ind = env.get_expert_index(t) + prev_bquat = env.prev_bquat + prev_qpos = env.prev_qpos + # learner + cur_qpos = env.data.qpos.copy() + cur_qvel = get_qvel_fd(prev_qpos, cur_qpos, env.dt, cfg.obs_coord) + cur_rlinv_local = cur_qvel[:3] + cur_rangv = cur_qvel[3:6] + cur_rq_rmh = de_heading(cur_qpos[3:7]) + cur_ee = env.get_ee_pos(cfg.obs_coord) + cur_bquat = env.get_body_quat() + cur_bangvel = get_angvel_fd(prev_bquat, cur_bquat, env.dt) + # expert + e_qpos = env.get_expert_attr('qpos', ind) + e_rlinv_local = env.get_expert_attr('rlinv_local', ind) + e_rangv = env.get_expert_attr('rangv', ind) + e_rq_rmh = env.get_expert_attr('rq_rmh', ind) + e_ee = env.get_expert_attr('ee_pos', ind) + e_bquat = env.get_expert_attr('bquat', ind) + e_bangvel = env.get_expert_attr('bangvel', ind) + # pose reward + pose_diff = multi_quat_norm(multi_quat_diff(cur_bquat[4:], e_bquat[4:])) # ignore root + pose_diff *= cfg.b_diffw + pose_dist = np.linalg.norm(pose_diff) + pose_reward = math.exp(-k_p * (pose_dist ** 2)) + # velocity reward + vel_dist = np.linalg.norm(cur_bangvel[3:] - e_bangvel[3:], ord=v_ord) # ignore root + vel_reward = math.exp(-k_v * (vel_dist ** 2)) + # ee reward + ee_dist = np.linalg.norm(cur_ee - e_ee) + ee_reward = math.exp(-k_e * (ee_dist ** 2)) + # root position reward + root_height_dist = cur_qpos[2] - e_qpos[2] + root_quat_dist = multi_quat_norm(multi_quat_diff(cur_rq_rmh, e_rq_rmh))[0] + root_pose_reward = math.exp(-k_rh * (root_height_dist ** 2) - k_rq * (root_quat_dist ** 2)) + # root velocity reward + root_linv_dist = np.linalg.norm(cur_rlinv_local - e_rlinv_local) + root_angv_dist = np.linalg.norm(cur_rangv - e_rangv) + root_vel_reward = math.exp(-k_rl * (root_linv_dist ** 2) - k_ra * (root_angv_dist ** 2)) + # overall reward + reward = w_p * pose_reward + w_v * vel_reward + w_e * ee_reward + w_rp * root_pose_reward + w_rv * root_vel_reward + reward /= w_p + w_v + w_e + w_rp + w_rv + if ws.get('decay', False): + reward *= 1.0 - t / cfg.env_episode_len + if info['end']: + reward += env.end_reward + return reward, np.array([pose_reward, vel_reward, ee_reward, root_pose_reward, root_vel_reward]) + + +def constant_reward(env, state, action, info): + reward = 1.0 + if info['end']: + reward += env.end_reward + return 1.0, np.zeros(1) + + +def pose_dist_reward(env, state, action, info): + pose_dist = env.get_pose_dist() + reward = 5.0 - 3.0 * pose_dist + if info['end']: + reward += env.end_reward + return reward, np.array([pose_dist]) + + +reward_func = {'quat_v3': quat_space_reward_v3, + 'constant': constant_reward, + 'pose_dist': pose_dist_reward} + diff --git a/ego_pose/core/trajbatch_ego.py b/ego_pose/core/trajbatch_ego.py new file mode 100644 index 0000000..390fc86 --- /dev/null +++ b/ego_pose/core/trajbatch_ego.py @@ -0,0 +1,10 @@ +from core import TrajBatch +import numpy as np + + +class TrajBatchEgo(TrajBatch): + + def __init__(self, memory_list): + super().__init__(memory_list) + self.v_metas = np.stack(next(self.batch)) + diff --git a/ego_pose/data_process/gen_cnn_feature.py b/ego_pose/data_process/gen_cnn_feature.py new file mode 100644 index 0000000..b8b3ae7 --- /dev/null +++ b/ego_pose/data_process/gen_cnn_feature.py @@ -0,0 +1,70 @@ +import argparse +import os +import sys +import pickle +import math +import datetime +import numpy as np +sys.path.append(os.getcwd()) + +from utils import * +from models.video_reg_net import * +from ego_pose.utils.statereg_dataset import Dataset +from ego_pose.utils.statereg_config import Config + + +parser = argparse.ArgumentParser() +parser.add_argument('--cfg', default=None) +parser.add_argument('--meta-id', default=None) +parser.add_argument('--out-id', default=None) +parser.add_argument('--gpu-index', type=int, default=0) +parser.add_argument('--iter', type=int, default=100) +args = parser.parse_args() + +cfg = Config(args.cfg, create_dirs=False) + + +"""setup""" +dtype = torch.float64 +torch.set_default_dtype(dtype) +device = torch.device('cuda', index=args.gpu_index) if torch.cuda.is_available() else torch.device('cpu') +if torch.cuda.is_available(): + torch.cuda.set_device(args.gpu_index) +torch.set_grad_enabled(False) +logger = create_logger(os.path.join(cfg.log_dir, 'gen_cnn_feature.txt')) + +cp_path = '%s/iter_%04d.p' % (cfg.model_dir, args.iter) +logger.info('loading model from checkpoint: %s' % cp_path) +model_cp, meta = pickle.load(open(cp_path, "rb")) + +meta_id = cfg.meta_id if args.meta_id is None else args.meta_id +dataset = Dataset(meta_id, 'all', cfg.fr_num, 'iter', False, 0) +dataset.set_mean_std(meta['mean'], meta['std']) +state_net = VideoRegNet(dataset.mean.size, cfg.v_hdim, cfg.cnn_fdim, mlp_dim=cfg.mlp_dim, cnn_type=cfg.cnn_type, + v_net_type=cfg.v_net, v_net_param=cfg.v_net_param, causal=cfg.causal) +state_net.load_state_dict(model_cp['state_net_dict']) +state_net.to(device) +to_test(state_net) + + +num_sample = 0 +take = dataset.takes[0] +cnn_features = {} +feature_arr = [] +for of_np, _, _ in dataset: + of = tensor(of_np, dtype=dtype, device=device) + of = torch.cat((of, zeros(of.shape[:-1] + (1,), device=device)), dim=-1).permute(0, 3, 1, 2).unsqueeze(1).contiguous() + feature = state_net.get_cnn_feature(of).cpu().numpy() + feature_arr.append(feature) + num_sample += feature.shape[0] + # print(feature.shape, of_np.shape, next_take) + if dataset.cur_ind >= len(dataset.takes) or dataset.takes[dataset.cur_tid] != take: + cnn_features[take] = np.vstack(feature_arr) + feature_arr = [] + take = dataset.takes[dataset.cur_tid] + +logger.info('cfg: %s, iter: %d, total sample: %d, dataset length: %d' % (args.cfg, args.iter, num_sample, dataset.len)) + +meta = {'cfg': args.cfg, 'iter': args.iter, 'meta': meta_id, 'time': datetime.datetime.now()} +path = 'datasets/features/cnn_feat_%s.p' % args.out_id +pickle.dump((cnn_features, meta), open(path, 'wb')) diff --git a/ego_pose/data_process/gen_expert.py b/ego_pose/data_process/gen_expert.py new file mode 100644 index 0000000..72eaec2 --- /dev/null +++ b/ego_pose/data_process/gen_expert.py @@ -0,0 +1,102 @@ +import argparse +import os +import sys +import pickle +sys.path.append(os.getcwd()) + +from utils import * +from ego_pose.envs.humanoid_v1 import HumanoidEnv +from ego_pose.utils.statereg_dataset import Dataset +from ego_pose.utils.egomimic_config import Config as EgoConfig + +parser = argparse.ArgumentParser() +parser.add_argument('--meta-id', default=None) +parser.add_argument('--out-id', default=None) +args = parser.parse_args() + +cfg_dict = { + 'meta_id': args.meta_id, + 'mujoco_model': 'humanoid_1205_v1', + 'vis_model': 'humanoid_1205_vis', + 'obs_coord': 'heading', +} +cfg = EgoConfig(None, create_dirs=False, cfg_dict=cfg_dict) +env = HumanoidEnv(cfg) +dataset = Dataset(args.meta_id, 'all', 0, 'iter', False, 0) + + +def get_expert(expert_qpos, lb, ub): + expert = {'qpos': expert_qpos} + feat_keys = {'qvel', 'rlinv', 'rlinv_local', 'rangv', 'rq_rmh', + 'com', 'head_pos', 'obs', 'ee_pos', 'ee_wpos', 'bquat', 'bangvel'} + for key in feat_keys: + expert[key] = [] + + for i in range(expert_qpos.shape[0]): + qpos = expert_qpos[i] + # remove noisy hand data + qpos[slice(*env.body_qposaddr['LeftHand'])] = 0.0 + qpos[slice(*env.body_qposaddr['RightHand'])] = 0.0 + env.data.qpos[:] = qpos + env.sim.forward() + rq_rmh = de_heading(qpos[3:7]) + obs = env.get_obs() + ee_pos = env.get_ee_pos(env.cfg.obs_coord) + ee_wpos = env.get_ee_pos(None) + bquat = env.get_body_quat() + com = env.get_com() + head_pos = env.get_body_com('Head').copy() + if i > 0: + prev_qpos = expert_qpos[i - 1] + qvel = get_qvel_fd(prev_qpos, qpos, env.dt) + rlinv = qvel[:3].copy() + rlinv_local = transform_vec(qvel[:3].copy(), qpos[3:7], env.cfg.obs_coord) + rangv = qvel[3:6].copy() + expert['qvel'].append(qvel) + expert['rlinv'].append(rlinv) + expert['rlinv_local'].append(rlinv_local) + expert['rangv'].append(rangv) + expert['obs'].append(obs) + expert['ee_pos'].append(ee_pos) + expert['ee_wpos'].append(ee_wpos) + expert['bquat'].append(bquat) + expert['com'].append(com) + expert['head_pos'].append(head_pos) + expert['rq_rmh'].append(rq_rmh) + expert['qvel'].insert(0, expert['qvel'][0].copy()) + expert['rlinv'].insert(0, expert['rlinv'][0].copy()) + expert['rlinv_local'].insert(0, expert['rlinv_local'][0].copy()) + expert['rangv'].insert(0, expert['rangv'][0].copy()) + # get expert body quaternions + for i in range(expert_qpos.shape[0]): + if i > 0: + bangvel = get_angvel_fd(expert['bquat'][i - 1], expert['bquat'][i], env.dt) + expert['bangvel'].append(bangvel) + expert['bangvel'].insert(0, expert['bangvel'][0].copy()) + + expert['qpos'] = expert['qpos'][lb:ub, :] + for key in feat_keys: + expert[key] = np.vstack(expert[key][lb:ub]) + expert['len'] = expert['qpos'].shape[0] + expert['height_lb'] = expert['qpos'][:, 2].min() + expert['head_height_lb'] = expert['head_pos'][:, 2].min() + return expert + + +num_sample = 0 +expert_dict = {} +for i in range(len(dataset.takes)): + take = dataset.takes[i] + _, lb, ub = dataset.msync[take] + expert_qpos = dataset.orig_trajs[i] + expert = get_expert(expert_qpos, lb, ub) + expert_dict[take] = expert + num_sample += expert['len'] + print(take, expert['len'], expert['qvel'].min(), expert['qvel'].max(), expert['head_height_lb']) + +print('meta: %s, total sample: %d, dataset length: %d' % (args.meta_id, num_sample, dataset.len)) + +path = 'datasets/features/expert_%s.p' % args.out_id +pickle.dump(expert_dict, open(path, 'wb')) + + diff --git a/ego_pose/ego_forecast.py b/ego_pose/ego_forecast.py new file mode 100644 index 0000000..77f5f6d --- /dev/null +++ b/ego_pose/ego_forecast.py @@ -0,0 +1,162 @@ +import argparse +import os +import sys +import pickle +import time +sys.path.append(os.getcwd()) + +from utils import * +from core.policy_gaussian import PolicyGaussian +from core.critic import Value +from models.mlp import MLP +from models.video_forecast_net import VideoForecastNet +from ego_pose.envs.humanoid_v1 import HumanoidEnv +from ego_pose.core.agent_ego import AgentEgo +from ego_pose.utils.egoforecast_config import Config +from ego_pose.utils.egomimic_config import Config as EgoMimicConfig +from ego_pose.core.reward_function import reward_func + + +parser = argparse.ArgumentParser() +parser.add_argument('--cfg', default=None) +parser.add_argument('--render', action='store_true', default=False) +parser.add_argument('--num-threads', type=int, default=12) +parser.add_argument('--gpu-index', type=int, default=0) +parser.add_argument('--iter', type=int, default=0) +parser.add_argument('--show-noise', action='store_true', default=False) +args = parser.parse_args() +if args.render: + args.num_threads = 1 +cfg = Config(args.cfg, create_dirs=not (args.render or args.iter > 0)) + +dtype = torch.float64 +torch.set_default_dtype(dtype) +device = torch.device('cuda', index=args.gpu_index) if torch.cuda.is_available() else torch.device('cpu') +if torch.cuda.is_available(): + torch.cuda.set_device(args.gpu_index) +np.random.seed(cfg.seed) +torch.manual_seed(cfg.seed) +tb_logger = Logger(cfg.tb_dir) if not args.render else None +logger = create_logger(os.path.join(cfg.log_dir, 'log.txt'), file_handle=not args.render) + +"""environment""" +env = HumanoidEnv(cfg) +env.seed(cfg.seed) +env.load_experts(cfg.takes['train'], cfg.expert_feat_file, cfg.cnn_feat_file) +cnn_feat_dim = env.cnn_feat[0].shape[-1] +actuators = env.model.actuator_names +state_dim = env.observation_space.shape[0] +action_dim = env.action_space.shape[0] +running_state = ZFilter((state_dim,), clip=5) + +"""define actor and critic""" +policy_vs_net = VideoForecastNet(cnn_feat_dim, state_dim, cfg.policy_v_hdim, cfg.fr_margin, cfg.policy_v_net, + cfg.policy_v_net_param, cfg.policy_s_hdim, cfg.policy_s_net, cfg.policy_dyn_v) +value_vs_net = VideoForecastNet(cnn_feat_dim, state_dim, cfg.value_v_hdim, cfg.fr_margin, cfg.value_v_net, + cfg.value_v_net_param, cfg.value_s_hdim, cfg.value_s_net, cfg.value_dyn_v) +policy_net = PolicyGaussian(MLP(policy_vs_net.out_dim, cfg.policy_hsize, cfg.policy_htype), action_dim, + log_std=cfg.log_std, fix_std=cfg.fix_std) +value_net = Value(MLP(value_vs_net.out_dim, cfg.value_hsize, cfg.value_htype)) +if args.iter == 0: + em_cfg = EgoMimicConfig(cfg.ego_mimic_cfg) + cp_path = '%s/iter_%04d.p' % (em_cfg.model_dir, cfg.ego_mimic_iter) + logger.info('loading model from ego mimic checkpoint: %s' % cp_path) + model_cp = pickle.load(open(cp_path, "rb")) + if cfg.obs_phase or cfg.policy_s_net != 'id' or cfg.policy_v_hdim != em_cfg.policy_v_hdim: + filter_state_dict(model_cp['policy_dict'], {'net.affine_layers.0'}) + filter_state_dict(model_cp['value_dict'], {'net.affine_layers.0'}) + policy_net.load_state_dict(model_cp['policy_dict'], strict=False) + value_net.load_state_dict(model_cp['value_dict'], strict=False) +elif args.iter > 0: + cp_path = '%s/iter_%04d.p' % (cfg.model_dir, args.iter) + logger.info('loading model from checkpoint: %s' % cp_path) + model_cp = pickle.load(open(cp_path, "rb")) + policy_net.load_state_dict(model_cp['policy_dict']) + policy_vs_net.load_state_dict(model_cp['policy_vs_dict']) + value_net.load_state_dict(model_cp['value_dict']) + value_vs_net.load_state_dict(model_cp['value_vs_dict']) + running_state = model_cp['running_state'] +to_device(device, policy_net, value_net, policy_vs_net, value_vs_net) + +policy_params = list(policy_net.parameters()) + list(policy_vs_net.parameters()) +value_params = list(value_net.parameters()) + list(value_vs_net.parameters()) +if cfg.policy_optimizer == 'Adam': + optimizer_policy = torch.optim.Adam(policy_params, lr=cfg.policy_lr, weight_decay=cfg.policy_weightdecay) +else: + optimizer_policy = torch.optim.SGD(policy_params, lr=cfg.policy_lr, momentum=cfg.policy_momentum, weight_decay=cfg.policy_weightdecay) +if cfg.value_optimizer == 'Adam': + optimizer_value = torch.optim.Adam(value_params, lr=cfg.value_lr, weight_decay=cfg.value_weightdecay) +else: + optimizer_value = torch.optim.SGD(value_params, lr=cfg.value_lr, momentum=cfg.value_momentum, weight_decay=cfg.value_weightdecay) + +# reward functions +expert_reward = reward_func[cfg.reward_id] + +"""create agent""" +agent = AgentEgo(env=env, dtype=dtype, device=device, running_state=running_state, + custom_reward=expert_reward, mean_action=args.render and not args.show_noise, + render=args.render, num_threads=args.num_threads, + policy_net=policy_net, policy_vs_net=policy_vs_net, + value_net=value_net, value_vs_net=value_vs_net, + optimizer_policy=optimizer_policy, optimizer_value=optimizer_value, opt_num_epochs=cfg.num_optim_epoch, + gamma=cfg.gamma, tau=cfg.tau, clip_epsilon=cfg.clip_epsilon, + policy_grad_clip=[(policy_params, 40)]) + + +def pre_iter_update(i_iter): + cfg.update_adaptive_params(i_iter) + cfg.env_init_noise = cfg.adp_init_noise + agent.set_noise_rate(cfg.adp_noise_rate) + set_optimizer_lr(optimizer_policy, cfg.adp_policy_lr) + if cfg.fix_std: + policy_net.action_log_std.fill_(cfg.adp_log_std) + return + + +def main_loop(): + + if args.render: + pre_iter_update(args.iter) + agent.sample(1e8) + else: + for i_iter in range(args.iter, cfg.max_iter_num): + """generate multiple trajectories that reach the minimum batch_size""" + pre_iter_update(i_iter) + batch, log = agent.sample(cfg.min_batch_size) + if cfg.end_reward: + agent.env.end_reward = log.avg_c_reward * cfg.gamma / (1 - cfg.gamma) # set end reward + + """update networks""" + t0 = time.time() + agent.update_params(batch) + t1 = time.time() + + """logging""" + c_info = log.avg_c_info + logger.info( + '{}\tT_sample {:.2f}\tT_update {:.2f}\tR_avg {:.4f} {}' + '\tR_range ({:.4f}, {:.4f})\teps_len_avg {:.2f}' + .format(i_iter, log.sample_time, t1 - t0, log.avg_c_reward, + np.array2string(c_info, formatter={'all': lambda x: '%.4f' % x}, separator=','), + log.min_c_reward, log.max_c_reward, log.avg_episode_reward)) + + tb_logger.scalar_summary('total_reward', log.avg_c_reward, i_iter) + tb_logger.scalar_summary('episode_len', log.avg_episode_reward, i_iter) + for i in range(c_info.shape[0]): + tb_logger.scalar_summary('reward_%d' % i, c_info[i], i_iter) + + if cfg.save_model_interval > 0 and (i_iter+1) % cfg.save_model_interval == 0: + with to_cpu(policy_net, value_net, policy_vs_net, value_vs_net): + cp_path = '%s/iter_%04d.p' % (cfg.model_dir, i_iter + 1) + model_cp = {'policy_dict': policy_net.state_dict(), 'policy_vs_dict': policy_vs_net.state_dict(), + 'value_dict': value_net.state_dict(), 'value_vs_dict': value_vs_net.state_dict(), + 'running_state': running_state} + pickle.dump(model_cp, open(cp_path, 'wb')) + + """clean up gpu memory""" + torch.cuda.empty_cache() + + logger.info('training done!') + + +main_loop() diff --git a/ego_pose/ego_forecast_eval.py b/ego_pose/ego_forecast_eval.py new file mode 100644 index 0000000..5437f7c --- /dev/null +++ b/ego_pose/ego_forecast_eval.py @@ -0,0 +1,278 @@ +import argparse +import os +import sys +import pickle +import time +sys.path.append(os.getcwd()) + +from utils import * +from core.policy_gaussian import PolicyGaussian +from core.critic import Value +from models.mlp import MLP +from models.video_forecast_net import VideoForecastNet +from models.video_reg_net import VideoRegNet +from envs.visual.humanoid_vis import HumanoidVisEnv +from ego_pose.envs.humanoid_v1 import HumanoidEnv +from ego_pose.utils.egoforecast_config import Config +from ego_pose.utils.egomimic_config import Config as EgoMimicConfig +from ego_pose.utils.tools import * +from ego_pose.core.reward_function import reward_func + + +parser = argparse.ArgumentParser() +parser.add_argument('--cfg', default=None) +parser.add_argument('--render', action='store_true', default=False) +parser.add_argument('--iter', type=int, default=0) +parser.add_argument('--expert-ind', type=int, default=0) +parser.add_argument('--start-ind', type=int, default=None) +parser.add_argument('--data', default='test') +parser.add_argument('--show-noise', action='store_true', default=False) +parser.add_argument('--verbose', action='store_true', default=False) +parser.add_argument('--gt-init', action='store_true', default=False) +parser.add_argument('--mode', default='save') + +args = parser.parse_args() +cfg = Config(args.cfg, create_dirs=False) +if hasattr(cfg, 'random_cur_t'): + cfg.random_cur_t = False + +dtype = torch.float64 +torch.set_default_dtype(dtype) +np.random.seed(cfg.seed) +torch.manual_seed(cfg.seed) +torch.set_grad_enabled(False) +logger = create_logger(os.path.join(cfg.log_dir, 'log_eval.txt')) + +"""environment""" +env = HumanoidEnv(cfg) +env.load_experts(cfg.takes[args.data], cfg.expert_feat_file, cfg.cnn_feat_file) +env_vis = HumanoidVisEnv(cfg.vis_model_file, 10) +env.seed(cfg.seed) +cnn_feat_dim = env.cnn_feat[0].shape[-1] +actuators = env.model.actuator_names +state_dim = env.observation_space.shape[0] +action_dim = env.action_space.shape[0] +body_qposaddr = get_body_qposaddr(env.model) + +"""load policy net""" +policy_vs_net = VideoForecastNet(cnn_feat_dim, state_dim, cfg.policy_v_hdim, cfg.fr_margin, cfg.policy_v_net, + cfg.policy_v_net_param, cfg.policy_s_hdim, cfg.policy_s_net) +value_vs_net = VideoForecastNet(cnn_feat_dim, state_dim, cfg.value_v_hdim, cfg.fr_margin, cfg.value_v_net, + cfg.value_v_net_param, cfg.value_s_hdim, cfg.value_s_net) +policy_net = PolicyGaussian(MLP(policy_vs_net.out_dim, cfg.policy_hsize, cfg.policy_htype), action_dim, + log_std=cfg.log_std, fix_std=cfg.fix_std) +value_net = Value(MLP(value_vs_net.out_dim, cfg.value_hsize, cfg.value_htype)) +cp_path = '%s/iter_%04d.p' % (cfg.model_dir, args.iter) +logger.info('loading policy net from checkpoint: %s' % cp_path) +model_cp = pickle.load(open(cp_path, "rb")) +policy_net.load_state_dict(model_cp['policy_dict']) +policy_vs_net.load_state_dict(model_cp['policy_vs_dict']) +value_net.load_state_dict(model_cp['value_dict']) +value_vs_net.load_state_dict(model_cp['value_vs_dict']) +running_state = model_cp['running_state'] +value_stat = RunningStat(1) +to_test(policy_vs_net, policy_net, value_vs_net, value_net) + +"""load ego mimic results""" +em_cfg = EgoMimicConfig(cfg.ego_mimic_cfg) +em_res_path = '%s/iter_%04d_%s.p' % (em_cfg.result_dir, cfg.ego_mimic_iter, args.data) +em_res, em_meta = pickle.load(open(em_res_path, 'rb')) + +# reward functions +expert_reward = reward_func[cfg.reward_id] + + +def render(qpos, epos, hide_expert=False): + env_vis.data.qpos[:env.model.nq] = qpos + env_vis.data.qpos[env.model.nq:] = epos + env_vis.data.qpos[env.model.nq] += 1.0 + if hide_expert: + env_vis.data.qpos[env.model.nq + 1] += 100.0 + env_vis.sim_forward() + env_vis.render() + + +def eval_expert(expert_ind, start_ind, test_len): + + take = env.expert_list[expert_ind] + traj_pred = [] + traj_orig = [] + reward_episode = 0 + env.set_fix_sampling(expert_ind, start_ind, test_len) + state = env.reset() + cnn_feat = tensor(env.get_episode_cnn_feat()) + policy_vs_net.initialize(cnn_feat) + value_vs_net.initialize(cnn_feat) + + # use prediction from ego pose estimation + if not args.gt_init: + em_offset = em_cfg.fr_margin + state_pred = em_res['traj_pred'][take][max(0, start_ind - cfg.fr_margin - em_offset): start_ind + test_len - em_offset] + vel_pred = em_res['vel_pred'][take][max(0, start_ind - cfg.fr_margin - em_offset): start_ind + test_len - em_offset] + miss_len = cfg.fr_margin + test_len - state_pred.shape[0] + if start_ind - cfg.fr_margin - em_offset >= 0: + ref_qpos = env.get_expert_attr('qpos', env.get_expert_index(-cfg.fr_margin)).copy() + state_pred, vel_pred = sync_traj(state_pred, vel_pred, ref_qpos) + ind = cfg.fr_margin - miss_len + qpos = state_pred[ind].copy() + qvel = vel_pred[ind].copy() + env.set_state(qpos, qvel) + state = env.get_obs() + + if running_state is not None: + state = running_state(state, update=False) + + for t in range(-cfg.fr_margin, 0): + ind = env.get_expert_index(t) + epos = env.get_expert_attr('qpos', ind).copy() + if args.gt_init or t + cfg.fr_margin < miss_len: + qpos = epos.copy() + else: + qpos = state_pred[t + cfg.fr_margin - miss_len] + traj_pred.append(qpos.copy()) + traj_orig.append(epos.copy()) + if args.render: + render(qpos, epos, False) + if args.render: + env_vis.viewer._paused = True + + fail = False + for t in range(test_len): + + ind = env.get_expert_index(t) + epos = env.get_expert_attr('qpos', ind).copy() + qpos = env.data.qpos + traj_pred.append(qpos.copy()) + traj_orig.append(epos.copy()) + + if args.render: + render(qpos, epos) + + """learner policy""" + state_var = tensor(state, dtype=dtype).unsqueeze(0) + policy_vs_out = policy_vs_net(state_var) + value_vs_out = value_vs_net(state_var) + value = value_net(value_vs_out).item() + value_stat.push(np.array([value])) + + action = policy_net.select_action(policy_vs_out, mean_action=not args.show_noise)[0].numpy() + next_state, reward, done, info = env.step(action) + if running_state is not None: + next_state = running_state(next_state, update=False) + + if args.verbose: + reward, cinfo = reward_func[cfg.reward_id](env, state, action, info) + logger.info("{} {:.2f} {} {:.2f}".format(t, reward, np.array2string(cinfo, formatter={'all': lambda x: '%.4f' % x}), value)) + + reward_episode += reward + + if info['fail']: + fail = True + + state = next_state + + if fail: + logger.info('fail - expert_ind: %d, start_ind %d' % (expert_ind, start_ind)) + + return np.vstack(traj_pred), np.vstack(traj_orig) + + +if args.mode == 'save': + args.render = False + test_len = cfg.env_episode_len + traj_pred_dict = {} + traj_orig_dict = {} + for i, take in enumerate(env.expert_list): + logger.info('Testing on expert trajectory %s' % take) + take_len = env.cnn_feat[i].shape[0] + traj_pred_arr = [] + traj_orig_arr = [] + start_ind = cfg.fr_margin + while start_ind + test_len <= take_len: + # logger.info('%d %d %d' % (start_ind, start_ind + test_len, take_len)) + traj_pred, traj_orig = eval_expert(i, start_ind, test_len) + traj_pred_arr.append(traj_pred) + traj_orig_arr.append(traj_orig) + start_ind += cfg.fr_margin + traj_pred_dict[take] = np.stack(traj_pred_arr, axis=0) + traj_orig_dict[take] = np.stack(traj_orig_arr, axis=0) + logger.info('%s %s' % (traj_pred_dict[take].shape, traj_orig_dict[take].shape)) + results = {'traj_pred': traj_pred_dict, 'traj_orig': traj_orig_dict} + meta = {'algo': 'ego_forecast'} + res_path = '%s/iter_%04d_%s%s.p' % (cfg.result_dir, args.iter, args.data, '_gt' if args.gt_init else '') + pickle.dump((results, meta), open(res_path, 'wb')) + logger.info('saved results to %s' % res_path) + +elif args.mode == 'vis': + args.render = True + + def key_callback(key, action, mods): + global T, fr, paused, start, reverse, expert_ind + + if action != glfw.RELEASE: + return False + elif not start: + if key == glfw.KEY_D: + T *= 1.5 + elif key == glfw.KEY_F: + T = max(1, T / 1.5) + elif key == glfw.KEY_R: + start = True + elif key == glfw.KEY_Q: + fr = cfg.fr_margin + expert_ind = (expert_ind - 1) % len(env.expert_list) + elif key == glfw.KEY_E: + fr = cfg.fr_margin + expert_ind = (expert_ind + 1) % len(env.expert_list) + elif key == glfw.KEY_W: + fr = cfg.fr_margin + elif key == glfw.KEY_S: + reverse = not reverse + elif key == glfw.KEY_RIGHT: + if fr < env.cnn_feat[expert_ind].shape[0] - 1 - cfg.env_episode_len: + fr += 1 + elif key == glfw.KEY_LEFT: + if fr > cfg.fr_margin: + fr -= 1 + elif key == glfw.KEY_UP: + if fr < env.cnn_feat[expert_ind].shape[0] - 5 - cfg.env_episode_len: + fr += 5 + elif key == glfw.KEY_DOWN: + if fr >= cfg.fr_margin + 5: + fr -= 5 + elif key == glfw.KEY_SPACE: + paused = not paused + else: + return False + return True + + return False + + T = 1 + start = False + paused = True + reverse = False + env_vis.set_custom_key_callback(key_callback) + expert_ind = args.expert_ind + fr = cfg.fr_margin if args.start_ind is None else args.start_ind + + while True: + t = 0 + paused = True + while not start: + if t >= math.floor(T): + if not reverse and fr < env.cnn_feat[expert_ind].shape[0] - 1 - cfg.env_episode_len: + fr += 1 + elif reverse and fr > cfg.fr_margin: + fr -= 1 + t = 0 + + qpos = env.expert_arr[expert_ind]['qpos'][fr].copy() + render(qpos, qpos, True) + if not paused: + t += 1 + + eval_expert(expert_ind, start_ind=fr, test_len=cfg.env_episode_len) + fr += 30 + start = False + diff --git a/ego_pose/ego_forecast_eval_wild.py b/ego_pose/ego_forecast_eval_wild.py new file mode 100644 index 0000000..88559dc --- /dev/null +++ b/ego_pose/ego_forecast_eval_wild.py @@ -0,0 +1,253 @@ +import argparse +import os +import sys +import pickle +import time +sys.path.append(os.getcwd()) + +from utils import * +from core.policy_gaussian import PolicyGaussian +from core.critic import Value +from models.mlp import MLP +from models.video_forecast_net import VideoForecastNet +from models.video_reg_net import VideoRegNet +from envs.visual.humanoid_vis import HumanoidVisEnv +from ego_pose.envs.humanoid_v1 import HumanoidEnv +from ego_pose.utils.egoforecast_config import Config +from ego_pose.utils.egomimic_config import Config as EgoMimicConfig +from ego_pose.utils.tools import * + + +parser = argparse.ArgumentParser() +parser.add_argument('--vis-model', default='humanoid_1205_vis_forecast_v1') +parser.add_argument('--cfg', default=None) +parser.add_argument('--render', action='store_true', default=False) +parser.add_argument('--iter', type=int, default=0) +parser.add_argument('--test-feat', default=None) +parser.add_argument('--test-ind', type=int, default=0) +parser.add_argument('--start-ind', type=int, default=None) +parser.add_argument('--show-noise', action='store_true', default=False) +parser.add_argument('--mode', default='save') + +args = parser.parse_args() +cfg = Config(args.cfg, create_dirs=False) +if hasattr(cfg, 'random_cur_t'): + cfg.random_cur_t = False + +dtype = torch.float64 +torch.set_default_dtype(dtype) +np.random.seed(cfg.seed) +torch.manual_seed(cfg.seed) +torch.set_grad_enabled(False) +logger = create_logger(os.path.join(cfg.log_dir, 'log_eval_wild.txt')) + +"""test data""" +cnn_feat_file = '%s/features/cnn_feat_%s.p' % (cfg.data_dir, args.test_feat) +cnn_feat_dict, _ = pickle.load(open(cnn_feat_file, 'rb')) +takes = list(cnn_feat_dict.keys()) + +"""environment""" +env = HumanoidEnv(cfg) +env_vis = HumanoidVisEnv('assets/mujoco_models/%s.xml' % args.vis_model, 10) +env.seed(cfg.seed) +cnn_feat_dim = cnn_feat_dict[takes[0]].shape[-1] +actuators = env.model.actuator_names +state_dim = env.observation_space.shape[0] +action_dim = env.action_space.shape[0] +body_qposaddr = get_body_qposaddr(env.model) + +"""load policy net""" +policy_vs_net = VideoForecastNet(cnn_feat_dim, state_dim, cfg.policy_v_hdim, cfg.fr_margin, cfg.policy_v_net, + cfg.policy_v_net_param, cfg.policy_s_hdim, cfg.policy_s_net) +value_vs_net = VideoForecastNet(cnn_feat_dim, state_dim, cfg.value_v_hdim, cfg.fr_margin, cfg.value_v_net, + cfg.value_v_net_param, cfg.value_s_hdim, cfg.value_s_net) +policy_net = PolicyGaussian(MLP(policy_vs_net.out_dim, cfg.policy_hsize, cfg.policy_htype), action_dim, + log_std=cfg.log_std, fix_std=cfg.fix_std) +value_net = Value(MLP(value_vs_net.out_dim, cfg.value_hsize, cfg.value_htype)) +cp_path = '%s/iter_%04d.p' % (cfg.model_dir, args.iter) +logger.info('loading policy net from checkpoint: %s' % cp_path) +model_cp = pickle.load(open(cp_path, "rb")) +policy_net.load_state_dict(model_cp['policy_dict']) +policy_vs_net.load_state_dict(model_cp['policy_vs_dict']) +value_net.load_state_dict(model_cp['value_dict']) +value_vs_net.load_state_dict(model_cp['value_vs_dict']) +running_state = model_cp['running_state'] +value_stat = RunningStat(1) +to_test(policy_vs_net, policy_net, value_vs_net, value_net) + +"""load ego mimic results""" +em_cfg = EgoMimicConfig(cfg.ego_mimic_cfg) +em_res_path = '%s/iter_%04d_%s.p' % (em_cfg.result_dir, cfg.ego_mimic_iter, args.test_feat) +em_res, em_meta = pickle.load(open(em_res_path, 'rb')) + + +def render(qpos, epos, hide_expert=False): + env_vis.data.qpos[:env.model.nq] = qpos + env_vis.data.qpos[env.model.nq:] = epos + env_vis.data.qpos[env.model.nq] += 1.0 + if hide_expert: + env_vis.data.qpos[env.model.nq + 1] += 100.0 + env_vis.sim_forward() + env_vis.render() + + +def eval_expert(test_ind, start_ind, test_len): + + take = takes[test_ind] + traj_pred = [] + cnn_feat = tensor(cnn_feat_dict[take][start_ind - cfg.fr_margin: start_ind + test_len]) + policy_vs_net.initialize(cnn_feat) + value_vs_net.initialize(cnn_feat) + + # use prediction from ego pose estimation + em_offset = em_cfg.fr_margin + assert start_ind >= cfg.fr_margin + em_offset + state_pred = em_res['traj_pred'][take][start_ind - cfg.fr_margin - em_offset: start_ind + test_len - em_offset] + vel_pred = em_res['vel_pred'][take][start_ind - cfg.fr_margin - em_offset: start_ind + test_len - em_offset] + ind = cfg.fr_margin + qpos = state_pred[ind].copy() + qvel = vel_pred[ind].copy() + env.set_state(qpos, qvel) + state = env.get_obs() + + if running_state is not None: + state = running_state(state, update=False) + + for t in range(-cfg.fr_margin, 0): + qpos = state_pred[t + cfg.fr_margin] + traj_pred.append(qpos.copy()) + if args.render: + render(qpos, qpos, False) + if args.render: + env_vis.viewer._paused = True + + fail = False + for t in range(test_len): + qpos = env.data.qpos + traj_pred.append(qpos.copy()) + + if args.render: + epos = state_pred[min(t + cfg.fr_margin, state_pred.shape[0])] + render(qpos, epos) + + """learner policy""" + state_var = tensor(state, dtype=dtype).unsqueeze(0) + policy_vs_out = policy_vs_net(state_var) + value_vs_out = value_vs_net(state_var) + value = value_net(value_vs_out).item() + value_stat.push(np.array([value])) + + action = policy_net.select_action(policy_vs_out, mean_action=not args.show_noise)[0].numpy() + next_state, reward, done, info = env.step(action) + if running_state is not None: + next_state = running_state(next_state, update=False) + + if info['fail']: + fail = True + + state = next_state + + if fail: + logger.info('fail - test_ind: %d, start_ind %d' % (test_ind, start_ind)) + + return np.vstack(traj_pred) + + +if args.mode == 'save': + test_len = cfg.env_episode_len + traj_pred_dict = {} + traj_orig_dict = {} + for i, take in enumerate(takes): + logger.info('Testing on expert trajectory %s' % take) + take_len = cnn_feat_dict[take].shape[0] + traj_pred_arr = [] + start_ind = cfg.fr_margin + em_cfg.fr_margin + while start_ind + test_len <= take_len: + # logger.info('%d %d %d' % (start_ind, start_ind + test_len, take_len)) + traj_pred = eval_expert(i, start_ind, test_len) + traj_pred_arr.append(traj_pred) + start_ind += cfg.fr_margin + traj_pred_dict[take] = np.stack(traj_pred_arr, axis=0) + logger.info('%s' % str(traj_pred_dict[take].shape)) + results = {'traj_pred': traj_pred_dict} + meta = {'algo': 'ego_forecast'} + res_path = '%s/iter_%04d_%s.p' % (cfg.result_dir, args.iter, args.test_feat) + pickle.dump((results, meta), open(res_path, 'wb')) + logger.info('saved results to %s' % res_path) + +elif args.mode == 'vis': + args.render = True + + def key_callback(key, action, mods): + global T, fr, paused, start, reverse, test_ind + + if action != glfw.RELEASE: + return False + elif not start: + if key == glfw.KEY_D: + T *= 1.5 + elif key == glfw.KEY_F: + T = max(1, T / 1.5) + elif key == glfw.KEY_R: + start = True + elif key == glfw.KEY_Q: + fr = cfg.fr_margin + test_ind = (test_ind - 1) % len(env.expert_list) + elif key == glfw.KEY_E: + fr = cfg.fr_margin + test_ind = (test_ind + 1) % len(env.expert_list) + elif key == glfw.KEY_W: + fr = cfg.fr_margin + em_cfg.fr_margin + elif key == glfw.KEY_S: + reverse = not reverse + elif key == glfw.KEY_RIGHT: + if fr < cnn_feat_dict[takes[test_ind]].shape[0] - em_cfg.fr_margin - 1 - cfg.env_episode_len: + fr += 1 + elif key == glfw.KEY_LEFT: + if fr > cfg.fr_margin + em_cfg.fr_margin: + fr -= 1 + elif key == glfw.KEY_UP: + if fr < cnn_feat_dict[takes[test_ind]].shape[0] - em_cfg.fr_margin - 5 - cfg.env_episode_len: + fr += 5 + elif key == glfw.KEY_DOWN: + if fr >= cfg.fr_margin + em_cfg.fr_margin + 5: + fr -= 5 + elif key == glfw.KEY_SPACE: + paused = not paused + else: + return False + return True + + return False + + T = 1 + start = False + paused = True + reverse = False + env_vis.set_custom_key_callback(key_callback) + test_ind = args.test_ind + fr = cfg.fr_margin + em_cfg.fr_margin if args.start_ind is None else args.start_ind + qpos = em_res['traj_pred'][takes[test_ind]][fr - em_cfg.fr_margin] + render(qpos, qpos, True) + + while True: + t = 0 + paused = True + while not start: + if t >= math.floor(T): + if not reverse and fr < cnn_feat_dict[takes[test_ind]].shape[0] - 1 - cfg.env_episode_len: + fr += 1 + elif reverse and fr > cfg.fr_margin: + fr -= 1 + t = 0 + + print(fr - em_cfg.fr_margin) + qpos = em_res['traj_pred'][takes[test_ind]][fr - em_cfg.fr_margin] + render(qpos, qpos, True) + if not paused: + t += 1 + + eval_expert(test_ind, start_ind=fr, test_len=cfg.env_episode_len) + fr += 30 + start = False + diff --git a/ego_pose/ego_mimic.py b/ego_pose/ego_mimic.py new file mode 100644 index 0000000..1ffff35 --- /dev/null +++ b/ego_pose/ego_mimic.py @@ -0,0 +1,147 @@ +import argparse +import os +import sys +import pickle +import time +sys.path.append(os.getcwd()) + +from utils import * +from core.policy_gaussian import PolicyGaussian +from core.critic import Value +from models.mlp import MLP +from models.video_state_net import VideoStateNet +from ego_pose.envs.humanoid_v1 import HumanoidEnv +from ego_pose.core.agent_ego import AgentEgo +from ego_pose.utils.egomimic_config import Config +from ego_pose.core.reward_function import reward_func + + +parser = argparse.ArgumentParser() +parser.add_argument('--cfg', default=None) +parser.add_argument('--render', action='store_true', default=False) +parser.add_argument('--num-threads', type=int, default=12) +parser.add_argument('--gpu-index', type=int, default=0) +parser.add_argument('--iter', type=int, default=0) +parser.add_argument('--show-noise', action='store_true', default=False) +args = parser.parse_args() +if args.render: + args.num_threads = 1 +cfg = Config(args.cfg, create_dirs=not (args.render or args.iter > 0)) + +dtype = torch.float64 +torch.set_default_dtype(dtype) +device = torch.device('cuda', index=args.gpu_index) if torch.cuda.is_available() else torch.device('cpu') +if torch.cuda.is_available(): + torch.cuda.set_device(args.gpu_index) +np.random.seed(cfg.seed) +torch.manual_seed(cfg.seed) +tb_logger = Logger(cfg.tb_dir) if not args.render else None +logger = create_logger(os.path.join(cfg.log_dir, 'log.txt'), file_handle=not args.render) + +"""environment""" +env = HumanoidEnv(cfg) +env.seed(cfg.seed) +env.load_experts(cfg.takes['train'], cfg.expert_feat_file, cfg.cnn_feat_file) +cnn_feat_dim = env.cnn_feat[0].shape[-1] +actuators = env.model.actuator_names +state_dim = env.observation_space.shape[0] +action_dim = env.action_space.shape[0] +running_state = ZFilter((state_dim,), clip=5) + +"""define actor and critic""" +policy_vs_net = VideoStateNet(cnn_feat_dim, cfg.policy_v_hdim, cfg.fr_margin, cfg.policy_v_net, cfg.policy_v_net_param, cfg.causal) +value_vs_net = VideoStateNet(cnn_feat_dim, cfg.value_v_hdim, cfg.fr_margin, cfg.value_v_net, cfg.value_v_net_param, cfg.causal) +policy_net = PolicyGaussian(MLP(state_dim + cfg.policy_v_hdim, cfg.policy_hsize, cfg.policy_htype), action_dim, + log_std=cfg.log_std, fix_std=cfg.fix_std) +value_net = Value(MLP(state_dim + cfg.value_v_hdim, cfg.value_hsize, cfg.value_htype)) +if args.iter > 0: + cp_path = '%s/iter_%04d.p' % (cfg.model_dir, args.iter) + logger.info('loading model from checkpoint: %s' % cp_path) + model_cp = pickle.load(open(cp_path, "rb")) + policy_net.load_state_dict(model_cp['policy_dict']) + policy_vs_net.load_state_dict(model_cp['policy_vs_dict']) + value_net.load_state_dict(model_cp['value_dict']) + value_vs_net.load_state_dict(model_cp['value_vs_dict']) + running_state = model_cp['running_state'] +to_device(device, policy_net, value_net, policy_vs_net, value_vs_net) + +policy_params = list(policy_net.parameters()) + list(policy_vs_net.parameters()) +value_params = list(value_net.parameters()) + list(value_vs_net.parameters()) +if cfg.policy_optimizer == 'Adam': + optimizer_policy = torch.optim.Adam(policy_params, lr=cfg.policy_lr, weight_decay=cfg.policy_weightdecay) +else: + optimizer_policy = torch.optim.SGD(policy_params, lr=cfg.policy_lr, momentum=cfg.policy_momentum, weight_decay=cfg.policy_weightdecay) +if cfg.value_optimizer == 'Adam': + optimizer_value = torch.optim.Adam(value_params, lr=cfg.value_lr, weight_decay=cfg.value_weightdecay) +else: + optimizer_value = torch.optim.SGD(value_params, lr=cfg.value_lr, momentum=cfg.value_momentum, weight_decay=cfg.value_weightdecay) + +# reward functions +expert_reward = reward_func[cfg.reward_id] + +"""create agent""" +agent = AgentEgo(env=env, dtype=dtype, device=device, running_state=running_state, + custom_reward=expert_reward, mean_action=args.render and not args.show_noise, + render=args.render, num_threads=args.num_threads, + policy_net=policy_net, policy_vs_net=policy_vs_net, + value_net=value_net, value_vs_net=value_vs_net, + optimizer_policy=optimizer_policy, optimizer_value=optimizer_value, opt_num_epochs=cfg.num_optim_epoch, + gamma=cfg.gamma, tau=cfg.tau, clip_epsilon=cfg.clip_epsilon, + policy_grad_clip=[(policy_params, 40)]) + + +def pre_iter_update(i_iter): + cfg.update_adaptive_params(i_iter) + agent.set_noise_rate(cfg.adp_noise_rate) + set_optimizer_lr(optimizer_policy, cfg.adp_policy_lr) + if cfg.fix_std: + policy_net.action_log_std.fill_(cfg.adp_log_std) + return + + +def main_loop(): + + if args.render: + pre_iter_update(args.iter) + agent.sample(1e8) + else: + for i_iter in range(args.iter, cfg.max_iter_num): + """generate multiple trajectories that reach the minimum batch_size""" + pre_iter_update(i_iter) + batch, log = agent.sample(cfg.min_batch_size) + agent.env.end_reward = log.avg_c_reward * cfg.gamma / (1 - cfg.gamma) # set end reward + + """update networks""" + t0 = time.time() + agent.update_params(batch) + t1 = time.time() + + """logging""" + c_info = log.avg_c_info + logger.info( + '{}\tT_sample {:.2f}\tT_update {:.2f}\tR_avg {:.4f} {}' + '\tR_range ({:.4f}, {:.4f})\teps_len_avg {:.2f}' + .format(i_iter, log.sample_time, t1 - t0, log.avg_c_reward, + np.array2string(c_info, formatter={'all': lambda x: '%.4f' % x}, separator=','), + log.min_c_reward, log.max_c_reward, log.avg_episode_reward)) + + tb_logger.scalar_summary('total_reward', log.avg_c_reward, i_iter) + tb_logger.scalar_summary('episode_len', log.avg_episode_reward, i_iter) + for i in range(c_info.shape[0]): + tb_logger.scalar_summary('reward_%d' % i, c_info[i], i_iter) + + if cfg.save_model_interval > 0 and (i_iter+1) % cfg.save_model_interval == 0: + with to_cpu(policy_net, value_net, policy_vs_net, value_vs_net): + cp_path = '%s/iter_%04d.p' % (cfg.model_dir, i_iter + 1) + model_cp = {'policy_dict': policy_net.state_dict(), 'policy_vs_dict': policy_vs_net.state_dict(), + 'value_dict': value_net.state_dict(), 'value_vs_dict': value_vs_net.state_dict(), + 'running_state': running_state} + pickle.dump(model_cp, open(cp_path, 'wb')) + + """clean up gpu memory""" + torch.cuda.empty_cache() + + logger.info('training done!') + + +main_loop() diff --git a/ego_pose/ego_mimic_eval.py b/ego_pose/ego_mimic_eval.py new file mode 100644 index 0000000..c718405 --- /dev/null +++ b/ego_pose/ego_mimic_eval.py @@ -0,0 +1,199 @@ +import argparse +import os +import sys +import pickle +import time +sys.path.append(os.getcwd()) + +from utils import * +from core.policy_gaussian import PolicyGaussian +from core.critic import Value +from models.mlp import MLP +from models.video_state_net import VideoStateNet +from models.video_reg_net import VideoRegNet +from envs.visual.humanoid_vis import HumanoidVisEnv +from ego_pose.envs.humanoid_v1 import HumanoidEnv +from ego_pose.utils.egomimic_config import Config +from ego_pose.core.reward_function import reward_func + + +parser = argparse.ArgumentParser() +parser.add_argument('--cfg', default=None) +parser.add_argument('--render', action='store_true', default=False) +parser.add_argument('--iter', type=int, default=0) +parser.add_argument('--expert-ind', type=int, default=-1) +parser.add_argument('--sync', action='store_true', default=False) +parser.add_argument('--causal', action='store_true', default=False) +parser.add_argument('--data', default='test') +parser.add_argument('--show-noise', action='store_true', default=False) +parser.add_argument('--fail-safe', default='valuefs') +args = parser.parse_args() +cfg = Config(args.cfg, create_dirs=False) + +dtype = torch.float64 +torch.set_default_dtype(dtype) +np.random.seed(cfg.seed) +torch.manual_seed(cfg.seed) +torch.set_grad_enabled(False) +logger = create_logger(os.path.join(cfg.log_dir, 'log_eval.txt')) + +"""environment""" +env = HumanoidEnv(cfg) +env.load_experts(cfg.takes[args.data], cfg.expert_feat_file, cfg.cnn_feat_file) +env_vis = HumanoidVisEnv(cfg.vis_model_file, 10) +env.seed(cfg.seed) +epos = None +cnn_feat_dim = env.cnn_feat[0].shape[-1] +actuators = env.model.actuator_names +state_dim = env.observation_space.shape[0] +action_dim = env.action_space.shape[0] +body_qposaddr = get_body_qposaddr(env.model) +if args.fail_safe == 'naivefs': + env.set_fix_head_lb(0.3) + +"""load policy net""" +policy_vs_net = VideoStateNet(cnn_feat_dim, cfg.policy_v_hdim, cfg.fr_margin, cfg.policy_v_net, cfg.policy_v_net_param, cfg.causal) +value_vs_net = VideoStateNet(cnn_feat_dim, cfg.value_v_hdim, cfg.fr_margin, cfg.value_v_net, cfg.value_v_net_param, cfg.causal) +policy_net = PolicyGaussian(MLP(state_dim + cfg.policy_v_hdim, cfg.policy_hsize, cfg.policy_htype), action_dim, + log_std=cfg.log_std, fix_std=cfg.fix_std) +value_net = Value(MLP(state_dim + cfg.value_v_hdim, cfg.value_hsize, cfg.value_htype)) +cp_path = '%s/iter_%04d.p' % (cfg.model_dir, args.iter) +logger.info('loading policy net from checkpoint: %s' % cp_path) +model_cp = pickle.load(open(cp_path, "rb")) +policy_net.load_state_dict(model_cp['policy_dict']) +policy_vs_net.load_state_dict(model_cp['policy_vs_dict']) +value_net.load_state_dict(model_cp['value_dict']) +value_vs_net.load_state_dict(model_cp['value_vs_dict']) +running_state = model_cp['running_state'] +value_stat = RunningStat(1) + +"""load state net""" +cp_path = cfg.state_net_model +logger.info('loading state net from checkpoint: %s' % cp_path) +model_cp, meta = pickle.load(open(cp_path, "rb")) +state_net_mean, state_net_std, state_net_cfg = meta['mean'], meta['std'], meta['cfg'] +state_net = VideoRegNet(state_net_mean.size, state_net_cfg.v_hdim, cnn_feat_dim, no_cnn=True, + cnn_type=state_net_cfg.cnn_type, mlp_dim=state_net_cfg.mlp_dim, v_net_type=state_net_cfg.v_net, + v_net_param=state_net_cfg.v_net_param, causal=state_net_cfg.causal) +state_net.load_state_dict(model_cp['state_net_dict']) +to_test(policy_vs_net, policy_net, value_vs_net, value_net, state_net) + +# reward functions +expert_reward = reward_func[cfg.reward_id] + + +def render(): + env_vis.data.qpos[:env.model.nq] = env.data.qpos.copy() + env_vis.data.qpos[env.model.nq:] = epos + env_vis.data.qpos[env.model.nq] += 1.0 + env_vis.sim_forward() + env_vis.render() + + +def reset_env_state(state, ref_qpos): + qpos = ref_qpos.copy() + qpos[2:] = state[:qpos.size - 2] + qvel = state[qpos.size - 2:] + align_human_state(qpos, qvel, ref_qpos) + env.set_state(qpos, qvel) + return env.get_obs() + + +def eval_expert(expert_ind): + global epos + + expert_name = env.expert_list[expert_ind] + logger.info('Testing on expert trajectory %s' % expert_name) + + traj_pred = [] + traj_orig = [] + vel_pred = [] + num_reset = 0 + reward_episode = 0 + data_len = env.cnn_feat[expert_ind].shape[0] + test_len = data_len - 2 * cfg.fr_margin + env.set_fix_sampling(expert_ind, cfg.fr_margin, test_len) + + state = env.reset() + cnn_feat = tensor(env.get_episode_cnn_feat()) + policy_vs_net.initialize(cnn_feat) + value_vs_net.initialize(cnn_feat) + state_pred = state_net(cnn_feat.unsqueeze(1)).squeeze(1)[cfg.fr_margin: -cfg.fr_margin, :].numpy() + state_pred = state_pred * state_net_std[None, :] + state_net_mean[None, :] + + state = reset_env_state(state_pred[0, :], env.data.qpos) + if running_state is not None: + state = running_state(state, update=False) + + for t in range(test_len): + + ind = env.get_expert_index(t) + epos = env.get_expert_attr('qpos', ind).copy() + traj_pred.append(env.data.qpos.copy()) + traj_orig.append(epos.copy()) + vel_pred.append(env.data.qvel.copy()) + + if args.sync: + epos[:3] = quat_mul_vec(env.expert['rel_heading'], epos[:3] - env.expert['start_pos']) + env.expert['sim_pos'] + epos[3:7] = quaternion_multiply(env.expert['rel_heading'], epos[3:7]) + + if args.render: + render() + + if args.causal: + policy_vs_net.initialize(cnn_feat[:t + 2*cfg.fr_margin + 1]) + policy_vs_net.t = t + + """learner policy""" + state_var = tensor(state, dtype=dtype).unsqueeze(0) + policy_vs_out = policy_vs_net(state_var) + value_vs_out = value_vs_net(state_var) + value = value_net(value_vs_out).item() + value_stat.push(np.array([value])) + + action = policy_net.select_action(policy_vs_out, mean_action=not args.show_noise)[0].numpy() + next_state, reward, done, info = env.step(action) + if running_state is not None: + next_state = running_state(next_state, update=False) + + reward, cinfo = reward_func[cfg.reward_id](env, state, action, info) + logger.debug("{} {:.2f} {} {:.2f}".format(t, reward, np.array2string(cinfo, formatter={'all': lambda x: '%.4f' % x}), value)) + + reward_episode += reward + + if info['end']: + break + + if args.fail_safe == 'valuefs' and value < 0.6 * value_stat.mean[0] or \ + args.fail_safe == 'naivefs' and info['fail']: + logger.info('reset state!') + num_reset += 1 + state = reset_env_state(state_pred[t+1, :], env.data.qpos) + if running_state is not None: + state = running_state(state, update=False) + else: + state = next_state + + return np.vstack(traj_pred), np.vstack(traj_orig), np.vstack(vel_pred), num_reset + + +# if expert_ind is defined, then keeping visualizing this trajectory +if args.expert_ind >= 0: + for i in range(100): + eval_expert(args.expert_ind) +else: + traj_pred = {} + traj_orig = {} + vel_pred = {} + num_reset = 0 + for i, take in enumerate(env.expert_list): + traj_pred[take], traj_orig[take], vel_pred[take], t_num_reset = eval_expert(i) + num_reset += t_num_reset + results = {'traj_pred': traj_pred, 'traj_orig': traj_orig, 'vel_pred': vel_pred} + meta = {'algo': 'ego_mimic', 'num_reset': num_reset} + fs_tag = '' if args.fail_safe == 'valuefs' else '_' + args.fail_safe + c_tag = '_causal' if args.causal else '' + res_path = '%s/iter_%04d_%s%s%s.p' % (cfg.result_dir, args.iter, args.data, fs_tag, c_tag) + pickle.dump((results, meta), open(res_path, 'wb')) + logger.info('num reset: %d' % num_reset) + logger.info('saved results to %s' % res_path) diff --git a/ego_pose/ego_mimic_eval_wild.py b/ego_pose/ego_mimic_eval_wild.py new file mode 100644 index 0000000..dd04428 --- /dev/null +++ b/ego_pose/ego_mimic_eval_wild.py @@ -0,0 +1,156 @@ +import argparse +import os +import sys +import pickle +import time +sys.path.append(os.getcwd()) + +from utils import * +from core.policy_gaussian import PolicyGaussian +from core.critic import Value +from models.mlp import MLP +from models.video_state_net import VideoStateNet +from models.video_reg_net import VideoRegNet +from envs.visual.humanoid_vis import HumanoidVisEnv +from ego_pose.envs.humanoid_v1 import HumanoidEnv +from ego_pose.utils.egomimic_config import Config + + +parser = argparse.ArgumentParser() +parser.add_argument('--cfg', default=None) +parser.add_argument('--render', action='store_true', default=False) +parser.add_argument('--iter', type=int, default=0) +parser.add_argument('--test-ind', type=int, default=-1) +parser.add_argument('--test-feat', default=None) +parser.add_argument('--show-noise', action='store_true', default=False) +args = parser.parse_args() +cfg = Config(args.cfg, create_dirs=False) + +dtype = torch.float64 +torch.set_default_dtype(dtype) +np.random.seed(cfg.seed) +torch.manual_seed(cfg.seed) +torch.set_grad_enabled(False) +logger = create_logger(os.path.join(cfg.log_dir, 'log_eval_wild.txt')) + +"""test data""" +cnn_feat_file = '%s/features/cnn_feat_%s.p' % (cfg.data_dir, args.test_feat) +cnn_feat_dict, _ = pickle.load(open(cnn_feat_file, 'rb')) +takes = list(cnn_feat_dict.keys()) + +"""environment""" +env = HumanoidEnv(cfg) +env_vis = HumanoidVisEnv(cfg.mujoco_model_file, 10) +env.seed(cfg.seed) +cnn_feat_dim = cnn_feat_dict[takes[0]].shape[-1] +actuators = env.model.actuator_names +state_dim = env.observation_space.shape[0] +action_dim = env.action_space.shape[0] +body_qposaddr = get_body_qposaddr(env.model) + +"""load policy net""" +policy_vs_net = VideoStateNet(cnn_feat_dim, cfg.policy_v_hdim, cfg.fr_margin, cfg.policy_v_net, cfg.policy_v_net_param, cfg.causal) +value_vs_net = VideoStateNet(cnn_feat_dim, cfg.value_v_hdim, cfg.fr_margin, cfg.value_v_net, cfg.value_v_net_param, cfg.causal) +policy_net = PolicyGaussian(MLP(state_dim + cfg.policy_v_hdim, cfg.policy_hsize, cfg.policy_htype), action_dim, + log_std=cfg.log_std, fix_std=cfg.fix_std) +value_net = Value(MLP(state_dim + cfg.value_v_hdim, cfg.value_hsize, cfg.value_htype)) +cp_path = '%s/iter_%04d.p' % (cfg.model_dir, args.iter) +logger.info('loading policy net from checkpoint: %s' % cp_path) +model_cp = pickle.load(open(cp_path, "rb")) +policy_net.load_state_dict(model_cp['policy_dict']) +policy_vs_net.load_state_dict(model_cp['policy_vs_dict']) +value_net.load_state_dict(model_cp['value_dict']) +value_vs_net.load_state_dict(model_cp['value_vs_dict']) +running_state = model_cp['running_state'] +value_stat = RunningStat(1) + +"""load state net""" +cp_path = cfg.state_net_model +logger.info('loading state net from checkpoint: %s' % cp_path) +model_cp, meta = pickle.load(open(cp_path, "rb")) +state_net_mean, state_net_std, state_net_cfg = meta['mean'], meta['std'], meta['cfg'] +state_net = VideoRegNet(state_net_mean.size, state_net_cfg.v_hdim, cnn_feat_dim, no_cnn=True, + cnn_type=state_net_cfg.cnn_type, mlp_dim=state_net_cfg.mlp_dim, v_net_type=state_net_cfg.v_net, + v_net_param=state_net_cfg.v_net_param, causal=state_net_cfg.causal) +state_net.load_state_dict(model_cp['state_net_dict']) +to_test(policy_vs_net, policy_net, value_vs_net, value_net, state_net) + + +def render(): + env_vis.data.qpos[:] = env.data.qpos + env_vis.sim_forward() + env_vis.render() + + +def reset_env_state(state, ref_qpos): + qpos = ref_qpos.copy() + qpos[2:] = state[:qpos.size - 2] + qvel = state[qpos.size - 2:] + align_human_state(qpos, qvel, ref_qpos) + env.set_state(qpos, qvel) + return env.get_obs() + + +def eval_take(take): + logger.info('Testing on %s' % take) + + traj_pred = [] + vel_pred = [] + cnn_feat = cnn_feat_dict[take] + data_len = cnn_feat.shape[0] + test_len = data_len - 2 * cfg.fr_margin + state = env.reset() + cnn_feat = tensor(cnn_feat) + policy_vs_net.initialize(cnn_feat) + value_vs_net.initialize(cnn_feat) + state_pred = state_net(cnn_feat.unsqueeze(1)).squeeze(1)[cfg.fr_margin: -cfg.fr_margin, :].numpy() + state_pred = state_pred * state_net_std[None, :] + state_net_mean[None, :] + + state = reset_env_state(state_pred[0, :], env.data.qpos) + if running_state is not None: + state = running_state(state, update=False) + + for t in range(test_len): + traj_pred.append(env.data.qpos.copy()) + vel_pred.append(env.data.qvel.copy()) + + if args.render: + render() + + """learner policy""" + state_var = tensor(state, dtype=dtype).unsqueeze(0) + policy_vs_out = policy_vs_net(state_var) + value_vs_out = value_vs_net(state_var) + value = value_net(value_vs_out).item() + value_stat.push(np.array([value])) + + action = policy_net.select_action(policy_vs_out, mean_action=not args.show_noise)[0].numpy() + next_state, reward, done, info = env.step(action) + if running_state is not None: + next_state = running_state(next_state, update=False) + + if value < 0.6 * value_stat.mean[0]: + logger.info('reset state!') + state = reset_env_state(state_pred[t+1, :], env.data.qpos) + if running_state is not None: + state = running_state(state, update=False) + else: + state = next_state + + return np.vstack(traj_pred), np.vstack(vel_pred) + + +# if test_ind is defined, then keeping visualizing this trajectory +if args.test_ind >= 0: + for i in range(100): + eval_take(takes[args.test_ind]) +else: + traj_pred = {} + vel_pred = {} + for take in takes: + traj_pred[take], vel_pred[take] = eval_take(take) + results = {'traj_pred': traj_pred, 'vel_pred': vel_pred} + meta = {'algo': 'ego_mimic'} + res_path = '%s/iter_%04d_%s.p' % (cfg.result_dir, args.iter, args.test_feat) + pickle.dump((results, meta), open(res_path, 'wb')) + logger.info('saved results to %s' % res_path) diff --git a/ego_pose/envs/humanoid_v1.py b/ego_pose/envs/humanoid_v1.py new file mode 100644 index 0000000..d4bdcdb --- /dev/null +++ b/ego_pose/envs/humanoid_v1.py @@ -0,0 +1,295 @@ +import numpy as np +from envs.common import mujoco_env +from gym import spaces +from utils import * +from utils.transformation import quaternion_from_euler, rotation_from_quaternion, quaternion_about_axis +from mujoco_py import functions as mjf +import pickle +import time +import cv2 as cv +from scipy.linalg import cho_solve, cho_factor + + +class HumanoidEnv(mujoco_env.MujocoEnv): + + def __init__(self, cfg): + mujoco_env.MujocoEnv.__init__(self, cfg.mujoco_model_file, 15) + self.cfg = cfg + # visualization + self.save_video = False + self.video_res = (224, 224) + self.video_dir = cfg.video_dir + self.set_cam_first = set() + self.subsample_rate = 1 + # env specific + self.end_reward = 0.0 + self.start_ind = 0 + self.body_qposaddr = get_body_qposaddr(self.model) + self.bquat = self.get_body_quat() + self.prev_bquat = None + self.set_spaces() + self.set_model_params() + # expert + self.expert_ind = None + self.expert_id = None + self.expert_list = None # name only + self.expert_arr = None # store actual experts + self.expert = None + self.cnn_feat = None + # fixed sampling + self.fix_expert_ind = None + self.fix_start_ind = None + self.fix_len = None + self.fix_start_state = None + self.fix_cnn_feat = None + self.fix_head_lb = None + + def load_experts(self, expert_list, expert_feat_file, cnn_feat_file): + self.expert_ind = 0 + self.expert_list = expert_list + expert_dict = pickle.load(open(expert_feat_file, 'rb')) + self.expert_arr = [expert_dict[x] for x in self.expert_list] + self.set_expert(0) + cnn_feat_dict, _ = pickle.load(open(cnn_feat_file, 'rb')) + self.cnn_feat = [cnn_feat_dict[x] for x in self.expert_list] + + def set_model_params(self): + if self.cfg.action_type == 'torque' and hasattr(self.cfg, 'j_stiff'): + self.model.jnt_stiffness[1:] = self.cfg.j_stiff + self.model.dof_damping[6:] = self.cfg.j_damp + + def set_spaces(self): + bounds = self.model.actuator_ctrlrange.copy() + self.action_space = spaces.Box(low=bounds[:, 0], high=bounds[:, 1], dtype=np.float32) + self.obs_dim = self.get_obs().size + high = np.inf * np.ones(self.obs_dim) + low = -high + self.observation_space = spaces.Box(low, high, dtype=np.float32) + + def get_obs(self): + obs = self.get_full_obs() + return obs + + def get_full_obs(self): + data = self.data + qpos = data.qpos.copy() + qvel = data.qvel.copy() + # transform velocity + qvel[:3] = transform_vec(qvel[:3], qpos[3:7], self.cfg.obs_coord).ravel() + obs = [] + # pos + if self.cfg.obs_heading: + obs.append(np.array([get_heading(qpos[3:7])])) + if self.cfg.root_deheading: + qpos[3:7] = de_heading(qpos[3:7]) + obs.append(qpos[2:]) + # vel + if self.cfg.obs_vel == 'root': + obs.append(qvel[:6]) + elif self.cfg.obs_vel == 'full': + obs.append(qvel) + # phase + if hasattr(self.cfg, 'obs_phase') and self.cfg.obs_phase: + phase = min(self.cur_t / self.cfg.env_episode_len, 1.0) + obs.append(np.array([phase])) + obs = np.concatenate(obs) + return obs + + def get_ee_pos(self, transform): + data = self.data + ee_name = ['LeftFoot', 'RightFoot', 'LeftHand', 'RightHand', 'Head'] + ee_pos = [] + root_pos = data.qpos[:3] + root_q = data.qpos[3:7].copy() + for name in ee_name: + bone_id = self.model._body_name2id[name] + bone_vec = self.data.body_xpos[bone_id] + if transform is not None: + bone_vec = bone_vec - root_pos + bone_vec = transform_vec(bone_vec, root_q, transform) + ee_pos.append(bone_vec) + return np.concatenate(ee_pos) + + def get_body_quat(self): + qpos = self.data.qpos.copy() + body_quat = [qpos[3:7]] + for body in self.model.body_names[1:]: + if body == 'Hips': + continue + start, end = self.body_qposaddr[body] + euler = np.zeros(3) + euler[:end - start] = qpos[start:end] + quat = quaternion_from_euler(euler[0], euler[1], euler[2]) + body_quat.append(quat) + body_quat = np.concatenate(body_quat) + return body_quat + + def get_com(self): + return self.data.subtree_com[0, :].copy() + + def compute_desired_accel(self, qpos_err, qvel_err): + dt = self.model.opt.timestep + nv = self.model.nv + M = np.zeros(nv * nv) + mjf.mj_fullM(self.model, M, self.data.qM) + M.resize(self.model.nv, self.model.nv) + C = self.data.qfrc_bias.copy() + k_p = np.zeros(nv) + k_d = np.zeros(nv) + k_p[6:] = self.cfg.jkp + k_d[6:] = self.cfg.jkd + K_p = np.diag(k_p) + K_d = np.diag(k_d) + q_accel = cho_solve(cho_factor(M + K_d*dt), -C[:, None] - K_p.dot(qpos_err[:, None]) - K_d.dot(qvel_err[:, None])) + return q_accel.squeeze() + + def compute_torque(self, ctrl): + cfg = self.cfg + dt = self.model.opt.timestep + qpos = self.data.qpos.copy() + qvel = self.data.qvel.copy() + qpos_err = np.concatenate((np.zeros(6), qpos[7:] - ctrl)) + qvel_err = qvel + q_accel = self.compute_desired_accel(qpos_err, qvel_err) + qvel_err += q_accel * dt + torque = -cfg.jkp * qpos_err[6:] - cfg.jkd * qvel_err[6:] + return torque + + def do_simulation(self, action, n_frames): + t0 = time.time() + cfg = self.cfg + for i in range(n_frames): + if self.save_video and i % self.subsample_rate == 0: + img = self.render('image', *self.video_res) + vind = self.cur_t * (n_frames // self.subsample_rate) + i // self.subsample_rate + cv.imwrite('%s/%04d.png' % (self.video_dir, vind), img) + + ctrl = cfg.a_ref + action * cfg.a_scale + if cfg.action_type == 'position': + torque = self.compute_torque(ctrl) + elif cfg.action_type == 'torque': + torque = ctrl + torque = np.clip(torque, -cfg.torque_lim, cfg.torque_lim) + self.data.ctrl[:] = torque + self.sim.step() + + if self.viewer is not None: + self.viewer.sim_time = time.time() - t0 + + def step(self, a): + cfg = self.cfg + # record prev state + self.prev_qpos = self.data.qpos.copy() + self.prev_qvel = self.data.qvel.copy() + self.prev_bquat = self.bquat.copy() + # do simulation + self.do_simulation(a, self.frame_skip) + self.cur_t += 1 + self.bquat = self.get_body_quat() + self.sync_expert() + # get obs + head_pos = self.get_body_com('Head') + reward = 1.0 + if self.fix_head_lb is not None: + fail = head_pos[2] < self.fix_head_lb + else: + fail = self.expert is not None and head_pos[2] < self.expert['head_height_lb'] - 0.1 + end = self.cur_t >= (cfg.env_episode_len if self.fix_len is None else self.fix_len) + done = fail or end + return self.get_obs(), reward, done, {'fail': fail, 'end': end} + + def reset_model(self): + if self.fix_start_state is not None: + init_pose = self.fix_start_state[:self.model.nq] + init_vel = self.fix_start_state[self.model.nq:] + self.set_state(init_pose, init_vel) + elif self.expert_list is not None: + cfg = self.cfg + fr_margin = cfg.fr_margin + # sample expert + expert_ind = self.np_random.randint(len(self.expert_list)) if self.fix_expert_ind is None else self.fix_expert_ind + self.set_expert(expert_ind) + # sample start frame + if self.fix_start_ind is None: + ind = 0 if cfg.env_start_first else self.np_random.randint(fr_margin, self.expert['len'] - cfg.env_episode_len - fr_margin) + else: + ind = self.fix_start_ind + self.start_ind = ind + if hasattr(cfg, 'random_cur_t') and cfg.random_cur_t: + self.cur_t = np.random.randint(cfg.env_episode_len) + ind += self.cur_t + init_pose = self.expert['qpos'][ind, :].copy() + init_vel = self.expert['qvel'][ind, :].copy() + init_pose[7:] += self.np_random.normal(loc=0.0, scale=cfg.env_init_noise, size=self.model.nq - 7) + self.set_state(init_pose, init_vel) + self.bquat = self.get_body_quat() + self.sync_expert() + else: + init_pose = self.data.qpos + init_pose[2] += 1.0 + self.set_state(init_pose, self.data.qvel) + return self.get_obs() + + def viewer_setup(self, mode): + self.viewer.cam.trackbodyid = 1 + self.viewer.cam.lookat[:2] = self.data.qpos[0:2] + if mode not in self.set_cam_first: + self.viewer.video_fps = 33 + self.viewer.frame_skip = self.frame_skip + self.viewer.cam.distance = self.model.stat.extent * 1.2 + self.viewer.cam.elevation = -20 + self.viewer.cam.azimuth = 45 + self.set_cam_first.add(mode) + + def set_fix_sampling(self, expert_ind=None, start_ind=None, len=None, start_state=None, cnn_feat=None): + self.fix_expert_ind = expert_ind + self.fix_start_ind = start_ind + self.fix_len = len + self.fix_start_state = start_state + self.fix_cnn_feat = cnn_feat + + def set_fix_head_lb(self, fix_head_lb=None): + self.fix_head_lb = fix_head_lb + + def sync_expert(self): + if self.expert is not None and self.cur_t % self.cfg.sync_exp_interval == 0: + expert = self.expert + ind = self.get_expert_index(self.cur_t) + e_qpos = self.get_expert_attr('qpos', ind).copy() + expert['rel_heading'] = quaternion_multiply(get_heading_q(self.data.qpos[3:7]), + quaternion_inverse(get_heading_q(e_qpos[3:7]))) + expert['start_pos'] = e_qpos[:3] + expert['sim_pos'] = np.concatenate((self.data.qpos[:2], np.array([e_qpos[2]]))) + + def set_expert(self, expert_ind): + self.expert_ind = expert_ind + self.expert_id = self.expert_list[expert_ind] + self.expert = self.expert_arr[expert_ind] + + def get_expert_index(self, t): + return self.start_ind + t + + def get_expert_attr(self, attr, ind): + return self.expert[attr][ind, :] + + def get_pose_dist(self): + ind = self.get_expert_index(self.cur_t) + qpos_e = self.expert['qpos'][ind, :] + qpos_g = self.data.qpos + diff = qpos_e - qpos_g + return np.linalg.norm(diff[2:]) + + def get_pose_diff(self): + ind = self.get_expert_index(self.cur_t) + qpos_e = self.expert['qpos'][ind, :] + qpos_g = self.data.qpos + diff = qpos_e - qpos_g + return np.abs(diff[2:]) + + def get_episode_cnn_feat(self): + fm = self.cfg.fr_margin + num_fr = self.cfg.env_episode_len if self.fix_len is None else self.fix_len + return self.cnn_feat[self.expert_ind][self.start_ind - fm: self.start_ind + num_fr + fm, :] \ + if self.fix_cnn_feat is None else self.fix_cnn_feat + + diff --git a/ego_pose/eval_forecast.py b/ego_pose/eval_forecast.py new file mode 100644 index 0000000..e28695d --- /dev/null +++ b/ego_pose/eval_forecast.py @@ -0,0 +1,260 @@ +import argparse +import os +import sys +import pickle +import math +import time +import numpy as np +sys.path.append(os.getcwd()) + +from ego_pose.utils.metrics import * +from ego_pose.utils.tools import * +from ego_pose.utils.egoforecast_config import Config +from ego_pose.envs.humanoid_v1 import HumanoidEnv +from envs.visual.humanoid_vis import HumanoidVisEnv + + +parser = argparse.ArgumentParser() +parser.add_argument('--vis-model', default='humanoid_1205_vis_ghost_v1') +parser.add_argument('--multi-vis-model', default='humanoid_1205_vis_forecast_v1') +parser.add_argument('--multi', action='store_true', default=False) +parser.add_argument('--egoforecast-cfg', default='subject_03') +parser.add_argument('--egoforecast-iter', type=int, default=3000) +parser.add_argument('--suffix', default='') +parser.add_argument('--data', default='test') +parser.add_argument('--mode', default='vis') +args = parser.parse_args() + + +def compute_metrics(results, algo, horizon, verbose=True): + if results is None: + return + + if verbose: + print('=' * 10 + ' %s ' % algo + '=' * 10) + g_pose_dist = 0 + g_vel_dist = 0 + g_smoothness = 0 + traj_orig = results['traj_orig'] + traj_pred = results['traj_pred'] + + for take in traj_pred.keys(): + t_pose_dist = 0 + t_vel_dist = 0 + t_smoothness = 0 + for i in range(traj_orig[take].shape[0]): + traj = traj_pred[take][i, cfg.fr_margin:cfg.fr_margin + horizon, :] + traj_gt = traj_orig[take][i, cfg.fr_margin:cfg.fr_margin + horizon, :] + # compute gt stats + angs_gt = get_joint_angles(traj_gt) + vels_gt = get_joint_vels(traj_gt, dt) + # compute pred stats + angs = get_joint_angles(traj) + vels = get_joint_vels(traj, dt) + accels = get_joint_accels(vels, dt) + # compute metrics + pose_dist = get_mean_dist(angs, angs_gt) + vel_dist = get_mean_dist(vels, vels_gt) # exclude root + smoothness = get_mean_abs(accels) # exclude root + t_pose_dist += pose_dist + t_vel_dist += vel_dist + t_smoothness += smoothness + + t_pose_dist /= traj_orig[take].shape[0] + t_vel_dist /= traj_orig[take].shape[0] + t_smoothness /= traj_orig[take].shape[0] + + if verbose: + print('%s - horizon: %d, pose dist: %.4f, vel dist: %.4f, accels: %.4f' % + (take, horizon, t_pose_dist, t_vel_dist, t_smoothness)) + + g_pose_dist += t_pose_dist + g_vel_dist += t_vel_dist + g_smoothness += t_smoothness + + g_pose_dist /= len(traj_pred) + g_vel_dist /= len(traj_pred) + g_smoothness /= len(traj_pred) + + if verbose: + print('-' * 60) + print('all - horizon: %d, pose dist: %.4f, vel dist: %.4f, accels: %.4f' % + (horizon, g_pose_dist, g_vel_dist, g_smoothness)) + print('-' * 60 + '\n') + + return g_pose_dist, g_vel_dist, g_smoothness + + +def compute_err_vs_h(results, algo, horizon, step=10): + errors = [] + for h in range(step, horizon, step): + err, _, _ = compute_metrics(results, algo, h, False) + errors.append(err) + errors = np.array(errors) + print('-' * 60) + print(algo) + print(np.array2string(errors, formatter={'all': lambda x: '%.4f' % x}, separator=', ')) + print('-' * 60 + '\n') + return errors + + +cfg = Config(args.egoforecast_cfg, False) +dt = 1 / 30.0 + +res_base_dir = 'results' +ef_res_path = '%s/egoforecast/%s/results/iter_%04d_%s%s.p' % (res_base_dir, args.egoforecast_cfg, args.egoforecast_iter, args.data, args.suffix) +ef_res, ef_meta = pickle.load(open(ef_res_path, 'rb')) if args.egoforecast_cfg is not None else (None, None) +# some mocap data we captured have very noisy hands, +# we set hand angles to 0 during training and we do not use it in evaluation +remove_noisy_hands(ef_res) + +if args.mode == 'stats': + # compute_err_vs_h(ef_res, 'ego forecast', 95) + compute_metrics(ef_res, 'ego forecast', 30) + compute_metrics(ef_res, 'ego forecast', 90) + +elif args.mode == 'vis': + """visualization""" + + def key_callback(key, action, mods): + global T, fr, paused, stop, reverse, algo_ind, take_ind, traj_ind, ss_ind, show_gt, mfr_int + + if action != glfw.RELEASE: + return False + if key == glfw.KEY_D: + T *= 1.5 + elif key == glfw.KEY_F: + T = max(1, T / 1.5) + elif key == glfw.KEY_R: + stop = True + elif key == glfw.KEY_W: + fr = 0 + update_pose() + elif key == glfw.KEY_S: + reverse = not reverse + elif key == glfw.KEY_Z: + fr = 0 + traj_ind = 0 + take_ind = (take_ind - 1) % len(takes) + load_take() + update_pose() + elif key == glfw.KEY_C: + fr = 0 + traj_ind = 0 + take_ind = (take_ind + 1) % len(takes) + load_take() + update_pose() + elif key == glfw.KEY_Q: + fr = 0 + traj_ind = (traj_ind - 1) % traj_orig.shape[0] + update_pose() + elif key == glfw.KEY_E: + fr = 0 + traj_ind = (traj_ind + 1) % traj_orig.shape[0] + update_pose() + elif key == glfw.KEY_X: + save_screen_shots(env_vis.viewer.window, 'out/%04d.png' % ss_ind) + ss_ind += 1 + elif glfw.KEY_1 <= key < glfw.KEY_1 + len(algos): + algo_ind = key - glfw.KEY_1 + load_take() + update_pose() + elif key == glfw.KEY_0: + show_gt = not show_gt + update_pose() + elif key == glfw.KEY_MINUS: + mfr_int -= 1 + update_pose() + elif key == glfw.KEY_EQUAL: + mfr_int += 1 + update_pose() + elif key == glfw.KEY_RIGHT: + if fr < traj_orig.shape[1] - 1: + fr += 1 + update_pose() + elif key == glfw.KEY_LEFT: + if fr > 0: + fr -= 1 + update_pose() + elif key == glfw.KEY_SPACE: + paused = not paused + else: + return False + + return True + + def update_pose(): + print('take ind: %d, traj ind: %d, mfr int: %d' % (take_ind, traj_ind, mfr_int)) + if args.multi: + fr_s = cfg.fr_margin + nq = 59 + traj = traj_orig if show_gt else traj_pred + num_model = env_vis.model.nq // nq + hq = get_heading_q(traj_orig[traj_ind, fr_s, 3:7]) + rel_q = quaternion_multiply(hq, quaternion_inverse(get_heading_q(traj[traj_ind, fr_s, 3:7]))) + vec = quat_mul_vec(hq, np.array([0, -1, 0]))[:2] + for i in range(num_model): + fr_m = max(0, min(fr_s + (i-3) * mfr_int, traj.shape[1] - 1)) + env_vis.data.qpos[i*nq: (i + 1) * nq] = traj[traj_ind, fr_m, :] + env_vis.data.qpos[i*nq + 3: i * nq + 7] = quaternion_multiply(rel_q, traj[traj_ind, fr_m, 3:7]) + env_vis.data.qpos[i*nq: i * nq + 2] = traj_orig[traj_ind, fr_s, :2] + vec * 0.8 * i + else: + nq = env_vis.model.nq // 2 + traj = traj_orig if show_gt else traj_pred + if fr >= cfg.fr_margin: + env_vis.data.qpos[:nq] = traj[traj_ind, fr, :] + env_vis.data.qpos[nq:] = traj[traj_ind, cfg.fr_margin, :] + else: + env_vis.data.qpos[:nq] = traj[traj_ind, fr, :] + env_vis.data.qpos[nq:] = traj[traj_ind, fr, :] + env_vis.sim_forward() + + def load_take(): + global traj_pred, traj_orig + algo_res, algo = algos[algo_ind] + if algo_res is None: + return + take = takes[take_ind] + print('%s ---------- %s' % (algo, take)) + traj_pred = algo_res['traj_pred'][take] + traj_orig = algo_res['traj_orig'][take] + + traj_pred = None + traj_orig = None + vis_model_file = 'assets/mujoco_models/%s.xml' % (args.multi_vis_model if args.multi else args.vis_model) + env_vis = HumanoidVisEnv(vis_model_file, 1, focus=not args.multi) + env_vis.set_custom_key_callback(key_callback) + takes = cfg.takes[args.data] + algos = [(ef_res, 'ego forecast')] + algo_ind = 0 + take_ind = 0 + traj_ind = 0 + ss_ind = 0 + mfr_int = 10 + show_gt = False + load_take() + + """render or select part of the clip""" + T = 10 + fr = 0 + paused = False + stop = False + reverse = False + + update_pose() + t = 0 + while not stop: + if t >= math.floor(T): + if not reverse and fr < traj_orig.shape[1] - 1: + fr += 1 + update_pose() + elif reverse and fr > 0: + fr -= 1 + update_pose() + t = 0 + + env_vis.render() + + if not paused: + t += 1 + diff --git a/ego_pose/eval_forecast_wild.py b/ego_pose/eval_forecast_wild.py new file mode 100644 index 0000000..1a1e4e1 --- /dev/null +++ b/ego_pose/eval_forecast_wild.py @@ -0,0 +1,294 @@ +import argparse +import os +import sys +import pickle +import math +import time +import numpy as np +import yaml +import glob +sys.path.append(os.getcwd()) + +from ego_pose.utils.metrics import * +from ego_pose.utils.tools import * +from ego_pose.utils.pose2d import Pose2DContext +from ego_pose.utils.egoforecast_config import Config +from ego_pose.envs.humanoid_v1 import HumanoidEnv +from envs.visual.humanoid_vis import HumanoidVisEnv + + +parser = argparse.ArgumentParser() +parser.add_argument('--vis-model', default='humanoid_1205_vis_ghost_v1') +parser.add_argument('--multi-vis-model', default='humanoid_1205_vis_blank_v1') +parser.add_argument('--multi', action='store_true', default=False) +parser.add_argument('--egoforecast-cfg', default='cross_01') +parser.add_argument('--egoforecast-iter', type=int, default=6000) +parser.add_argument('--horizon', type=int, default=30) +parser.add_argument('--data', default='wild_01') +parser.add_argument('--take-ind', type=int, default=0) +parser.add_argument('--mode', default='vis') +parser.add_argument('--tpv', action='store_true', default=True) +parser.add_argument('--stats-vis', action='store_true', default=False) +args = parser.parse_args() + + +cfg = Config(args.egoforecast_cfg, False) +dt = 1 / 30.0 +data_dir = 'datasets' +meta = yaml.load(open('%s/meta/meta_%s.yml' % (data_dir, args.data))) +pose_ctx = Pose2DContext(cfg) + +res_base_dir = 'results' +ef_res_path = '%s/egoforecast/%s/results/iter_%04d_%s.p' % (res_base_dir, args.egoforecast_cfg, args.egoforecast_iter, args.data) +ef_res, ef_meta = pickle.load(open(ef_res_path, 'rb')) if args.egoforecast_cfg is not None else (None, None) +# some mocap data we captured have very noisy hands, +# we set hand angles to 0 during training and we do not use it in evaluation +remove_noisy_hands(ef_res) + +if args.mode == 'stats': + + def get_kp_dist(traj, take, start_fr): + pose_dist = 0 + traj_ub = meta['traj_ub'].get(take, None) + tpv_offset = meta['tpv_offset'].get(take, cfg.fr_margin) + flip = meta['tpv_flip'].get(take, False) + fr_num = traj.shape[0] + valid_num = 0 + for fr in range(fr_num): + if traj_ub is not None and start_fr + fr >= traj_ub: + break + gt_fr = start_fr + fr + tpv_offset + gt_file = '%s/tpv/poses/%s/%05d_keypoints.json' % (data_dir, take, gt_fr) + gt_p = pose_ctx.load_gt_pose(gt_file) + if not pose_ctx.check_gt(gt_p): + # print('invalid frame: %s, %d' % (take, fr)) + continue + valid_num += 1 + qpos = traj[fr, :] + p = pose_ctx.align_qpos(qpos, gt_p, flip=flip) + dist = pose_ctx.get_pose_dist(p, gt_p) + pose_dist += dist + if args.stats_vis: + img = cv2.imread('%s/tpv/s_frames/%s/%05d.jpg' % (data_dir, take, gt_fr)) + pose_ctx.draw_pose(img, p * 0.25, flip=flip) + cv2.imshow('', img) + cv2.waitKey(1) + pose_dist /= valid_num + return pose_dist + + def compute_metrics(results, algo, horizon): + if results is None: + return + + print('=' * 10 + ' %s ' % algo + '=' * 10) + g_pose_dist = 0 + g_smoothness = 0 + traj_pred = results['traj_pred'] + + for take in traj_pred.keys(): + t_pose_dist = 0 + t_smoothness = 0 + for i in range(traj_pred[take].shape[0]): + traj = traj_pred[take][i, cfg.fr_margin:cfg.fr_margin + horizon, :] + kp_dist = get_kp_dist(traj, take, (i + 1) * cfg.fr_margin) + # compute pred stats + vels = get_joint_vels(traj, dt) + accels = get_joint_accels(vels, dt) + # compute metrics + smoothness = get_mean_abs(accels) # exclude root + t_pose_dist += kp_dist + t_smoothness += smoothness + + t_pose_dist /= traj_pred[take].shape[0] + t_smoothness /= traj_pred[take].shape[0] + + print('%s - pose dist: %.4f, accels: %.4f' % (take, t_pose_dist, t_smoothness)) + + g_pose_dist += t_pose_dist + g_smoothness += t_smoothness + + g_pose_dist /= len(traj_pred) + g_smoothness /= len(traj_pred) + + print('-' * 60) + print('all - pose dist: %.4f, accels: %.4f' % (g_pose_dist, g_smoothness)) + print('-' * 60 + '\n') + + compute_metrics(ef_res, 'ego forecast', args.horizon) + +elif args.mode == 'vis': + """visualization""" + + def key_callback(key, action, mods): + global T, fr, paused, stop, reverse, algo_ind, take_ind, traj_ind, ss_ind, mfr_int + + if action != glfw.RELEASE: + return False + if key == glfw.KEY_D: + T *= 1.5 + elif key == glfw.KEY_F: + T = max(1, T / 1.5) + elif key == glfw.KEY_R: + stop = True + elif key == glfw.KEY_W: + fr = 0 + update_all() + elif key == glfw.KEY_S: + reverse = not reverse + elif key == glfw.KEY_Z: + fr = 0 + traj_ind = 0 + take_ind = (take_ind - 1) % len(takes) + load_take() + update_all() + elif key == glfw.KEY_C: + fr = 0 + traj_ind = 0 + take_ind = (take_ind + 1) % len(takes) + load_take() + update_all() + elif key == glfw.KEY_Q: + fr = 0 + traj_ind = (traj_ind - 1) % traj_pred.shape[0] + update_pose() + elif key == glfw.KEY_E: + fr = 0 + traj_ind = (traj_ind + 1) % traj_pred.shape[0] + update_pose() + elif key == glfw.KEY_X: + save_screen_shots(env_vis.viewer.window, 'out/%04d.png' % ss_ind) + ss_ind += 1 + elif key == glfw.KEY_V: + env_vis.focus = not env_vis.focus + elif glfw.KEY_1 <= key < glfw.KEY_1 + len(algos): + algo_ind = key - glfw.KEY_1 + load_take(False) + update_all() + elif key == glfw.KEY_MINUS: + mfr_int -= 1 + update_pose() + elif key == glfw.KEY_EQUAL: + mfr_int += 1 + update_pose() + elif key == glfw.KEY_RIGHT: + if fr < traj_pred.shape[1] - 1 and fr < len(traj_fpv) - 1: + fr += 1 + update_all() + elif key == glfw.KEY_LEFT: + if fr > 0: + fr -= 1 + update_all() + elif key == glfw.KEY_SPACE: + paused = not paused + else: + return False + + return True + + def update_pose(): + print('take_ind: %d, traj ind: %d, tpv fr: %d, mfr int: %d' % (take_ind, traj_ind, traj_ind * cfg.fr_margin + tpv_offset, mfr_int)) + if args.multi: + fr_s = cfg.fr_margin + nq = 59 + num_model = env_vis.model.nq // nq + vec = np.array([0, -1]) + for i in range(num_model): + fr_m = max(0, min(fr_s + (i - 3) * mfr_int, traj_pred.shape[1] - 1)) + env_vis.data.qpos[i * nq: (i + 1) * nq] = traj_pred[traj_ind, fr_m, :] + env_vis.data.qpos[i * nq + 3: i * nq + 7] = de_heading(traj_pred[traj_ind, fr_m, 3:7]) + env_vis.data.qpos[i * nq: i * nq + 2] = traj_pred[traj_ind, fr_s, :2] + vec * 0.8 * i + else: + nq = env_vis.model.nq // 2 + if fr >= cfg.fr_margin: + env_vis.data.qpos[:nq] = traj_pred[traj_ind, fr, :] + env_vis.data.qpos[nq:] = traj_pred[traj_ind, cfg.fr_margin, :] + else: + env_vis.data.qpos[:nq] = traj_pred[traj_ind, fr, :] + env_vis.data.qpos[nq:] = traj_pred[traj_ind, fr, :] + env_vis.sim_forward() + + + def update_images(): + cv2.imshow('FPV', traj_fpv[traj_ind * cfg.fr_margin + min(fr, cfg.fr_margin) + 10]) + if args.tpv: + cv2.imshow('TPV', traj_tpv[traj_ind * cfg.fr_margin + min(fr, cfg.fr_margin) + tpv_offset]) + + + def update_all(): + update_pose() + update_images() + + + def load_take(load_images=True): + global traj_pred, traj_fpv, traj_tpv, tpv_offset + algo_res, algo = algos[algo_ind] + if algo_res is None: + return + take = takes[take_ind] + traj_pred = algo_res['traj_pred'][take] + traj_ub = meta['traj_ub'].get(take, traj_pred.shape[0]) + traj_pred = traj_pred[:traj_ub] + if load_images: + frame_folder = 'datasets/fpv_frames/%s' % take + frame_files = glob.glob(os.path.join(frame_folder, '*.png')) + frame_files.sort() + traj_fpv = [cv2.imread(file) for file in frame_files] + if args.tpv: + frame_folder = 'datasets/tpv/s_frames/%s' % take + frame_files = glob.glob(os.path.join(frame_folder, '*.jpg')) + frame_files.sort() + traj_tpv = [cv2.imread(file) for file in frame_files] + tpv_offset = meta['tpv_offset'][take] + print(len(traj_fpv), len(traj_tpv) if args.tpv else 0, traj_pred.shape[0]) + print('%s ---------- %s' % (algo, take)) + + traj_pred = None + traj_fpv = None + traj_tpv = None + vis_model_file = 'assets/mujoco_models/%s.xml' % (args.multi_vis_model if args.multi else args.vis_model) + env_vis = HumanoidVisEnv(vis_model_file, 1) + env_vis.set_custom_key_callback(key_callback) + takes = list(ef_res['traj_pred'].keys()) + algos = [(ef_res, 'ego forecast')] + algo_ind = 0 + take_ind = args.take_ind + traj_ind = 0 + tpv_offset = 0 + ss_ind = 0 + mfr_int = 10 + load_take() + + """render or select part of the clip""" + cv2.namedWindow('FPV') + cv2.moveWindow('FPV', 150, 400) + cv2.namedWindow('TPV') + glfw.set_window_size(env_vis.viewer.window, 1000, 960) + glfw.set_window_pos(env_vis.viewer.window, 500, 0) + env_vis.viewer._hide_overlay = True + env_vis.viewer.cam.distance = 6 + T = 10 + fr = 0 + paused = False + stop = False + reverse = False + + update_all() + t = 0 + while not stop: + if t >= math.floor(T): + if not reverse and fr < traj_pred.shape[1] - 1: + fr += 1 + update_all() + elif reverse and fr > 0: + fr -= 1 + update_all() + t = 0 + + heading = get_heading(traj_pred[traj_ind, 0, 3:7]) + flip = meta['tpv_flip'].get(takes[take_ind], False) + env_vis.viewer.cam.azimuth = math.degrees(heading) + (0 if flip else 180) + env_vis.render() + + if not paused: + t += 1 + diff --git a/ego_pose/eval_pose.py b/ego_pose/eval_pose.py new file mode 100644 index 0000000..99741b6 --- /dev/null +++ b/ego_pose/eval_pose.py @@ -0,0 +1,217 @@ +import argparse +import os +import sys +import pickle +import math +import time +import numpy as np +sys.path.append(os.getcwd()) + +from ego_pose.utils.metrics import * +from ego_pose.utils.tools import * +from ego_pose.utils.egomimic_config import Config +from ego_pose.envs.humanoid_v1 import HumanoidEnv +from envs.visual.humanoid_vis import HumanoidVisEnv + + +parser = argparse.ArgumentParser() +parser.add_argument('--vis-model', default='humanoid_1205_vis_double_v1') +parser.add_argument('--multi-vis-model', default='humanoid_1205_vis_estimate_v1') +parser.add_argument('--multi', action='store_true', default=False) +parser.add_argument('--egomimic-cfg', default='subject_03') +parser.add_argument('--statereg-cfg', default='subject_03') +parser.add_argument('--egomimic-iter', type=int, default=3000) +parser.add_argument('--statereg-iter', type=int, default=100) +parser.add_argument('--egomimic-tag', default='') +parser.add_argument('--data', default='test') +parser.add_argument('--mode', default='vis') +args = parser.parse_args() + + +def compute_metrics(results, meta, algo): + if results is None: + return + + print('=' * 10 + ' %s ' % algo + '=' * 10) + g_pose_dist = 0 + g_vel_dist = 0 + g_smoothness = 0 + traj_orig = results['traj_orig'] + traj_pred = results['traj_pred'] + + for take in traj_pred.keys(): + traj = traj_pred[take] + traj_gt = traj_orig[take] + # compute gt stats + angs_gt = get_joint_angles(traj_gt) + vels_gt = get_joint_vels(traj_gt, dt) + # compute pred stats + angs = get_joint_angles(traj) + vels = get_joint_vels(traj, dt) + accels = get_joint_accels(vels, dt) + # compute metrics + pose_dist = get_mean_dist(angs, angs_gt) + vel_dist = get_mean_dist(vels, vels_gt) + smoothness = get_mean_abs(accels) + g_pose_dist += pose_dist + g_vel_dist += vel_dist + g_smoothness += smoothness + + # print('%s - pose dist: %.4f, vel dist: %.4f, accels: %.4f' % + # (take, pose_dist, vel_dist, smoothness)) + + g_pose_dist /= len(traj_pred) + g_vel_dist /= len(traj_pred) + g_smoothness /= len(traj_pred) + + print('-' * 60) + print('all - pose dist: %.4f, vel dist: %.4f, accels: %.4f' % (g_pose_dist, g_vel_dist, g_smoothness)) + print('-' * 60 + '\n') + + +cfg = Config(args.egomimic_cfg) +dt = 1 / 30.0 + +res_base_dir = 'results' +em_res_path = '%s/egomimic/%s/results/iter_%04d_%s%s.p' % (res_base_dir, args.egomimic_cfg, args.egomimic_iter, args.data, args.egomimic_tag) +sr_res_path = '%s/statereg/%s/results/iter_%04d_%s.p' % (res_base_dir, args.statereg_cfg, args.statereg_iter, args.data) +em_res, em_meta = pickle.load(open(em_res_path, 'rb')) if args.egomimic_cfg is not None else (None, None) +sr_res, sr_meta = pickle.load(open(sr_res_path, 'rb')) if args.statereg_cfg is not None else (None, None) +# some mocap data we captured have very noisy hands, +# we set hand angles to 0 during training and we do not use it in evaluation +remove_noisy_hands(em_res) +remove_noisy_hands(sr_res) + +if args.mode == 'stats': + compute_metrics(em_res, em_meta, 'ego mimic') + compute_metrics(sr_res, sr_meta, 'state reg') + +elif args.mode == 'vis': + """visualization""" + + def key_callback(key, action, mods): + global T, fr, paused, stop, reverse, algo_ind, take_ind, ss_ind, show_gt, mfr_int + + if action != glfw.RELEASE: + return False + if key == glfw.KEY_D: + T *= 1.5 + elif key == glfw.KEY_F: + T = max(1, T / 1.5) + elif key == glfw.KEY_R: + stop = True + elif key == glfw.KEY_W: + fr = 0 + update_pose() + elif key == glfw.KEY_S: + reverse = not reverse + elif key == glfw.KEY_Z: + fr = 0 + take_ind = (take_ind - 1) % len(takes) + load_take() + update_pose() + elif key == glfw.KEY_C: + fr = 0 + take_ind = (take_ind + 1) % len(takes) + load_take() + update_pose() + elif key == glfw.KEY_X: + save_screen_shots(env_vis.viewer.window, 'out/%04d.png' % ss_ind) + ss_ind += 1 + elif glfw.KEY_1 <= key < glfw.KEY_1 + len(algos): + algo_ind = key - glfw.KEY_1 + load_take() + update_pose() + elif key == glfw.KEY_0: + show_gt = not show_gt + update_pose() + elif key == glfw.KEY_MINUS: + mfr_int -= 1 + update_pose() + elif key == glfw.KEY_EQUAL: + mfr_int += 1 + update_pose() + elif key == glfw.KEY_RIGHT: + if fr < traj_orig.shape[0] - 1: + fr += 1 + update_pose() + elif key == glfw.KEY_LEFT: + if fr > 0: + fr -= 1 + update_pose() + elif key == glfw.KEY_SPACE: + paused = not paused + else: + return False + + return True + + def update_pose(): + print('take_ind: %d, fr: %d, mfr int: %d' % (take_ind, fr, mfr_int)) + if args.multi: + nq = 59 + traj = traj_orig if show_gt else traj_pred + num_model = env_vis.model.nq // nq + hq = get_heading_q(traj_orig[fr, 3:7]) + rel_q = quaternion_multiply(hq, quaternion_inverse(get_heading_q(traj[fr, 3:7]))) + vec = quat_mul_vec(hq, np.array([0, -1, 0]))[:2] + for i in range(num_model): + fr_m = min(fr + i * mfr_int, traj.shape[0] - 1) + env_vis.data.qpos[i*nq: (i+1)*nq] = traj[fr_m, :] + env_vis.data.qpos[i*nq + 3: i*nq + 7] = quaternion_multiply(rel_q, traj[fr_m, 3:7]) + env_vis.data.qpos[i*nq: i*nq + 2] = traj_orig[fr, :2] + vec * 0.8 * i + else: + nq = env_vis.model.nq // 2 + env_vis.data.qpos[:nq] = traj_pred[fr, :] + env_vis.data.qpos[nq:] = traj_orig[fr, :] + # add x offset + env_vis.data.qpos[nq] += 1.0 + env_vis.sim_forward() + + def load_take(): + global traj_pred, traj_orig + algo_res, algo = algos[algo_ind] + if algo_res is None: + return + take = takes[take_ind] + print('%s ---------- %s' % (algo, take)) + traj_pred = algo_res['traj_pred'][take] + traj_orig = algo_res['traj_orig'][take] + + traj_pred = None + traj_orig = None + vis_model_file = 'assets/mujoco_models/%s.xml' % (args.multi_vis_model if args.multi else args.vis_model) + env_vis = HumanoidVisEnv(vis_model_file, 1, focus=not args.multi) + env_vis.set_custom_key_callback(key_callback) + takes = cfg.takes[args.data] + algos = [(em_res, 'ego mimic'), (sr_res, 'state reg')] + algo_ind = 0 + take_ind = 0 + ss_ind = 0 + mfr_int = 10 + show_gt = False + load_take() + + """render or select part of the clip""" + T = 10 + fr = 0 + paused = False + stop = False + reverse = False + + update_pose() + t = 0 + while not stop: + if t >= math.floor(T): + if not reverse and fr < traj_orig.shape[0] - 1: + fr += 1 + update_pose() + elif reverse and fr > 0: + fr -= 1 + update_pose() + t = 0 + + env_vis.render() + if not paused: + t += 1 + diff --git a/ego_pose/eval_pose_wild.py b/ego_pose/eval_pose_wild.py new file mode 100644 index 0000000..3270c08 --- /dev/null +++ b/ego_pose/eval_pose_wild.py @@ -0,0 +1,273 @@ +import argparse +import os +import sys +import pickle +import math +import time +import glob +import cv2 +import yaml +import numpy as np +sys.path.append(os.getcwd()) + +from ego_pose.utils.metrics import * +from ego_pose.utils.egomimic_config import Config +from ego_pose.utils.pose2d import Pose2DContext +from envs.visual.humanoid_vis import HumanoidVisEnv + + +parser = argparse.ArgumentParser() +parser.add_argument('--vis-model', default='humanoid_1205_vis_single_v1') +parser.add_argument('--multi-vis-model', default='humanoid_1205_vis_estimate_v1') +parser.add_argument('--multi', action='store_true', default=False) +parser.add_argument('--egomimic-cfg', default='cross_01') +parser.add_argument('--statereg-cfg', default='cross_01') +parser.add_argument('--egomimic-iter', type=int, default=6000) +parser.add_argument('--statereg-iter', type=int, default=100) +parser.add_argument('--data', default='wild_01') +parser.add_argument('--take-ind', type=int, default=0) +parser.add_argument('--mode', default='vis') +parser.add_argument('--tpv', action='store_true', default=True) +parser.add_argument('--stats-vis', action='store_true', default=False) +args = parser.parse_args() + + +cfg = Config(args.egomimic_cfg, False) +dt = 1 / 30.0 +data_dir = 'datasets' +meta = yaml.load(open('%s/meta/meta_%s.yml' % (data_dir, args.data))) + +res_base_dir = 'results' +em_res_path = '%s/egomimic/%s/results/iter_%04d_%s.p' % (res_base_dir, args.egomimic_cfg, args.egomimic_iter, args.data) +sr_res_path = '%s/statereg/%s/results/iter_%04d_%s.p' % (res_base_dir, args.statereg_cfg, args.statereg_iter, args.data) +em_res, em_meta = pickle.load(open(em_res_path, 'rb')) if args.egomimic_cfg is not None else (None, None) +sr_res, sr_meta = pickle.load(open(sr_res_path, 'rb')) if args.statereg_cfg is not None else (None, None) +takes = list(em_res['traj_pred'].keys()) + +if args.mode == 'stats': + pose_ctx = Pose2DContext(cfg) + + def eval_take(res, take): + pose_dist = 0 + traj_pred = res['traj_pred'][take] + traj_ub = meta['traj_ub'].get(take, traj_pred.shape[0]) + traj_pred = traj_pred[:traj_ub] + tpv_offset = meta['tpv_offset'].get(take, cfg.fr_margin) + flip = meta['tpv_flip'].get(take, False) + fr_num = traj_pred.shape[0] + valid_num = 0 + for fr in range(max(0, -tpv_offset), fr_num): + gt_fr = fr + tpv_offset + gt_file = '%s/tpv/poses/%s/%05d_keypoints.json' % (data_dir, take, gt_fr) + gt_p = pose_ctx.load_gt_pose(gt_file) + if not pose_ctx.check_gt(gt_p): + # print('invalid frame: %s, %d' % (take, fr)) + continue + valid_num += 1 + qpos = traj_pred[fr, :] + p = pose_ctx.align_qpos(qpos, gt_p, flip=flip) + dist = pose_ctx.get_pose_dist(p, gt_p) + pose_dist += dist + if args.stats_vis: + img = cv2.imread('%s/tpv/s_frames/%s/%05d.jpg' % (data_dir, take, gt_fr)) + pose_ctx.draw_pose(img, p * 0.25, flip=flip) + cv2.imshow('', img) + cv2.waitKey(1) + pose_dist /= valid_num + vels = get_joint_vels(traj_pred, dt) + accels = get_joint_accels(vels, dt) + smoothness = get_mean_abs(accels) + return pose_dist, smoothness + + def compute_metrics(res, algo): + if res is None: + return + + print('=' * 10 + ' %s ' % algo + '=' * 10) + g_pose_dist = 0 + g_smoothness = 0 + for take in takes: + pose_dist, smoothness = eval_take(res, take) + g_pose_dist += pose_dist + g_smoothness += smoothness + print('%s - pose dist: %.4f, accels: %.4f' % (take, pose_dist, smoothness)) + g_pose_dist /= len(takes) + g_smoothness /= len(takes) + print('-' * 60) + print('all - pose dist: %.4f, accels: %.4f' % (g_pose_dist, g_smoothness)) + print('-' * 60 + '\n') + + compute_metrics(em_res, 'ego mimic') + compute_metrics(sr_res, 'state reg') + +elif args.mode == 'vis': + + def key_callback(key, action, mods): + global T, fr, paused, stop, reverse, algo_ind, take_ind, tpv_offset, mfr_int + + if action != glfw.RELEASE: + return False + if key == glfw.KEY_D: + T *= 1.5 + elif key == glfw.KEY_F: + T = max(1, T / 1.5) + elif key == glfw.KEY_R: + stop = True + elif key == glfw.KEY_W: + fr = 0 + update_all() + elif key == glfw.KEY_S: + reverse = not reverse + elif key == glfw.KEY_Q: + tpv_offset -= 1 + update_images() + print('tpv offset: %d' % tpv_offset) + elif key == glfw.KEY_E: + tpv_offset += 1 + update_images() + print('tpv offset: %d' % tpv_offset) + elif key == glfw.KEY_Z: + fr = 0 + take_ind = (take_ind - 1) % len(takes) + load_take() + update_all() + elif key == glfw.KEY_C: + fr = 0 + take_ind = (take_ind + 1) % len(takes) + load_take() + update_all() + elif key == glfw.KEY_V: + env_vis.focus = not env_vis.focus + elif glfw.KEY_1 <= key < glfw.KEY_1 + len(algos): + algo_ind = key - glfw.KEY_1 + load_take(False) + update_all() + elif key == glfw.KEY_MINUS: + mfr_int -= 1 + update_pose() + elif key == glfw.KEY_EQUAL: + mfr_int += 1 + update_pose() + elif key == glfw.KEY_RIGHT: + if fr < traj_pred.shape[0] - 1 and fr < len(traj_fpv) - 1: + fr += 1 + update_all() + elif key == glfw.KEY_LEFT: + if fr > 0: + fr -= 1 + update_all() + elif key == glfw.KEY_SPACE: + paused = not paused + else: + return False + + return True + + + def update_pose(): + print('take_ind: %d, fr: %d, tpv fr: %d, mfr int: %d' % (take_ind, fr, fr + tpv_offset, mfr_int)) + if args.multi: + nq = 59 + num_model = env_vis.model.nq // nq + vec = np.array([0, -1]) + for i in range(num_model): + fr_m = min(fr + i * mfr_int, traj_pred.shape[0] - 1) + env_vis.data.qpos[i * nq: (i + 1) * nq] = traj_pred[fr_m, :] + env_vis.data.qpos[i * nq + 3: i * nq + 7] = de_heading(traj_pred[fr_m, 3:7]) + env_vis.data.qpos[i * nq: i * nq + 2] = traj_pred[fr, :2] + vec * 0.8 * i + else: + env_vis.data.qpos[:] = traj_pred[fr, :] + env_vis.sim_forward() + + + def update_images(): + cv2.imshow('FPV', traj_fpv[fr]) + if args.tpv: + cv2.imshow('TPV', traj_tpv[fr + tpv_offset]) + + + def update_all(): + update_pose() + update_images() + + + def load_take(load_images=True): + global traj_pred, traj_fpv, traj_tpv, tpv_offset, h_arr, flip + algo_res, algo = algos[algo_ind] + if algo_res is None: + return + take = takes[take_ind] + traj_pred = algo_res['traj_pred'][take] + traj_ub = meta['traj_ub'].get(take, traj_pred.shape[0]) + traj_pred = traj_pred[:traj_ub] + h_arr = [] + for fr in range(0, traj_pred.shape[0]): + heading = get_heading(traj_pred[fr, 3:7]) + if heading > np.pi: + heading -= 2 * np.pi + elif heading < -np.pi: + heading += 2 * np.pi + h_arr.append(heading) + flip = meta['tpv_flip'].get(take, False) + + if load_images: + frame_folder = 'datasets/fpv_frames/%s' % take + frame_files = glob.glob(os.path.join(frame_folder, '*.png')) + frame_files.sort() + traj_fpv = [cv2.imread(file) for file in frame_files[cfg.fr_margin:-cfg.fr_margin]] + if args.tpv: + frame_folder = 'datasets/tpv/s_frames/%s' % take + frame_files = glob.glob(os.path.join(frame_folder, '*.jpg')) + frame_files.sort() + traj_tpv = [cv2.imread(file) for file in frame_files] + tpv_offset = meta['tpv_offset'].get(take, cfg.fr_margin) + print(len(traj_fpv), len(traj_tpv) if args.tpv else 0, traj_pred.shape[0]) + print('%s ---------- %s' % (algo, take)) + + + traj_pred = None + traj_fpv = None + traj_tpv = None + h_arr = [] + flip = False + vis_model_file = 'assets/mujoco_models/%s.xml' % (args.multi_vis_model if args.multi else args.vis_model) + env_vis = HumanoidVisEnv(vis_model_file, 1) + env_vis.set_custom_key_callback(key_callback) + algos = [(em_res, 'ego mimic'), (sr_res, 'state reg')] + algo_ind = 0 + take_ind = args.take_ind + tpv_offset = 0 + mfr_int = 30 + load_take() + + """render or select part of the clip""" + cv2.namedWindow('FPV') + cv2.moveWindow('FPV', 150, 400) + cv2.namedWindow('TPV') + glfw.set_window_size(env_vis.viewer.window, 1000, 960) + glfw.set_window_pos(env_vis.viewer.window, 500, 0) + env_vis.viewer._hide_overlay = True + T = 10 + fr = 0 + paused = False + stop = False + reverse = False + + update_all() + t = 0 + while not stop: + if t >= math.floor(T): + if not reverse and fr < traj_pred.shape[0] - 1: + fr += 1 + update_all() + elif reverse and fr > 0: + fr -= 1 + update_all() + t = 0 + + env_vis.viewer.cam.azimuth = math.degrees(sum(h_arr[fr - 30:fr]) / 30 + if fr > 30 else sum(h_arr[:30]) / 30) + (0 if flip else 180) + env_vis.render() + if not paused: + t += 1 + diff --git a/ego_pose/state_reg.py b/ego_pose/state_reg.py new file mode 100644 index 0000000..9fef73e --- /dev/null +++ b/ego_pose/state_reg.py @@ -0,0 +1,184 @@ +import argparse +import os +import sys +import pickle +import time +import torch +import numpy as np +sys.path.append(os.getcwd()) + +from utils import * +from models.video_reg_net import * +from ego_pose.utils.statereg_dataset import Dataset +from ego_pose.utils.statereg_config import Config + + +parser = argparse.ArgumentParser() +parser.add_argument('--cfg', default=None) +parser.add_argument('--mode', default='train') +parser.add_argument('--data', default=None) +parser.add_argument('--test-feat', default=None) +parser.add_argument('--gpu-index', type=int, default=0) +parser.add_argument('--iter', type=int, default=0) +args = parser.parse_args() +if args.data is None: + args.data = args.mode if args.mode in {'train', 'test'} else 'train' + +cfg = Config(args.cfg, create_dirs=(args.iter == 0)) + +"""setup""" +dtype = torch.float64 +torch.set_default_dtype(dtype) +device = torch.device('cuda', index=args.gpu_index) if torch.cuda.is_available() else torch.device('cpu') +if torch.cuda.is_available(): + torch.cuda.set_device(args.gpu_index) +np.random.seed(cfg.seed) +torch.manual_seed(cfg.seed) +tb_logger = Logger(cfg.tb_dir) +logger = create_logger(os.path.join(cfg.log_dir, 'log.txt')) + +dataset = Dataset(cfg.meta_id, args.data, cfg.fr_num, cfg.iter_method, cfg.shuffle, 2*cfg.fr_margin, cfg.num_sample) + +"""networks""" +# if only predicting pose, then the out dim will be n_pose + 6 (include root velocities) +state_dim = (dataset.traj_dim - 1) // 2 + 6 if cfg.pose_only else dataset.traj_dim +no_cnn = (args.mode == 'save_inf' or args.test_feat is not None) +state_net = VideoRegNet(state_dim, cfg.v_hdim, cfg.cnn_fdim, no_cnn=no_cnn, cnn_type=cfg.cnn_type, + mlp_dim=cfg.mlp_dim, v_net_type=cfg.v_net, v_net_param=cfg.v_net_param, causal=cfg.causal) +if args.iter > 0: + cp_path = '%s/iter_%04d.p' % (cfg.model_dir, args.iter) + logger.info('loading model from checkpoint: %s' % cp_path) + model_cp, meta = pickle.load(open(cp_path, "rb")) + if args.data != 'train': + dataset.set_mean_std(meta['mean'], meta['std']) + state_net.load_state_dict(model_cp['state_net_dict'], strict=not no_cnn) +if args.mode != 'save_inf': + state_net.to(device) +optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, state_net.parameters()), lr=cfg.lr) +fr_margin = cfg.fr_margin + +if args.mode == 'train': + to_train(state_net) + for i_epoch in range(args.iter, cfg.num_epoch): + t0 = time.time() + epoch_num_sample = 0 + epoch_loss = 0 + for of_np, traj_np, _ in dataset: + num = traj_np.shape[0] - 2 * fr_margin + of = tensor(of_np, dtype=dtype, device=device) + of = torch.cat((of, zeros(of.shape[:-1] + (1,), device=device)), dim=-1).permute(0, 3, 1, 2).unsqueeze(1).contiguous() + state_gt = tensor(traj_np[fr_margin: -fr_margin, :state_dim], dtype=dtype, device=device) + state_pred = state_net(of)[fr_margin: -fr_margin, :] + # compute loss + loss = (state_gt - state_pred).pow(2).sum(dim=1).mean() + optimizer.zero_grad() + loss.backward() + optimizer.step() + # logging + epoch_loss += loss.cpu().item() * num + epoch_num_sample += num + + """clean up gpu memory""" + del of + torch.cuda.empty_cache() + + epoch_loss /= epoch_num_sample + logger.info('epoch {:4d} time {:.2f} nsample {} loss {:.4f} ' + .format(i_epoch, time.time() - t0, epoch_num_sample, epoch_loss)) + tb_logger.scalar_summary('loss', epoch_loss, i_epoch) + + with to_cpu(state_net): + if cfg.save_model_interval > 0 and (i_epoch + 1) % cfg.save_model_interval == 0: + cp_path = '%s/iter_%04d.p' % (cfg.model_dir, i_epoch + 1) + model_cp = {'state_net_dict': state_net.state_dict()} + meta = {'mean': dataset.mean, 'std': dataset.std} + pickle.dump((model_cp, meta), open(cp_path, 'wb')) + +elif args.mode == 'test': + + def get_traj_from_state_pred(state_pred, init_pos, init_heading, dt): + nv = (dataset.traj_dim + 1) // 2 + nq = nv + 1 + pos = init_pos.copy() + heading = init_heading.copy() + traj_pred = [] + for i in range(state_pred.shape[0]): + qpos = np.concatenate((pos, state_pred[i, :nq-2])) + qvel = state_pred[i, nq-2:] + qpos[3:7] = quaternion_multiply(heading, qpos[3:7]) + linv = quat_mul_vec(heading, qvel[:3]) + angv = quat_mul_vec(qpos[3:7], qvel[3:6]) + pos += linv[:2] * dt + new_q = quaternion_multiply(quat_from_expmap(angv * dt), qpos[3:7]) + heading = get_heading_q(new_q) + traj_pred.append(qpos) + traj_qpos = np.vstack(traj_pred) + return traj_qpos + + to_test(state_net) + dataset.iter_method = 'iter' + dataset.shuffle = False + torch.set_grad_enabled(False) + epoch_num_sample = 0 + epoch_loss = 0 + state_pred_arr = [] + traj_orig_arr = [] + res_pred = {} + res_orig = {} + meta = {} + if args.test_feat is None: + take = dataset.takes[0] + for of_np, traj_np, traj_orig_np in dataset: + num = traj_np.shape[0] - 2 * fr_margin + of = tensor(of_np, dtype=dtype, device=device) + of = torch.cat((of, zeros(of.shape[:-1] + (1,), device=device)), dim=-1).permute(0, 3, 1, 2).unsqueeze(1).contiguous() + state_gt = tensor(traj_np[fr_margin: -fr_margin, :state_dim], dtype=dtype, device=device) + state_pred = state_net(of)[fr_margin: -fr_margin, :] + # logging + loss = (state_gt - state_pred).pow(2).sum(dim=1).mean() + state_pred = state_pred.cpu().numpy() * dataset.std[None, :state_dim] + dataset.mean[None, :state_dim] + traj_orig = traj_orig_np[fr_margin: -fr_margin, :] + state_pred_arr.append(state_pred) + traj_orig_arr.append(traj_orig) + epoch_loss += loss.cpu().item() * num + epoch_num_sample += num + # save if datset moved on to next take + if dataset.cur_ind >= len(dataset.takes) or dataset.takes[dataset.cur_tid] != take: + state_pred = np.vstack(state_pred_arr) + traj_orig = np.vstack(traj_orig_arr) + init_pos = traj_orig[0, :2] + init_heading = get_heading_q(traj_orig[0, 3:7]) + traj_pred = get_traj_from_state_pred(state_pred, init_pos, init_heading, dataset.dt) + res_pred[take] = traj_pred + res_orig[take] = np.vstack(traj_orig_arr) + state_pred_arr, traj_orig_arr = [], [] + take = dataset.takes[dataset.cur_tid] + epoch_loss /= epoch_num_sample + results = {'traj_pred': res_pred, 'traj_orig': res_orig} + meta['algo'] = 'state_reg' + meta['num_sample'] = epoch_num_sample + meta['epoch_loss'] = epoch_loss + res_path = '%s/iter_%04d_%s.p' % (cfg.result_dir, args.iter, args.data) + else: + cnn_feat_file = '%s/features/cnn_feat_%s.p' % (dataset.base_folder, args.test_feat) + cnn_feat_dict, _ = pickle.load(open(cnn_feat_file, 'rb')) + for take, cnn_feat in cnn_feat_dict.items(): + cnn_feat = tensor(cnn_feat, dtype=dtype, device=device) + state_pred = state_net(cnn_feat.unsqueeze(1)).squeeze(1)[cfg.fr_margin: -cfg.fr_margin, :].cpu().numpy() + state_pred = state_pred * dataset.std[None, :state_dim] + dataset.mean[None, :state_dim] + traj_pred = get_traj_from_state_pred(state_pred, np.zeros(2), np.array([1, 0, 0, 0]), dataset.dt) + res_pred[take] = traj_pred + epoch_num_sample += state_pred.shape[0] + results = {'traj_pred': res_pred} + meta['algo'] = 'state_reg' + meta['num_sample'] = epoch_num_sample + res_path = '%s/iter_%04d_%s.p' % (cfg.result_dir, args.iter, args.test_feat) + pickle.dump((results, meta), open(res_path, 'wb')) + logger.info('nsample {} loss {:.4f} '.format(epoch_num_sample, epoch_loss)) + logger.info('saved results to %s' % res_path) + +elif args.mode == 'save_inf': + cp_path = '%s/iter_%04d_inf.p' % (cfg.model_dir, args.iter) + model_cp = {'state_net_dict': state_net.state_dict()} + meta = {'mean': dataset.mean, 'std': dataset.std, 'cfg': cfg} + pickle.dump((model_cp, meta), open(cp_path, 'wb')) diff --git a/ego_pose/utils/egoforecast_config.py b/ego_pose/utils/egoforecast_config.py new file mode 100644 index 0000000..a44c444 --- /dev/null +++ b/ego_pose/utils/egoforecast_config.py @@ -0,0 +1,138 @@ +import yaml +import os +import numpy as np +from utils import recreate_dirs + + +class Config: + + def __init__(self, cfg_id, create_dirs=False): + self.id = cfg_id + cfg_name = 'config/egoforecast/%s.yml' % cfg_id + if not os.path.exists(cfg_name): + print("Config file doesn't exist: %s" % cfg_name) + exit(0) + cfg = yaml.load(open(cfg_name, 'r')) + + # create dirs + self.base_dir = 'results' + self.cfg_dir = '%s/egoforecast/%s' % (self.base_dir, cfg_id) + self.model_dir = '%s/models' % self.cfg_dir + self.result_dir = '%s/results' % self.cfg_dir + self.log_dir = '%s/log' % self.cfg_dir + self.tb_dir = '%s/tb' % self.cfg_dir + os.makedirs(self.model_dir, exist_ok=True) + os.makedirs(self.result_dir, exist_ok=True) + if create_dirs: + recreate_dirs(self.log_dir, self.tb_dir) + + # data + self.meta_id = cfg['meta_id'] + self.data_dir = 'datasets' + self.meta = yaml.load(open('%s/meta/%s.yml' % (self.data_dir, self.meta_id), 'r')) + self.takes = {x: self.meta[x] for x in ['train', 'test']} + self.expert_feat_file = '%s/features/expert_%s.p' % (self.data_dir, cfg['expert_feat']) if 'expert_feat' in cfg else None + self.cnn_feat_file = '%s/features/cnn_feat_%s.p' % (self.data_dir, cfg['cnn_feat']) if 'cnn_feat' in cfg else None + self.fr_margin = cfg.get('fr_margin', 10) + + # ego mimic + self.ego_mimic_cfg = cfg.get('ego_mimic_cfg', None) + self.ego_mimic_iter = cfg.get('ego_mimic_iter', None) + + # training config + self.gamma = cfg.get('gamma', 0.95) + self.tau = cfg.get('tau', 0.95) + self.policy_htype = cfg.get('policy_htype', 'relu') + self.policy_hsize = cfg.get('policy_hsize', [300, 200]) + self.policy_v_hdim = cfg.get('policy_v_hdim', 128) + self.policy_v_net = cfg.get('policy_v_net', 'lstm') + self.policy_v_net_param = cfg.get('policy_v_net_param', None) + self.policy_s_net = cfg.get('policy_s_net', 'id') + self.policy_s_hdim = cfg.get('policy_s_hdim', None) + self.policy_dyn_v = cfg.get('policy_dyn_v', False) + self.policy_optimizer = cfg.get('policy_optimizer', 'Adam') + self.policy_lr = cfg.get('policy_lr', 5e-5) + self.policy_momentum = cfg.get('policy_momentum', 0.0) + self.policy_weightdecay = cfg.get('policy_weightdecay', 0.0) + self.value_htype = cfg.get('value_htype', 'relu') + self.value_hsize = cfg.get('value_hsize', [300, 200]) + self.value_v_hdim = cfg.get('value_v_hdim', 128) + self.value_v_net = cfg.get('value_v_net', 'lstm') + self.value_v_net_param = cfg.get('value_v_net_param', None) + self.value_s_net = cfg.get('value_s_net', 'id') + self.value_s_hdim = cfg.get('value_s_hdim', None) + self.value_dyn_v = cfg.get('value_dyn_v', False) + self.value_optimizer = cfg.get('value_optimizer', 'Adam') + self.value_lr = cfg.get('value_lr', 3e-4) + self.value_momentum = cfg.get('value_momentum', 0.0) + self.value_weightdecay = cfg.get('value_weightdecay', 0.0) + self.adv_clip = cfg.get('adv_clip', np.inf) + self.clip_epsilon = cfg.get('clip_epsilon', 0.2) + self.log_std = cfg.get('log_std', -2.3) + self.fix_std = cfg.get('fix_std', False) + self.num_optim_epoch = cfg.get('num_optim_epoch', 10) + self.min_batch_size = cfg.get('min_batch_size', 50000) + self.max_iter_num = cfg.get('max_iter_num', 1000) + self.seed = cfg.get('seed', 1) + self.save_model_interval = cfg.get('save_model_interval', 100) + self.reward_id = cfg.get('reward_id', 'quat') + self.reward_weights = cfg.get('reward_weights', None) + self.end_reward = cfg.get('end_reward', True) + + # adaptive parameters + self.adp_iter_cp = np.array(cfg.get('adp_iter_cp', [0])) + self.adp_noise_rate_cp = np.array(cfg.get('adp_noise_rate_cp', [1.0])) + self.adp_noise_rate_cp = np.pad(self.adp_noise_rate_cp, (0, self.adp_iter_cp.size - self.adp_noise_rate_cp.size), 'edge') + self.adp_log_std_cp = np.array(cfg.get('adp_log_std_cp', [self.log_std])) + self.adp_log_std_cp = np.pad(self.adp_log_std_cp, (0, self.adp_iter_cp.size - self.adp_log_std_cp.size), 'edge') + self.adp_policy_lr_cp = np.array(cfg.get('adp_policy_lr_cp', [self.policy_lr])) + self.adp_policy_lr_cp = np.pad(self.adp_policy_lr_cp, (0, self.adp_iter_cp.size - self.adp_policy_lr_cp.size), 'edge') + self.adp_init_noise_cp = np.array(cfg.get('adp_init_noise_cp', [0.0])) + self.adp_init_noise_cp = np.pad(self.adp_init_noise_cp, (0, self.adp_iter_cp.size - self.adp_init_noise_cp.size), 'edge') + self.adp_noise_rate = None + self.adp_log_std = None + self.adp_policy_lr = None + self.adp_init_noise = None + + # env config + self.mujoco_model_file = '%s/assets/mujoco_models/%s.xml' % (os.getcwd(), cfg['mujoco_model']) + self.vis_model_file = '%s/assets/mujoco_models/%s.xml' % (os.getcwd(), cfg['vis_model']) + self.env_start_first = cfg.get('env_start_first', False) + self.env_init_noise = cfg.get('env_init_noise', 0.0) + self.env_episode_len = cfg.get('env_episode_len', 200) + self.obs_type = cfg.get('obs_type', 'full') + self.obs_coord = cfg.get('obs_coord', 'heading') + self.obs_heading = cfg.get('obs_heading', False) + self.obs_vel = cfg.get('obs_vel', 'full') + self.obs_phase = cfg.get('obs_phase', False) + self.random_cur_t = cfg.get('random_cur_t', False) + self.root_deheading = cfg.get('root_deheading', True) + self.sync_exp_interval = cfg.get('sync_exp_interval', 100) + self.action_type = cfg.get('action_type', 'position') + + # joint param + if 'joint_params' in cfg: + jparam = zip(*cfg['joint_params']) + jparam = [np.array(p) for p in jparam] + self.jkp, self.jkd, self.a_ref, self.a_scale, self.torque_lim = jparam[1:6] + self.a_ref = np.deg2rad(self.a_ref) + jkp_multiplier = cfg.get('jkp_multiplier', 1.0) + jkd_multiplier = cfg.get('jkd_multiplier', jkp_multiplier) + self.jkp *= jkp_multiplier + self.jkd *= jkd_multiplier + + # body param + if 'body_params' in cfg: + bparam = zip(*cfg['body_params']) + bparam = [np.array(p) for p in bparam] + self.b_diffw = bparam[1] + + def update_adaptive_params(self, i_iter): + cp = self.adp_iter_cp + ind = np.where(i_iter >= cp)[0][-1] + nind = ind + int(ind < len(cp) - 1) + t = (i_iter - self.adp_iter_cp[ind]) / (cp[nind] - cp[ind]) if nind > ind else 0.0 + self.adp_noise_rate = self.adp_noise_rate_cp[ind] * (1-t) + self.adp_noise_rate_cp[nind] * t + self.adp_log_std = self.adp_log_std_cp[ind] * (1-t) + self.adp_log_std_cp[nind] * t + self.adp_policy_lr = self.adp_policy_lr_cp[ind] * (1-t) + self.adp_policy_lr_cp[nind] * t + self.adp_init_noise = self.adp_init_noise_cp[ind] * (1-t) + self.adp_init_noise_cp[nind] * t diff --git a/ego_pose/utils/egomimic_config.py b/ego_pose/utils/egomimic_config.py new file mode 100644 index 0000000..b21acba --- /dev/null +++ b/ego_pose/utils/egomimic_config.py @@ -0,0 +1,131 @@ +import yaml +import os +import numpy as np +from utils import recreate_dirs + + +class Config: + + def __init__(self, cfg_id=None, create_dirs=False, cfg_dict=None): + self.id = cfg_id + if cfg_dict is not None: + cfg = cfg_dict + else: + cfg_name = 'config/egomimic/%s.yml' % cfg_id + if not os.path.exists(cfg_name): + print("Config file doesn't exist: %s" % cfg_name) + exit(0) + cfg = yaml.load(open(cfg_name, 'r')) + + # create dirs + self.base_dir = 'results' + self.cfg_dir = '%s/egomimic/%s' % (self.base_dir, cfg_id) + self.model_dir = '%s/models' % self.cfg_dir + self.result_dir = '%s/results' % self.cfg_dir + self.log_dir = '%s/log' % self.cfg_dir + self.tb_dir = '%s/tb' % self.cfg_dir + os.makedirs(self.model_dir, exist_ok=True) + os.makedirs(self.result_dir, exist_ok=True) + if create_dirs: + recreate_dirs(self.log_dir, self.tb_dir) + + # data + self.meta_id = cfg['meta_id'] + self.data_dir = 'datasets' + self.meta = yaml.load(open('%s/meta/%s.yml' % (self.data_dir, self.meta_id), 'r')) + self.takes = {x: self.meta[x] for x in ['train', 'test']} + self.expert_feat_file = '%s/features/expert_%s.p' % (self.data_dir, cfg['expert_feat']) if 'expert_feat' in cfg else None + self.cnn_feat_file = '%s/features/cnn_feat_%s.p' % (self.data_dir, cfg['cnn_feat']) if 'cnn_feat' in cfg else None + self.fr_margin = cfg.get('fr_margin', 10) + + # state net + self.state_net_cfg = cfg.get('state_net_cfg', None) + self.state_net_iter = cfg.get('state_net_iter', None) + if self.state_net_cfg is not None: + self.state_net_model = '%s/statereg/%s/models/iter_%04d_inf.p' % (self.base_dir, self.state_net_cfg, self.state_net_iter) + + # training config + self.gamma = cfg.get('gamma', 0.95) + self.tau = cfg.get('tau', 0.95) + self.causal = cfg.get('causal', False) + self.policy_htype = cfg.get('policy_htype', 'relu') + self.policy_hsize = cfg.get('policy_hsize', [300, 200]) + self.policy_v_hdim = cfg.get('policy_v_hdim', 128) + self.policy_v_net = cfg.get('policy_v_net', 'lstm') + self.policy_v_net_param = cfg.get('policy_v_net_param', None) + self.policy_optimizer = cfg.get('policy_optimizer', 'Adam') + self.policy_lr = cfg.get('policy_lr', 5e-5) + self.policy_momentum = cfg.get('policy_momentum', 0.0) + self.policy_weightdecay = cfg.get('policy_weightdecay', 0.0) + self.value_htype = cfg.get('value_htype', 'relu') + self.value_hsize = cfg.get('value_hsize', [300, 200]) + self.value_v_hdim = cfg.get('value_v_hdim', 128) + self.value_v_net = cfg.get('value_v_net', 'lstm') + self.value_v_net_param = cfg.get('value_v_net_param', None) + self.value_optimizer = cfg.get('value_optimizer', 'Adam') + self.value_lr = cfg.get('value_lr', 3e-4) + self.value_momentum = cfg.get('value_momentum', 0.0) + self.value_weightdecay = cfg.get('value_weightdecay', 0.0) + self.adv_clip = cfg.get('adv_clip', np.inf) + self.clip_epsilon = cfg.get('clip_epsilon', 0.2) + self.log_std = cfg.get('log_std', -2.3) + self.fix_std = cfg.get('fix_std', False) + self.num_optim_epoch = cfg.get('num_optim_epoch', 10) + self.min_batch_size = cfg.get('min_batch_size', 50000) + self.max_iter_num = cfg.get('max_iter_num', 1000) + self.seed = cfg.get('seed', 1) + self.save_model_interval = cfg.get('save_model_interval', 100) + self.reward_id = cfg.get('reward_id', 'quat') + self.reward_weights = cfg.get('reward_weights', None) + + # adaptive parameters + self.adp_iter_cp = np.array(cfg.get('adp_iter_cp', [0])) + self.adp_noise_rate_cp = np.array(cfg.get('adp_noise_rate_cp', [1.0])) + self.adp_noise_rate_cp = np.pad(self.adp_noise_rate_cp, (0, self.adp_iter_cp.size - self.adp_noise_rate_cp.size), 'edge') + self.adp_log_std_cp = np.array(cfg.get('adp_log_std_cp', [self.log_std])) + self.adp_log_std_cp = np.pad(self.adp_log_std_cp, (0, self.adp_iter_cp.size - self.adp_log_std_cp.size), 'edge') + self.adp_policy_lr_cp = np.array(cfg.get('adp_policy_lr_cp', [self.policy_lr])) + self.adp_policy_lr_cp = np.pad(self.adp_policy_lr_cp, (0, self.adp_iter_cp.size - self.adp_policy_lr_cp.size), 'edge') + self.adp_noise_rate = None + self.adp_log_std = None + self.adp_policy_lr = None + + # env config + self.mujoco_model_file = '%s/assets/mujoco_models/%s.xml' % (os.getcwd(), cfg['mujoco_model']) + self.vis_model_file = '%s/assets/mujoco_models/%s.xml' % (os.getcwd(), cfg['vis_model']) + self.env_start_first = cfg.get('env_start_first', False) + self.env_init_noise = cfg.get('env_init_noise', 0.0) + self.env_episode_len = cfg.get('env_episode_len', 200) + self.obs_type = cfg.get('obs_type', 'full') + self.obs_coord = cfg.get('obs_coord', 'heading') + self.obs_heading = cfg.get('obs_heading', False) + self.obs_vel = cfg.get('obs_vel', 'full') + self.root_deheading = cfg.get('root_deheading', True) + self.sync_exp_interval = cfg.get('sync_exp_interval', 100) + self.action_type = cfg.get('action_type', 'position') + + # joint param + if 'joint_params' in cfg: + jparam = zip(*cfg['joint_params']) + jparam = [np.array(p) for p in jparam] + self.jkp, self.jkd, self.a_ref, self.a_scale, self.torque_lim = jparam[1:6] + self.a_ref = np.deg2rad(self.a_ref) + jkp_multiplier = cfg.get('jkp_multiplier', 1.0) + jkd_multiplier = cfg.get('jkd_multiplier', jkp_multiplier) + self.jkp *= jkp_multiplier + self.jkd *= jkd_multiplier + + # body param + if 'body_params' in cfg: + bparam = zip(*cfg['body_params']) + bparam = [np.array(p) for p in bparam] + self.b_diffw = bparam[1] + + def update_adaptive_params(self, i_iter): + cp = self.adp_iter_cp + ind = np.where(i_iter >= cp)[0][-1] + nind = ind + int(ind < len(cp) - 1) + t = (i_iter - self.adp_iter_cp[ind]) / (cp[nind] - cp[ind]) if nind > ind else 0.0 + self.adp_noise_rate = self.adp_noise_rate_cp[ind] * (1-t) + self.adp_noise_rate_cp[nind] * t + self.adp_log_std = self.adp_log_std_cp[ind] * (1-t) + self.adp_log_std_cp[nind] * t + self.adp_policy_lr = self.adp_policy_lr_cp[ind] * (1-t) + self.adp_policy_lr_cp[nind] * t diff --git a/ego_pose/utils/metrics.py b/ego_pose/utils/metrics.py new file mode 100644 index 0000000..f516c0b --- /dev/null +++ b/ego_pose/utils/metrics.py @@ -0,0 +1,38 @@ +from utils import * +from utils.transformation import euler_from_quaternion + + +def get_joint_angles(poses): + root_angs = [] + for pose in poses: + root_euler = np.array(euler_from_quaternion(pose[3:7])) + root_euler[2] = 0.0 + root_angs.append(root_euler) + root_angs = np.vstack(root_angs) + angles = np.hstack((root_angs, poses[:, 7:])) + return angles + + +def get_joint_vels(poses, dt): + vels = [] + for i in range(poses.shape[0] - 1): + v = get_qvel_fd(poses[i], poses[i+1], dt, 'heading') + vels.append(v) + vels = np.vstack(vels) + return vels + + +def get_joint_accels(vels, dt): + accels = np.diff(vels, axis=0) / dt + accels = np.vstack(accels) + return accels + + +def get_mean_dist(x, y): + return np.linalg.norm(x - y, axis=1).mean() + + +def get_mean_abs(x): + return np.abs(x).mean() + + diff --git a/ego_pose/utils/pose2d.py b/ego_pose/utils/pose2d.py new file mode 100644 index 0000000..0e41718 --- /dev/null +++ b/ego_pose/utils/pose2d.py @@ -0,0 +1,149 @@ +import json +import numpy as np +import math +import cv2 +from ego_pose.envs.humanoid_v1 import HumanoidEnv + + +class Pose2DContext: + + def __init__(self, cfg): + self.env = HumanoidEnv(cfg) + self.body2id = self.env.model._body_name2id + self.body_names = self.env.model.body_names[1:] + self.body_set = {'LeftForeArm', 'RightForeArm', 'LeftHand', 'RightHand', 'LeftArm', 'RightArm', + 'LeftUpLeg', 'RightUpLeg', 'LeftLeg', 'RightLeg', 'LeftFoot', 'RightFoot'} + self.nbody = len(self.body_set) + self.body_filter = np.zeros((len(self.body_names),), dtype=bool) + for body in self.body2id.keys(): + if body in self.body_set: + self.body_filter[self.body2id[body]-1] = True + self.body_names = [self.body_names[i] for i in range(len(self.body_filter)) if self.body_filter[i]] + self.body2id = {body: i for i, body in enumerate(self.body_names)} + + self.conn = [('RightUpLeg', 'RightArm', (255, 255, 0)), + ('RightArm', 'RightForeArm', (255, 191, 0)), + ('RightForeArm', 'RightHand', (255, 191, 0)), + ('RightUpLeg', 'RightLeg', (255, 64, 0.0)), + ('RightLeg', 'RightFoot', (255, 64, 0.0)), + ('LeftUpLeg', 'LeftArm', (0, 255, 128)), + ('LeftArm', 'LeftForeArm', (0, 255, 255)), + ('LeftForeArm', 'LeftHand', (0, 255, 255)), + ('LeftUpLeg', 'LeftLeg', (0, 64, 255)), + ('LeftLeg', 'LeftFoot', (0, 64, 255))] + + self.joints_map = [(2, self.body2id['RightArm']), + (3, self.body2id['RightForeArm']), + (4, self.body2id['RightHand']), + (5, self.body2id['LeftArm']), + (6, self.body2id['LeftForeArm']), + (7, self.body2id['LeftHand']), + (9, self.body2id['RightUpLeg']), + (10, self.body2id['RightLeg']), + (11, self.body2id['RightFoot']), + (12, self.body2id['LeftUpLeg']), + (13, self.body2id['LeftLeg']), + (14, self.body2id['LeftFoot'])] + + def draw_pose(self, img, pose, flip=False): + conn = self.conn + if flip: + conn = self.conn[5:] + self.conn[:5] + for b1, b2, c in conn: + p1 = pose[self.body2id[b1], :2] + p2 = pose[self.body2id[b2], :2] + self.draw_bone(img, p1, p2, c) + for x in self.body_set: + e = pose[self.body2id[x], :2] + cv2.circle(img, (int(e[0]), int(e[1])), 1, (0, 0, 255), -1) + + def draw_bone(self, img, p1, p2, c): + center = (int((p1[0]+p2[0])/2), int((p1[1]+p2[1])/2)) + angle = int(math.atan2(p2[1]-p1[1], p2[0]-p1[0]) / np.pi * 180) + axes = (int(np.linalg.norm(p2-p1)/2), 1) + cv2.ellipse(img, center, axes, angle, 0, 360, c, -1) + # cv2.line(img, (int(p[i1, 0]), int(p[i1, 1])), (int(p[i2, 0]), int(p[i2, 1])), (c[2], c[1], c[0]), thickness=5) + + def load_gt_pose(self, filename): + data = json.load(open(filename)) + keypoints = data['people'][0]['pose_keypoints_2d'] + p = np.zeros((self.nbody, 3)) + for i1, i2 in self.joints_map: + p[i2, :] = keypoints[3*i1: 3*i1 + 3] + return p + + def check_gt(self, gt_pose): + return gt_pose[self.body2id['LeftUpLeg'], 2] > 0.1 or gt_pose[self.body2id['RightUpLeg'], 2] > 0.1 + + def get_pose_dist(self, p, gt_p): + body2id = self.body2id + if gt_p[body2id['LeftArm'], 2] > 0.1 and gt_p[body2id['LeftUpLeg'], 2] > 0.1: + kp1 = 'LeftArm' + kp2 = 'LeftUpLeg' + else: + kp1 = 'RightArm' + kp2 = 'RightUpLeg' + scale = 0.5 / abs(gt_p[body2id[kp1], 1] - gt_p[body2id[kp2], 1]) + + dist = 0 + num = 0 + for i in range(gt_p.shape[0]): + if gt_p[i, 2] > 0.1: + dist += np.linalg.norm(gt_p[i, :2] - p[i, :]) * scale + num += 1 + dist /= num + return dist + + def project_qpos(self, qpos, flip): + self.env.data.qpos[:] = qpos + self.env.sim.forward() + pose_3d = np.vstack(self.env.data.body_xpos[1:]) + pose_3d = pose_3d[self.body_filter, :] + body2id = self.body2id + + """make projection matrix""" + vp = (pose_3d[body2id['LeftUpLeg'], :] + pose_3d[body2id['RightUpLeg'], :]) * 0.5 + v = pose_3d[body2id['RightUpLeg'], :] - pose_3d[body2id['LeftUpLeg'], :] + if flip: + v *= -1 + v[2] = 0 + v /= np.linalg.norm(v) + x = v + z = np.array([0, 0, 1]) + y = np.cross(z, x) + # R, t transfrom camera coordinate to world coordiante + R = np.hstack((-y[:, None], z[:, None], x[:, None])) + t = vp - 10 * x + t = t[:, None] + E = np.hstack((R.T, -R.T.dot(t))) + + p = np.hstack((pose_3d, np.ones((pose_3d.shape[0], 1)))).dot(E.T) + p = p[:, :2] / p[:, [2]] + p[:, 1] *= -1 + return p + + def align_qpos(self, qpos, gt_p, scale=None, flip=False): + body2id = self.body2id + p = self.project_qpos(qpos, flip) + base = np.zeros((1, 2)) + n = 0 + if gt_p[body2id['LeftUpLeg'], 2] > 0.1: + base += gt_p[[body2id['LeftUpLeg']], :2] + n += 1 + if gt_p[body2id['RightUpLeg'], 2] > 0.1: + base += gt_p[[body2id['RightUpLeg']], :2] + n += 1 + base /= n + + if scale is None: + if gt_p[body2id['LeftLeg'], 2] > 0.1 and gt_p[body2id['LeftUpLeg'], 2] > 0.1: + kp1 = 'LeftLeg' + kp2 = 'LeftUpLeg' + else: + kp1 = 'RightLeg' + kp2 = 'RightUpLeg' + scale = np.linalg.norm(gt_p[body2id[kp1]] - gt_p[body2id[kp2]]) / np.linalg.norm(p[body2id[kp1]] - p[body2id[kp2]]) + + p = p * scale + base + return p + diff --git a/ego_pose/utils/statereg_config.py b/ego_pose/utils/statereg_config.py new file mode 100644 index 0000000..e5fdee5 --- /dev/null +++ b/ego_pose/utils/statereg_config.py @@ -0,0 +1,50 @@ +import yaml +import os +from utils import recreate_dirs + + +class Config: + + def __init__(self, cfg_id, create_dirs=False): + self.id = cfg_id + cfg_name = 'config/statereg/%s.yml' % cfg_id + if not os.path.exists(cfg_name): + print("Config file doesn't exist: %s" % cfg_name) + exit(0) + cfg = yaml.load(open(cfg_name, 'r')) + + # create dirs + self.base_dir = 'results' + self.cfg_dir = '%s/statereg/%s' % (self.base_dir, cfg_id) + self.model_dir = '%s/models' % self.cfg_dir + self.result_dir = '%s/results' % self.cfg_dir + self.log_dir = '%s/log' % self.cfg_dir + self.tb_dir = '%s/tb' % self.cfg_dir + os.makedirs(self.model_dir, exist_ok=True) + os.makedirs(self.result_dir, exist_ok=True) + if create_dirs: + recreate_dirs(self.log_dir, self.tb_dir) + + # training config + self.meta_id = cfg['meta_id'] + self.seed = cfg['seed'] + self.fr_num = cfg['fr_num'] + self.v_net = cfg.get('v_net', 'lstm') + self.v_net_param = cfg.get('v_net_param', None) + self.v_hdim = cfg['v_hdim'] + self.mlp_dim = cfg['mlp_dim'] + self.cnn_fdim = cfg['cnn_fdim'] + self.lr = cfg['lr'] + self.num_epoch = cfg['num_epoch'] + self.iter_method = cfg['iter_method'] + self.shuffle = cfg.get('shuffle', False) + self.num_sample = cfg.get('num_sample', 20000) + self.save_model_interval = cfg['save_model_interval'] + self.fr_margin = cfg['fr_margin'] + self.pose_only = cfg.get('pose_only', False) + self.causal = cfg.get('causal', False) + self.cnn_type = cfg.get('cnn_type', 'resnet') + + # misc config + self.humanoid_model = cfg['humanoid_model'] + self.vis_model = cfg['vis_model'] diff --git a/ego_pose/utils/statereg_dataset.py b/ego_pose/utils/statereg_dataset.py new file mode 100644 index 0000000..530ce0c --- /dev/null +++ b/ego_pose/utils/statereg_dataset.py @@ -0,0 +1,163 @@ +import numpy as np +import os +import yaml +import math +from utils import get_qvel_fd, de_heading + + +class Dataset: + + def __init__(self, meta_id, mode, fr_num, iter_method='iter', shuffle=False, overlap=0, num_sample=20000): + self.meta_id = meta_id + self.mode = mode + self.fr_num = fr_num + self.iter_method = iter_method + self.shuffle = shuffle + self.overlap = overlap + self.num_sample = num_sample + self.base_folder = 'datasets' + self.of_folder = os.path.join(self.base_folder, 'fpv_of') + self.traj_folder = os.path.join(self.base_folder, 'traj') + meta_file = '%s/meta/%s.yml' % (self.base_folder, meta_id) + self.meta = yaml.load(open(meta_file, 'r')) + self.no_traj = self.meta.get('no_traj', False) + self.msync = self.meta['video_mocap_sync'] + self.dt = 1 / self.meta['capture']['fps'] + # get take names + if mode == 'all': + self.takes = self.meta['train'] + self.meta['test'] + else: + self.takes = self.meta[mode] + # get dataset len + self.len = np.sum([self.msync[x][2] - self.msync[x][1] for x in self.takes]) + # preload trajectories + if self.no_traj: + self.trajs = None + self.orig_trajs = None + self.norm_trajs = None + else: + self.trajs = [] + self.orig_trajs = [] + for i, take in enumerate(self.takes): + traj_file = '%s/%s_traj.p' % (self.traj_folder, take) + orig_traj = np.load(traj_file) + # remove noisy hand pose + orig_traj[:, 32:35] = 0.0 + orig_traj[:, 42:45] = 0.0 + traj_pos = self.get_traj_pos(orig_traj) + traj_vel = self.get_traj_vel(orig_traj) + traj = np.hstack((traj_pos, traj_vel)) + self.trajs.append(traj) + self.orig_trajs.append(orig_traj) + if mode == 'train': + all_traj = np.vstack(self.trajs) + self.mean = np.mean(all_traj, axis=0) + self.std = np.std(all_traj, axis=0) + self.norm_trajs = self.normalize_traj() + else: + self.mean, self.std, self.norm_trajs = None, None, None + self.traj_dim = self.trajs[0].shape[1] + # iterator specific + self.sample_count = None + self.take_indices = None + self.cur_ind = None + self.cur_tid = None + self.cur_fr = None + self.fr_lb = None + self.fr_ub = None + self.im_offset = None + + def __iter__(self): + if self.iter_method == 'sample': + self.sample_count = 0 + elif self.iter_method == 'iter': + self.cur_ind = -1 + self.take_indices = np.arange(len(self.takes)) + if self.shuffle: + np.random.shuffle(self.take_indices) + self.__next_take() + return self + + def __next_take(self): + self.cur_ind = self.cur_ind + 1 + if self.cur_ind < len(self.take_indices): + self.cur_tid = self.take_indices[self.cur_ind] + self.im_offset, self.fr_lb, self.fr_ub = self.msync[self.takes[self.cur_tid]] + self.cur_fr = self.fr_lb + + def __next__(self): + if self.iter_method == 'sample': + if self.sample_count >= self.num_sample: + raise StopIteration + self.sample_count += self.fr_num - self.overlap + return self.sample() + elif self.iter_method == 'iter': + if self.cur_ind >= len(self.takes): + raise StopIteration + fr_start = self.cur_fr + fr_end = self.cur_fr + self.fr_num if self.cur_fr + self.fr_num + 30 < self.fr_ub else self.fr_ub + # print(self.cur_ind, self.cur_tid, fr_start, fr_end) + of = self.load_of(self.cur_tid, fr_start + self.im_offset, fr_end + self.im_offset) + if self.no_traj: + norm_traj, orig_traj = None, None + else: + norm_traj = self.norm_trajs[self.cur_tid][fr_start: fr_end] + orig_traj = self.orig_trajs[self.cur_tid][fr_start: fr_end] + self.cur_fr = fr_end - self.overlap + if fr_end == self.fr_ub: + self.__next_take() + return of, norm_traj, orig_traj + + def get_traj_pos(self, orig_traj): + traj_pos = orig_traj[:, 2:].copy() + for i in range(traj_pos.shape[0]): + traj_pos[i, 1:5] = de_heading(traj_pos[i, 1:5]) + return traj_pos + + def get_traj_vel(self, orig_traj): + traj_vel = [] + for i in range(orig_traj.shape[0] - 1): + vel = get_qvel_fd(orig_traj[i, :], orig_traj[i + 1, :], self.dt, 'heading') + traj_vel.append(vel) + traj_vel.append(traj_vel[-1].copy()) + traj_vel = np.vstack(traj_vel) + return traj_vel + + def set_mean_std(self, mean, std): + self.mean, self.std = mean, std + if not self.no_traj: + self.norm_trajs = self.normalize_traj() + + def normalize_traj(self): + norm_trajs = [] + for traj in self.trajs: + norm_traj = (traj - self.mean[None, :]) / (self.std[None, :] + 1e-8) + norm_trajs.append(norm_traj) + return norm_trajs + + def sample(self): + take_ind = np.random.randint(len(self.takes)) + im_offset, fr_lb, fr_ub = self.msync[self.takes[take_ind]] + fr_start = np.random.randint(fr_lb, fr_ub - self.fr_num) + fr_end = fr_start + self.fr_num + of = self.load_of(take_ind, fr_start + im_offset, fr_end + im_offset) + if self.no_traj: + norm_traj, orig_traj = None, None + else: + norm_traj = self.norm_trajs[take_ind][fr_start: fr_end] + orig_traj = self.orig_trajs[take_ind][fr_start: fr_end] + return of, norm_traj, orig_traj + + def load_of(self, take_ind, start, end): + take_folder = '%s/%s' % (self.of_folder, self.takes[take_ind]) + of = [] + for i in range(start, end): + of_file = '%s/%05d.npy' % (take_folder, i) + of_i = np.load(of_file) + of.append(of_i) + of = np.stack(of) + return of + + + + diff --git a/ego_pose/utils/tools.py b/ego_pose/utils/tools.py new file mode 100644 index 0000000..af4f438 --- /dev/null +++ b/ego_pose/utils/tools.py @@ -0,0 +1,40 @@ +import numpy as np +from utils.math import * + + +def normalize_traj(qpos_traj, qvel_traj): + new_qpos_traj = [] + new_qvel_traj = [] + for qpos, qvel in zip(qpos_traj, qvel_traj): + new_qpos = qpos.copy() + new_qvel = qvel.copy() + new_qvel[:3] = transform_vec(qvel[:3], qpos[3:7], 'heading') + new_qpos[3:7] = de_heading(qpos[3:7]) + new_qpos_traj.append(new_qpos) + new_qvel_traj.append(new_qvel) + return np.vstack(new_qpos_traj), np.vstack(new_qvel_traj) + + +def sync_traj(qpos_traj, qvel_traj, ref_qpos): + new_qpos_traj = [] + new_qvel_traj = [] + rel_heading = quaternion_multiply(get_heading_q(ref_qpos[3:7]), quaternion_inverse(get_heading_q(qpos_traj[0, 3:7]))) + ref_pos = ref_qpos[:3] + start_pos = np.concatenate((qpos_traj[0, :2], ref_pos[[2]])) + for qpos, qvel in zip(qpos_traj, qvel_traj): + new_qpos = qpos.copy() + new_qvel = qvel.copy() + new_qpos[:2] = quat_mul_vec(rel_heading, qpos[:3] - start_pos)[:2] + ref_pos[:2] + new_qpos[3:7] = quaternion_multiply(rel_heading, qpos[3:7]) + new_qvel[:3] = quat_mul_vec(rel_heading, qvel[:3]) + new_qpos_traj.append(new_qpos) + new_qvel_traj.append(new_qvel) + return np.vstack(new_qpos_traj), np.vstack(new_qvel_traj) + + +def remove_noisy_hands(results): + for traj in results.values(): + for take in traj.keys(): + traj[take][..., 32:35] = 0 + traj[take][..., 42:45] = 0 + return diff --git a/envs/common/mjviewer.py b/envs/common/mjviewer.py new file mode 100644 index 0000000..3a78aee --- /dev/null +++ b/envs/common/mjviewer.py @@ -0,0 +1,384 @@ +from threading import Lock +import glfw +from mujoco_py.builder import cymj +from mujoco_py.generated import const +import time +import copy +from multiprocessing import Process, Queue +from mujoco_py.utils import rec_copy, rec_assign +import numpy as np +import imageio + + +class MjViewerBasic(cymj.MjRenderContextWindow): + """ + A simple display GUI showing the scene of an :class:`.MjSim` with a mouse-movable camera. + :class:`.MjViewer` extends this class to provide more sophisticated playback and interaction controls. + Parameters + ---------- + sim : :class:`.MjSim` + The simulator to display. + """ + + def __init__(self, sim): + super().__init__(sim) + + self._gui_lock = Lock() + self._button_left_pressed = False + self._button_right_pressed = False + self._last_mouse_x = 0 + self._last_mouse_y = 0 + + framebuffer_width, _ = glfw.get_framebuffer_size(self.window) + window_width, _ = glfw.get_window_size(self.window) + self._scale = framebuffer_width * 1.0 / window_width + + glfw.set_cursor_pos_callback(self.window, self._cursor_pos_callback) + glfw.set_mouse_button_callback( + self.window, self._mouse_button_callback) + glfw.set_scroll_callback(self.window, self._scroll_callback) + glfw.set_key_callback(self.window, self.key_callback) + + def render(self): + """ + Render the current simulation state to the screen or off-screen buffer. + Call this in your main loop. + """ + if self.window is None: + return + elif glfw.window_should_close(self.window): + exit(0) + + with self._gui_lock: + super().render() + + glfw.poll_events() + + def key_callback(self, window, key, scancode, action, mods): + if action == glfw.RELEASE and key == glfw.KEY_ESCAPE: + print("Pressed ESC") + print("Quitting.") + exit(0) + + def _cursor_pos_callback(self, window, xpos, ypos): + if not (self._button_left_pressed or self._button_right_pressed): + return + + # Determine whether to move, zoom or rotate view + mod_shift = ( + glfw.get_key(window, glfw.KEY_LEFT_SHIFT) == glfw.PRESS or + glfw.get_key(window, glfw.KEY_RIGHT_SHIFT) == glfw.PRESS) + if self._button_right_pressed: + action = const.MOUSE_MOVE_H if mod_shift else const.MOUSE_MOVE_V + elif self._button_left_pressed: + action = const.MOUSE_ROTATE_H if mod_shift else const.MOUSE_ROTATE_V + else: + action = const.MOUSE_ZOOM + + # Determine + dx = int(self._scale * xpos) - self._last_mouse_x + dy = int(self._scale * ypos) - self._last_mouse_y + width, height = glfw.get_framebuffer_size(window) + + with self._gui_lock: + self.move_camera(action, dx / height, dy / height) + + self._last_mouse_x = int(self._scale * xpos) + self._last_mouse_y = int(self._scale * ypos) + + def _mouse_button_callback(self, window, button, act, mods): + self._button_left_pressed = ( + glfw.get_mouse_button(window, glfw.MOUSE_BUTTON_LEFT) == glfw.PRESS) + self._button_right_pressed = ( + glfw.get_mouse_button(window, glfw.MOUSE_BUTTON_RIGHT) == glfw.PRESS) + + x, y = glfw.get_cursor_pos(window) + self._last_mouse_x = int(self._scale * x) + self._last_mouse_y = int(self._scale * y) + + def _scroll_callback(self, window, x_offset, y_offset): + with self._gui_lock: + self.move_camera(const.MOUSE_ZOOM, 0, -0.05 * y_offset) + + +class MjViewer(MjViewerBasic): + """ + Extends :class:`.MjViewerBasic` to add video recording, interactive time and interaction controls. + The key bindings are as follows: + - TAB: Switch between MuJoCo cameras. + - H: Toggle hiding all GUI components. + - SPACE: Pause/unpause the simulation. + - RIGHT: Advance simulation by one step. + - V: Start/stop video recording. + - T: Capture screenshot. + - I: Drop into ``ipdb`` debugger. + - S/F: Decrease/Increase simulation playback speed. + - C: Toggle visualization of contact forces (off by default). + - D: Enable/disable frame skipping when rendering lags behind real time. + - R: Toggle transparency of geoms. + - M: Toggle display of mocap bodies. + - 0-4: Toggle display of geomgroups + Parameters + ---------- + sim : :class:`.MjSim` + The simulator to display. + """ + + def __init__(self, sim): + super().__init__(sim) + + self._ncam = sim.model.ncam + self._paused = False # is viewer paused. + # should we advance viewer just by one step. + self._advance_by_one_step = False + + # Vars for recording video + self._record_video = False + self._video_queue = Queue() + self._video_idx = 0 + self._video_path = "/tmp/video_%07d.mp4" + + # vars for capturing screen + self._image_idx = 0 + self._image_path = "/tmp/frame_%07d.png" + + # run_speed = x1, means running real time, x2 means fast-forward times + # two. + self._run_speed = 1.0 + self._loop_count = 0 + self._render_every_frame = False + + self._show_mocap = True # Show / hide mocap bodies. + self._transparent = False # Make everything transparent. + + # this variable is estamated as a running average. + self._time_per_render = 1 / 60.0 + self._hide_overlay = False # hide the entire overlay. + self._user_overlay = {} + + self.video_fps = 30 + self.frame_skip = 1 + self.sim_time = 0 + self.custom_key_callback = None + + def render(self): + """ + Render the current simulation state to the screen or off-screen buffer. + Call this in your main loop. + """ + + def render_inner_loop(self): + render_start = time.time() + + self._overlay.clear() + if not self._hide_overlay: + for k, v in self._user_overlay.items(): + self._overlay[k] = copy.deepcopy(v) + self._create_full_overlay() + + super().render() + if self._record_video: + frame = self._read_pixels_as_in_window() + self._video_queue.put(frame) + else: + self._time_per_render = 0.9 * self._time_per_render + \ + 0.1 * (time.time() - render_start) + + self._user_overlay = copy.deepcopy(self._overlay) + # Render the same frame if paused. + if self._paused: + while self._paused: + render_inner_loop(self) + if self._advance_by_one_step: + self._advance_by_one_step = False + break + else: + # inner_loop runs "_loop_count" times in expectation (where "_loop_count" is a float). + # Therefore, frames are displayed in the real-time. + self._loop_count += (self.sim.model.opt.timestep * self.frame_skip - self.sim_time) / \ + (self._time_per_render * self._run_speed) + if self._render_every_frame: + self._loop_count = 1 + while self._loop_count > 0: + render_inner_loop(self) + self._loop_count -= 1 + # Markers and overlay are regenerated in every pass. + self._markers[:] = [] + self._overlay.clear() + + def _read_pixels_as_in_window(self): + # Reads pixels with markers and overlay from the same camera as screen. + resolution = glfw.get_framebuffer_size( + self.sim._render_context_window.window) + + resolution = np.array(resolution) + resolution = resolution * min(1000 / np.min(resolution), 1) + resolution = resolution.astype(np.int32) + resolution -= resolution % 16 + if self.sim._render_context_offscreen is None: + self.sim.render(resolution[0], resolution[1]) + offscreen_ctx = self.sim._render_context_offscreen + window_ctx = self.sim._render_context_window + # Save markers and overlay from offscreen. + saved = [copy.deepcopy(offscreen_ctx._markers), + copy.deepcopy(offscreen_ctx._overlay), + rec_copy(offscreen_ctx.cam)] + # Copy markers and overlay from window. + offscreen_ctx._markers[:] = window_ctx._markers[:] + offscreen_ctx._overlay.clear() + offscreen_ctx._overlay.update(window_ctx._overlay) + rec_assign(offscreen_ctx.cam, rec_copy(window_ctx.cam)) + + img = self.sim.render(*resolution) + img = img[::-1, :, :] # Rendered images are upside-down. + # Restore markers and overlay to offscreen. + offscreen_ctx._markers[:] = saved[0][:] + offscreen_ctx._overlay.clear() + offscreen_ctx._overlay.update(saved[1]) + rec_assign(offscreen_ctx.cam, saved[2]) + return img + + def _create_full_overlay(self): + if self._render_every_frame: + self.add_overlay(const.GRID_TOPLEFT, "", "") + else: + self.add_overlay(const.GRID_TOPLEFT, "Run speed = %.3f x real time" % + self._run_speed, "[S]lower, [F]aster") + self.add_overlay( + const.GRID_TOPLEFT, "Ren[d]er every frame", "Off" if self._render_every_frame else "On") + self.add_overlay(const.GRID_TOPLEFT, "Switch camera (#cams = %d)" % (self._ncam + 1), + "[Tab] (camera ID = %d)" % self.cam.fixedcamid) + self.add_overlay(const.GRID_TOPLEFT, "[C]ontact forces", "Off" if self.vopt.flags[ + 10] == 1 else "On") + self.add_overlay( + const.GRID_TOPLEFT, "Referenc[e] frames", "Off" if self.vopt.frame == 1 else "On") + self.add_overlay(const.GRID_TOPLEFT, + "T[r]ansparent", "On" if self._transparent else "Off") + self.add_overlay( + const.GRID_TOPLEFT, "Display [M]ocap bodies", "On" if self._show_mocap else "Off") + if self._paused is not None: + if not self._paused: + self.add_overlay(const.GRID_TOPLEFT, "Stop", "[Space]") + else: + self.add_overlay(const.GRID_TOPLEFT, "Start", "[Space]") + self.add_overlay(const.GRID_TOPLEFT, + "Advance simulation by one step", "[right arrow]") + self.add_overlay(const.GRID_TOPLEFT, "[H]ide Menu", "") + if self._record_video: + ndots = int(7 * (time.time() % 1)) + dots = ("." * ndots) + (" " * (6 - ndots)) + self.add_overlay(const.GRID_TOPLEFT, + "Record [V]ideo (On) " + dots, "") + else: + self.add_overlay(const.GRID_TOPLEFT, "Record [V]ideo (Off) ", "") + if self._video_idx > 0: + fname = self._video_path % (self._video_idx - 1) + self.add_overlay(const.GRID_TOPLEFT, " saved as %s" % fname, "") + + self.add_overlay(const.GRID_TOPLEFT, "Cap[t]ure frame", "") + if self._image_idx > 0: + fname = self._image_path % (self._image_idx - 1) + self.add_overlay(const.GRID_TOPLEFT, " saved as %s" % fname, "") + self.add_overlay(const.GRID_TOPLEFT, "Start [i]pdb", "") + if self._record_video: + extra = " (while video is not recorded)" + else: + extra = "" + self.add_overlay(const.GRID_BOTTOMLEFT, "FPS", "%d%s" % + (1 / self._time_per_render, extra)) + self.add_overlay(const.GRID_BOTTOMLEFT, "Solver iterations", str( + self.sim.data.solver_iter + 1)) + step = round(self.sim.data.time / self.sim.model.opt.timestep) + self.add_overlay(const.GRID_BOTTOMRIGHT, "Step", str(step)) + self.add_overlay(const.GRID_TOPLEFT, "Toggle geomgroup visibility", "0-4") + + def key_callback(self, window, key, scancode, action, mods): + if self.custom_key_callback is not None: + res = self.custom_key_callback(key, action, mods) + if res: + return + + if action != glfw.RELEASE: + return + elif key == glfw.KEY_TAB: # Switches cameras. + self.cam.fixedcamid += 1 + self.cam.type = const.CAMERA_FIXED + if self.cam.fixedcamid >= self._ncam: + self.cam.fixedcamid = -1 + self.cam.type = const.CAMERA_FREE + elif key == glfw.KEY_H: # hides all overlay. + self._hide_overlay = not self._hide_overlay + elif key == glfw.KEY_SPACE and self._paused is not None: # stops simulation. + self._paused = not self._paused + # Advances simulation by one step. + elif key == glfw.KEY_RIGHT and self._paused is not None: + self._advance_by_one_step = True + self._paused = True + elif key == glfw.KEY_V or \ + (key == glfw.KEY_ESCAPE and self._record_video): # Records video. Trigers with V or if in progress by ESC. + self._record_video = not self._record_video + if self._record_video: + fps = self.video_fps # (1 / self._time_per_render) + self._video_process = Process(target=save_video, + args=(self._video_queue, self._video_path % self._video_idx, fps)) + self._video_process.start() + if not self._record_video: + self._video_queue.put(None) + self._video_process.join() + self._video_idx += 1 + elif key == glfw.KEY_T: # capture screenshot + img = self._read_pixels_as_in_window() + imageio.imwrite(self._image_path % self._image_idx, img) + self._image_idx += 1 + elif key == glfw.KEY_I: # drops in debugger. + print('You can access the simulator by self.sim') + import ipdb + ipdb.set_trace() + elif key == glfw.KEY_S: # Slows down simulation. + self._run_speed /= 2.0 + elif key == glfw.KEY_F: # Speeds up simulation. + self._run_speed *= 2.0 + elif key == glfw.KEY_C: # Displays contact forces. + vopt = self.vopt + vopt.flags[10] = vopt.flags[11] = not vopt.flags[10] + elif key == glfw.KEY_D: # turn off / turn on rendering every frame. + self._render_every_frame = not self._render_every_frame + elif key == glfw.KEY_E: + vopt = self.vopt + vopt.frame = 1 - vopt.frame + elif key == glfw.KEY_R: # makes everything little bit transparent. + self._transparent = not self._transparent + if self._transparent: + self.sim.model.geom_rgba[:, 3] /= 5.0 + else: + self.sim.model.geom_rgba[:, 3] *= 5.0 + elif key == glfw.KEY_M: # Shows / hides mocap bodies + self._show_mocap = not self._show_mocap + for body_idx1, val in enumerate(self.sim.model.body_mocapid): + if val != -1: + for geom_idx, body_idx2 in enumerate(self.sim.model.geom_bodyid): + if body_idx1 == body_idx2: + if not self._show_mocap: + # Store transparency for later to show it. + self.sim.extras[ + geom_idx] = self.sim.model.geom_rgba[geom_idx, 3] + self.sim.model.geom_rgba[geom_idx, 3] = 0 + else: + self.sim.model.geom_rgba[ + geom_idx, 3] = self.sim.extras[geom_idx] + elif key in (glfw.KEY_0, glfw.KEY_1, glfw.KEY_2, glfw.KEY_3, glfw.KEY_4): + self.vopt.geomgroup[key - glfw.KEY_0] ^= 1 + super().key_callback(window, key, scancode, action, mods) + +# Separate Process to save video. This way visualization is +# less slowed down. + + +def save_video(queue, filename, fps): + writer = imageio.get_writer(filename, fps=fps) + while True: + frame = queue.get() + if frame is None: + break + writer.append_data(frame) + writer.close() diff --git a/envs/common/mujoco_env.py b/envs/common/mujoco_env.py new file mode 100644 index 0000000..ddbcadf --- /dev/null +++ b/envs/common/mujoco_env.py @@ -0,0 +1,149 @@ +import os + +from gym import error, spaces +from gym.utils import seeding +import numpy as np +from os import path +import gym +import mujoco_py +from envs.common.mjviewer import MjViewer + +DEFAULT_SIZE = 500 + + +class MujocoEnv: + """Superclass for all MuJoCo environments. + """ + + def __init__(self, fullpath, frame_skip): + if not path.exists(fullpath): + raise IOError("File %s does not exist" % fullpath) + self.frame_skip = frame_skip + self.model = mujoco_py.load_model_from_path(fullpath) + self.sim = mujoco_py.MjSim(self.model) + self.data = self.sim.data + self.viewer = None + self._viewers = {} + self.obs_dim = None + self.action_space = None + self.observation_space = None + self.np_random = None + self.cur_t = 0 # number of steps taken + + self.metadata = { + 'render.modes': ['human', 'rgb_array'], + 'video.frames_per_second': int(np.round(1.0 / self.dt)) + } + + self.init_qpos = self.sim.data.qpos.ravel().copy() + self.init_qvel = self.sim.data.qvel.ravel().copy() + self.prev_qpos = None + self.prev_qvel = None + self.seed() + + def set_spaces(self): + observation, _reward, done, _info = self.step(np.zeros(self.model.nu)) + self.obs_dim = observation.size + bounds = self.model.actuator_ctrlrange.copy() + low = bounds[:, 0] + high = bounds[:, 1] + self.action_space = spaces.Box(low=low, high=high, dtype=np.float32) + high = np.inf * np.ones(self.obs_dim) + low = -high + self.observation_space = spaces.Box(low, high, dtype=np.float32) + + def seed(self, seed=None): + self.np_random, seed = seeding.np_random(seed) + return [seed] + + # methods to override: + # ---------------------------- + def step(self, action): + """ + Step the environment forward. + """ + raise NotImplementedError + + def reset_model(self): + """ + Reset the robot degrees of freedom (qpos and qvel). + Implement this in each subclass. + """ + raise NotImplementedError + + def viewer_setup(self, mode): + """ + This method is called when the viewer is initialized and after every reset + Optionally implement this method, if you need to tinker with camera position + and so forth. + """ + pass + + # ----------------------------- + + def reset(self): + self.sim.reset() + self.cur_t = 0 + ob = self.reset_model() + old_viewer = self.viewer + for mode, v in self._viewers.items(): + self.viewer = v + self.viewer_setup(mode) + self.viewer = old_viewer + return ob + + def set_state(self, qpos, qvel): + assert qpos.shape == (self.model.nq,) and qvel.shape == (self.model.nv,) + old_state = self.sim.get_state() + new_state = mujoco_py.MjSimState(old_state.time, qpos, qvel, + old_state.act, old_state.udd_state) + self.sim.set_state(new_state) + self.sim.forward() + + @property + def dt(self): + return self.model.opt.timestep * self.frame_skip + + def do_simulation(self, ctrl, n_frames): + self.sim.data.ctrl[:] = ctrl + for _ in range(n_frames): + self.sim.step() + + def render(self, mode='human', width=DEFAULT_SIZE, height=DEFAULT_SIZE): + if mode == 'image': + self._get_viewer(mode).render(width, height) + # window size used for old mujoco-py: + data = self._get_viewer(mode).read_pixels(width, height, depth=False) + # original image is upside-down, so flip it, and the image format is BGR for OpenCV + return data[::-1, :, [2, 1, 0]] + elif mode == 'human': + self._get_viewer(mode).render() + + def close(self): + if self.viewer is not None: + # self.viewer.finish() + self.viewer = None + self._viewers = {} + + def _get_viewer(self, mode): + self.viewer = self._viewers.get(mode) + if self.viewer is None: + if mode == 'human': + self.viewer = MjViewer(self.sim) + elif mode == 'image': + self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, 0) + self._viewers[mode] = self.viewer + self.viewer_setup(mode) + return self.viewer + + def set_custom_key_callback(self, key_func): + self._get_viewer('human').custom_key_callback = key_func + + def get_body_com(self, body_name): + return self.data.get_body_xpos(body_name) + + def state_vector(self): + return np.concatenate([ + self.sim.data.qpos.flat, + self.sim.data.qvel.flat + ]) diff --git a/envs/visual/humanoid_vis.py b/envs/visual/humanoid_vis.py new file mode 100644 index 0000000..4c9b656 --- /dev/null +++ b/envs/visual/humanoid_vis.py @@ -0,0 +1,36 @@ +import numpy as np +from envs.common import mujoco_env + + +class HumanoidVisEnv(mujoco_env.MujocoEnv): + def __init__(self, vis_model_file, nframes=6, focus=True): + mujoco_env.MujocoEnv.__init__(self, vis_model_file, nframes) + self.set_cam_first = set() + self.focus = focus + + def step(self, a): + return np.zeros((10, 1)), 0, False, dict() + + def reset_model(self): + c = 0 + self.set_state( + self.np_random.uniform(low=-c, high=c, size=self.model.nq), + self.np_random.uniform(low=-c, high=c, size=self.model.nv) + ) + return None + + def sim_forward(self): + self.sim.forward() + + def viewer_setup(self, mode): + self.viewer.cam.trackbodyid = 1 + if self.focus: + self.viewer.cam.lookat[:2] = self.data.qpos[:2] + self.viewer.cam.lookat[2] = 0.8 + if mode not in self.set_cam_first: + self.viewer.video_fps = 30 + self.viewer.frame_skip = self.frame_skip + self.viewer.cam.distance = self.model.stat.extent * 1.5 + self.viewer.cam.elevation = -10 + self.viewer.cam.azimuth = 45 + self.set_cam_first.add(mode) diff --git a/mocap/pose.py b/mocap/pose.py new file mode 100644 index 0000000..79880c7 --- /dev/null +++ b/mocap/pose.py @@ -0,0 +1,93 @@ +import numpy as np +import math +from bvh import Bvh +from utils.transformation import quaternion_slerp, quaternion_from_euler, euler_from_quaternion + + +def load_amc_file(fname, scale): + + with open(fname) as f: + content = f.readlines() + + bone_addr = dict() + poses = [] + cur_pos = None + fr = 1 + for line in content: + line_words = line.split() + cmd = line_words[0] + if cmd == str(fr): + if cur_pos: + poses.append(np.array(cur_pos)) + cur_pos = [] + fr += 1 + elif cur_pos is not None: + start_ind = len(cur_pos) + if cmd == 'root': + cur_pos += [float(word)*scale for word in line_words[1:4]] + cur_pos += [math.radians(float(word)) for word in line_words[4:]] + elif cmd == 'lfoot' or cmd == 'rfoot': + cur_pos += reversed([math.radians(float(word)) for word in line_words[1:]]) + if len(cur_pos) < 3: + cur_pos.insert(-1, 0.0) + else: + cur_pos += reversed([math.radians(float(word)) for word in line_words[1:]]) + if fr == 2: + end_ind = len(cur_pos) + bone_addr[cmd] = (start_ind, end_ind) + + if cur_pos: + poses.append(np.array(cur_pos)) + poses = np.vstack(poses) + return poses, bone_addr + + +def load_bvh_file(fname, skeleton): + with open(fname) as f: + mocap = Bvh(f.read()) + + # build bone_addr + bone_addr = dict() + start_ind = 0 + for bone in skeleton.bones: + end_ind = start_ind + len(bone.channels) + bone_addr[bone.name] = (start_ind, end_ind) + start_ind = end_ind + dof_num = start_ind + + poses = np.zeros((mocap.nframes, dof_num)) + for i in range(mocap.nframes): + for bone in skeleton.bones: + trans = np.array(mocap.frame_joint_channels(i, bone.name, bone.channels)) + if bone == skeleton.root: + trans[:3] *= skeleton.len_scale + trans[3:6] = np.deg2rad(trans[3:6]) + else: + trans = np.deg2rad(trans) + start_ind, end_ind = bone_addr[bone.name] + poses[i, start_ind:end_ind] = trans + + return poses, bone_addr + + +def lin_interp(pose1, pose2, t): + pose_t = (1 - t) * pose1 + t * pose2 + return pose_t + + +def interpolated_traj(poses, sample_t=0.030, mocap_fr=120): + N = poses.shape[0] + T = float(N-1)/mocap_fr + num = int(math.floor(T/sample_t)) + sampling_times = np.arange(num+1)*sample_t*mocap_fr + + poses_samp = [] + for t in sampling_times: + start = int(math.floor(t)) + end = min(int(math.ceil(t)), poses.shape[0] - 1) + poses_samp.append(lin_interp(poses[start, :], poses[end, :], t-math.floor(t))) + poses_samp = np.vstack(poses_samp) + + return poses_samp + + diff --git a/mocap/skeleton.py b/mocap/skeleton.py new file mode 100644 index 0000000..e026712 --- /dev/null +++ b/mocap/skeleton.py @@ -0,0 +1,256 @@ +from utils.transformation import euler_matrix +from lxml.etree import XMLParser, parse, ElementTree, Element, SubElement +import math +import re +from bvh import Bvh +import numpy as np + + +class Bone: + + def __init__(self): + # original bone info + self.id = None + self.name = None + self.orient = np.identity(3) + self.dof_index = [] + self.channels = [] # bvh only + self.lb = [] + self.ub = [] + self.parent = None + self.child = [] + + # asf specific + self.dir = np.zeros(3) + self.len = 0 + # bvh specific + self.offset = np.zeros(3) + + # inferred info + self.pos = np.zeros(3) + self.end = np.zeros(3) + + +class Skeleton: + + def __init__(self): + self.bones = [] + self.name2bone = {} + self.mass_scale = 1.0 + self.len_scale = 1.0 + self.dof_name = ['x', 'y', 'z'] + self.root = None + + def load_from_asf(self, fname, swap_axes=False): + with open(fname) as f: + content = f.readlines() + + dof_ind = {'rx': 0, 'ry': 1, 'rz': 2} + phase = 0 + lastbone = None + setting_limit = False + + for line in content: + line_words = line.split() + cmd = line_words[0] + if cmd == ':root': + phase = 0 + self.root = Bone() + self.root.id = 0 + self.root.name = 'root' + self.name2bone['root'] = self.root + self.bones.append(self.root) + continue + if cmd == ':bonedata': + phase = 1 + continue + if cmd == ':hierarchy': + phase = 2 + continue + + if phase == 0: + if cmd == 'mass': + self.mass_scale = float(line_words[1]) + elif cmd == 'length': + self.len_scale = 1 / float(line_words[1]) * 0.0254 + if phase == 1: + if cmd == 'begin': + lastbone = Bone() + self.bones.append(lastbone) + elif cmd == 'end': + setting_limit = False + lastbone = None + elif cmd == 'id': + lastbone.id = len(self.bones) - 1 + elif cmd == 'name': + lastbone.name = line_words[1] + self.name2bone[lastbone.name] = lastbone + elif cmd == 'direction': + for i in range(3): + lastbone.dir[i] = float(line_words[i+1]) + if swap_axes: + lastbone.dir[1], lastbone.dir[2] = -lastbone.dir[2], lastbone.dir[1] + elif cmd == 'length': + lastbone.len = float(line_words[1]) * self.len_scale + elif cmd == 'axis': + args = [math.radians(float(word)) for word in line_words[1:4]] + lastbone.orient = euler_matrix(*args, "sxyz")[:3, :3] + if swap_axes: + orient = lastbone.orient.copy() + lastbone.orient[1, :], lastbone.orient[2, :] = -orient[2, :], orient[1, :] + elif cmd == 'dof': + for word in reversed(line_words[1:]): + if word in dof_ind: + ind = dof_ind[word] + lastbone.dof_index.append(ind) + elif cmd == 'limits' or setting_limit: + lastbone.lb.append(float(re.sub('[(]', ' ', line_words[1 - setting_limit]))) + lastbone.ub.append(float(re.sub('[)]', ' ', line_words[2 - setting_limit]))) + setting_limit = True + + if phase == 2: + if cmd != "begin" and cmd != "end": + bone_p = self.name2bone[line_words[0]] + for child_name in line_words[1:]: + bone_c = self.name2bone[child_name] + bone_p.child.append(bone_c) + bone_c.parent = bone_p + + self.forward_asf(self.root) + + def forward_asf(self, bone): + if bone.parent: + bone.pos = bone.parent.end + bone.end = bone.pos + bone.dir * bone.len + for bone_c in bone.child: + self.forward_asf(bone_c) + + def load_from_bvh(self, fname, exclude_bones=None, spec_channels=None): + if exclude_bones is None: + exclude_bones = {} + if spec_channels is None: + spec_channels = dict() + with open(fname) as f: + mocap = Bvh(f.read()) + + joint_names = list(filter(lambda x: all([t not in x for t in exclude_bones]), mocap.get_joints_names())) + dof_ind = {'x': 0, 'y': 1, 'z': 2} + self.len_scale = 0.0254 + self.root = Bone() + self.root.id = 0 + self.root.name = joint_names[0] + self.root.channels = mocap.joint_channels(self.root.name) + self.name2bone[self.root.name] = self.root + self.bones.append(self.root) + for i, joint in enumerate(joint_names[1:]): + bone = Bone() + bone.id = i + 1 + bone.name = joint + bone.channels = spec_channels[joint] if joint in spec_channels.keys() else mocap.joint_channels(joint) + bone.dof_index = [dof_ind[x[0].lower()] for x in bone.channels] + bone.offset = np.array(mocap.joint_offset(joint)) * self.len_scale + bone.lb = [-180.0] * 3 + bone.ub = [180.0] * 3 + self.bones.append(bone) + self.name2bone[joint] = bone + + for bone in self.bones[1:]: + parent_name = mocap.joint_parent(bone.name).name + if parent_name in self.name2bone.keys(): + bone_p = self.name2bone[parent_name] + bone_p.child.append(bone) + bone.parent = bone_p + + self.forward_bvh(self.root) + for bone in self.bones: + if len(bone.child) == 0: + bone.end = bone.pos + np.array([float(x) for x in mocap.get_joint(bone.name).children[-1]['OFFSET']]) * self.len_scale + else: + bone.end = sum([bone_c.pos for bone_c in bone.child]) / len(bone.child) + + def forward_bvh(self, bone): + if bone.parent: + bone.pos = bone.parent.pos + bone.offset + else: + bone.pos = bone.offset + for bone_c in bone.child: + self.forward_bvh(bone_c) + + def write_xml(self, fname, template_fname): + parser = XMLParser(remove_blank_text=True) + tree = parse(template_fname, parser=parser) + worldbody = tree.getroot().find('worldbody') + self.write_xml_bodynode(self.root, worldbody) + + # create actuators + actuators = tree.getroot().find('actuator') + joints = worldbody.findall('.//joint') + for joint in joints[1:]: + name = joint.attrib['name'] + attr = dict() + attr['name'] = name + attr['joint'] = name + attr['gear'] = '1' + SubElement(actuators, 'motor', attr) + + tree.write(fname, pretty_print=True) + + def write_xml_bodynode(self, bone, parent_node): + attr = dict() + attr['name'] = bone.name + attr['user'] = '{0:.4f} {1:.4f} {2:.4f}'.format(*bone.end) + node = SubElement(parent_node, 'body', attr) + + # write joints + if bone.parent is None: + j_attr = dict() + j_attr['name'] = bone.name + j_attr['pos'] = '{0:.4f} {1:.4f} {2:.4f}'.format(*bone.pos) + j_attr['limited'] = 'false' + j_attr['type'] = 'free' + j_attr['armature'] = '0' + j_attr['damping'] = '0' + j_attr['stiffness'] = '0' + SubElement(node, 'joint', j_attr) + else: + for i in range(len(bone.dof_index)): + ind = bone.dof_index[i] + axis = bone.orient[:, ind] + j_attr = dict() + j_attr['name'] = bone.name + '_' + self.dof_name[ind] + j_attr['type'] = 'hinge' + j_attr['pos'] = '{0:.4f} {1:.4f} {2:.4f}'.format(*bone.pos) + j_attr['axis'] = '{0:.4f} {1:.4f} {2:.4f}'.format(*axis) + if i < len(bone.lb): + j_attr['range'] = '{0:.4f} {1:.4f}'.format(bone.lb[i], bone.ub[i]) + else: + j_attr['range'] = '-180.0 180.0' + SubElement(node, 'joint', j_attr) + + # write geometry + if bone.parent is None: + g_attr = dict() + g_attr['size'] = '0.03' + g_attr['type'] = 'sphere' + g_attr['pos'] = '{0:.4f} {1:.4f} {2:.4f}'.format(*bone.pos) + SubElement(node, 'geom', g_attr) + else: + e1 = bone.pos.copy() + e2 = bone.end.copy() + v = e2 - e1 + if np.linalg.norm(v) > 1e-6: + v /= np.linalg.norm(v) + else: + v = np.array([0.0, 0.0, 0.2]) + e1 += v * 0.02 + e2 -= v * 0.02 + g_attr = dict() + g_attr['size'] = '0.03' + g_attr['type'] = 'capsule' + g_attr['fromto'] = '{0:.4f} {1:.4f} {2:.4f} {3:.4f} {4:.4f} {5:.4f}'.format(*np.concatenate([e1, e2])) + SubElement(node, 'geom', g_attr) + + # write child bones + for bone_c in bone.child: + self.write_xml_bodynode(bone_c, node) + diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000..0e4c909 --- /dev/null +++ b/models/__init__.py @@ -0,0 +1,4 @@ +from models.mlp import MLP +from models.rnn import RNN +from models.tcn import TemporalConvNet +from models.resnet import ResNet diff --git a/models/mlp.py b/models/mlp.py new file mode 100644 index 0000000..ef79eb6 --- /dev/null +++ b/models/mlp.py @@ -0,0 +1,25 @@ +import torch.nn as nn +import torch + + +class MLP(nn.Module): + def __init__(self, input_dim, hidden_dims=(128, 128), activation='tanh'): + super().__init__() + if activation == 'tanh': + self.activation = torch.tanh + elif activation == 'relu': + self.activation = torch.relu + elif activation == 'sigmoid': + self.activation = torch.sigmoid + + self.out_dim = hidden_dims[-1] + self.affine_layers = nn.ModuleList() + last_dim = input_dim + for nh in hidden_dims: + self.affine_layers.append(nn.Linear(last_dim, nh)) + last_dim = nh + + def forward(self, x): + for affine in self.affine_layers: + x = self.activation(affine(x)) + return x diff --git a/models/mobile_net.py b/models/mobile_net.py new file mode 100644 index 0000000..0cde640 --- /dev/null +++ b/models/mobile_net.py @@ -0,0 +1,63 @@ +from torch import nn +from utils.torch import * + + +class MobileNet(nn.Module): + def __init__(self, out_dim): + super().__init__() + self.out_dim = out_dim + + def conv_bn(inp, oup, stride): + return nn.Sequential( + nn.Conv2d(inp, oup, 3, stride, 1, bias=False), + nn.BatchNorm2d(oup), + nn.ReLU(inplace=True) + ) + + def conv_dw(inp, oup, stride): + return nn.Sequential( + nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False), + nn.BatchNorm2d(inp), + nn.ReLU(inplace=True), + + nn.Conv2d(inp, oup, 1, 1, 0, bias=False), + nn.BatchNorm2d(oup), + nn.ReLU(inplace=True), + ) + + self.model = nn.Sequential( + conv_bn(3, 32, 2), + conv_dw(32, 64, 1), + conv_dw(64, 128, 2), + conv_dw(128, 128, 1), + conv_dw(128, 256, 2), + conv_dw(256, 256, 1), + conv_dw(256, 512, 2), + conv_dw(512, 512, 1), + conv_dw(512, 512, 1), + conv_dw(512, 512, 1), + conv_dw(512, 512, 1), + conv_dw(512, 512, 1), + conv_dw(512, 1024, 2), + conv_dw(1024, 1024, 1), + nn.AvgPool2d(7), + ) + self.fc = nn.Linear(1024, out_dim) + + def forward(self, x): + x = self.model(x) + x = x.view(-1, 1024) + x = self.fc(x) + return x + + +if __name__ == '__main__': + import time + torch.set_grad_enabled(False) + net = MobileNet(128) + input = ones(1, 3, 224, 224) + for i in range(10): + t0 = time.time() + out = net(input) + print(time.time() - t0) + print(out.shape) diff --git a/models/resnet.py b/models/resnet.py new file mode 100644 index 0000000..72a8907 --- /dev/null +++ b/models/resnet.py @@ -0,0 +1,28 @@ +from torch import nn +from torchvision import models +from utils.torch import * + + +class ResNet(nn.Module): + + def __init__(self, out_dim, fix_params=False): + super().__init__() + self.out_dim = out_dim + self.resnet = models.resnet18(pretrained=True) + if fix_params: + for param in self.resnet.parameters(): + param.requires_grad = False + self.resnet.fc = nn.Linear(self.resnet.fc.in_features, out_dim) + + def forward(self, x): + return self.resnet(x) + + +if __name__ == '__main__': + import time + net = ResNet(128) + t0 = time.time() + input = ones(1, 3, 224, 224) + out = net(input) + print(time.time() - t0) + print(out.shape) diff --git a/models/rnn.py b/models/rnn.py new file mode 100644 index 0000000..8198227 --- /dev/null +++ b/models/rnn.py @@ -0,0 +1,68 @@ +import torch.nn as nn +from utils.torch import * + + +class RNN(nn.Module): + def __init__(self, input_dim, out_dim, cell_type='lstm', bi_dir=False): + super().__init__() + self.input_dim = input_dim + self.out_dim = out_dim + self.cell_type = cell_type + self.bi_dir = bi_dir + self.mode = 'batch' + rnn_cls = nn.LSTMCell if cell_type == 'lstm' else nn.GRUCell + hidden_dim = out_dim // 2 if bi_dir else out_dim + self.rnn_f = rnn_cls(self.input_dim, hidden_dim) + if bi_dir: + self.rnn_b = rnn_cls(self.input_dim, hidden_dim) + self.hx, self.cx = None, None + + def set_mode(self, mode): + self.mode = mode + + def initialize(self, batch_size=1): + if self.mode == 'step': + self.hx = zeros((batch_size, self.rnn_f.hidden_size)) + if self.cell_type == 'lstm': + self.cx = zeros((batch_size, self.rnn_f.hidden_size)) + + def forward(self, x): + if self.mode == 'step': + self.hx, self.cx = batch_to(x.device, self.hx, self.cx) + if self.cell_type == 'lstm': + self.hx, self.cx = self.rnn_f(x, (self.hx, self.cx)) + else: + self.hx = self.rnn_f(x, self.hx) + rnn_out = self.hx + else: + rnn_out_f = self.batch_forward(x) + if not self.bi_dir: + return rnn_out_f + rnn_out_b = self.batch_forward(x, reverse=True) + rnn_out = torch.cat((rnn_out_f, rnn_out_b), 2) + return rnn_out + + def batch_forward(self, x, reverse=False): + rnn = self.rnn_b if reverse else self.rnn_f + rnn_out = [] + hx = zeros((x.size(1), rnn.hidden_size), device=x.device) + if self.cell_type == 'lstm': + cx = zeros((x.size(1), rnn.hidden_size), device=x.device) + ind = reversed(range(x.size(0))) if reverse else range(x.size(0)) + for t in ind: + if self.cell_type == 'lstm': + hx, cx = rnn(x[t, ...], (hx, cx)) + else: + hx = rnn(x[t, ...], hx) + rnn_out.append(hx.unsqueeze(0)) + if reverse: + rnn_out.reverse() + rnn_out = torch.cat(rnn_out, 0) + return rnn_out + + +if __name__ == '__main__': + rnn = RNN(12, 24, 'gru', bi_dir=True) + input = zeros(5, 3, 12) + out = rnn(input) + print(out.shape) diff --git a/models/tcn.py b/models/tcn.py new file mode 100644 index 0000000..b76a6a0 --- /dev/null +++ b/models/tcn.py @@ -0,0 +1,78 @@ +import torch.nn as nn +from torch.nn.utils import weight_norm +from utils.torch import * + + +class Chomp1d(nn.Module): + def __init__(self, chomp_size): + super().__init__() + self.chomp_size = chomp_size + + def forward(self, x): + return x[:, :, :-self.chomp_size].contiguous() + + +class TemporalBlock(nn.Module): + def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation, dropout, causal): + super().__init__() + padding = (kernel_size - 1) * dilation // (1 if causal else 2) + modules = [] + self.conv1 = weight_norm(nn.Conv1d(n_inputs, n_outputs, kernel_size, stride=stride, padding=padding, dilation=dilation)) + modules.append(self.conv1) + if causal: + modules.append(Chomp1d(padding)) + modules.append(nn.ReLU()) + if dropout > 0: + modules.append(nn.Dropout(dropout)) + + self.conv2 = weight_norm(nn.Conv1d(n_outputs, n_outputs, kernel_size, stride=stride, padding=padding, dilation=dilation)) + modules.append(self.conv2) + if causal: + modules.append(Chomp1d(padding)) + modules.append(nn.ReLU()) + if dropout > 0: + modules.append(nn.Dropout(dropout)) + + self.net = nn.Sequential(*modules) + + self.downsample = nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None + self.relu = nn.ReLU() + self.init_weights() + + def init_weights(self): + self.conv1.weight.data.normal_(0, 0.01) + self.conv2.weight.data.normal_(0, 0.01) + if self.downsample is not None: + self.downsample.weight.data.normal_(0, 0.01) + + def forward(self, x): + out = self.net(x) + res = x if self.downsample is None else self.downsample(x) + return self.relu(out + res) + + +class TemporalConvNet(nn.Module): + def __init__(self, num_inputs, num_channels, kernel_size=3, dropout=0.2, causal=False): + super().__init__() + assert kernel_size % 2 == 1 + layers = [] + num_levels = len(num_channels) + for i in range(num_levels): + dilation_size = 2 ** i + in_channels = num_inputs if i == 0 else num_channels[i-1] + out_channels = num_channels[i] + layers += [TemporalBlock(in_channels, out_channels, kernel_size, stride=1, dilation=dilation_size, + dropout=dropout, causal=causal)] + + self.network = nn.Sequential(*layers) + + def forward(self, x): + return self.network(x) + + +if __name__ == '__main__': + tcn = TemporalConvNet(4, [1, 2, 8], kernel_size=3, causal=False) + input = zeros(3, 4, 80) + out = tcn(input) + print(tcn) + print(out.shape) diff --git a/models/video_forecast_net.py b/models/video_forecast_net.py new file mode 100644 index 0000000..28f3785 --- /dev/null +++ b/models/video_forecast_net.py @@ -0,0 +1,121 @@ +import torch.nn as nn +from models.tcn import TemporalConvNet +from models.rnn import RNN +from utils.torch import * + + +class VideoForecastNet(nn.Module): + def __init__(self, cnn_feat_dim, state_dim, v_hdim=128, v_margin=10, v_net_type='lstm', v_net_param=None, + s_hdim=None, s_net_type='id', dynamic_v=False): + super().__init__() + s_hdim = state_dim if s_hdim is None else s_hdim + self.mode = 'test' + self.cnn_feat_dim = cnn_feat_dim + self.state_dim = state_dim + self.v_net_type = v_net_type + self.v_hdim = v_hdim + self.v_margin = v_margin + self.s_net_type = s_net_type + self.s_hdim = s_hdim + self.dynamic_v = dynamic_v + self.out_dim = v_hdim + s_hdim + + if v_net_type == 'lstm': + self.v_net = RNN(cnn_feat_dim, v_hdim, v_net_type, bi_dir=False) + elif v_net_type == 'tcn': + if v_net_param is None: + v_net_param = {} + tcn_size = v_net_param.get('size', [64, 128]) + dropout = v_net_param.get('dropout', 0.2) + kernel_size = v_net_param.get('kernel_size', 3) + assert tcn_size[-1] == v_hdim + self.v_net = TemporalConvNet(cnn_feat_dim, tcn_size, kernel_size=kernel_size, dropout=dropout, causal=True) + + if s_net_type == 'lstm': + self.s_net = RNN(state_dim, s_hdim, s_net_type, bi_dir=False) + + self.v_out = None + self.t = 0 + # training only + self.indices = None + self.s_scatter_indices = None + self.s_gather_indices = None + self.v_gather_indices = None + self.cnn_feat_ctx = None + self.num_episode = None + self.max_episode_len = None + self.set_mode('test') + + def set_mode(self, mode): + self.mode = mode + if self.s_net_type == 'lstm': + if mode == 'train': + self.s_net.set_mode('batch') + else: + self.s_net.set_mode('step') + + def initialize(self, x): + if self.mode == 'test': + self.v_out = self.forward_v_net(x.unsqueeze(1)[:self.v_margin])[-1] + if self.s_net_type == 'lstm': + self.s_net.initialize() + self.t = 0 + elif self.mode == 'train': + masks, cnn_feat, v_metas = x + device, dtype = masks.device, masks.dtype + end_indice = np.where(masks.cpu().numpy() == 0)[0] + v_metas = v_metas[end_indice, :] + num_episode = len(end_indice) + end_indice = np.insert(end_indice, 0, -1) + max_episode_len = int(np.diff(end_indice).max()) + self.num_episode = num_episode + self.max_episode_len = max_episode_len + self.indices = np.arange(masks.shape[0]) + for i in range(1, num_episode): + start_index = end_indice[i] + 1 + end_index = end_indice[i + 1] + 1 + self.indices[start_index:end_index] += i * max_episode_len - start_index + self.cnn_feat_ctx = np.zeros((self.v_margin + max_episode_len if + self.dynamic_v else self.v_margin, num_episode, self.cnn_feat_dim)) + for i in range(num_episode): + exp_ind, start_ind = v_metas[i, :] + self.cnn_feat_ctx[:self.v_margin, i, :] = cnn_feat[exp_ind][start_ind - self.v_margin: start_ind] + self.cnn_feat_ctx = tensor(self.cnn_feat_ctx, dtype=dtype, device=device) + self.s_scatter_indices = LongTensor(np.tile(self.indices[:, None], (1, self.state_dim))).to(device) + self.s_gather_indices = LongTensor(np.tile(self.indices[:, None], (1, self.s_hdim))).to(device) + self.v_gather_indices = LongTensor(np.tile(self.indices[:, None], (1, self.v_hdim))).to(device) + + def forward(self, x): + if self.mode == 'test': + if self.s_net_type == 'lstm': + x = self.s_net(x) + x = torch.cat((self.v_out, x), dim=1) + self.t += 1 + elif self.mode == 'train': + if self.dynamic_v: + v_ctx = self.forward_v_net(self.cnn_feat_ctx)[self.v_margin:] + else: + v_ctx = self.forward_v_net(self.cnn_feat_ctx)[[-1]] + v_ctx = v_ctx.repeat(self.max_episode_len, 1, 1) + v_ctx = v_ctx.transpose(0, 1).contiguous().view(-1, self.v_hdim) + v_out = torch.gather(v_ctx, 0, self.v_gather_indices) + if self.s_net_type == 'lstm': + s_ctx = zeros((self.num_episode * self.max_episode_len, self.state_dim), device=x.device) + s_ctx.scatter_(0, self.s_scatter_indices, x) + s_ctx = s_ctx.view(-1, self.max_episode_len, self.state_dim).transpose(0, 1).contiguous() + s_ctx = self.s_net(s_ctx).transpose(0, 1).contiguous().view(-1, self.s_hdim) + s_out = torch.gather(s_ctx, 0, self.s_gather_indices) + else: + s_out = x + x = torch.cat((v_out, s_out), dim=1) + return x + + def forward_v_net(self, x): + if self.v_net_type == 'tcn': + x = x.permute(1, 2, 0).contiguous() + x = self.v_net(x) + if self.v_net_type == 'tcn': + x = x.permute(2, 0, 1).contiguous() + return x + + diff --git a/models/video_reg_net.py b/models/video_reg_net.py new file mode 100644 index 0000000..553792a --- /dev/null +++ b/models/video_reg_net.py @@ -0,0 +1,66 @@ +from utils.torch import * +from torch import nn +from models.resnet import ResNet +from models.tcn import TemporalConvNet +from models.rnn import RNN +from models.mlp import MLP +from models.mobile_net import MobileNet + + +class VideoRegNet(nn.Module): + + def __init__(self, out_dim, v_hdim, cnn_fdim, no_cnn=False, frame_shape=(3, 224, 224), mlp_dim=(300, 200), + cnn_type='resnet', v_net_type='lstm', v_net_param=None, causal=False): + super().__init__() + self.out_dim = out_dim + self.cnn_fdim = cnn_fdim + self.v_hdim = v_hdim + self.no_cnn = no_cnn + self.frame_shape = frame_shape + if no_cnn: + self.cnn = None + elif cnn_type == 'resnet': + self.cnn = ResNet(cnn_fdim) + elif cnn_type == 'mobile': + self.cnn = MobileNet(cnn_fdim) + + self.v_net_type = v_net_type + if v_net_type == 'lstm': + self.v_net = RNN(cnn_fdim, v_hdim, v_net_type, bi_dir=not causal) + elif v_net_type == 'tcn': + if v_net_param is None: + v_net_param = {} + tcn_size = v_net_param.get('size', [64, 128]) + dropout = v_net_param.get('dropout', 0.2) + kernel_size = v_net_param.get('kernel_size', 3) + assert tcn_size[-1] == v_hdim + self.v_net = TemporalConvNet(cnn_fdim, tcn_size, kernel_size=kernel_size, dropout=dropout, causal=causal) + self.mlp = MLP(v_hdim, mlp_dim, 'relu') + self.linear = nn.Linear(self.mlp.out_dim, out_dim) + + def forward_v_net(self, x): + if self.v_net_type == 'tcn': + x = x.permute(1, 2, 0).contiguous() + x = self.v_net(x) + if self.v_net_type == 'tcn': + x = x.permute(2, 0, 1).contiguous() + return x + + def forward(self, x): + # CNN + if self.cnn is not None: + x = self.cnn(x.view((-1,) + self.frame_shape)).view(-1, x.size(1), self.cnn_fdim) + x = self.forward_v_net(x).view(-1, self.v_hdim) + x = self.mlp(x) + x = self.linear(x) + return x + + def get_cnn_feature(self, x): + return self.cnn(x.view((-1,) + self.frame_shape)) + + +if __name__ == '__main__': + net = VideoRegNet(64, 128, 128, v_net_type='tcn', cnn_type='mobile') + input = ones(32, 1, 3, 224, 224) + out = net(input) + print(out.shape) diff --git a/models/video_state_net.py b/models/video_state_net.py new file mode 100644 index 0000000..a941346 --- /dev/null +++ b/models/video_state_net.py @@ -0,0 +1,79 @@ +import torch.nn as nn +from models.tcn import TemporalConvNet +from models.rnn import RNN +from utils.torch import * + + +class VideoStateNet(nn.Module): + def __init__(self, cnn_feat_dim, v_hdim=128, v_margin=10, v_net_type='lstm', v_net_param=None, causal=False): + super().__init__() + self.mode = 'test' + self.cnn_feat_dim = cnn_feat_dim + self.v_net_type = v_net_type + self.v_hdim = v_hdim + self.v_margin = v_margin + if v_net_type == 'lstm': + self.v_net = RNN(cnn_feat_dim, v_hdim, v_net_type, bi_dir=not causal) + elif v_net_type == 'tcn': + if v_net_param is None: + v_net_param = {} + tcn_size = v_net_param.get('size', [64, 128]) + dropout = v_net_param.get('dropout', 0.2) + kernel_size = v_net_param.get('kernel_size', 3) + assert tcn_size[-1] == v_hdim + self.v_net = TemporalConvNet(cnn_feat_dim, tcn_size, kernel_size=kernel_size, dropout=dropout, causal=causal) + self.v_out = None + self.t = 0 + # training only + self.indices = None + self.scatter_indices = None + self.gather_indices = None + self.cnn_feat_ctx = None + + def set_mode(self, mode): + self.mode = mode + + def initialize(self, x): + if self.mode == 'test': + self.v_out = self.forward_v_net(x.unsqueeze(1)).squeeze(1)[self.v_margin:-self.v_margin] + self.t = 0 + elif self.mode == 'train': + masks, cnn_feat, v_metas = x + device, dtype = masks.device, masks.dtype + end_indice = np.where(masks.cpu().numpy() == 0)[0] + v_metas = v_metas[end_indice, :] + num_episode = len(end_indice) + end_indice = np.insert(end_indice, 0, -1) + max_episode_len = int(np.diff(end_indice).max()) + self.indices = np.arange(masks.shape[0]) + for i in range(1, num_episode): + start_index = end_indice[i] + 1 + end_index = end_indice[i + 1] + 1 + self.indices[start_index:end_index] += i * max_episode_len - start_index + self.cnn_feat_ctx = np.zeros((max_episode_len + 2*self.v_margin, num_episode, self.cnn_feat_dim)) + for i in range(num_episode): + exp_ind, start_ind = v_metas[i, :] + self.cnn_feat_ctx[:, i, :] = cnn_feat[exp_ind][start_ind - self.v_margin: start_ind + max_episode_len + self.v_margin] + self.cnn_feat_ctx = tensor(self.cnn_feat_ctx, dtype=dtype, device=device) + self.scatter_indices = LongTensor(np.tile(self.indices[:, None], (1, self.cnn_feat_dim))).to(device) + self.gather_indices = LongTensor(np.tile(self.indices[:, None], (1, self.v_hdim))).to(device) + + def forward(self, x): + if self.mode == 'test': + x = torch.cat((self.v_out[[self.t], :], x), dim=1) + self.t += 1 + elif self.mode == 'train': + v_ctx = self.forward_v_net(self.cnn_feat_ctx)[self.v_margin:-self.v_margin] + v_ctx = v_ctx.transpose(0, 1).contiguous().view(-1, self.v_hdim) + v_out = torch.gather(v_ctx, 0, self.gather_indices) + x = torch.cat((v_out, x), dim=1) + return x + + def forward_v_net(self, x): + if self.v_net_type == 'tcn': + x = x.permute(1, 2, 0).contiguous() + x = self.v_net(x) + if self.v_net_type == 'tcn': + x = x.permute(2, 0, 1).contiguous() + return x + diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..4007647 --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1,7 @@ +from utils.memory import * +from utils.zfilter import * +from utils.torch import * +from utils.math import * +from utils.tools import * +from utils.logger import * +from utils.tb_logger import * diff --git a/utils/logger.py b/utils/logger.py new file mode 100644 index 0000000..d81f600 --- /dev/null +++ b/utils/logger.py @@ -0,0 +1,27 @@ +import logging +import os + + +def create_logger(filename, file_handle=True): + # create logger + logger = logging.getLogger(filename) + logger.propagate = False + logger.setLevel(logging.DEBUG) + # create console handler with a higher log level + ch = logging.StreamHandler() + ch.setLevel(logging.INFO) + stream_formatter = logging.Formatter('%(message)s') + ch.setFormatter(stream_formatter) + logger.addHandler(ch) + + if file_handle: + # create file handler which logs even debug messages + os.makedirs(os.path.dirname(filename), exist_ok=True) + fh = logging.FileHandler(filename, mode='a') + fh.setLevel(logging.DEBUG) + file_formatter = logging.Formatter('[%(asctime)s] %(message)s') + fh.setFormatter(file_formatter) + logger.addHandler(fh) + + return logger + diff --git a/utils/math.py b/utils/math.py new file mode 100644 index 0000000..236a97e --- /dev/null +++ b/utils/math.py @@ -0,0 +1,121 @@ +import torch +import math +import numpy as np +from utils.transformation import quaternion_matrix, quaternion_about_axis,\ + quaternion_inverse, quaternion_multiply, rotation_from_quaternion, rotation_from_matrix + + +def normal_entropy(std): + var = std.pow(2) + entropy = 0.5 + 0.5 * torch.log(2 * var * math.pi) + return entropy.sum(1, keepdim=True) + + +def normal_log_density(x, mean, log_std, std): + var = std.pow(2) + log_density = -(x - mean).pow(2) / (2 * var) - 0.5 * math.log(2 * math.pi) - log_std + return log_density.sum(1, keepdim=True) + + +def get_qvel_fd(cur_qpos, next_qpos, dt, transform=None): + v = (next_qpos[:3] - cur_qpos[:3]) / dt + qrel = quaternion_multiply(next_qpos[3:7], quaternion_inverse(cur_qpos[3:7])) + axis, angle = rotation_from_quaternion(qrel, True) + if angle > np.pi: + angle -= 2 * np.pi + elif angle < -np.pi: + angle += 2 * np.pi + rv = (axis * angle) / dt + rv = transform_vec(rv, cur_qpos[3:7], 'root') # angular velocity is in root coord + qvel = (next_qpos[7:] - cur_qpos[7:]) / dt + qvel = np.concatenate((v, rv, qvel)) + if transform is not None: + v = transform_vec(v, cur_qpos[3:7], transform) + qvel[:3] = v + return qvel + + +def get_angvel_fd(prev_bquat, cur_bquat, dt): + q_diff = multi_quat_diff(cur_bquat, prev_bquat) + n_joint = q_diff.shape[0] // 4 + body_angvel = np.zeros(n_joint * 3) + for i in range(n_joint): + body_angvel[3*i: 3*i + 3] = rotation_from_quaternion(q_diff[4*i: 4*i + 4]) / dt + return body_angvel + + +def transform_vec(v, q, trans='root'): + if trans == 'root': + rot = quaternion_matrix(q)[:3, :3] + elif trans == 'heading': + hq = q.copy() + hq[1] = 0 + hq[2] = 0 + hq /= np.linalg.norm(hq) + rot = quaternion_matrix(hq)[:3, :3] + else: + assert False + v = rot.T.dot(v[:, None]).ravel() + return v + + +def get_heading_q(q): + hq = q.copy() + hq[1] = 0 + hq[2] = 0 + hq /= np.linalg.norm(hq) + return hq + + +def get_heading(q): + hq = q.copy() + hq[1] = 0 + hq[2] = 0 + if hq[3] < 0: + hq *= -1 + hq /= np.linalg.norm(hq) + return 2 * math.acos(hq[0]) + + +def de_heading(q): + return quaternion_multiply(quaternion_inverse(get_heading_q(q)), q) + + +def multi_quat_diff(nq1, nq0): + """return the relative quaternions q1-q0 of N joints""" + + nq_diff = np.zeros_like(nq0) + for i in range(nq1.shape[0] // 4): + ind = slice(4*i, 4*i + 4) + q1 = nq1[ind] + q0 = nq0[ind] + nq_diff[ind] = quaternion_multiply(q1, quaternion_inverse(q0)) + return nq_diff + + +def multi_quat_norm(nq): + """return the scalar rotation of a N joints""" + + nq_norm = np.arccos(np.clip(nq[::4], -1.0, 1.0)) + return nq_norm + + +def quat_mul_vec(q, v): + return quaternion_matrix(q)[:3, :3].dot(v[:, None]).ravel() + + +def quat_to_bullet(q): + return np.array([q[1], q[2], q[3], q[0]]) + + +def quat_from_bullet(q): + return np.array([q[3], q[0], q[1], q[2]]) + + +def quat_from_expmap(e): + angle = np.linalg.norm(e) + if angle < 1e-12: + axis = np.array([1.0, 0.0, 0.0]) + else: + axis = e / angle + return quaternion_about_axis(angle, axis) diff --git a/utils/memory.py b/utils/memory.py new file mode 100644 index 0000000..6f9a31e --- /dev/null +++ b/utils/memory.py @@ -0,0 +1,24 @@ +import random + + +class Memory(object): + def __init__(self): + self.memory = [] + + def push(self, *args): + """Saves a tuple.""" + self.memory.append([*args]) + + def sample(self, batch_size=None): + if batch_size is None: + return self.memory + else: + random_batch = random.sample(self.memory, batch_size) + return random_batch + + def append(self, new_memory): + self.memory += new_memory.memory + + def __len__(self): + return len(self.memory) + diff --git a/utils/tb_logger.py b/utils/tb_logger.py new file mode 100644 index 0000000..0b60b5b --- /dev/null +++ b/utils/tb_logger.py @@ -0,0 +1,116 @@ +""" +File: logger.py +Modified by: Senthil Purushwalkam +Code referenced from https://gist.github.com/gyglim/1f8dfb1b5c82627ae3efcfbbadb9f514 +Email: spurushwandrewcmuedu +Github: https://github.com/senthilps8 +Description: +""" + +import tensorflow as tf +from torch.autograd import Variable +import numpy as np +import scipy.misc +import os +import torch +from os import path + +try: + from StringIO import StringIO # Python 2.7 +except ImportError: + from io import BytesIO # Python 3.x + + +class Logger(object): + + def __init__(self, log_dir, name=None): + """Create a summary writer logging to log_dir.""" + self.name = name + if name is not None: + try: + os.makedirs(os.path.join(log_dir, name)) + except: + pass + self.writer = tf.summary.FileWriter(os.path.join(log_dir, name), + filename_suffix=name) + else: + self.writer = tf.summary.FileWriter(log_dir) + + def scalar_summary(self, tag, value, step): + """Log a scalar variable.""" + summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)]) + self.writer.add_summary(summary, step) + + def image_summary(self, tag, images, step): + """Log a list of images.""" + + img_summaries = [] + for i, img in enumerate(images): + # Write the image to a string + try: + s = StringIO() + except: + s = BytesIO() + scipy.misc.toimage(img).save(s, format="png") + + # Create an Image object + img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(), + height=img.shape[0], + width=img.shape[1]) + # Create a Summary value + img_summaries.append(tf.Summary.Value(tag='%s/%d' % (tag, i), image=img_sum)) + + # Create and write Summary + summary = tf.Summary(value=img_summaries) + self.writer.add_summary(summary, step) + + def histo_summary(self, tag, values, step, bins=1000): + """Log a histogram of the tensor of values.""" + + # Create a histogram using numpy + counts, bin_edges = np.histogram(values, bins=bins) + + # Fill the fields of the histogram proto + hist = tf.HistogramProto() + hist.min = float(np.min(values)) + hist.max = float(np.max(values)) + hist.num = int(np.prod(values.shape)) + hist.sum = float(np.sum(values)) + hist.sum_squares = float(np.sum(values ** 2)) + + # Drop the start of the first bin + bin_edges = bin_edges[1:] + + # Add bin edges and counts + for edge in bin_edges: + hist.bucket_limit.append(edge) + for c in counts: + hist.bucket.append(c) + + # Create and write Summary + summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)]) + self.writer.add_summary(summary, step) + self.writer.flush() + + def to_np(self, x): + return x.data.cpu().numpy() + + def to_var(self, x): + if torch.cuda.is_available(): + x = x.cuda() + return Variable(x) + + def model_param_histo_summary(self, model, step): + """log histogram summary of model's parameters + and parameter gradients + """ + for tag, value in model.named_parameters(): + if value.grad is None: + continue + tag = tag.replace('.', '/') + tag = self.name + '/' + tag + self.histo_summary(tag, self.to_np(value), step) + self.histo_summary(tag + '/grad', self.to_np(value.grad), step) + + + diff --git a/utils/tools.py b/utils/tools.py new file mode 100644 index 0000000..966a615 --- /dev/null +++ b/utils/tools.py @@ -0,0 +1,75 @@ +import numpy as np +import os +import shutil +from os import path +from os import listdir +from PIL import Image +from OpenGL import GL +from gym.envs.mujoco.mujoco_env import MujocoEnv +from utils.math import * +import glfw +import cv2 + + +def assets_dir(): + return path.abspath(path.join(path.dirname(path.abspath(__file__)), '../assets')) + + +def out_dir(): + return path.abspath(path.join(path.dirname(path.abspath(__file__)), '../out')) + + +def log_dir(): + return path.abspath(path.join(path.dirname(path.abspath(__file__)), '../logs')) + + +def recreate_dirs(*dirs): + for d in dirs: + if os.path.exists(d): + shutil.rmtree(d) + os.makedirs(d) + + +def load_img(path): + # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835) + with open(path, 'rb') as f: + I = Image.open(f) + img = I.resize((224, 224), Image.ANTIALIAS).convert('RGB') + return img + + +def save_screen_shots(window, file_name, transparent=False): + import pyautogui + xpos, ypos = glfw.get_window_pos(window) + width, height = glfw.get_window_size(window) + image = pyautogui.screenshot(region=(xpos*2, ypos*2, width*2, height*2)) + image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGRA if transparent else cv2.COLOR_RGB2BGR) + if transparent: + image[np.all(image >= [240, 240, 240, 240], axis=2)] = [255, 255, 255, 0] + cv2.imwrite(file_name, image) + + +"""mujoco helper""" + + +def get_body_qposaddr(model): + body_qposaddr = dict() + for i, body_name in enumerate(model.body_names): + start_joint = model.body_jntadr[i] + if start_joint < 0: + continue + end_joint = start_joint + model.body_jntnum[i] + start_qposaddr = model.jnt_qposadr[start_joint] + if end_joint < len(model.jnt_qposadr): + end_qposaddr = model.jnt_qposadr[end_joint] + else: + end_qposaddr = model.nq + body_qposaddr[body_name] = (start_qposaddr, end_qposaddr) + return body_qposaddr + + +def align_human_state(qpos, qvel, ref_qpos): + qpos[:2] = ref_qpos[:2] + hq = get_heading_q(ref_qpos[3:7]) + qpos[3:7] = quaternion_multiply(hq, qpos[3:7]) + qvel[:3] = quat_mul_vec(hq, qvel[:3]) diff --git a/utils/torch.py b/utils/torch.py new file mode 100644 index 0000000..3fbfd50 --- /dev/null +++ b/utils/torch.py @@ -0,0 +1,158 @@ +import torch +import numpy as np + +tensor = torch.tensor +DoubleTensor = torch.DoubleTensor +FloatTensor = torch.FloatTensor +LongTensor = torch.LongTensor +ByteTensor = torch.ByteTensor +ones = torch.ones +zeros = torch.zeros + + +class to_cpu: + + def __init__(self, *models): + self.models = list(filter(lambda x: x is not None, models)) + self.prev_devices = [x.device if hasattr(x, 'device') else next(x.parameters()).device for x in self.models] + for x in self.models: + x.to(torch.device('cpu')) + + def __enter__(self): + pass + + def __exit__(self, *args): + for x, device in zip(self.models, self.prev_devices): + x.to(device) + return False + + +class to_device: + + def __init__(self, device, *models): + self.models = list(filter(lambda x: x is not None, models)) + self.prev_devices = [x.device if hasattr(x, 'device') else next(x.parameters()).device for x in self.models] + for x in self.models: + x.to(device) + + def __enter__(self): + pass + + def __exit__(self, *args): + for x, device in zip(self.models, self.prev_devices): + x.to(device) + return False + + +class to_test: + + def __init__(self, *models): + self.models = list(filter(lambda x: x is not None, models)) + self.prev_modes = [x.training for x in self.models] + for x in self.models: + x.train(False) + + def __enter__(self): + pass + + def __exit__(self, *args): + for x, mode in zip(self.models, self.prev_modes): + x.train(mode) + return False + + +class to_train: + + def __init__(self, *models): + self.models = list(filter(lambda x: x is not None, models)) + self.prev_modes = [x.training for x in self.models] + for x in self.models: + x.train(True) + + def __enter__(self): + pass + + def __exit__(self, *args): + for x, mode in zip(self.models, self.prev_modes): + x.train(mode) + return False + + +def batch_to(dst, *args): + return [x.to(dst) for x in args if x is not None] + + +def get_flat_params_from(models): + if not hasattr(models, '__iter__'): + models = (models, ) + params = [] + for model in models: + for param in model.parameters(): + params.append(param.data.view(-1)) + + flat_params = torch.cat(params) + return flat_params + + +def set_flat_params_to(model, flat_params): + prev_ind = 0 + for param in model.parameters(): + flat_size = int(np.prod(list(param.size()))) + param.data.copy_( + flat_params[prev_ind:prev_ind + flat_size].view(param.size())) + prev_ind += flat_size + + +def get_flat_grad_from(inputs, grad_grad=False): + grads = [] + for param in inputs: + if grad_grad: + grads.append(param.grad.grad.view(-1)) + else: + if param.grad is None: + grads.append(zeros(param.view(-1).shape)) + else: + grads.append(param.grad.view(-1)) + + flat_grad = torch.cat(grads) + return flat_grad + + +def compute_flat_grad(output, inputs, filter_input_ids=set(), retain_graph=False, create_graph=False): + if create_graph: + retain_graph = True + + inputs = list(inputs) + params = [] + for i, param in enumerate(inputs): + if i not in filter_input_ids: + params.append(param) + + grads = torch.autograd.grad(output, params, retain_graph=retain_graph, create_graph=create_graph) + + j = 0 + out_grads = [] + for i, param in enumerate(inputs): + if i in filter_input_ids: + out_grads.append(zeros(param.view(-1).shape)) + else: + out_grads.append(grads[j].view(-1)) + j += 1 + grads = torch.cat(out_grads) + + for param in params: + param.grad = None + return grads + + +def set_optimizer_lr(optimizer, lr): + for param_group in optimizer.param_groups: + param_group['lr'] = lr + + +def filter_state_dict(state_dict, filter_keys): + for key in list(state_dict.keys()): + for f_key in filter_keys: + if f_key in key: + del state_dict[key] + break diff --git a/utils/transformation.py b/utils/transformation.py new file mode 100644 index 0000000..b4a74a5 --- /dev/null +++ b/utils/transformation.py @@ -0,0 +1,1944 @@ +# -*- coding: utf-8 -*- +# transformations.py + +# Copyright (c) 2006-2017, Christoph Gohlke +# Copyright (c) 2006-2017, The Regents of the University of California +# Produced at the Laboratory for Fluorescence Dynamics +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of the copyright holders nor the names of any +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. + +"""Homogeneous Transformation Matrices and Quaternions. + +A library for calculating 4x4 matrices for translating, rotating, reflecting, +scaling, shearing, projecting, orthogonalizing, and superimposing arrays of +3D homogeneous coordinates as well as for converting between rotation matrices, +Euler angles, and quaternions. Also includes an Arcball control object and +functions to decompose transformation matrices. + +:Author: + `Christoph Gohlke `_ + +:Organization: + Laboratory for Fluorescence Dynamics, University of California, Irvine + +:Version: 2017.02.17 + +Requirements +------------ +* `CPython 2.7 or 3.5 `_ +* `Numpy 1.11 `_ +* `Transformations.c 2017.02.17 `_ + (recommended for speedup of some functions) + +Notes +----- +The API is not stable yet and is expected to change between revisions. + +This Python code is not optimized for speed. Refer to the transformations.c +module for a faster implementation of some functions. + +Documentation in HTML format can be generated with epydoc. + +Matrices (M) can be inverted using numpy.linalg.inv(M), be concatenated using +numpy.dot(M0, M1), or transform homogeneous coordinate arrays (v) using +numpy.dot(M, v) for shape (4, \*) column vectors, respectively +numpy.dot(v, M.T) for shape (\*, 4) row vectors ("array of points"). + +This module follows the "column vectors on the right" and "row major storage" +(C contiguous) conventions. The translation components are in the right column +of the transformation matrix, i.e. M[:3, 3]. +The transpose of the transformation matrices may have to be used to interface +with other graphics systems, e.g. with OpenGL's glMultMatrixd(). See also [16]. + +Calculations are carried out with numpy.float64 precision. + +Vector, point, quaternion, and matrix function arguments are expected to be +"array like", i.e. tuple, list, or numpy arrays. + +Return types are numpy arrays unless specified otherwise. + +Angles are in radians unless specified otherwise. + +Quaternions w+ix+jy+kz are represented as [w, x, y, z]. + +A triple of Euler angles can be applied/interpreted in 24 ways, which can +be specified using a 4 character string or encoded 4-tuple: + + *Axes 4-string*: e.g. 'sxyz' or 'ryxy' + + - first character : rotations are applied to 's'tatic or 'r'otating frame + - remaining characters : successive rotation axis 'x', 'y', or 'z' + + *Axes 4-tuple*: e.g. (0, 0, 0, 0) or (1, 1, 1, 1) + + - inner axis: code of axis ('x':0, 'y':1, 'z':2) of rightmost matrix. + - parity : even (0) if inner axis 'x' is followed by 'y', 'y' is followed + by 'z', or 'z' is followed by 'x'. Otherwise odd (1). + - repetition : first and last axis are same (1) or different (0). + - frame : rotations are applied to static (0) or rotating (1) frame. + +Other Python packages and modules for 3D transformations and quaternions: + +* `Transforms3d `_ + includes most code of this module. +* `Blender.mathutils `_ +* `numpy-dtypes `_ + +References +---------- +(1) Matrices and transformations. Ronald Goldman. + In "Graphics Gems I", pp 472-475. Morgan Kaufmann, 1990. +(2) More matrices and transformations: shear and pseudo-perspective. + Ronald Goldman. In "Graphics Gems II", pp 320-323. Morgan Kaufmann, 1991. +(3) Decomposing a matrix into simple transformations. Spencer Thomas. + In "Graphics Gems II", pp 320-323. Morgan Kaufmann, 1991. +(4) Recovering the data from the transformation matrix. Ronald Goldman. + In "Graphics Gems II", pp 324-331. Morgan Kaufmann, 1991. +(5) Euler angle conversion. Ken Shoemake. + In "Graphics Gems IV", pp 222-229. Morgan Kaufmann, 1994. +(6) Arcball rotation control. Ken Shoemake. + In "Graphics Gems IV", pp 175-192. Morgan Kaufmann, 1994. +(7) Representing attitude: Euler angles, unit quaternions, and rotation + vectors. James Diebel. 2006. +(8) A discussion of the solution for the best rotation to relate two sets + of vectors. W Kabsch. Acta Cryst. 1978. A34, 827-828. +(9) Closed-form solution of absolute orientation using unit quaternions. + BKP Horn. J Opt Soc Am A. 1987. 4(4):629-642. +(10) Quaternions. Ken Shoemake. + http://www.sfu.ca/~jwa3/cmpt461/files/quatut.pdf +(11) From quaternion to matrix and back. JMP van Waveren. 2005. + http://www.intel.com/cd/ids/developer/asmo-na/eng/293748.htm +(12) Uniform random rotations. Ken Shoemake. + In "Graphics Gems III", pp 124-132. Morgan Kaufmann, 1992. +(13) Quaternion in molecular modeling. CFF Karney. + J Mol Graph Mod, 25(5):595-604 +(14) New method for extracting the quaternion from a rotation matrix. + Itzhack Y Bar-Itzhack, J Guid Contr Dynam. 2000. 23(6): 1085-1087. +(15) Multiple View Geometry in Computer Vision. Hartley and Zissermann. + Cambridge University Press; 2nd Ed. 2004. Chapter 4, Algorithm 4.7, p 130. +(16) Column Vectors vs. Row Vectors. + http://steve.hollasch.net/cgindex/math/matrix/column-vec.html + +Examples +-------- +>>> alpha, beta, gamma = 0.123, -1.234, 2.345 +>>> origin, xaxis, yaxis, zaxis = [0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1] +>>> I = identity_matrix() +>>> Rx = rotation_matrix(alpha, xaxis) +>>> Ry = rotation_matrix(beta, yaxis) +>>> Rz = rotation_matrix(gamma, zaxis) +>>> R = concatenate_matrices(Rx, Ry, Rz) +>>> euler = euler_from_matrix(R, 'rxyz') +>>> numpy.allclose([alpha, beta, gamma], euler) +True +>>> Re = euler_matrix(alpha, beta, gamma, 'rxyz') +>>> is_same_transform(R, Re) +True +>>> al, be, ga = euler_from_matrix(Re, 'rxyz') +>>> is_same_transform(Re, euler_matrix(al, be, ga, 'rxyz')) +True +>>> qx = quaternion_about_axis(alpha, xaxis) +>>> qy = quaternion_about_axis(beta, yaxis) +>>> qz = quaternion_about_axis(gamma, zaxis) +>>> q = quaternion_multiply(qx, qy) +>>> q = quaternion_multiply(q, qz) +>>> Rq = quaternion_matrix(q) +>>> is_same_transform(R, Rq) +True +>>> S = scale_matrix(1.23, origin) +>>> T = translation_matrix([1, 2, 3]) +>>> Z = shear_matrix(beta, xaxis, origin, zaxis) +>>> R = random_rotation_matrix(numpy.random.rand(3)) +>>> M = concatenate_matrices(T, R, Z, S) +>>> scale, shear, angles, trans, persp = decompose_matrix(M) +>>> numpy.allclose(scale, 1.23) +True +>>> numpy.allclose(trans, [1, 2, 3]) +True +>>> numpy.allclose(shear, [0, math.tan(beta), 0]) +True +>>> is_same_transform(R, euler_matrix(axes='sxyz', *angles)) +True +>>> M1 = compose_matrix(scale, shear, angles, trans, persp) +>>> is_same_transform(M, M1) +True +>>> v0, v1 = random_vector(3), random_vector(3) +>>> M = rotation_matrix(angle_between_vectors(v0, v1), vector_product(v0, v1)) +>>> v2 = numpy.dot(v0, M[:3,:3].T) +>>> numpy.allclose(unit_vector(v1), unit_vector(v2)) +True + +""" + +from __future__ import division, print_function + +import math + +import numpy + +__version__ = '2017.02.17' +__docformat__ = 'restructuredtext en' +__all__ = () + + +def identity_matrix(): + """Return 4x4 identity/unit matrix. + + >>> I = identity_matrix() + >>> numpy.allclose(I, numpy.dot(I, I)) + True + >>> numpy.sum(I), numpy.trace(I) + (4.0, 4.0) + >>> numpy.allclose(I, numpy.identity(4)) + True + + """ + return numpy.identity(4) + + +def translation_matrix(direction): + """Return matrix to translate by direction vector. + + >>> v = numpy.random.random(3) - 0.5 + >>> numpy.allclose(v, translation_matrix(v)[:3, 3]) + True + + """ + M = numpy.identity(4) + M[:3, 3] = direction[:3] + return M + + +def translation_from_matrix(matrix): + """Return translation vector from translation matrix. + + >>> v0 = numpy.random.random(3) - 0.5 + >>> v1 = translation_from_matrix(translation_matrix(v0)) + >>> numpy.allclose(v0, v1) + True + + """ + return numpy.array(matrix, copy=False)[:3, 3].copy() + + +def reflection_matrix(point, normal): + """Return matrix to mirror at plane defined by point and normal vector. + + >>> v0 = numpy.random.random(4) - 0.5 + >>> v0[3] = 1. + >>> v1 = numpy.random.random(3) - 0.5 + >>> R = reflection_matrix(v0, v1) + >>> numpy.allclose(2, numpy.trace(R)) + True + >>> numpy.allclose(v0, numpy.dot(R, v0)) + True + >>> v2 = v0.copy() + >>> v2[:3] += v1 + >>> v3 = v0.copy() + >>> v2[:3] -= v1 + >>> numpy.allclose(v2, numpy.dot(R, v3)) + True + + """ + normal = unit_vector(normal[:3]) + M = numpy.identity(4) + M[:3, :3] -= 2.0 * numpy.outer(normal, normal) + M[:3, 3] = (2.0 * numpy.dot(point[:3], normal)) * normal + return M + + +def reflection_from_matrix(matrix): + """Return mirror plane point and normal vector from reflection matrix. + + >>> v0 = numpy.random.random(3) - 0.5 + >>> v1 = numpy.random.random(3) - 0.5 + >>> M0 = reflection_matrix(v0, v1) + >>> point, normal = reflection_from_matrix(M0) + >>> M1 = reflection_matrix(point, normal) + >>> is_same_transform(M0, M1) + True + + """ + M = numpy.array(matrix, dtype=numpy.float64, copy=False) + # normal: unit eigenvector corresponding to eigenvalue -1 + w, V = numpy.linalg.eig(M[:3, :3]) + i = numpy.where(abs(numpy.real(w) + 1.0) < 1e-8)[0] + if not len(i): + raise ValueError("no unit eigenvector corresponding to eigenvalue -1") + normal = numpy.real(V[:, i[0]]).squeeze() + # point: any unit eigenvector corresponding to eigenvalue 1 + w, V = numpy.linalg.eig(M) + i = numpy.where(abs(numpy.real(w) - 1.0) < 1e-8)[0] + if not len(i): + raise ValueError("no unit eigenvector corresponding to eigenvalue 1") + point = numpy.real(V[:, i[-1]]).squeeze() + point /= point[3] + return point, normal + + +def rotation_matrix(angle, direction, point=None): + """Return matrix to rotate about axis defined by point and direction. + + >>> R = rotation_matrix(math.pi/2, [0, 0, 1], [1, 0, 0]) + >>> numpy.allclose(numpy.dot(R, [0, 0, 0, 1]), [1, -1, 0, 1]) + True + >>> angle = (random.random() - 0.5) * (2*math.pi) + >>> direc = numpy.random.random(3) - 0.5 + >>> point = numpy.random.random(3) - 0.5 + >>> R0 = rotation_matrix(angle, direc, point) + >>> R1 = rotation_matrix(angle-2*math.pi, direc, point) + >>> is_same_transform(R0, R1) + True + >>> R0 = rotation_matrix(angle, direc, point) + >>> R1 = rotation_matrix(-angle, -direc, point) + >>> is_same_transform(R0, R1) + True + >>> I = numpy.identity(4, numpy.float64) + >>> numpy.allclose(I, rotation_matrix(math.pi*2, direc)) + True + >>> numpy.allclose(2, numpy.trace(rotation_matrix(math.pi/2, + ... direc, point))) + True + + """ + sina = math.sin(angle) + cosa = math.cos(angle) + direction = unit_vector(direction[:3]) + # rotation matrix around unit vector + R = numpy.diag([cosa, cosa, cosa]) + R += numpy.outer(direction, direction) * (1.0 - cosa) + direction *= sina + R += numpy.array([[ 0.0, -direction[2], direction[1]], + [ direction[2], 0.0, -direction[0]], + [-direction[1], direction[0], 0.0]]) + M = numpy.identity(4) + M[:3, :3] = R + if point is not None: + # rotation not around origin + point = numpy.array(point[:3], dtype=numpy.float64, copy=False) + M[:3, 3] = point - numpy.dot(R, point) + return M + +import numpy as np + + +def rotation_from_quaternion(quaternion, separate=False): + if 1 - quaternion[0] < 1e-8: + axis = np.array([1.0, 0.0, 0.0]) + angle = 0.0 + else: + s = math.sqrt(1 - quaternion[0]*quaternion[0]) + axis = quaternion[1:4] / s + angle = 2 * math.acos(quaternion[0]) + return (axis, angle) if separate else axis * angle + + +def rotation_from_matrix(matrix): + """Return rotation angle and axis from rotation matrix. + + >>> angle = (random.random() - 0.5) * (2*math.pi) + >>> direc = numpy.random.random(3) - 0.5 + >>> point = numpy.random.random(3) - 0.5 + >>> R0 = rotation_matrix(angle, direc, point) + >>> angle, direc, point = rotation_from_matrix(R0) + >>> R1 = rotation_matrix(angle, direc, point) + >>> is_same_transform(R0, R1) + True + + """ + R = numpy.array(matrix, dtype=numpy.float64, copy=False) + R33 = R[:3, :3] + # direction: unit eigenvector of R33 corresponding to eigenvalue of 1 + w, W = numpy.linalg.eig(R33.T) + i = numpy.where(abs(numpy.real(w) - 1.0) < 1e-8)[0] + if not len(i): + raise ValueError("no unit eigenvector corresponding to eigenvalue 1") + direction = numpy.real(W[:, i[-1]]).squeeze() + # point: unit eigenvector of R33 corresponding to eigenvalue of 1 + w, Q = numpy.linalg.eig(R) + i = numpy.where(abs(numpy.real(w) - 1.0) < 1e-8)[0] + if not len(i): + raise ValueError("no unit eigenvector corresponding to eigenvalue 1") + point = numpy.real(Q[:, i[-1]]).squeeze() + point /= point[3] + # rotation angle depending on direction + cosa = (numpy.trace(R33) - 1.0) / 2.0 + if abs(direction[2]) > 1e-8: + sina = (R[1, 0] + (cosa-1.0)*direction[0]*direction[1]) / direction[2] + elif abs(direction[1]) > 1e-8: + sina = (R[0, 2] + (cosa-1.0)*direction[0]*direction[2]) / direction[1] + else: + sina = (R[2, 1] + (cosa-1.0)*direction[1]*direction[2]) / direction[0] + angle = math.atan2(sina, cosa) + return angle, direction, point + + +def scale_matrix(factor, origin=None, direction=None): + """Return matrix to scale by factor around origin in direction. + + Use factor -1 for point symmetry. + + >>> v = (numpy.random.rand(4, 5) - 0.5) * 20 + >>> v[3] = 1 + >>> S = scale_matrix(-1.234) + >>> numpy.allclose(numpy.dot(S, v)[:3], -1.234*v[:3]) + True + >>> factor = random.random() * 10 - 5 + >>> origin = numpy.random.random(3) - 0.5 + >>> direct = numpy.random.random(3) - 0.5 + >>> S = scale_matrix(factor, origin) + >>> S = scale_matrix(factor, origin, direct) + + """ + if direction is None: + # uniform scaling + M = numpy.diag([factor, factor, factor, 1.0]) + if origin is not None: + M[:3, 3] = origin[:3] + M[:3, 3] *= 1.0 - factor + else: + # nonuniform scaling + direction = unit_vector(direction[:3]) + factor = 1.0 - factor + M = numpy.identity(4) + M[:3, :3] -= factor * numpy.outer(direction, direction) + if origin is not None: + M[:3, 3] = (factor * numpy.dot(origin[:3], direction)) * direction + return M + + +def scale_from_matrix(matrix): + """Return scaling factor, origin and direction from scaling matrix. + + >>> factor = random.random() * 10 - 5 + >>> origin = numpy.random.random(3) - 0.5 + >>> direct = numpy.random.random(3) - 0.5 + >>> S0 = scale_matrix(factor, origin) + >>> factor, origin, direction = scale_from_matrix(S0) + >>> S1 = scale_matrix(factor, origin, direction) + >>> is_same_transform(S0, S1) + True + >>> S0 = scale_matrix(factor, origin, direct) + >>> factor, origin, direction = scale_from_matrix(S0) + >>> S1 = scale_matrix(factor, origin, direction) + >>> is_same_transform(S0, S1) + True + + """ + M = numpy.array(matrix, dtype=numpy.float64, copy=False) + M33 = M[:3, :3] + factor = numpy.trace(M33) - 2.0 + try: + # direction: unit eigenvector corresponding to eigenvalue factor + w, V = numpy.linalg.eig(M33) + i = numpy.where(abs(numpy.real(w) - factor) < 1e-8)[0][0] + direction = numpy.real(V[:, i]).squeeze() + direction /= vector_norm(direction) + except IndexError: + # uniform scaling + factor = (factor + 2.0) / 3.0 + direction = None + # origin: any eigenvector corresponding to eigenvalue 1 + w, V = numpy.linalg.eig(M) + i = numpy.where(abs(numpy.real(w) - 1.0) < 1e-8)[0] + if not len(i): + raise ValueError("no eigenvector corresponding to eigenvalue 1") + origin = numpy.real(V[:, i[-1]]).squeeze() + origin /= origin[3] + return factor, origin, direction + + +def projection_matrix(point, normal, direction=None, + perspective=None, pseudo=False): + """Return matrix to project onto plane defined by point and normal. + + Using either perspective point, projection direction, or none of both. + + If pseudo is True, perspective projections will preserve relative depth + such that Perspective = dot(Orthogonal, PseudoPerspective). + + >>> P = projection_matrix([0, 0, 0], [1, 0, 0]) + >>> numpy.allclose(P[1:, 1:], numpy.identity(4)[1:, 1:]) + True + >>> point = numpy.random.random(3) - 0.5 + >>> normal = numpy.random.random(3) - 0.5 + >>> direct = numpy.random.random(3) - 0.5 + >>> persp = numpy.random.random(3) - 0.5 + >>> P0 = projection_matrix(point, normal) + >>> P1 = projection_matrix(point, normal, direction=direct) + >>> P2 = projection_matrix(point, normal, perspective=persp) + >>> P3 = projection_matrix(point, normal, perspective=persp, pseudo=True) + >>> is_same_transform(P2, numpy.dot(P0, P3)) + True + >>> P = projection_matrix([3, 0, 0], [1, 1, 0], [1, 0, 0]) + >>> v0 = (numpy.random.rand(4, 5) - 0.5) * 20 + >>> v0[3] = 1 + >>> v1 = numpy.dot(P, v0) + >>> numpy.allclose(v1[1], v0[1]) + True + >>> numpy.allclose(v1[0], 3-v1[1]) + True + + """ + M = numpy.identity(4) + point = numpy.array(point[:3], dtype=numpy.float64, copy=False) + normal = unit_vector(normal[:3]) + if perspective is not None: + # perspective projection + perspective = numpy.array(perspective[:3], dtype=numpy.float64, + copy=False) + M[0, 0] = M[1, 1] = M[2, 2] = numpy.dot(perspective-point, normal) + M[:3, :3] -= numpy.outer(perspective, normal) + if pseudo: + # preserve relative depth + M[:3, :3] -= numpy.outer(normal, normal) + M[:3, 3] = numpy.dot(point, normal) * (perspective+normal) + else: + M[:3, 3] = numpy.dot(point, normal) * perspective + M[3, :3] = -normal + M[3, 3] = numpy.dot(perspective, normal) + elif direction is not None: + # parallel projection + direction = numpy.array(direction[:3], dtype=numpy.float64, copy=False) + scale = numpy.dot(direction, normal) + M[:3, :3] -= numpy.outer(direction, normal) / scale + M[:3, 3] = direction * (numpy.dot(point, normal) / scale) + else: + # orthogonal projection + M[:3, :3] -= numpy.outer(normal, normal) + M[:3, 3] = numpy.dot(point, normal) * normal + return M + + +def projection_from_matrix(matrix, pseudo=False): + """Return projection plane and perspective point from projection matrix. + + Return values are same as arguments for projection_matrix function: + point, normal, direction, perspective, and pseudo. + + >>> point = numpy.random.random(3) - 0.5 + >>> normal = numpy.random.random(3) - 0.5 + >>> direct = numpy.random.random(3) - 0.5 + >>> persp = numpy.random.random(3) - 0.5 + >>> P0 = projection_matrix(point, normal) + >>> result = projection_from_matrix(P0) + >>> P1 = projection_matrix(*result) + >>> is_same_transform(P0, P1) + True + >>> P0 = projection_matrix(point, normal, direct) + >>> result = projection_from_matrix(P0) + >>> P1 = projection_matrix(*result) + >>> is_same_transform(P0, P1) + True + >>> P0 = projection_matrix(point, normal, perspective=persp, pseudo=False) + >>> result = projection_from_matrix(P0, pseudo=False) + >>> P1 = projection_matrix(*result) + >>> is_same_transform(P0, P1) + True + >>> P0 = projection_matrix(point, normal, perspective=persp, pseudo=True) + >>> result = projection_from_matrix(P0, pseudo=True) + >>> P1 = projection_matrix(*result) + >>> is_same_transform(P0, P1) + True + + """ + M = numpy.array(matrix, dtype=numpy.float64, copy=False) + M33 = M[:3, :3] + w, V = numpy.linalg.eig(M) + i = numpy.where(abs(numpy.real(w) - 1.0) < 1e-8)[0] + if not pseudo and len(i): + # point: any eigenvector corresponding to eigenvalue 1 + point = numpy.real(V[:, i[-1]]).squeeze() + point /= point[3] + # direction: unit eigenvector corresponding to eigenvalue 0 + w, V = numpy.linalg.eig(M33) + i = numpy.where(abs(numpy.real(w)) < 1e-8)[0] + if not len(i): + raise ValueError("no eigenvector corresponding to eigenvalue 0") + direction = numpy.real(V[:, i[0]]).squeeze() + direction /= vector_norm(direction) + # normal: unit eigenvector of M33.T corresponding to eigenvalue 0 + w, V = numpy.linalg.eig(M33.T) + i = numpy.where(abs(numpy.real(w)) < 1e-8)[0] + if len(i): + # parallel projection + normal = numpy.real(V[:, i[0]]).squeeze() + normal /= vector_norm(normal) + return point, normal, direction, None, False + else: + # orthogonal projection, where normal equals direction vector + return point, direction, None, None, False + else: + # perspective projection + i = numpy.where(abs(numpy.real(w)) > 1e-8)[0] + if not len(i): + raise ValueError( + "no eigenvector not corresponding to eigenvalue 0") + point = numpy.real(V[:, i[-1]]).squeeze() + point /= point[3] + normal = - M[3, :3] + perspective = M[:3, 3] / numpy.dot(point[:3], normal) + if pseudo: + perspective -= normal + return point, normal, None, perspective, pseudo + + +def clip_matrix(left, right, bottom, top, near, far, perspective=False): + """Return matrix to obtain normalized device coordinates from frustum. + + The frustum bounds are axis-aligned along x (left, right), + y (bottom, top) and z (near, far). + + Normalized device coordinates are in range [-1, 1] if coordinates are + inside the frustum. + + If perspective is True the frustum is a truncated pyramid with the + perspective point at origin and direction along z axis, otherwise an + orthographic canonical view volume (a box). + + Homogeneous coordinates transformed by the perspective clip matrix + need to be dehomogenized (divided by w coordinate). + + >>> frustum = numpy.random.rand(6) + >>> frustum[1] += frustum[0] + >>> frustum[3] += frustum[2] + >>> frustum[5] += frustum[4] + >>> M = clip_matrix(perspective=False, *frustum) + >>> numpy.dot(M, [frustum[0], frustum[2], frustum[4], 1]) + array([-1., -1., -1., 1.]) + >>> numpy.dot(M, [frustum[1], frustum[3], frustum[5], 1]) + array([ 1., 1., 1., 1.]) + >>> M = clip_matrix(perspective=True, *frustum) + >>> v = numpy.dot(M, [frustum[0], frustum[2], frustum[4], 1]) + >>> v / v[3] + array([-1., -1., -1., 1.]) + >>> v = numpy.dot(M, [frustum[1], frustum[3], frustum[4], 1]) + >>> v / v[3] + array([ 1., 1., -1., 1.]) + + """ + if left >= right or bottom >= top or near >= far: + raise ValueError("invalid frustum") + if perspective: + if near <= _EPS: + raise ValueError("invalid frustum: near <= 0") + t = 2.0 * near + M = [[t/(left-right), 0.0, (right+left)/(right-left), 0.0], + [0.0, t/(bottom-top), (top+bottom)/(top-bottom), 0.0], + [0.0, 0.0, (far+near)/(near-far), t*far/(far-near)], + [0.0, 0.0, -1.0, 0.0]] + else: + M = [[2.0/(right-left), 0.0, 0.0, (right+left)/(left-right)], + [0.0, 2.0/(top-bottom), 0.0, (top+bottom)/(bottom-top)], + [0.0, 0.0, 2.0/(far-near), (far+near)/(near-far)], + [0.0, 0.0, 0.0, 1.0]] + return numpy.array(M) + + +def shear_matrix(angle, direction, point, normal): + """Return matrix to shear by angle along direction vector on shear plane. + + The shear plane is defined by a point and normal vector. The direction + vector must be orthogonal to the plane's normal vector. + + A point P is transformed by the shear matrix into P" such that + the vector P-P" is parallel to the direction vector and its extent is + given by the angle of P-P'-P", where P' is the orthogonal projection + of P onto the shear plane. + + >>> angle = (random.random() - 0.5) * 4*math.pi + >>> direct = numpy.random.random(3) - 0.5 + >>> point = numpy.random.random(3) - 0.5 + >>> normal = numpy.cross(direct, numpy.random.random(3)) + >>> S = shear_matrix(angle, direct, point, normal) + >>> numpy.allclose(1, numpy.linalg.det(S)) + True + + """ + normal = unit_vector(normal[:3]) + direction = unit_vector(direction[:3]) + if abs(numpy.dot(normal, direction)) > 1e-6: + raise ValueError("direction and normal vectors are not orthogonal") + angle = math.tan(angle) + M = numpy.identity(4) + M[:3, :3] += angle * numpy.outer(direction, normal) + M[:3, 3] = -angle * numpy.dot(point[:3], normal) * direction + return M + + +def shear_from_matrix(matrix): + """Return shear angle, direction and plane from shear matrix. + + >>> angle = (random.random() - 0.5) * 4*math.pi + >>> direct = numpy.random.random(3) - 0.5 + >>> point = numpy.random.random(3) - 0.5 + >>> normal = numpy.cross(direct, numpy.random.random(3)) + >>> S0 = shear_matrix(angle, direct, point, normal) + >>> angle, direct, point, normal = shear_from_matrix(S0) + >>> S1 = shear_matrix(angle, direct, point, normal) + >>> is_same_transform(S0, S1) + True + + """ + M = numpy.array(matrix, dtype=numpy.float64, copy=False) + M33 = M[:3, :3] + # normal: cross independent eigenvectors corresponding to the eigenvalue 1 + w, V = numpy.linalg.eig(M33) + i = numpy.where(abs(numpy.real(w) - 1.0) < 1e-4)[0] + if len(i) < 2: + raise ValueError("no two linear independent eigenvectors found %s" % w) + V = numpy.real(V[:, i]).squeeze().T + lenorm = -1.0 + for i0, i1 in ((0, 1), (0, 2), (1, 2)): + n = numpy.cross(V[i0], V[i1]) + w = vector_norm(n) + if w > lenorm: + lenorm = w + normal = n + normal /= lenorm + # direction and angle + direction = numpy.dot(M33 - numpy.identity(3), normal) + angle = vector_norm(direction) + direction /= angle + angle = math.atan(angle) + # point: eigenvector corresponding to eigenvalue 1 + w, V = numpy.linalg.eig(M) + i = numpy.where(abs(numpy.real(w) - 1.0) < 1e-8)[0] + if not len(i): + raise ValueError("no eigenvector corresponding to eigenvalue 1") + point = numpy.real(V[:, i[-1]]).squeeze() + point /= point[3] + return angle, direction, point, normal + + +def decompose_matrix(matrix): + """Return sequence of transformations from transformation matrix. + + matrix : array_like + Non-degenerative homogeneous transformation matrix + + Return tuple of: + scale : vector of 3 scaling factors + shear : list of shear factors for x-y, x-z, y-z axes + angles : list of Euler angles about static x, y, z axes + translate : translation vector along x, y, z axes + perspective : perspective partition of matrix + + Raise ValueError if matrix is of wrong type or degenerative. + + >>> T0 = translation_matrix([1, 2, 3]) + >>> scale, shear, angles, trans, persp = decompose_matrix(T0) + >>> T1 = translation_matrix(trans) + >>> numpy.allclose(T0, T1) + True + >>> S = scale_matrix(0.123) + >>> scale, shear, angles, trans, persp = decompose_matrix(S) + >>> scale[0] + 0.123 + >>> R0 = euler_matrix(1, 2, 3) + >>> scale, shear, angles, trans, persp = decompose_matrix(R0) + >>> R1 = euler_matrix(*angles) + >>> numpy.allclose(R0, R1) + True + + """ + M = numpy.array(matrix, dtype=numpy.float64, copy=True).T + if abs(M[3, 3]) < _EPS: + raise ValueError("M[3, 3] is zero") + M /= M[3, 3] + P = M.copy() + P[:, 3] = 0.0, 0.0, 0.0, 1.0 + if not numpy.linalg.det(P): + raise ValueError("matrix is singular") + + scale = numpy.zeros((3, )) + shear = [0.0, 0.0, 0.0] + angles = [0.0, 0.0, 0.0] + + if any(abs(M[:3, 3]) > _EPS): + perspective = numpy.dot(M[:, 3], numpy.linalg.inv(P.T)) + M[:, 3] = 0.0, 0.0, 0.0, 1.0 + else: + perspective = numpy.array([0.0, 0.0, 0.0, 1.0]) + + translate = M[3, :3].copy() + M[3, :3] = 0.0 + + row = M[:3, :3].copy() + scale[0] = vector_norm(row[0]) + row[0] /= scale[0] + shear[0] = numpy.dot(row[0], row[1]) + row[1] -= row[0] * shear[0] + scale[1] = vector_norm(row[1]) + row[1] /= scale[1] + shear[0] /= scale[1] + shear[1] = numpy.dot(row[0], row[2]) + row[2] -= row[0] * shear[1] + shear[2] = numpy.dot(row[1], row[2]) + row[2] -= row[1] * shear[2] + scale[2] = vector_norm(row[2]) + row[2] /= scale[2] + shear[1:] /= scale[2] + + if numpy.dot(row[0], numpy.cross(row[1], row[2])) < 0: + numpy.negative(scale, scale) + numpy.negative(row, row) + + angles[1] = math.asin(-row[0, 2]) + if math.cos(angles[1]): + angles[0] = math.atan2(row[1, 2], row[2, 2]) + angles[2] = math.atan2(row[0, 1], row[0, 0]) + else: + # angles[0] = math.atan2(row[1, 0], row[1, 1]) + angles[0] = math.atan2(-row[2, 1], row[1, 1]) + angles[2] = 0.0 + + return scale, shear, angles, translate, perspective + + +def compose_matrix(scale=None, shear=None, angles=None, translate=None, + perspective=None): + """Return transformation matrix from sequence of transformations. + + This is the inverse of the decompose_matrix function. + + Sequence of transformations: + scale : vector of 3 scaling factors + shear : list of shear factors for x-y, x-z, y-z axes + angles : list of Euler angles about static x, y, z axes + translate : translation vector along x, y, z axes + perspective : perspective partition of matrix + + >>> scale = numpy.random.random(3) - 0.5 + >>> shear = numpy.random.random(3) - 0.5 + >>> angles = (numpy.random.random(3) - 0.5) * (2*math.pi) + >>> trans = numpy.random.random(3) - 0.5 + >>> persp = numpy.random.random(4) - 0.5 + >>> M0 = compose_matrix(scale, shear, angles, trans, persp) + >>> result = decompose_matrix(M0) + >>> M1 = compose_matrix(*result) + >>> is_same_transform(M0, M1) + True + + """ + M = numpy.identity(4) + if perspective is not None: + P = numpy.identity(4) + P[3, :] = perspective[:4] + M = numpy.dot(M, P) + if translate is not None: + T = numpy.identity(4) + T[:3, 3] = translate[:3] + M = numpy.dot(M, T) + if angles is not None: + R = euler_matrix(angles[0], angles[1], angles[2], 'sxyz') + M = numpy.dot(M, R) + if shear is not None: + Z = numpy.identity(4) + Z[1, 2] = shear[2] + Z[0, 2] = shear[1] + Z[0, 1] = shear[0] + M = numpy.dot(M, Z) + if scale is not None: + S = numpy.identity(4) + S[0, 0] = scale[0] + S[1, 1] = scale[1] + S[2, 2] = scale[2] + M = numpy.dot(M, S) + M /= M[3, 3] + return M + + +def orthogonalization_matrix(lengths, angles): + """Return orthogonalization matrix for crystallographic cell coordinates. + + Angles are expected in degrees. + + The de-orthogonalization matrix is the inverse. + + >>> O = orthogonalization_matrix([10, 10, 10], [90, 90, 90]) + >>> numpy.allclose(O[:3, :3], numpy.identity(3, float) * 10) + True + >>> O = orthogonalization_matrix([9.8, 12.0, 15.5], [87.2, 80.7, 69.7]) + >>> numpy.allclose(numpy.sum(O), 43.063229) + True + + """ + a, b, c = lengths + angles = numpy.radians(angles) + sina, sinb, _ = numpy.sin(angles) + cosa, cosb, cosg = numpy.cos(angles) + co = (cosa * cosb - cosg) / (sina * sinb) + return numpy.array([ + [ a*sinb*math.sqrt(1.0-co*co), 0.0, 0.0, 0.0], + [-a*sinb*co, b*sina, 0.0, 0.0], + [ a*cosb, b*cosa, c, 0.0], + [ 0.0, 0.0, 0.0, 1.0]]) + + +def affine_matrix_from_points(v0, v1, shear=True, scale=True, usesvd=True): + """Return affine transform matrix to register two point sets. + + v0 and v1 are shape (ndims, \*) arrays of at least ndims non-homogeneous + coordinates, where ndims is the dimensionality of the coordinate space. + + If shear is False, a similarity transformation matrix is returned. + If also scale is False, a rigid/Euclidean transformation matrix + is returned. + + By default the algorithm by Hartley and Zissermann [15] is used. + If usesvd is True, similarity and Euclidean transformation matrices + are calculated by minimizing the weighted sum of squared deviations + (RMSD) according to the algorithm by Kabsch [8]. + Otherwise, and if ndims is 3, the quaternion based algorithm by Horn [9] + is used, which is slower when using this Python implementation. + + The returned matrix performs rotation, translation and uniform scaling + (if specified). + + >>> v0 = [[0, 1031, 1031, 0], [0, 0, 1600, 1600]] + >>> v1 = [[675, 826, 826, 677], [55, 52, 281, 277]] + >>> affine_matrix_from_points(v0, v1) + array([[ 0.14549, 0.00062, 675.50008], + [ 0.00048, 0.14094, 53.24971], + [ 0. , 0. , 1. ]]) + >>> T = translation_matrix(numpy.random.random(3)-0.5) + >>> R = random_rotation_matrix(numpy.random.random(3)) + >>> S = scale_matrix(random.random()) + >>> M = concatenate_matrices(T, R, S) + >>> v0 = (numpy.random.rand(4, 100) - 0.5) * 20 + >>> v0[3] = 1 + >>> v1 = numpy.dot(M, v0) + >>> v0[:3] += numpy.random.normal(0, 1e-8, 300).reshape(3, -1) + >>> M = affine_matrix_from_points(v0[:3], v1[:3]) + >>> numpy.allclose(v1, numpy.dot(M, v0)) + True + + More examples in superimposition_matrix() + + """ + v0 = numpy.array(v0, dtype=numpy.float64, copy=True) + v1 = numpy.array(v1, dtype=numpy.float64, copy=True) + + ndims = v0.shape[0] + if ndims < 2 or v0.shape[1] < ndims or v0.shape != v1.shape: + raise ValueError("input arrays are of wrong shape or type") + + # move centroids to origin + t0 = -numpy.mean(v0, axis=1) + M0 = numpy.identity(ndims+1) + M0[:ndims, ndims] = t0 + v0 += t0.reshape(ndims, 1) + t1 = -numpy.mean(v1, axis=1) + M1 = numpy.identity(ndims+1) + M1[:ndims, ndims] = t1 + v1 += t1.reshape(ndims, 1) + + if shear: + # Affine transformation + A = numpy.concatenate((v0, v1), axis=0) + u, s, vh = numpy.linalg.svd(A.T) + vh = vh[:ndims].T + B = vh[:ndims] + C = vh[ndims:2*ndims] + t = numpy.dot(C, numpy.linalg.pinv(B)) + t = numpy.concatenate((t, numpy.zeros((ndims, 1))), axis=1) + M = numpy.vstack((t, ((0.0,)*ndims) + (1.0,))) + elif usesvd or ndims != 3: + # Rigid transformation via SVD of covariance matrix + u, s, vh = numpy.linalg.svd(numpy.dot(v1, v0.T)) + # rotation matrix from SVD orthonormal bases + R = numpy.dot(u, vh) + if numpy.linalg.det(R) < 0.0: + # R does not constitute right handed system + R -= numpy.outer(u[:, ndims-1], vh[ndims-1, :]*2.0) + s[-1] *= -1.0 + # homogeneous transformation matrix + M = numpy.identity(ndims+1) + M[:ndims, :ndims] = R + else: + # Rigid transformation matrix via quaternion + # compute symmetric matrix N + xx, yy, zz = numpy.sum(v0 * v1, axis=1) + xy, yz, zx = numpy.sum(v0 * numpy.roll(v1, -1, axis=0), axis=1) + xz, yx, zy = numpy.sum(v0 * numpy.roll(v1, -2, axis=0), axis=1) + N = [[xx+yy+zz, 0.0, 0.0, 0.0], + [yz-zy, xx-yy-zz, 0.0, 0.0], + [zx-xz, xy+yx, yy-xx-zz, 0.0], + [xy-yx, zx+xz, yz+zy, zz-xx-yy]] + # quaternion: eigenvector corresponding to most positive eigenvalue + w, V = numpy.linalg.eigh(N) + q = V[:, numpy.argmax(w)] + q /= vector_norm(q) # unit quaternion + # homogeneous transformation matrix + M = quaternion_matrix(q) + + if scale and not shear: + # Affine transformation; scale is ratio of RMS deviations from centroid + v0 *= v0 + v1 *= v1 + M[:ndims, :ndims] *= math.sqrt(numpy.sum(v1) / numpy.sum(v0)) + + # move centroids back + M = numpy.dot(numpy.linalg.inv(M1), numpy.dot(M, M0)) + M /= M[ndims, ndims] + return M + + +def superimposition_matrix(v0, v1, scale=False, usesvd=True): + """Return matrix to transform given 3D point set into second point set. + + v0 and v1 are shape (3, \*) or (4, \*) arrays of at least 3 points. + + The parameters scale and usesvd are explained in the more general + affine_matrix_from_points function. + + The returned matrix is a similarity or Euclidean transformation matrix. + This function has a fast C implementation in transformations.c. + + >>> v0 = numpy.random.rand(3, 10) + >>> M = superimposition_matrix(v0, v0) + >>> numpy.allclose(M, numpy.identity(4)) + True + >>> R = random_rotation_matrix(numpy.random.random(3)) + >>> v0 = [[1,0,0], [0,1,0], [0,0,1], [1,1,1]] + >>> v1 = numpy.dot(R, v0) + >>> M = superimposition_matrix(v0, v1) + >>> numpy.allclose(v1, numpy.dot(M, v0)) + True + >>> v0 = (numpy.random.rand(4, 100) - 0.5) * 20 + >>> v0[3] = 1 + >>> v1 = numpy.dot(R, v0) + >>> M = superimposition_matrix(v0, v1) + >>> numpy.allclose(v1, numpy.dot(M, v0)) + True + >>> S = scale_matrix(random.random()) + >>> T = translation_matrix(numpy.random.random(3)-0.5) + >>> M = concatenate_matrices(T, R, S) + >>> v1 = numpy.dot(M, v0) + >>> v0[:3] += numpy.random.normal(0, 1e-9, 300).reshape(3, -1) + >>> M = superimposition_matrix(v0, v1, scale=True) + >>> numpy.allclose(v1, numpy.dot(M, v0)) + True + >>> M = superimposition_matrix(v0, v1, scale=True, usesvd=False) + >>> numpy.allclose(v1, numpy.dot(M, v0)) + True + >>> v = numpy.empty((4, 100, 3)) + >>> v[:, :, 0] = v0 + >>> M = superimposition_matrix(v0, v1, scale=True, usesvd=False) + >>> numpy.allclose(v1, numpy.dot(M, v[:, :, 0])) + True + + """ + v0 = numpy.array(v0, dtype=numpy.float64, copy=False)[:3] + v1 = numpy.array(v1, dtype=numpy.float64, copy=False)[:3] + return affine_matrix_from_points(v0, v1, shear=False, + scale=scale, usesvd=usesvd) + + +def euler_matrix(ai, aj, ak, axes='sxyz'): + """Return homogeneous rotation matrix from Euler angles and axis sequence. + + ai, aj, ak : Euler's roll, pitch and yaw angles + axes : One of 24 axis sequences as string or encoded tuple + + >>> R = euler_matrix(1, 2, 3, 'syxz') + >>> numpy.allclose(numpy.sum(R[0]), -1.34786452) + True + >>> R = euler_matrix(1, 2, 3, (0, 1, 0, 1)) + >>> numpy.allclose(numpy.sum(R[0]), -0.383436184) + True + >>> ai, aj, ak = (4*math.pi) * (numpy.random.random(3) - 0.5) + >>> for axes in _AXES2TUPLE.keys(): + ... R = euler_matrix(ai, aj, ak, axes) + >>> for axes in _TUPLE2AXES.keys(): + ... R = euler_matrix(ai, aj, ak, axes) + + """ + try: + firstaxis, parity, repetition, frame = _AXES2TUPLE[axes] + except (AttributeError, KeyError): + _TUPLE2AXES[axes] # validation + firstaxis, parity, repetition, frame = axes + + i = firstaxis + j = _NEXT_AXIS[i+parity] + k = _NEXT_AXIS[i-parity+1] + + if frame: + ai, ak = ak, ai + if parity: + ai, aj, ak = -ai, -aj, -ak + + si, sj, sk = math.sin(ai), math.sin(aj), math.sin(ak) + ci, cj, ck = math.cos(ai), math.cos(aj), math.cos(ak) + cc, cs = ci*ck, ci*sk + sc, ss = si*ck, si*sk + + M = numpy.identity(4) + if repetition: + M[i, i] = cj + M[i, j] = sj*si + M[i, k] = sj*ci + M[j, i] = sj*sk + M[j, j] = -cj*ss+cc + M[j, k] = -cj*cs-sc + M[k, i] = -sj*ck + M[k, j] = cj*sc+cs + M[k, k] = cj*cc-ss + else: + M[i, i] = cj*ck + M[i, j] = sj*sc-cs + M[i, k] = sj*cc+ss + M[j, i] = cj*sk + M[j, j] = sj*ss+cc + M[j, k] = sj*cs-sc + M[k, i] = -sj + M[k, j] = cj*si + M[k, k] = cj*ci + return M + + +def euler_from_matrix(matrix, axes='sxyz'): + """Return Euler angles from rotation matrix for specified axis sequence. + + axes : One of 24 axis sequences as string or encoded tuple + + Note that many Euler angle triplets can describe one matrix. + + >>> R0 = euler_matrix(1, 2, 3, 'syxz') + >>> al, be, ga = euler_from_matrix(R0, 'syxz') + >>> R1 = euler_matrix(al, be, ga, 'syxz') + >>> numpy.allclose(R0, R1) + True + >>> angles = (4*math.pi) * (numpy.random.random(3) - 0.5) + >>> for axes in _AXES2TUPLE.keys(): + ... R0 = euler_matrix(axes=axes, *angles) + ... R1 = euler_matrix(axes=axes, *euler_from_matrix(R0, axes)) + ... if not numpy.allclose(R0, R1): print(axes, "failed") + + """ + try: + firstaxis, parity, repetition, frame = _AXES2TUPLE[axes.lower()] + except (AttributeError, KeyError): + _TUPLE2AXES[axes] # validation + firstaxis, parity, repetition, frame = axes + + i = firstaxis + j = _NEXT_AXIS[i+parity] + k = _NEXT_AXIS[i-parity+1] + + M = numpy.array(matrix, dtype=numpy.float64, copy=False)[:3, :3] + if repetition: + sy = math.sqrt(M[i, j]*M[i, j] + M[i, k]*M[i, k]) + if sy > _EPS: + ax = math.atan2( M[i, j], M[i, k]) + ay = math.atan2( sy, M[i, i]) + az = math.atan2( M[j, i], -M[k, i]) + else: + ax = math.atan2(-M[j, k], M[j, j]) + ay = math.atan2( sy, M[i, i]) + az = 0.0 + else: + cy = math.sqrt(M[i, i]*M[i, i] + M[j, i]*M[j, i]) + if cy > _EPS: + ax = math.atan2( M[k, j], M[k, k]) + ay = math.atan2(-M[k, i], cy) + az = math.atan2( M[j, i], M[i, i]) + else: + ax = math.atan2(-M[j, k], M[j, j]) + ay = math.atan2(-M[k, i], cy) + az = 0.0 + + if parity: + ax, ay, az = -ax, -ay, -az + if frame: + ax, az = az, ax + return ax, ay, az + + +def euler_from_quaternion(quaternion, axes='sxyz'): + """Return Euler angles from quaternion for specified axis sequence. + + >>> angles = euler_from_quaternion([0.99810947, 0.06146124, 0, 0]) + >>> numpy.allclose(angles, [0.123, 0, 0]) + True + + """ + return euler_from_matrix(quaternion_matrix(quaternion), axes) + + +def quaternion_from_euler(ai, aj, ak, axes='sxyz'): + """Return quaternion from Euler angles and axis sequence. + + ai, aj, ak : Euler's roll, pitch and yaw angles + axes : One of 24 axis sequences as string or encoded tuple + + >>> q = quaternion_from_euler(1, 2, 3, 'ryxz') + >>> numpy.allclose(q, [0.435953, 0.310622, -0.718287, 0.444435]) + True + + """ + try: + firstaxis, parity, repetition, frame = _AXES2TUPLE[axes.lower()] + except (AttributeError, KeyError): + _TUPLE2AXES[axes] # validation + firstaxis, parity, repetition, frame = axes + + i = firstaxis + 1 + j = _NEXT_AXIS[i+parity-1] + 1 + k = _NEXT_AXIS[i-parity] + 1 + + if frame: + ai, ak = ak, ai + if parity: + aj = -aj + + ai /= 2.0 + aj /= 2.0 + ak /= 2.0 + ci = math.cos(ai) + si = math.sin(ai) + cj = math.cos(aj) + sj = math.sin(aj) + ck = math.cos(ak) + sk = math.sin(ak) + cc = ci*ck + cs = ci*sk + sc = si*ck + ss = si*sk + + q = numpy.empty((4, )) + if repetition: + q[0] = cj*(cc - ss) + q[i] = cj*(cs + sc) + q[j] = sj*(cc + ss) + q[k] = sj*(cs - sc) + else: + q[0] = cj*cc + sj*ss + q[i] = cj*sc - sj*cs + q[j] = cj*ss + sj*cc + q[k] = cj*cs - sj*sc + if parity: + q[j] *= -1.0 + + return q + + +def quaternion_about_axis(angle, axis): + """Return quaternion for rotation about axis. + + >>> q = quaternion_about_axis(0.123, [1, 0, 0]) + >>> numpy.allclose(q, [0.99810947, 0.06146124, 0, 0]) + True + + """ + q = numpy.array([0.0, axis[0], axis[1], axis[2]]) + qlen = vector_norm(q) + if qlen > _EPS: + q *= math.sin(angle/2.0) / qlen + q[0] = math.cos(angle/2.0) + return q + + +def quaternion_matrix(quaternion): + """Return homogeneous rotation matrix from quaternion. + + >>> M = quaternion_matrix([0.99810947, 0.06146124, 0, 0]) + >>> numpy.allclose(M, rotation_matrix(0.123, [1, 0, 0])) + True + >>> M = quaternion_matrix([1, 0, 0, 0]) + >>> numpy.allclose(M, numpy.identity(4)) + True + >>> M = quaternion_matrix([0, 1, 0, 0]) + >>> numpy.allclose(M, numpy.diag([1, -1, -1, 1])) + True + + """ + q = numpy.array(quaternion, dtype=numpy.float64, copy=True) + n = numpy.dot(q, q) + if n < _EPS: + return numpy.identity(4) + q *= math.sqrt(2.0 / n) + q = numpy.outer(q, q) + return numpy.array([ + [1.0-q[2, 2]-q[3, 3], q[1, 2]-q[3, 0], q[1, 3]+q[2, 0], 0.0], + [ q[1, 2]+q[3, 0], 1.0-q[1, 1]-q[3, 3], q[2, 3]-q[1, 0], 0.0], + [ q[1, 3]-q[2, 0], q[2, 3]+q[1, 0], 1.0-q[1, 1]-q[2, 2], 0.0], + [ 0.0, 0.0, 0.0, 1.0]]) + + +def quaternion_from_matrix(matrix, isprecise=False): + """Return quaternion from rotation matrix. + + If isprecise is True, the input matrix is assumed to be a precise rotation + matrix and a faster algorithm is used. + + >>> q = quaternion_from_matrix(numpy.identity(4), True) + >>> numpy.allclose(q, [1, 0, 0, 0]) + True + >>> q = quaternion_from_matrix(numpy.diag([1, -1, -1, 1])) + >>> numpy.allclose(q, [0, 1, 0, 0]) or numpy.allclose(q, [0, -1, 0, 0]) + True + >>> R = rotation_matrix(0.123, (1, 2, 3)) + >>> q = quaternion_from_matrix(R, True) + >>> numpy.allclose(q, [0.9981095, 0.0164262, 0.0328524, 0.0492786]) + True + >>> R = [[-0.545, 0.797, 0.260, 0], [0.733, 0.603, -0.313, 0], + ... [-0.407, 0.021, -0.913, 0], [0, 0, 0, 1]] + >>> q = quaternion_from_matrix(R) + >>> numpy.allclose(q, [0.19069, 0.43736, 0.87485, -0.083611]) + True + >>> R = [[0.395, 0.362, 0.843, 0], [-0.626, 0.796, -0.056, 0], + ... [-0.677, -0.498, 0.529, 0], [0, 0, 0, 1]] + >>> q = quaternion_from_matrix(R) + >>> numpy.allclose(q, [0.82336615, -0.13610694, 0.46344705, -0.29792603]) + True + >>> R = random_rotation_matrix() + >>> q = quaternion_from_matrix(R) + >>> is_same_transform(R, quaternion_matrix(q)) + True + >>> is_same_quaternion(quaternion_from_matrix(R, isprecise=False), + ... quaternion_from_matrix(R, isprecise=True)) + True + >>> R = euler_matrix(0.0, 0.0, numpy.pi/2.0) + >>> is_same_quaternion(quaternion_from_matrix(R, isprecise=False), + ... quaternion_from_matrix(R, isprecise=True)) + True + + """ + M = numpy.array(matrix, dtype=numpy.float64, copy=False)[:4, :4] + if isprecise: + q = numpy.empty((4, )) + t = numpy.trace(M) + if t > M[3, 3]: + q[0] = t + q[3] = M[1, 0] - M[0, 1] + q[2] = M[0, 2] - M[2, 0] + q[1] = M[2, 1] - M[1, 2] + else: + i, j, k = 0, 1, 2 + if M[1, 1] > M[0, 0]: + i, j, k = 1, 2, 0 + if M[2, 2] > M[i, i]: + i, j, k = 2, 0, 1 + t = M[i, i] - (M[j, j] + M[k, k]) + M[3, 3] + q[i] = t + q[j] = M[i, j] + M[j, i] + q[k] = M[k, i] + M[i, k] + q[3] = M[k, j] - M[j, k] + q = q[[3, 0, 1, 2]] + q *= 0.5 / math.sqrt(t * M[3, 3]) + else: + m00 = M[0, 0] + m01 = M[0, 1] + m02 = M[0, 2] + m10 = M[1, 0] + m11 = M[1, 1] + m12 = M[1, 2] + m20 = M[2, 0] + m21 = M[2, 1] + m22 = M[2, 2] + # symmetric matrix K + K = numpy.array([[m00-m11-m22, 0.0, 0.0, 0.0], + [m01+m10, m11-m00-m22, 0.0, 0.0], + [m02+m20, m12+m21, m22-m00-m11, 0.0], + [m21-m12, m02-m20, m10-m01, m00+m11+m22]]) + K /= 3.0 + # quaternion is eigenvector of K that corresponds to largest eigenvalue + w, V = numpy.linalg.eigh(K) + q = V[[3, 0, 1, 2], numpy.argmax(w)] + if q[0] < 0.0: + numpy.negative(q, q) + return q + + +def quaternion_multiply(quaternion1, quaternion0): + """Return multiplication of two quaternions. + + >>> q = quaternion_multiply([4, 1, -2, 3], [8, -5, 6, 7]) + >>> numpy.allclose(q, [28, -44, -14, 48]) + True + + """ + w0, x0, y0, z0 = quaternion0 + w1, x1, y1, z1 = quaternion1 + return numpy.array([ + -x1*x0 - y1*y0 - z1*z0 + w1*w0, + x1*w0 + y1*z0 - z1*y0 + w1*x0, + -x1*z0 + y1*w0 + z1*x0 + w1*y0, + x1*y0 - y1*x0 + z1*w0 + w1*z0], dtype=numpy.float64) + + +def quaternion_conjugate(quaternion): + """Return conjugate of quaternion. + + >>> q0 = random_quaternion() + >>> q1 = quaternion_conjugate(q0) + >>> q1[0] == q0[0] and all(q1[1:] == -q0[1:]) + True + + """ + q = numpy.array(quaternion, dtype=numpy.float64, copy=True) + numpy.negative(q[1:], q[1:]) + return q + + +def quaternion_inverse(quaternion): + """Return inverse of quaternion. + + >>> q0 = random_quaternion() + >>> q1 = quaternion_inverse(q0) + >>> numpy.allclose(quaternion_multiply(q0, q1), [1, 0, 0, 0]) + True + + """ + q = numpy.array(quaternion, dtype=numpy.float64, copy=True) + numpy.negative(q[1:], q[1:]) + return q / numpy.dot(q, q) + + +def quaternion_real(quaternion): + """Return real part of quaternion. + + >>> quaternion_real([3, 0, 1, 2]) + 3.0 + + """ + return float(quaternion[0]) + + +def quaternion_imag(quaternion): + """Return imaginary part of quaternion. + + >>> quaternion_imag([3, 0, 1, 2]) + array([ 0., 1., 2.]) + + """ + return numpy.array(quaternion[1:4], dtype=numpy.float64, copy=True) + + +def quaternion_slerp(quat0, quat1, fraction, spin=0, shortestpath=True): + """Return spherical linear interpolation between two quaternions. + + >>> q0 = random_quaternion() + >>> q1 = random_quaternion() + >>> q = quaternion_slerp(q0, q1, 0) + >>> numpy.allclose(q, q0) + True + >>> q = quaternion_slerp(q0, q1, 1, 1) + >>> numpy.allclose(q, q1) + True + >>> q = quaternion_slerp(q0, q1, 0.5) + >>> angle = math.acos(numpy.dot(q0, q)) + >>> numpy.allclose(2, math.acos(numpy.dot(q0, q1)) / angle) or \ + numpy.allclose(2, math.acos(-numpy.dot(q0, q1)) / angle) + True + + """ + q0 = unit_vector(quat0[:4]) + q1 = unit_vector(quat1[:4]) + if fraction == 0.0: + return q0 + elif fraction == 1.0: + return q1 + d = numpy.dot(q0, q1) + if abs(abs(d) - 1.0) < _EPS: + return q0 + if shortestpath and d < 0.0: + # invert rotation + d = -d + numpy.negative(q1, q1) + angle = math.acos(d) + spin * math.pi + if abs(angle) < _EPS: + return q0 + isin = 1.0 / math.sin(angle) + q0 *= math.sin((1.0 - fraction) * angle) * isin + q1 *= math.sin(fraction * angle) * isin + q0 += q1 + return q0 + + +def random_quaternion(rand=None): + """Return uniform random unit quaternion. + + rand: array like or None + Three independent random variables that are uniformly distributed + between 0 and 1. + + >>> q = random_quaternion() + >>> numpy.allclose(1, vector_norm(q)) + True + >>> q = random_quaternion(numpy.random.random(3)) + >>> len(q.shape), q.shape[0]==4 + (1, True) + + """ + if rand is None: + rand = numpy.random.rand(3) + else: + assert len(rand) == 3 + r1 = numpy.sqrt(1.0 - rand[0]) + r2 = numpy.sqrt(rand[0]) + pi2 = math.pi * 2.0 + t1 = pi2 * rand[1] + t2 = pi2 * rand[2] + return numpy.array([numpy.cos(t2)*r2, numpy.sin(t1)*r1, + numpy.cos(t1)*r1, numpy.sin(t2)*r2]) + + +def random_rotation_matrix(rand=None): + """Return uniform random rotation matrix. + + rand: array like + Three independent random variables that are uniformly distributed + between 0 and 1 for each returned quaternion. + + >>> R = random_rotation_matrix() + >>> numpy.allclose(numpy.dot(R.T, R), numpy.identity(4)) + True + + """ + return quaternion_matrix(random_quaternion(rand)) + + +class Arcball(object): + """Virtual Trackball Control. + + >>> ball = Arcball() + >>> ball = Arcball(initial=numpy.identity(4)) + >>> ball.place([320, 320], 320) + >>> ball.down([500, 250]) + >>> ball.drag([475, 275]) + >>> R = ball.matrix() + >>> numpy.allclose(numpy.sum(R), 3.90583455) + True + >>> ball = Arcball(initial=[1, 0, 0, 0]) + >>> ball.place([320, 320], 320) + >>> ball.setaxes([1, 1, 0], [-1, 1, 0]) + >>> ball.constrain = True + >>> ball.down([400, 200]) + >>> ball.drag([200, 400]) + >>> R = ball.matrix() + >>> numpy.allclose(numpy.sum(R), 0.2055924) + True + >>> ball.next() + + """ + def __init__(self, initial=None): + """Initialize virtual trackball control. + + initial : quaternion or rotation matrix + + """ + self._axis = None + self._axes = None + self._radius = 1.0 + self._center = [0.0, 0.0] + self._vdown = numpy.array([0.0, 0.0, 1.0]) + self._constrain = False + if initial is None: + self._qdown = numpy.array([1.0, 0.0, 0.0, 0.0]) + else: + initial = numpy.array(initial, dtype=numpy.float64) + if initial.shape == (4, 4): + self._qdown = quaternion_from_matrix(initial) + elif initial.shape == (4, ): + initial /= vector_norm(initial) + self._qdown = initial + else: + raise ValueError("initial not a quaternion or matrix") + self._qnow = self._qpre = self._qdown + + def place(self, center, radius): + """Place Arcball, e.g. when window size changes. + + center : sequence[2] + Window coordinates of trackball center. + radius : float + Radius of trackball in window coordinates. + + """ + self._radius = float(radius) + self._center[0] = center[0] + self._center[1] = center[1] + + def setaxes(self, *axes): + """Set axes to constrain rotations.""" + if axes is None: + self._axes = None + else: + self._axes = [unit_vector(axis) for axis in axes] + + @property + def constrain(self): + """Return state of constrain to axis mode.""" + return self._constrain + + @constrain.setter + def constrain(self, value): + """Set state of constrain to axis mode.""" + self._constrain = bool(value) + + def down(self, point): + """Set initial cursor window coordinates and pick constrain-axis.""" + self._vdown = arcball_map_to_sphere(point, self._center, self._radius) + self._qdown = self._qpre = self._qnow + if self._constrain and self._axes is not None: + self._axis = arcball_nearest_axis(self._vdown, self._axes) + self._vdown = arcball_constrain_to_axis(self._vdown, self._axis) + else: + self._axis = None + + def drag(self, point): + """Update current cursor window coordinates.""" + vnow = arcball_map_to_sphere(point, self._center, self._radius) + if self._axis is not None: + vnow = arcball_constrain_to_axis(vnow, self._axis) + self._qpre = self._qnow + t = numpy.cross(self._vdown, vnow) + if numpy.dot(t, t) < _EPS: + self._qnow = self._qdown + else: + q = [numpy.dot(self._vdown, vnow), t[0], t[1], t[2]] + self._qnow = quaternion_multiply(q, self._qdown) + + def next(self, acceleration=0.0): + """Continue rotation in direction of last drag.""" + q = quaternion_slerp(self._qpre, self._qnow, 2.0+acceleration, False) + self._qpre, self._qnow = self._qnow, q + + def matrix(self): + """Return homogeneous rotation matrix.""" + return quaternion_matrix(self._qnow) + + +def arcball_map_to_sphere(point, center, radius): + """Return unit sphere coordinates from window coordinates.""" + v0 = (point[0] - center[0]) / radius + v1 = (center[1] - point[1]) / radius + n = v0*v0 + v1*v1 + if n > 1.0: + # position outside of sphere + n = math.sqrt(n) + return numpy.array([v0/n, v1/n, 0.0]) + else: + return numpy.array([v0, v1, math.sqrt(1.0 - n)]) + + +def arcball_constrain_to_axis(point, axis): + """Return sphere point perpendicular to axis.""" + v = numpy.array(point, dtype=numpy.float64, copy=True) + a = numpy.array(axis, dtype=numpy.float64, copy=True) + v -= a * numpy.dot(a, v) # on plane + n = vector_norm(v) + if n > _EPS: + if v[2] < 0.0: + numpy.negative(v, v) + v /= n + return v + if a[2] == 1.0: + return numpy.array([1.0, 0.0, 0.0]) + return unit_vector([-a[1], a[0], 0.0]) + + +def arcball_nearest_axis(point, axes): + """Return axis, which arc is nearest to point.""" + point = numpy.array(point, dtype=numpy.float64, copy=False) + nearest = None + mx = -1.0 + for axis in axes: + t = numpy.dot(arcball_constrain_to_axis(point, axis), point) + if t > mx: + nearest = axis + mx = t + return nearest + + +# epsilon for testing whether a number is close to zero +_EPS = numpy.finfo(float).eps * 4.0 + +# axis sequences for Euler angles +_NEXT_AXIS = [1, 2, 0, 1] + +# map axes strings to/from tuples of inner axis, parity, repetition, frame +_AXES2TUPLE = { + 'sxyz': (0, 0, 0, 0), 'sxyx': (0, 0, 1, 0), 'sxzy': (0, 1, 0, 0), + 'sxzx': (0, 1, 1, 0), 'syzx': (1, 0, 0, 0), 'syzy': (1, 0, 1, 0), + 'syxz': (1, 1, 0, 0), 'syxy': (1, 1, 1, 0), 'szxy': (2, 0, 0, 0), + 'szxz': (2, 0, 1, 0), 'szyx': (2, 1, 0, 0), 'szyz': (2, 1, 1, 0), + 'rzyx': (0, 0, 0, 1), 'rxyx': (0, 0, 1, 1), 'ryzx': (0, 1, 0, 1), + 'rxzx': (0, 1, 1, 1), 'rxzy': (1, 0, 0, 1), 'ryzy': (1, 0, 1, 1), + 'rzxy': (1, 1, 0, 1), 'ryxy': (1, 1, 1, 1), 'ryxz': (2, 0, 0, 1), + 'rzxz': (2, 0, 1, 1), 'rxyz': (2, 1, 0, 1), 'rzyz': (2, 1, 1, 1)} + +_TUPLE2AXES = dict((v, k) for k, v in _AXES2TUPLE.items()) + + +def vector_norm(data, axis=None, out=None): + """Return length, i.e. Euclidean norm, of ndarray along axis. + + >>> v = numpy.random.random(3) + >>> n = vector_norm(v) + >>> numpy.allclose(n, numpy.linalg.norm(v)) + True + >>> v = numpy.random.rand(6, 5, 3) + >>> n = vector_norm(v, axis=-1) + >>> numpy.allclose(n, numpy.sqrt(numpy.sum(v*v, axis=2))) + True + >>> n = vector_norm(v, axis=1) + >>> numpy.allclose(n, numpy.sqrt(numpy.sum(v*v, axis=1))) + True + >>> v = numpy.random.rand(5, 4, 3) + >>> n = numpy.empty((5, 3)) + >>> vector_norm(v, axis=1, out=n) + >>> numpy.allclose(n, numpy.sqrt(numpy.sum(v*v, axis=1))) + True + >>> vector_norm([]) + 0.0 + >>> vector_norm([1]) + 1.0 + + """ + data = numpy.array(data, dtype=numpy.float64, copy=True) + if out is None: + if data.ndim == 1: + return math.sqrt(numpy.dot(data, data)) + data *= data + out = numpy.atleast_1d(numpy.sum(data, axis=axis)) + numpy.sqrt(out, out) + return out + else: + data *= data + numpy.sum(data, axis=axis, out=out) + numpy.sqrt(out, out) + + +def unit_vector(data, axis=None, out=None): + """Return ndarray normalized by length, i.e. Euclidean norm, along axis. + + >>> v0 = numpy.random.random(3) + >>> v1 = unit_vector(v0) + >>> numpy.allclose(v1, v0 / numpy.linalg.norm(v0)) + True + >>> v0 = numpy.random.rand(5, 4, 3) + >>> v1 = unit_vector(v0, axis=-1) + >>> v2 = v0 / numpy.expand_dims(numpy.sqrt(numpy.sum(v0*v0, axis=2)), 2) + >>> numpy.allclose(v1, v2) + True + >>> v1 = unit_vector(v0, axis=1) + >>> v2 = v0 / numpy.expand_dims(numpy.sqrt(numpy.sum(v0*v0, axis=1)), 1) + >>> numpy.allclose(v1, v2) + True + >>> v1 = numpy.empty((5, 4, 3)) + >>> unit_vector(v0, axis=1, out=v1) + >>> numpy.allclose(v1, v2) + True + >>> list(unit_vector([])) + [] + >>> list(unit_vector([1])) + [1.0] + + """ + if out is None: + data = numpy.array(data, dtype=numpy.float64, copy=True) + if data.ndim == 1: + data /= math.sqrt(numpy.dot(data, data)) + return data + else: + if out is not data: + out[:] = numpy.array(data, copy=False) + data = out + length = numpy.atleast_1d(numpy.sum(data*data, axis)) + numpy.sqrt(length, length) + if axis is not None: + length = numpy.expand_dims(length, axis) + data /= length + if out is None: + return data + + +def random_vector(size): + """Return array of random doubles in the half-open interval [0.0, 1.0). + + >>> v = random_vector(10000) + >>> numpy.all(v >= 0) and numpy.all(v < 1) + True + >>> v0 = random_vector(10) + >>> v1 = random_vector(10) + >>> numpy.any(v0 == v1) + False + + """ + return numpy.random.random(size) + + +def vector_product(v0, v1, axis=0): + """Return vector perpendicular to vectors. + + >>> v = vector_product([2, 0, 0], [0, 3, 0]) + >>> numpy.allclose(v, [0, 0, 6]) + True + >>> v0 = [[2, 0, 0, 2], [0, 2, 0, 2], [0, 0, 2, 2]] + >>> v1 = [[3], [0], [0]] + >>> v = vector_product(v0, v1) + >>> numpy.allclose(v, [[0, 0, 0, 0], [0, 0, 6, 6], [0, -6, 0, -6]]) + True + >>> v0 = [[2, 0, 0], [2, 0, 0], [0, 2, 0], [2, 0, 0]] + >>> v1 = [[0, 3, 0], [0, 0, 3], [0, 0, 3], [3, 3, 3]] + >>> v = vector_product(v0, v1, axis=1) + >>> numpy.allclose(v, [[0, 0, 6], [0, -6, 0], [6, 0, 0], [0, -6, 6]]) + True + + """ + return numpy.cross(v0, v1, axis=axis) + + +def angle_between_vectors(v0, v1, directed=True, axis=0): + """Return angle between vectors. + + If directed is False, the input vectors are interpreted as undirected axes, + i.e. the maximum angle is pi/2. + + >>> a = angle_between_vectors([1, -2, 3], [-1, 2, -3]) + >>> numpy.allclose(a, math.pi) + True + >>> a = angle_between_vectors([1, -2, 3], [-1, 2, -3], directed=False) + >>> numpy.allclose(a, 0) + True + >>> v0 = [[2, 0, 0, 2], [0, 2, 0, 2], [0, 0, 2, 2]] + >>> v1 = [[3], [0], [0]] + >>> a = angle_between_vectors(v0, v1) + >>> numpy.allclose(a, [0, 1.5708, 1.5708, 0.95532]) + True + >>> v0 = [[2, 0, 0], [2, 0, 0], [0, 2, 0], [2, 0, 0]] + >>> v1 = [[0, 3, 0], [0, 0, 3], [0, 0, 3], [3, 3, 3]] + >>> a = angle_between_vectors(v0, v1, axis=1) + >>> numpy.allclose(a, [1.5708, 1.5708, 1.5708, 0.95532]) + True + + """ + v0 = numpy.array(v0, dtype=numpy.float64, copy=False) + v1 = numpy.array(v1, dtype=numpy.float64, copy=False) + dot = numpy.sum(v0 * v1, axis=axis) + dot /= vector_norm(v0, axis=axis) * vector_norm(v1, axis=axis) + return numpy.arccos(dot if directed else numpy.fabs(dot)) + + +def inverse_matrix(matrix): + """Return inverse of square transformation matrix. + + >>> M0 = random_rotation_matrix() + >>> M1 = inverse_matrix(M0.T) + >>> numpy.allclose(M1, numpy.linalg.inv(M0.T)) + True + >>> for size in range(1, 7): + ... M0 = numpy.random.rand(size, size) + ... M1 = inverse_matrix(M0) + ... if not numpy.allclose(M1, numpy.linalg.inv(M0)): print(size) + + """ + return numpy.linalg.inv(matrix) + + +def concatenate_matrices(*matrices): + """Return concatenation of series of transformation matrices. + + >>> M = numpy.random.rand(16).reshape((4, 4)) - 0.5 + >>> numpy.allclose(M, concatenate_matrices(M)) + True + >>> numpy.allclose(numpy.dot(M, M.T), concatenate_matrices(M, M.T)) + True + + """ + M = numpy.identity(4) + for i in matrices: + M = numpy.dot(M, i) + return M + + +def is_same_transform(matrix0, matrix1): + """Return True if two matrices perform same transformation. + + >>> is_same_transform(numpy.identity(4), numpy.identity(4)) + True + >>> is_same_transform(numpy.identity(4), random_rotation_matrix()) + False + + """ + matrix0 = numpy.array(matrix0, dtype=numpy.float64, copy=True) + matrix0 /= matrix0[3, 3] + matrix1 = numpy.array(matrix1, dtype=numpy.float64, copy=True) + matrix1 /= matrix1[3, 3] + return numpy.allclose(matrix0, matrix1) + + +def is_same_quaternion(q0, q1): + """Return True if two quaternions are equal.""" + q0 = numpy.array(q0) + q1 = numpy.array(q1) + return numpy.allclose(q0, q1) or numpy.allclose(q0, -q1) + + +def _import_module(name, package=None, warn=True, prefix='_py_', ignore='_'): + """Try import all public attributes from module into global namespace. + + Existing attributes with name clashes are renamed with prefix. + Attributes starting with underscore are ignored by default. + + Return True on successful import. + + """ + import warnings + from importlib import import_module + try: + if not package: + module = import_module(name) + else: + module = import_module('.' + name, package=package) + except ImportError: + if warn: + warnings.warn("failed to import module %s" % name) + else: + for attr in dir(module): + if ignore and attr.startswith(ignore): + continue + if prefix: + if attr in globals(): + globals()[prefix + attr] = globals()[attr] + elif warn: + warnings.warn("no Python implementation of " + attr) + globals()[attr] = getattr(module, attr) + return True + + +_import_module('_transformations', warn=False) + +if __name__ == "__main__": + import doctest + import random # noqa: used in doctests + numpy.set_printoptions(suppress=True, precision=5) + doctest.testmod() \ No newline at end of file diff --git a/utils/zfilter.py b/utils/zfilter.py new file mode 100644 index 0000000..7d61338 --- /dev/null +++ b/utils/zfilter.py @@ -0,0 +1,74 @@ +import numpy as np + +# from https://github.com/joschu/modular_rl +# http://www.johndcook.com/blog/standard_deviation/ + + +class RunningStat(object): + def __init__(self, shape): + self._n = 0 + self._M = np.zeros(shape) + self._S = np.zeros(shape) + + def push(self, x): + x = np.asarray(x) + assert x.shape == self._M.shape + self._n += 1 + if self._n == 1: + self._M[...] = x + else: + oldM = self._M.copy() + self._M[...] = oldM + (x - oldM) / self._n + self._S[...] = self._S + (x - oldM) * (x - self._M) + + @property + def n(self): + return self._n + + @property + def mean(self): + return self._M + + @property + def var(self): + return self._S / (self._n - 1) if self._n > 1 else np.square(self._M) + + @property + def std(self): + return np.sqrt(self.var) + + @property + def shape(self): + return self._M.shape + + +class ZFilter: + """ + y = (x-mean)/std + using running estimates of mean,std + """ + + def __init__(self, shape, demean=True, destd=True, clip=10.0): + self.demean = demean + self.destd = destd + self.clip = clip + + self.rs = RunningStat(shape) + + def __call__(self, x, update=True): + if update: + self.rs.push(x) + if self.demean: + x = x - self.rs.mean + if self.destd: + x = x / (self.rs.std + 1e-8) + if self.clip: + x = np.clip(x, -self.clip, self.clip) + return x + + def set_mean_std(self, mean, std, n): + self.rs._n = n + self.rs._M[...] = mean + self.rs._S[...] = std + +