├── .gitignore ├── COPYING ├── data ├── 2Dletters │ ├── A.npy │ ├── A_2.npy │ ├── B.npy │ ├── B_2.npy │ ├── C.npy │ ├── C_2.npy │ ├── D.npy │ ├── D_2.npy │ ├── E.npy │ ├── E_2.npy │ ├── F.npy │ ├── F_2.npy │ ├── G.npy │ ├── G_2.npy │ ├── H.npy │ ├── H_2.npy │ ├── I.npy │ ├── I_2.npy │ ├── J.npy │ ├── J_2.npy │ ├── K.npy │ ├── K_2.npy │ ├── L.npy │ ├── L_2.npy │ ├── M.npy │ ├── M_2.npy │ ├── N.npy │ ├── N_2.npy │ ├── O.npy │ ├── O_2.npy │ ├── P.npy │ ├── P_2.npy │ ├── Q.npy │ ├── Q_2.npy │ ├── R.npy │ ├── R_2.npy │ ├── S.npy │ ├── S_2.npy │ ├── T.npy │ ├── T_2.npy │ ├── U.npy │ ├── U_2.npy │ ├── V.npy │ ├── V_2.npy │ ├── W.npy │ ├── W_2.npy │ ├── X.npy │ ├── X_2.npy │ ├── Y.npy │ ├── Y_2.npy │ ├── Z.npy │ ├── Z_2.npy │ └── qx │ │ └── SV.npy └── urdf │ ├── panda_arm_gripper.urdf │ └── talos_reduced.urdf ├── notebooks ├── gamp │ ├── gamp_panda.ipynb │ ├── gamp_time_dependant_dynamics_learning_ensemble.ipynb │ └── gamp_time_dependant_mvn_discriminator.ipynb └── kinematic │ ├── talos_configuration_svi_hierarchy.ipynb │ └── test_franka_urdf.ipynb ├── publications ├── pignat_corl2020.pdf └── pignat_icra2020.pdf ├── readme.md ├── setup.py └── tf_robot_learning ├── __init__.py ├── control ├── __init__.py ├── lqr.py ├── rollout │ ├── __init__.py │ └── rollout.py └── utils.py ├── distributions ├── __init__.py ├── approx │ ├── __init__.py │ └── variational_gmm.py ├── mixture_models │ ├── __init__.py │ ├── gmm_ml.py │ └── moe.py ├── mvn.py ├── poe.py ├── promp.py └── soft_uniform.py ├── kinematic ├── __init__.py ├── chain.py ├── frame.py ├── joint.py ├── rotation.py ├── segment.py ├── urdf_utils.py └── utils │ ├── __init__.py │ ├── import_pykdl.py │ ├── layout.py │ ├── plot_utils.py │ ├── pykdl_utils.py │ ├── tf_utils.py │ └── urdf_parser_py │ ├── __init__.py │ ├── sdf.py │ ├── urdf.py │ └── xml_reflection │ ├── __init__.py │ ├── basics.py │ └── core.py ├── nn ├── __init__.py ├── density_mlp.py ├── inv_net.py └── mlp.py ├── planar_robots ├── __init__.py ├── bimanual_robot.py ├── n_joint.py ├── robot.py ├── three_joint.py └── two_joint.py ├── policy ├── __init__.py ├── poe_policy.py └── policy.py └── utils ├── __init__.py ├── basis_utils.py ├── data_utils.py ├── img_utils.py ├── param_utils.py ├── plot_utils.py └── tf_utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Created by .ignore support plugin (hsz.mobi) 2 | /tf_robot_learning.egg-info/ 3 | /notebooks/gamp/.ipynb_checkpoints/ 4 | /notebooks/kinematic/.ipynb_checkpoints/ 5 | -------------------------------------------------------------------------------- /data/2Dletters/A.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/A.npy -------------------------------------------------------------------------------- /data/2Dletters/A_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/A_2.npy -------------------------------------------------------------------------------- /data/2Dletters/B.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/B.npy -------------------------------------------------------------------------------- /data/2Dletters/B_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/B_2.npy -------------------------------------------------------------------------------- /data/2Dletters/C.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/C.npy -------------------------------------------------------------------------------- /data/2Dletters/C_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/C_2.npy -------------------------------------------------------------------------------- /data/2Dletters/D.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/D.npy -------------------------------------------------------------------------------- /data/2Dletters/D_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/D_2.npy -------------------------------------------------------------------------------- /data/2Dletters/E.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/E.npy -------------------------------------------------------------------------------- /data/2Dletters/E_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/E_2.npy -------------------------------------------------------------------------------- /data/2Dletters/F.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/F.npy -------------------------------------------------------------------------------- /data/2Dletters/F_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/F_2.npy -------------------------------------------------------------------------------- /data/2Dletters/G.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/G.npy -------------------------------------------------------------------------------- /data/2Dletters/G_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/G_2.npy -------------------------------------------------------------------------------- /data/2Dletters/H.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/H.npy -------------------------------------------------------------------------------- /data/2Dletters/H_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/H_2.npy -------------------------------------------------------------------------------- /data/2Dletters/I.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/I.npy -------------------------------------------------------------------------------- /data/2Dletters/I_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/I_2.npy -------------------------------------------------------------------------------- /data/2Dletters/J.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/J.npy -------------------------------------------------------------------------------- /data/2Dletters/J_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/J_2.npy -------------------------------------------------------------------------------- /data/2Dletters/K.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/K.npy -------------------------------------------------------------------------------- /data/2Dletters/K_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/K_2.npy -------------------------------------------------------------------------------- /data/2Dletters/L.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/L.npy -------------------------------------------------------------------------------- /data/2Dletters/L_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/L_2.npy -------------------------------------------------------------------------------- /data/2Dletters/M.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/M.npy -------------------------------------------------------------------------------- /data/2Dletters/M_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/M_2.npy -------------------------------------------------------------------------------- /data/2Dletters/N.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/N.npy -------------------------------------------------------------------------------- /data/2Dletters/N_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/N_2.npy -------------------------------------------------------------------------------- /data/2Dletters/O.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/O.npy -------------------------------------------------------------------------------- /data/2Dletters/O_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/O_2.npy -------------------------------------------------------------------------------- /data/2Dletters/P.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/P.npy -------------------------------------------------------------------------------- /data/2Dletters/P_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/P_2.npy -------------------------------------------------------------------------------- /data/2Dletters/Q.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/Q.npy -------------------------------------------------------------------------------- /data/2Dletters/Q_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/Q_2.npy -------------------------------------------------------------------------------- /data/2Dletters/R.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/R.npy -------------------------------------------------------------------------------- /data/2Dletters/R_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/R_2.npy -------------------------------------------------------------------------------- /data/2Dletters/S.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/S.npy -------------------------------------------------------------------------------- /data/2Dletters/S_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/S_2.npy -------------------------------------------------------------------------------- /data/2Dletters/T.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/T.npy -------------------------------------------------------------------------------- /data/2Dletters/T_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/T_2.npy -------------------------------------------------------------------------------- /data/2Dletters/U.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/U.npy -------------------------------------------------------------------------------- /data/2Dletters/U_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/U_2.npy -------------------------------------------------------------------------------- /data/2Dletters/V.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/V.npy -------------------------------------------------------------------------------- /data/2Dletters/V_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/V_2.npy -------------------------------------------------------------------------------- /data/2Dletters/W.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/W.npy -------------------------------------------------------------------------------- /data/2Dletters/W_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/W_2.npy -------------------------------------------------------------------------------- /data/2Dletters/X.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/X.npy -------------------------------------------------------------------------------- /data/2Dletters/X_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/X_2.npy -------------------------------------------------------------------------------- /data/2Dletters/Y.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/Y.npy -------------------------------------------------------------------------------- /data/2Dletters/Y_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/Y_2.npy -------------------------------------------------------------------------------- /data/2Dletters/Z.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/Z.npy -------------------------------------------------------------------------------- /data/2Dletters/Z_2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/Z_2.npy -------------------------------------------------------------------------------- /data/2Dletters/qx/SV.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/data/2Dletters/qx/SV.npy -------------------------------------------------------------------------------- /data/urdf/panda_arm_gripper.urdf: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | 102 | 103 | 104 | 105 | 106 | 107 | 108 | 109 | 110 | 111 | 112 | 113 | 114 | 115 | 116 | 117 | 118 | 119 | 120 | 121 | 122 | 123 | 124 | 125 | 126 | 127 | 128 | 129 | 130 | 131 | 132 | 133 | 134 | 135 | 136 | 137 | 138 | 139 | 140 | 141 | 142 | 143 | 144 | 145 | 146 | 147 | 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | 159 | 160 | 161 | 162 | 163 | 164 | 165 | 166 | 167 | 168 | 169 | 170 | 171 | 172 | 173 | 174 | 175 | 176 | 177 | 178 | 179 | 180 | 181 | 182 | 183 | 184 | 185 | 186 | 187 | 188 | 189 | 190 | 191 | 192 | 193 | 194 | 195 | 196 | 197 | 198 | 199 | 200 | 201 | 202 | 203 | 204 | 205 | 206 | 207 | 208 | 209 | 210 | 211 | 212 | 213 | 214 | 215 | 216 | 217 | 218 | 219 | 220 | 221 | 222 | 223 | 224 | 225 | 226 | 227 | 228 | 229 | 230 | 231 | 232 | 233 | 234 | 235 | 236 | 237 | 238 | 239 | 240 | 241 | 242 | 243 | 244 | 245 | 246 | 247 | 248 | 249 | 250 | 251 | 252 | 253 | 254 | 255 | 256 | 257 | 258 | 259 | 260 | 261 | 262 | 263 | 264 | 265 | 266 | 267 | 268 | 269 | 270 | 271 | 272 | 273 | 274 | 275 | 276 | 277 | 278 | 279 | 280 | 281 | 282 | 283 | 284 | 285 | 286 | 287 | 288 | 289 | 290 | 291 | 292 | -------------------------------------------------------------------------------- /publications/pignat_corl2020.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/publications/pignat_corl2020.pdf -------------------------------------------------------------------------------- /publications/pignat_icra2020.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/publications/pignat_icra2020.pdf -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # Tensorflow robot learning library 2 | 3 | Author : Emmanuel Pignat, emmanuel.pignat@gmail.com 4 | 5 | If you find these codes useful, you can acknowledge the authors by citing the following paper [(Click here to download pdf)](https://github.com/emmanuelpignat/tf_robot_learning/raw/master/publications/pignat_corl2020.pdf): 6 | 7 | @inproceedings{Pignat20CORL, 8 | author="Pignat, E. and Girgin, H. and Calinon, S.", 9 | title="Generative Adversarial Training of Product of Policies for Robust and Adaptive Movement Primitives", 10 | booktitle="Proc.\ Conference on Robot Learning ({CoRL})", 11 | year="2020", 12 | pages="" 13 | } 14 | 15 | ## How to install 16 | 17 | Is compatible with python 2 and 3, tensorflow 1 and 2 18 | 19 | git clone https://gitlab.idiap.ch/rli/tf_robot_learning.git 20 | 21 | cd tf_robot_learning 22 | 23 | pip install -e . 24 | 25 | ### Install jupyter kernels 26 | 27 | python2 -m pip install ipykernel 28 | python2 -m ipykernel install --user 29 | 30 | python3 -m pip install ipykernel 31 | python3 -m ipykernel install --user 32 | 33 | ## Notebooks 34 | 35 | cd [...]/tf_robot_learning/notebooks 36 | jupyter notebook 37 | 38 | Then navigate through folders and click on desired notebook. 39 | 40 | | Filename | Description | 41 | |----------|-------------| 42 | | gamp/gamp_time_dependant_mvn_discriminator.ipynb | Simplest example. | 43 | | gamp/gamp_time_dependant_dynamics_learning_ensemble.ipynb | Learning NN dynamics with ensemble network.| 44 | | gamp/gamp_panda.ipynb | Acceleration PoE policy on Panda robot.| 45 | 46 | 47 | 48 | ### URDF PARSER 49 | 50 | The folder *tf_robot_learning/kinematic/utils/urdf_parser_py* is adapted from https://github.com/ros/urdf_parser_py. 51 | The authors of this code are: 52 | 53 | - Thomas Moulard - urdfpy implementation, integration 54 | - David Lu - urdf_python implementation, integration 55 | - Kelsey Hawkins - urdf_parser_python implementation, integration 56 | - Antonio El Khoury - bugfixes 57 | - Eric Cousineau - reflection update -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from setuptools import setup, find_packages 21 | import sys 22 | 23 | 24 | # Setup for Python3 25 | setup(name='tf_robot_learning', 26 | version='0.1', 27 | description='Tensorflow robotic toolbox', 28 | url='', 29 | author='Emmanuel Pignat', 30 | author_email='emmanuel.pignat@gmail.com', 31 | license='MIT', 32 | packages=find_packages(), 33 | install_requires = ['numpy', 'matplotlib','jupyter', 'tensorflow', 'tensorflow_probability', 'pyyaml', 'lxml'], 34 | zip_safe=False) 35 | -------------------------------------------------------------------------------- /tf_robot_learning/__init__.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from . import control 21 | from . import kinematic 22 | from . import distributions 23 | from . import utils 24 | from .utils import param as p 25 | from . import planar_robots 26 | from . import nn 27 | from . import policy 28 | 29 | from .utils import tf_utils as tf 30 | 31 | datapath = __path__[0][:-17] + 'data' 32 | -------------------------------------------------------------------------------- /tf_robot_learning/control/__init__.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from . import utils 21 | from . import lqr 22 | from . import rollout -------------------------------------------------------------------------------- /tf_robot_learning/control/lqr.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | from tensorflow_probability import distributions as ds 22 | from .rollout import make_rollout_mvn, make_rollout_samples, make_rollout_autonomous_mvn 23 | 24 | def lqr_cost(Q, R, z=None, seq=None): 25 | def cost(x, u, Q, R, z=None, seq=None): 26 | """ 27 | 28 | :param x: [horizon, xi_dim] or [horizon, batch_size, xi_dim] 29 | :param u: [horizon, u_dim] or [horizon, batch_size, u_dim] 30 | :param Q: 31 | :param R: 32 | :param z: 33 | :param seq: 34 | :return: 35 | """ 36 | if x.shape.ndims == 2: 37 | if z is None: 38 | y = x 39 | elif z.shape.ndims == 2: 40 | if seq is not None: 41 | z = tf.gather(z, seq) 42 | y = z - x 43 | elif z.shape.ndims == 1: 44 | y = z[None] - x 45 | else: 46 | raise ValueError("Unknown target z") 47 | if x.shape.ndims == 3: 48 | if z is None: 49 | y = x 50 | elif z.shape.ndims == 2: 51 | if seq is not None: 52 | z = tf.gather(z, seq) 53 | y = z[:, None] - x 54 | elif z.shape.ndims == 1: 55 | y = z[None, None] - x 56 | else: 57 | raise ValueError("Unknown target z") 58 | 59 | if x.shape.ndims == 2: 60 | if Q.shape.ndims == 3: 61 | if seq is not None: 62 | Q = tf.gather(Q, seq) 63 | state_cost = tf.reduce_sum(y * tf.einsum('aij,aj->ai', Q, y)) 64 | else: 65 | state_cost = tf.reduce_sum(y * tf.einsum('ij,aj->ai', Q, y)) 66 | if x.shape.ndims == 3: 67 | if Q.shape.ndims == 3: 68 | if seq is not None: 69 | Q = tf.gather(Q, seq) 70 | state_cost = tf.reduce_sum(y * tf.einsum('aij,abj->abi', Q, y), axis=(0, 2)) 71 | else: 72 | state_cost = tf.reduce_sum(y * tf.einsum('ij,abj->abi', Q, y), axis=(0, 2)) 73 | 74 | if u.shape.ndims == 2: 75 | if R.shape.ndims == 3: 76 | control_cost = tf.reduce_sum(u * tf.einsum('aij,aj->ai', R, u)) 77 | else: 78 | control_cost = tf.reduce_sum(u * tf.einsum('ij,aj->ai', R, u)) 79 | if u.shape.ndims == 3: 80 | if R.shape.ndims == 3: 81 | control_cost = tf.reduce_sum(u * tf.einsum('aij,abj->abi', R, u), axis=(0, 2)) 82 | else: 83 | control_cost = tf.reduce_sum(u * tf.einsum('ij,abj->abi', R, u), axis=(0, 2)) 84 | 85 | return control_cost + state_cost 86 | 87 | return lambda x, u: cost(x, u, Q, R, z, seq) 88 | 89 | def lqr(A, B, Q, R, z=None, horizon=None, seq=None): 90 | """ 91 | 92 | http://web.mst.edu/~bohner/papers/tlqtots.pdf 93 | 94 | 95 | :param A: [x_dim, x_dim] 96 | :param B: [x_dim, u_dim] 97 | :param Q: [horizon, x_dim, x_dim] or [x_dim, x_dim] 98 | :param R: [horizon, u_dim, u_dim] or [u_dim, u_dim] 99 | :param z: [horizon, x_dim] or [x_dim] 100 | :param seq: 101 | :param horizon: int or tf.Tensor() 102 | :return: 103 | """ 104 | u_dim = B.shape[-1] 105 | x_dim = B.shape[0] 106 | 107 | if horizon is None: 108 | assert Q.shape.ndims == 3 and Q.shape[0] is not None, \ 109 | "If horizon is not specified, Q should be of rank 3 with first dimension specified" 110 | 111 | horizon = Q.shape[0] 112 | 113 | if R.shape.ndims == 3: 114 | get_R = lambda i: R[i] 115 | else: 116 | get_R = lambda i: R 117 | 118 | if Q.shape.ndims == 3: 119 | if seq is None: 120 | get_Q = lambda i: Q[i] 121 | else: 122 | get_Q = lambda i: Q[seq[i]] 123 | else: 124 | get_Q = lambda i: Q 125 | 126 | if z is None: 127 | get_z = lambda i: tf.zeros(x_dim) 128 | elif z.shape.ndims == 2: 129 | if seq is None: 130 | get_z = lambda i: z[i] 131 | else: 132 | get_z = lambda i: z[seq[i]] 133 | else: 134 | get_z = lambda i: z 135 | 136 | ### Init empty 137 | _Q = tf.zeros((0, u_dim, u_dim)) 138 | _K = tf.zeros((0, u_dim, x_dim)) 139 | _Kv = tf.zeros((0, u_dim, x_dim)) 140 | 141 | S = get_Q(-1)[None] 142 | v = tf.matmul(get_Q(-1), get_z(-1)[:, None])[:, 0][None] 143 | 144 | i0 = tf.constant(0) # counter 145 | 146 | c = lambda i, S, v, _Kv, _K, _Q: tf.less(i, horizon) # condition 147 | 148 | 149 | def pred(i, S, v, _Kv, _K, _Q): 150 | prev_Q = tf.linalg.inv(get_R(-i) + tf.matmul(tf.matmul(B, S[0], transpose_a=True), B)) 151 | prev_Kv = tf.matmul(prev_Q, B, transpose_b=True) # already wrong 152 | prev_K = tf.matmul(tf.matmul(prev_Kv, S[0]), A) # are already wrong 153 | 154 | AmBK = A - tf.matmul(B, prev_K) 155 | 156 | prev_S = tf.matmul(tf.matmul(A, S[0], transpose_a=True), AmBK) + get_Q(-i) 157 | prev_v = tf.matmul(AmBK, v[0][:, None], transpose_a=True)[:, 0] + \ 158 | tf.matmul(get_Q(-i), get_z(-i)[:, None])[:, 0] 159 | 160 | return tf.add(i, 1), tf.concat([prev_S[None], S], axis=0), tf.concat([prev_v[None], v], 161 | axis=0), \ 162 | tf.concat([prev_Kv[None], _Kv], axis=0), \ 163 | tf.concat([prev_K[None], _K], axis=0), tf.concat([prev_Q[None], _Q], axis=0) 164 | 165 | 166 | _, Ss, vs, _Kvs, _Ks, _Qs = tf.while_loop( 167 | c, pred, loop_vars=[i0, S, v, _Kv, _K, _Q], shape_invariants= 168 | [i0.get_shape(), tf.TensorShape([None, x_dim, x_dim]), 169 | tf.TensorShape([None, x_dim]), tf.TensorShape([None, u_dim, x_dim]), 170 | tf.TensorShape([None, u_dim, x_dim]), tf.TensorShape([None, u_dim, u_dim]) 171 | ] 172 | ) 173 | 174 | return Ss, vs, _Kvs, _Ks, _Qs 175 | 176 | class LQRPolicy(object): 177 | def __init__(self, A, B, Q, R, z=None, horizon=None, seq=None): 178 | self._Ss, self._vs, self._Kvs, self._Ks, self._Qs = lqr( 179 | A, B, Q, R, z=z, horizon=horizon, seq=seq) 180 | 181 | self.horizon = horizon 182 | self.A = A 183 | self.B = B 184 | 185 | 186 | def get_u_mvn(self, xi, i=0): 187 | xi_loc, xi_cov = xi 188 | 189 | u_loc = -tf.linalg.matvec(self._Ks[i], xi_loc) + \ 190 | tf.linalg.matvec(self._Kvs[i], self._vs[i + 1]) 191 | 192 | u_cov = self._Qs[i] + tf.linalg.matmul( 193 | tf.linalg.matmul(-self._Ks[i], xi_cov), -self._Ks[i], transpose_b=True) 194 | 195 | return u_loc, u_cov 196 | 197 | def f_mvn(self, xi, u, t=0): 198 | """ 199 | xi: [batch_size, xi_dim] 200 | u: [batch_size, u_dim] 201 | t : int 202 | return xi : [batch_size, xi_dim] 203 | """ 204 | u_loc, u_cov = u 205 | xi_loc, xi_cov = xi 206 | 207 | xi_n_loc = tf.linalg.matvec(self.A, xi_loc) + tf.linalg.matvec(self.B, u_loc) 208 | xi_n_cov = tf.linalg.matmul(tf.linalg.matmul(self.A, xi_cov), self.A, transpose_b=True) +\ 209 | tf.linalg.matmul(tf.linalg.matmul(self.B, u_cov), self.B , transpose_b=True) 210 | 211 | return xi_n_loc, xi_n_cov 212 | 213 | def make_rollout_mvn(self, p_xi0, return_ds=True): 214 | 215 | def f_mvn(xi, i=0): 216 | xi_loc, xi_cov = xi 217 | 218 | xi_n_loc = tf.linalg.matvec(self.A, xi_loc) + tf.linalg.matvec( 219 | self.B, -tf.linalg.matvec(self._Ks[i], xi_loc) + 220 | tf.linalg.matvec(self._Kvs[i], self._vs[i + 1])) 221 | 222 | C = self.A - tf.linalg.matmul(self.B, self._Ks[i]) 223 | xi_n_cov = tf.linalg.matmul(tf.linalg.matmul(C, xi_cov), C, transpose_b=True) +\ 224 | tf.linalg.matmul(tf.linalg.matmul(self.B, self._Qs[i]), self.B, transpose_b=True) 225 | 226 | return xi_n_loc, xi_n_cov 227 | 228 | return make_rollout_autonomous_mvn( 229 | p_xi0, f_mvn, T=self.horizon, return_ds=return_ds 230 | ) 231 | 232 | def make_rollout_samples(self, p_xi0, n=10): 233 | xi0_samples = p_xi0.sample(n) 234 | 235 | def f(xi, u): 236 | return tf.linalg.LinearOperatorFullMatrix(self.A).matvec(xi) + \ 237 | tf.linalg.LinearOperatorFullMatrix(self.B).matvec(u) 238 | 239 | return make_rollout_samples( 240 | xi0_samples, f, self.get_u_samples, self.B.shape[-1], T=self.horizon, 241 | batch_shape=n 242 | ) 243 | 244 | def make_rollout(self, p_xi0, n=1): 245 | xi0_samples = p_xi0 246 | 247 | def f(xi, u): 248 | return tf.linalg.LinearOperatorFullMatrix(self.A).matvec(xi) + \ 249 | tf.linalg.LinearOperatorFullMatrix(self.B).matvec(u) 250 | 251 | return make_rollout_samples( 252 | xi0_samples, f, self.get_u, self.B.shape[-1], T=self.horizon, 253 | batch_shape=n 254 | ) 255 | 256 | def get_u(self, xi, i=0): 257 | return -tf.einsum('ij,aj->ai', self._Ks[i], xi) + tf.linalg.matvec(self._Kvs[i], self._vs[i+1])[None] 258 | 259 | def log_prob(self, us, xis, i=None): 260 | """ 261 | 262 | :param xis: [horizon + 1, batch_size, xi_dim] 263 | :param us: [horizon, batch_size, xi_dim] 264 | :param i: 265 | :return: 266 | """ 267 | if i is not None: 268 | raise NotImplementedError 269 | 270 | # compute u for each timestep and batch b, as size [horzon, batch_size, u_dim] 271 | u_locs = -tf.einsum('aij,abj->abi', self._Ks[:self.horizon], xis[:self.horizon]) + \ 272 | tf.einsum('aij,aj->ai', self._Kvs[:self.horizon], self._vs[1:self.horizon+1])[:, None] 273 | 274 | u_covs = self._Qs[:self.horizon][:, None] 275 | 276 | return ds.MultivariateNormalFullCovariance(loc=u_locs, covariance_matrix=u_covs).log_prob(us) 277 | 278 | def get_u_samples(self, xi, i=0): 279 | loc = -tf.einsum('ij,aj->ai', self._Ks[i], xi) + tf.einsum('ij,j->i', self._Kvs[i], self._vs[i+1])[None] 280 | cov = self._Qs[i] 281 | return ds.MultivariateNormalFullCovariance(loc=loc, covariance_matrix=cov).sample(1)[0] 282 | 283 | def entropy(self): 284 | return tf.reduce_sum(ds.MultivariateNormalFullCovariance( 285 | covariance_matrix=self._Qs).entropy()) 286 | -------------------------------------------------------------------------------- /tf_robot_learning/control/rollout/__init__.py: -------------------------------------------------------------------------------- 1 | from .rollout import make_rollout_samples, make_multi_shooting_rollout_samples, \ 2 | make_rollout_samples_is, make_rollout_mvn, make_rollout_autonomous_mvn, make_multi_shooting_rollout_samples_nn -------------------------------------------------------------------------------- /tf_robot_learning/control/utils.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import numpy as np 21 | import tensorflow as tf 22 | from ..utils.basis_utils import build_fixed_psi 23 | 24 | def get_canonical(nb_dim, nb_deriv=2, dt=0.01, return_op=True, mass=1.): 25 | A1d = np.zeros((nb_deriv, nb_deriv)) 26 | 27 | for i in range(nb_deriv): 28 | A1d += 1. * np.diag(np.ones(nb_deriv - i), i) * np.power(dt, i) / np.math.factorial(i) 29 | 30 | if nb_deriv == 3: 31 | A1d[:1, 2] /= mass 32 | 33 | B1d = np.zeros((nb_deriv, 1)) 34 | for i in range(1, nb_deriv + 1): 35 | B1d[nb_deriv - i] = np.power(dt, i) / np.math.factorial(i) 36 | 37 | if nb_deriv == 2: 38 | B1d /= mass 39 | 40 | A, B = tf.constant(np.kron(A1d, np.eye(nb_dim)), dtype=tf.float32),\ 41 | tf.constant(np.kron(B1d, np.eye(nb_dim)), dtype=tf.float32) 42 | 43 | if return_op: 44 | return tf.linalg.LinearOperatorFullMatrix(A), tf.linalg.LinearOperatorFullMatrix(B) 45 | else: 46 | return A, B 47 | 48 | def get_perturbation_seq(ndim=1, p_pert=0.01, pert_range=0.1, batch_size=10, horizon=100): 49 | p_push_prior = tf.concat([ 50 | tf.log((1 - p_pert)) * tf.ones((batch_size, 1)), 51 | tf.log(p_pert) * tf.ones((batch_size, 1)) 52 | ], axis=1) 53 | 54 | push_seq = tf.cast(tf.random.categorical(p_push_prior, horizon), tf.float32)[:, :, None] * \ 55 | tf.random_normal(((batch_size, horizon, ndim)), 56 | tf.zeros((batch_size, horizon, ndim)), pert_range) 57 | 58 | return push_seq 59 | 60 | def get_push_seq(length=10, ndim=1, p_pert=0.01, pert_range=0.1, batch_size=10, horizon=100): 61 | k_basis = int(horizon / length * 2) 62 | 63 | _, h = build_fixed_psi(n_step=horizon, n_dim=ndim, n_state=k_basis, 64 | scale=.3 / k_basis) 65 | pert_seq = get_perturbation_seq( 66 | ndim=ndim, p_pert=p_pert, pert_range=pert_range, batch_size=batch_size, horizon=k_basis) 67 | 68 | push_seq = tf.reduce_sum(h[None, :, :, None] * pert_seq[:, None], axis=2) 69 | 70 | return push_seq 71 | -------------------------------------------------------------------------------- /tf_robot_learning/distributions/__init__.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from tensorflow_probability import distributions as _distributions 21 | 22 | from .poe import PoE 23 | 24 | from .mvn import MultivariateNormalFullCovarianceML, MultivariateNormalFullPrecision, \ 25 | MultivariateNormalIso, MultivariateNormalFullCovariance 26 | from .soft_uniform import SoftUniformNormalCdf, SoftUniform 27 | # import from tensorflow_probability 28 | from . import approx 29 | from .promp import ProMP, build_fixed_psi 30 | 31 | Categorical = _distributions.Categorical 32 | MultivariateNormalDiag = _distributions.MultivariateNormalDiag 33 | MultivariateNormalTriL = _distributions.MultivariateNormalTriL 34 | try: 35 | Wishart = _distributions.Wishart 36 | except: 37 | Wishart = None 38 | LogNormal = _distributions.LogNormal 39 | StudentT = _distributions.StudentT 40 | Normal = _distributions.Normal 41 | Uniform = _distributions.Uniform 42 | MixtureSameFamily = _distributions.MixtureSameFamily 43 | TransformedDistribution = _distributions.TransformedDistribution 44 | Mixture = _distributions.Mixture 45 | 46 | from .mixture_models import * -------------------------------------------------------------------------------- /tf_robot_learning/distributions/approx/__init__.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from .variational_gmm import VariationalGMM -------------------------------------------------------------------------------- /tf_robot_learning/distributions/approx/variational_gmm.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | import tensorflow.compat.v1 as tf1 22 | from ...utils.tf_utils import log_normalize 23 | from tensorflow_probability import distributions as _distributions 24 | 25 | 26 | class VariationalGMM(object): 27 | def __init__(self, log_unnormalized_prob, gmm=None, k=10, loc=0., std=1., ndim=None, loc_tril=None, 28 | samples=20, temp=1., cov_type='diag', loc_scale=1., priors_scale=1e1): 29 | """ 30 | 31 | :param log_unnormalized_prob: Unnormalized log density to estimate 32 | :type log_unnormalized_prob: a tensorflow function that takes [batch_size, ndim] 33 | as input and returns [batch_size] 34 | :param gmm: 35 | :param k: number of components for GMM approximation 36 | :param loc: for initialization, mean 37 | :param std: for initialization, standard deviation 38 | :param ndim: 39 | """ 40 | self.log_prob = log_unnormalized_prob 41 | self.ndim = ndim 42 | self.temp = temp 43 | 44 | if gmm is None: 45 | assert ndim is not None, "If no gmm is defined, should give the shape of x" 46 | 47 | if cov_type == 'diag': 48 | _log_priors_var = tf.Variable(1. / priors_scale * log_normalize(tf.ones(k))) 49 | log_priors = priors_scale * _log_priors_var 50 | 51 | if isinstance(loc, tf.Tensor) and loc.shape.ndims == 2: 52 | _locs_var = tf.Variable(1. / loc_scale * loc) 53 | locs = loc_scale * _locs_var 54 | else: 55 | _locs_var = tf.Variable( 56 | 1. / loc_scale * tf.random.normal((k, ndim), loc, std)) 57 | locs = loc_scale * _locs_var 58 | 59 | log_std_diags = tf.Variable(tf.log(std/k * tf.ones((k, ndim)))) 60 | 61 | self._opt_params = [_log_priors_var, _locs_var, log_std_diags] 62 | 63 | gmm = _distributions.MixtureSameFamily( 64 | mixture_distribution=_distributions.Categorical(logits=log_priors), 65 | components_distribution=_distributions.MultivariateNormalDiag( 66 | loc=locs, scale_diag=tf.math.exp(log_std_diags) 67 | ) 68 | ) 69 | 70 | elif cov_type == 'full': 71 | _log_priors_var = tf.Variable(1./priors_scale * log_normalize(tf.ones(k))) 72 | log_priors = priors_scale * _log_priors_var 73 | 74 | if isinstance(loc, tf.Tensor) and loc.shape.ndims == 2: 75 | _locs_var = tf.Variable(1. / loc_scale * loc) 76 | locs = loc_scale * _locs_var 77 | else: 78 | _locs_var = tf.Variable(1./loc_scale * tf.random.normal((k, ndim), loc, std)) 79 | locs = loc_scale * _locs_var 80 | 81 | 82 | loc_tril = loc_tril if loc_tril is not None else std/k 83 | # tril_cov = tf.Variable(loc_tril ** 2 * tf.eye(ndim, batch_shape=(k, ))) 84 | 85 | tril_cov = tf.Variable(tf1.log(loc_tril) * tf.eye(ndim, batch_shape=(k, ))) 86 | 87 | covariance = tf.linalg.expm(tril_cov + tf1.matrix_transpose(tril_cov)) 88 | # 89 | self._opt_params = [_log_priors_var, _locs_var, tril_cov] 90 | 91 | gmm = _distributions.MixtureSameFamily( 92 | mixture_distribution=_distributions.Categorical(logits=log_priors), 93 | components_distribution=_distributions.MultivariateNormalFullCovariance( 94 | loc=locs, covariance_matrix=covariance 95 | ) 96 | ) 97 | 98 | else: 99 | raise ValueError("Unrecognized covariance type") 100 | 101 | self.k = k 102 | self.num_samples = samples 103 | self.gmm = gmm 104 | 105 | @property 106 | def sample_shape(self): 107 | return (self.num_samples, ) 108 | 109 | @property 110 | def opt_params(self): 111 | """ 112 | Parameters to train 113 | :return: 114 | """ 115 | return self._opt_params 116 | 117 | @property 118 | def opt_params_wo_prior(self): 119 | return self._opt_params[1:] 120 | 121 | 122 | def mixture_elbo(self, *args): 123 | samples_conc = tf.reshape( 124 | tf.transpose(self.gmm.components_distribution.sample(self.sample_shape), perm=(1, 0, 2)) 125 | , (-1, self.ndim)) # [k * nsamples, ndim] 126 | 127 | log_qs = tf.reshape(self.gmm.log_prob(samples_conc), (self.k, self.num_samples)) 128 | log_ps = tf.reshape(self.temp * self.log_prob(samples_conc), (self.k, self.num_samples)) 129 | 130 | component_elbos = tf.reduce_mean(log_ps-log_qs, axis=1) 131 | 132 | return tf.reduce_mean(component_elbos * tf.exp(log_normalize(self.gmm.mixture_distribution.logits))) 133 | 134 | 135 | def mixture_elbo_cst_prior(self, *args): 136 | samples = tf.reshape(self.gmm.components_distribution.sample(self.sample_shape), (-1, self.ndim)) 137 | # samples = self.gmm.sample(self.sample_shape) 138 | 139 | log_ps = self.temp * self.log_prob(samples) 140 | log_qs = self.gmm.log_prob(samples) 141 | 142 | return tf.reduce_mean(log_ps - log_qs) 143 | 144 | 145 | @property 146 | def cost(self): 147 | return -self.mixture_elbo() 148 | 149 | @property 150 | def cost_cst_prior(self): 151 | return -self.mixture_elbo_cst_prior() -------------------------------------------------------------------------------- /tf_robot_learning/distributions/mixture_models/__init__.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from .gmm_ml import GaussianMixtureModelML, GaussianMixtureModelFromSK 21 | from .moe import MoE, MoEFromSkMixture -------------------------------------------------------------------------------- /tf_robot_learning/distributions/mixture_models/gmm_ml.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from tensorflow_probability import distributions as ds 21 | import tensorflow as tf 22 | from ...utils import tf_utils, param 23 | try: 24 | from sklearn.mixture import BayesianGaussianMixture, GaussianMixture 25 | except: 26 | print("Sklearn not installed, some features might not be usable") 27 | 28 | import numpy as np 29 | 30 | class EmptyMixture(object): 31 | def __init__(self): 32 | pass 33 | 34 | def filter_unused(mixture, default_cov=None): 35 | idx = mixture.degrees_of_freedom_ >= float( 36 | mixture.degrees_of_freedom_prior_) + 0.5 37 | 38 | _mixture = EmptyMixture() 39 | _mixture.means_ = mixture.means_[idx] 40 | _mixture.covariances_ = mixture.covariances_[idx] 41 | _mixture.weights_ = mixture.weights_[idx] 42 | 43 | _mixture.means_ = np.concatenate([ 44 | _mixture.means_, mixture.mean_prior_[None]], axis=0) 45 | 46 | _mixture.weights_ = np.concatenate([ 47 | _mixture.weights_/np.sum(_mixture.weights_, keepdims=True) * 0.99, 0.01 * np.ones(1)], axis=0) 48 | 49 | if default_cov is None: 50 | default_cov = np.cov(mixture.means_.T) * np.eye(mixture.means_.shape[-1]) 51 | 52 | _mixture.covariances_ = np.concatenate([ 53 | _mixture.covariances_, default_cov[None] 54 | ], axis=0) 55 | 56 | return _mixture 57 | 58 | class GaussianMixtureModelFromSK(ds.MixtureSameFamily): 59 | def __init__(self, mixture, marginal_slice=None, bayesian=False): 60 | 61 | self._logits = param.make_logits_from_value(mixture.weights_) 62 | 63 | 64 | if bayesian: 65 | mixture = filter_unused(mixture) 66 | 67 | if marginal_slice is not None: 68 | self._locs = param.make_loc_from_value( 69 | mixture.means_[:, marginal_slice]) 70 | self._covs = param.make_cov_from_value( 71 | mixture.covariances_[:, marginal_slice, marginal_slice]) 72 | else: 73 | self._locs = param.make_loc_from_value(mixture.means_) 74 | self._covs = param.make_cov_from_value(mixture.covariances_) 75 | 76 | self._priors = tf.math.exp(self._logits) 77 | 78 | ds.MixtureSameFamily.__init__( 79 | self, 80 | mixture_distribution=ds.Categorical(probs=self._priors), 81 | components_distribution=ds.MultivariateNormalFullCovariance( 82 | loc=self._locs, covariance_matrix=self._covs 83 | ) 84 | ) 85 | 86 | @property 87 | def variables(self): 88 | return [self._logits.variable, self._locs.variable, self._covs.variable] 89 | 90 | class GaussianMixtureModelML(ds.MixtureSameFamily): 91 | def __init__(self, priors, locs, covs, init_params='kmeans', warm_start=True, 92 | reg_cov=1e-6, bayesian=False): 93 | 94 | self._priors = priors 95 | self._locs = locs 96 | self._covs = covs 97 | 98 | self._covs_par = covs 99 | 100 | if tf.__version__[0] == '1': 101 | self._k = priors.shape[0].value 102 | self._n = locs.shape[-1].value 103 | else: 104 | self._k = priors.shape[0] 105 | self._n = locs.shape[-1] 106 | 107 | try: 108 | m = BayesianGaussianMixture if bayesian else GaussianMixture 109 | self._sk_gmm = m( 110 | self._k, 'full', n_init=1, init_params=init_params, warm_start=warm_start, reg_covar=reg_cov) 111 | except: 112 | print("Sklearn not install, cannot perform MLE of Gaussian mixture model") 113 | 114 | self._priors_ml = tf.compat.v1.placeholder(tf.float32, (self._k, )) 115 | self._locs_ml = tf.compat.v1.placeholder(tf.float32, (self._k, self._n)) 116 | self._covs_ml = tf.compat.v1.placeholder(tf.float32, (self._k, self._n, self._n)) 117 | 118 | self._ml_assign_op = [ 119 | self._priors.assign(self._priors_ml), 120 | self._locs.assign(self._locs_ml), 121 | self._covs.assign(self._covs_ml) 122 | ] 123 | 124 | ds.MixtureSameFamily.__init__( 125 | self, 126 | mixture_distribution=ds.Categorical(probs=self._priors), 127 | components_distribution=ds.MultivariateNormalFullCovariance( 128 | loc=self._locs, covariance_matrix=self._covs 129 | ) 130 | ) 131 | 132 | self._init = False 133 | 134 | def conditional_distribution(self, x, slice_in, slice_out): 135 | marginal_in = ds.MixtureSameFamily( 136 | mixture_distribution=ds.Categorical(probs=self._priors), 137 | components_distribution=ds.MultivariateNormalFullCovariance( 138 | loc=self._locs[:, slice_in], covariance_matrix=self._covs[:, slice_in, slice_in] 139 | )) 140 | 141 | p_k_in = ds.Categorical(logits=tf_utils.log_normalize( 142 | marginal_in.components_distribution.log_prob(x[:, None]) + 143 | marginal_in.mixture_distribution.logits[None], axis=1)) 144 | 145 | sigma_in_out = self._covs[:, slice_in, slice_out] 146 | inv_sigma_in_in = tf.linalg.inv(self._covs[:, slice_in, slice_in]) 147 | inv_sigma_out_in = tf.matmul(sigma_in_out, inv_sigma_in_in, transpose_a=True) 148 | 149 | A = inv_sigma_out_in 150 | b = self._locs[:, slice_out] - tf.matmul( 151 | inv_sigma_out_in, self._locs[:, slice_in, None])[:, :, 0] 152 | 153 | cov_est = (self._covs[:, slice_out, slice_out] - tf.matmul( 154 | inv_sigma_out_in, sigma_in_out)) 155 | 156 | ys = tf.einsum('aij,bj->abi', A, x) + b[:, None] 157 | 158 | p_out_in_k = ds.MultivariateNormalFullCovariance( 159 | tf.transpose(ys, perm=(1, 0, 2)), cov_est) 160 | 161 | return ds.MixtureSameFamily( 162 | mixture_distribution=p_k_in, 163 | components_distribution=p_out_in_k 164 | ) 165 | 166 | def ml(self, x, sess=None, alpha=1., warm_start=True, n_init=1, reg_diag=None, max_iter=100): 167 | if not self._init: alpha = 1. 168 | self._init = True 169 | 170 | if alpha != 1.: 171 | prev_x = self._sk_gmm.sample(int(x.shape[0]*(1.-alpha)))[0] 172 | x = np.concatenate([x, prev_x], axis=0) 173 | 174 | if sess is None: sess = tf.compat.v1.get_default_session() 175 | 176 | assert hasattr(self, '_sk_gmm'), ("Cannot perform MLE, sklearn not installed") 177 | 178 | self._sk_gmm.warm_start = warm_start 179 | self._sk_gmm.n_init = n_init 180 | self._sk_gmm.max_iter = max_iter 181 | self._sk_gmm.fit(x) 182 | 183 | _covs = self._sk_gmm.covariances_ 184 | 185 | if reg_diag is not None: 186 | _covs += np.diag(reg_diag)[None] ** 2 187 | 188 | feed_dict = { 189 | self._priors_ml: self._sk_gmm.weights_, 190 | self._locs_ml: self._sk_gmm.means_, 191 | self._covs_ml: _covs, 192 | } 193 | 194 | sess.run(self._ml_assign_op, feed_dict) 195 | 196 | 197 | -------------------------------------------------------------------------------- /tf_robot_learning/distributions/mixture_models/moe.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | from tensorflow_probability import distributions as ds 22 | from ...utils.tf_utils import log_normalize 23 | from ...utils import param 24 | 25 | class Gate(object): 26 | def __init__(self): 27 | pass 28 | 29 | @property 30 | def opt_params(self): 31 | raise NotImplementedError 32 | 33 | def conditional_mixture_distribution(self, x): 34 | """ 35 | x : [batch_shape, dim] 36 | return [batch_shape, nb_experts] 37 | """ 38 | raise NotImplementedError 39 | 40 | 41 | class Experts(object): 42 | def __init__(self): 43 | pass 44 | 45 | @property 46 | def opt_params(self): 47 | raise NotImplementedError 48 | 49 | @property 50 | def nb_experts(self): 51 | raise NotImplementedError 52 | 53 | @property 54 | def nb_dim(self): 55 | raise NotImplementedError 56 | 57 | 58 | def conditional_components_distribution(self, x): 59 | """ 60 | x : [batch_shape, dim] 61 | return distribution([batch_shape, nb_experts, dim_out]) 62 | """ 63 | raise NotImplementedError 64 | 65 | 66 | class LinearMVNExperts(Experts): 67 | def __init__(self, A, b, cov_tril): 68 | self._A = A 69 | self._b = b 70 | self._cov_tril = cov_tril 71 | self._covs = tf.linalg.LinearOperatorFullMatrix( 72 | self._cov_tril).matmul(self._cov_tril, adjoint_arg=True) 73 | 74 | @property 75 | def nb_dim(self): 76 | return self._b.shape[-1].value 77 | 78 | @property 79 | def nb_experts(self): 80 | return self._b.shape[0].value 81 | 82 | def conditional_components_distribution(self, x): 83 | ys = tf.einsum('aij,bj->abi', self._A, x) + self._b[:, None] 84 | 85 | return ds.MultivariateNormalTriL( 86 | tf.transpose(ys, perm=(1, 0, 2)), # One for each component. 87 | self._cov_tril) 88 | 89 | @property 90 | def opt_params(self): 91 | return [self._A, self._b, self._cov_tril] 92 | 93 | class MixtureGate(Gate): 94 | def __init__(self, mixture): 95 | self._mixture = mixture 96 | 97 | Gate.__init__(self) 98 | 99 | @property 100 | def opt_params(self): 101 | return self._mixture.opt_params 102 | 103 | @property 104 | def mixture(self): 105 | return self._mixture 106 | 107 | def conditional_mixture_distribution(self, x): 108 | # def logregexp(x, reg=1e-5, axis=0): 109 | # return [x, reg * tf.ones_like(x)] 110 | 111 | return ds.Categorical(logits=log_normalize( 112 | self._mixture.components_distribution.log_prob(x[:, None]) + 113 | self._mixture.mixture_distribution.logits[None], axis=1)) 114 | 115 | 116 | class MoE(object): 117 | def __init__(self, gate, experts, is_function=None): 118 | """ 119 | :type gate : Gate 120 | :type experts : Experts 121 | """ 122 | self._gate = gate 123 | self._experts = experts 124 | 125 | self._is_function = is_function 126 | 127 | @property 128 | def opt_params(self): 129 | return self._gate.opt_params + self._experts.opt_params 130 | 131 | @property 132 | def nb_experts(self): 133 | return self._experts.nb_experts 134 | 135 | def conditional_distribution(self, x): 136 | return ds.MixtureSameFamily( 137 | mixture_distribution=self._gate.conditional_mixture_distribution(x), 138 | components_distribution=self._experts.conditional_components_distribution(x) 139 | ) 140 | 141 | 142 | def sample_is(self, x, n=1): 143 | mixture_distribution, mixture_components = \ 144 | self._gate.conditional_mixture_distribution(x),\ 145 | self._experts.conditional_components_distribution(x) 146 | 147 | y = mixture_components.sample(n) 148 | # npdt = y.dtype.as_numpy_dtype 149 | 150 | is_logits = self._is_function(mixture_distribution.logits) 151 | is_mixture_distribution = ds.Categorical(logits=is_logits) 152 | idx = is_mixture_distribution.sample(n) 153 | 154 | # TODO check if we should not renormalize mixture.logits - tf.stop_... 155 | 156 | weights = tf.batch_gather(mixture_distribution.logits - tf.stop_gradient(is_logits), 157 | tf.transpose(idx)) 158 | # TODO check axis 159 | # weights = tf.batch_gather( 160 | # log_normalize(mixture_distribution.logits - tf.stop_gradient(is_logits), axis=1), 161 | # tf.transpose(idx)) 162 | 163 | if n == 1: 164 | return tf.batch_gather(y, idx[:, :, None])[0, :, 0], tf.transpose(weights)[0] 165 | else: 166 | return tf.batch_gather(y, idx[:, :, None])[:, :, 0], tf.transpose(weights) 167 | 168 | @property 169 | def gate(self): 170 | return self._gate 171 | 172 | @property 173 | def experts(self): 174 | return self._experts 175 | 176 | from .gmm_ml import GaussianMixtureModelFromSK, filter_unused 177 | import numpy as np 178 | 179 | class MoEFromSkMixture(MoE): 180 | def __init__( 181 | self, mixture, slice_in, slice_out, bayesian=False, default_cov=None, 182 | *args, **kwargs 183 | ): 184 | if bayesian: 185 | mixture = filter_unused(mixture, default_cov=default_cov) 186 | 187 | gate = MixtureGate(GaussianMixtureModelFromSK( 188 | mixture=mixture, marginal_slice=slice_in, bayesian=False)) 189 | 190 | sigma_in_out = mixture.covariances_[:, slice_in, slice_out] 191 | 192 | inv_sigma_in_in = np.linalg.inv( 193 | mixture.covariances_[:, slice_in, slice_in]) 194 | inv_sigma_out_in = np.einsum( 195 | 'aji,ajk->aik', sigma_in_out, inv_sigma_in_in) 196 | 197 | As = inv_sigma_out_in 198 | bs = mixture.means_[:, slice_out] - np.matmul( 199 | inv_sigma_out_in, mixture.means_[:, slice_in, None])[ 200 | :, :, 0] 201 | 202 | sigma_est = (mixture.covariances_[:, slice_out, slice_out] - np.matmul( 203 | inv_sigma_out_in, sigma_in_out)) 204 | 205 | As = param.make_loc_from_value(As) 206 | bs = param.make_loc_from_value(bs) 207 | cov_tril = param.make_cov_from_value(np.linalg.cholesky(sigma_est)) 208 | 209 | experts = LinearMVNExperts( 210 | A=As, b=bs, cov_tril=cov_tril 211 | ) 212 | 213 | MoE.__init__(self, gate, experts, *args, **kwargs) 214 | 215 | @property 216 | def variables(self): 217 | return [self._experts._A.variable, self._experts._b.variable, 218 | self._experts._cov_tril.variable] + self._gate.mixture.variables -------------------------------------------------------------------------------- /tf_robot_learning/distributions/poe.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | import tensorflow.compat.v1 as tf1 22 | from tensorflow_probability import distributions as ds 23 | from ..utils.tf_utils import batch_jacobians 24 | 25 | class PoE(ds.Distribution): 26 | def __init__(self, shape, experts, transfs, name='PoE', cost=None): 27 | """ 28 | 29 | :param shape: 30 | :param experts: 31 | :param transfs: a list of tensorflow function that goes from product space to expert space 32 | or a function f(x: tensor, i: index of transform) -> f_i(x) 33 | 34 | :param cost: additional cost [batch_size, n_dim] -> [batch_size, ] 35 | a function f(x) -> 36 | """ 37 | 38 | self._product_shape = shape 39 | self._experts = experts 40 | self._transfs = transfs 41 | self._laplace = None 42 | self._samples_approx = None 43 | 44 | self._cost = cost 45 | 46 | 47 | self.stepsize = tf.Variable(0.01) 48 | self._name = name 49 | 50 | self._x_ph = tf1.placeholder(tf.float32, (None, shape[0])) 51 | self._y = None 52 | self._jac = None 53 | 54 | def f(self, x, return_jac=False, sess=None, feed_dict={}): 55 | dim1 = False 56 | 57 | 58 | if self._y is None: 59 | self._y = tf.concat(self.get_transformed(self._x_ph), axis=1) 60 | if return_jac and self._jac is None: 61 | self._jac = batch_jacobians(self._y, self._x_ph) 62 | 63 | if x.ndim == 1: 64 | x = x[None] 65 | dim1 = True 66 | 67 | if not return_jac: 68 | feed_dict[self._x_ph] = x 69 | _y = self._y.eval(feed_dict) 70 | if dim1: return _y[0] 71 | else: return _y 72 | else: 73 | if sess is None: sess = tf.get_default_session() 74 | 75 | feed_dict[self._x_ph] = x 76 | _y, _jac = sess.run([self._y, self._jac], feed_dict) 77 | 78 | if dim1: return _y[0], _jac[0] 79 | else: return _y, _jac 80 | 81 | def J(self, x, feed_dict={}): 82 | dim1 = False 83 | 84 | if self._y is None: 85 | self._y = tf.concat(self.get_transformed(self._x_ph), axis=1) 86 | if self._jac is None: 87 | self._jac = batch_jacobians(self._y, self._x_ph) 88 | 89 | if x.ndim == 1: 90 | x = x[None] 91 | dim1 = True 92 | 93 | feed_dict[self._x_ph] = x 94 | _jac = self._jac.eval(feed_dict) 95 | if dim1: 96 | return _jac[0] 97 | else: 98 | return _jac 99 | 100 | def get_loc_prec(self): 101 | raise NotImplementedError 102 | # return tf.concat([exp.mean() for exp in self.experts], axis=0),\ 103 | # block_diagonal_different_sizes( 104 | # [tf.linalg.inv(exp.covariance()) for exp in self.experts]) 105 | 106 | @property 107 | def product_shape(self): 108 | return self._product_shape 109 | 110 | @property 111 | def experts(self): 112 | return self._experts 113 | 114 | @property 115 | def transfs(self): 116 | return self._transfs 117 | 118 | 119 | def _experts_probs(self, x): 120 | probs = [] 121 | 122 | for i, exp in enumerate(self.experts): 123 | if isinstance(self.transfs, list): 124 | if hasattr(exp, '_log_unnormalized_prob'): 125 | print('Using unnormalized prob for expert %d' % i) 126 | probs += [exp._log_unnormalized_prob(self.transfs[i](x))] 127 | else: 128 | probs += [exp.log_prob(self.transfs[i](x))] 129 | else: 130 | if hasattr(exp, '_log_unnormalized_prob'): 131 | probs += [exp._log_unnormalized_prob(self.transfs(x, i))] 132 | else: 133 | probs += [exp.log_prob(self.transfs(x, i))] 134 | 135 | return probs 136 | 137 | def _log_unnormalized_prob(self, x, wo_cost=False): 138 | if x.get_shape().ndims == 1: 139 | x = x[None] 140 | 141 | if wo_cost: 142 | return tf.reduce_sum(self._experts_probs(x), axis=0) 143 | else: 144 | cost = 0. if self._cost is None else self._cost(x) 145 | return tf.reduce_sum(self._experts_probs(x), axis=0) - cost 146 | 147 | @property 148 | def nb_experts(self): 149 | return len(self.experts) 150 | 151 | 152 | def get_ml_transformed(self, x): 153 | ys = self.get_transformed(x) 154 | 155 | return [[tf.reduce_mean(y, axis=0), reduce_cov(y, axis=0)] for y in ys] 156 | 157 | 158 | def get_transformed(self, x): 159 | return [f(x) for f in self.transfs] 160 | 161 | -------------------------------------------------------------------------------- /tf_robot_learning/distributions/promp.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | from tensorflow_probability import distributions as ds 22 | from ..utils.basis_utils import build_fixed_psi 23 | 24 | 25 | class ProMP(ds.MultivariateNormalFullCovariance): 26 | def __init__(self, 27 | loc=None, 28 | covariance_matrix=None, 29 | scale_obs=0.1, 30 | psi=None, 31 | fast_sample=True, 32 | validate_args=False, 33 | allow_nan_stats=True, 34 | name="ProMP"): 35 | 36 | self._psi = psi 37 | self._psi_op = tf.linalg.LinearOperatorFullMatrix(psi) 38 | self._loc_w = loc 39 | self._cov_w = covariance_matrix 40 | 41 | self._mvn_w = ds.MultivariateNormalFullCovariance( 42 | loc=self._loc_w, 43 | covariance_matrix=self._cov_w 44 | ) 45 | 46 | # (psi T cov psi ) T = (psi T cov )T 47 | 48 | # _loc = tf.linalg.matvec(psi, loc, transpose_a=True) 49 | # _cov = tf.linalg.matmul(tf.linalg.matmul(psi, covariance_matrix, transpose_a=True), 50 | # psi) + \ 51 | # tf.eye(psi.shape[1].value) * scale_obs ** 2 52 | 53 | _loc = self._psi_op.matvec(loc, adjoint=True) 54 | 55 | _cov = self._psi_op.matmul( 56 | self._psi_op.matmul(covariance_matrix, adjoint=True), adjoint=True, adjoint_arg=True) 57 | 58 | self._mvn_obs = ds.MultivariateNormalDiag( 59 | loc=tf.zeros_like(_loc), 60 | scale_diag=scale_obs * tf.ones(psi.shape[1].value), 61 | ) 62 | 63 | self._fast_sample = fast_sample 64 | 65 | if _cov.shape.ndims == 2: 66 | _cov += tf.eye(psi.shape[1].value) * scale_obs ** 2 67 | elif _cov.shape.ndims == 3: 68 | _cov += tf.eye(psi.shape[1].value)[None] * scale_obs ** 2 69 | else: 70 | raise NotImplementedError 71 | 72 | ds.MultivariateNormalFullCovariance.__init__( 73 | self, loc=_loc, covariance_matrix=_cov) 74 | 75 | def sample(self, sample_shape=(), seed=None, name="sample"): 76 | if self._fast_sample: 77 | _w_samples = self._mvn_w.sample(sample_shape, seed=seed, name=name + '_w') 78 | _obs_samples = self._mvn_obs.sample(sample_shape, seed=seed, name=name + '_obs') 79 | 80 | return self._psi_op.matvec(_w_samples, adjoint=True) + _obs_samples 81 | else: 82 | return ds.MultivariateNormalFullCovariance.sample( 83 | self, sample_shape=sample_shape, seed=seed, name=name) 84 | -------------------------------------------------------------------------------- /tf_robot_learning/distributions/soft_uniform.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | from tensorflow_probability import distributions as ds 22 | from tensorflow.python.framework import ops 23 | from tensorflow.python.ops import array_ops 24 | from tensorflow.python.ops import check_ops 25 | 26 | 27 | class SoftUniform(ds.Distribution): 28 | def __init__(self, 29 | low=0., 30 | high=1., 31 | temp=1e2, 32 | validate_args=False, 33 | allow_nan_stats=True, 34 | squared=True, 35 | activation=tf.nn.relu, 36 | name="SoftUniform"): 37 | parameters = locals() 38 | 39 | parameters = dict(locals()) 40 | 41 | self._activation = activation 42 | self._squared = squared 43 | 44 | with ops.name_scope(name, values=[low, high]) as name: 45 | with ops.control_dependencies([ 46 | check_ops.assert_less( 47 | low, high, 48 | message="uniform not defined when low >= high.") 49 | ] if validate_args else []): 50 | self._low = array_ops.identity(low, name="low") 51 | self._high = array_ops.identity(high, name="high") 52 | self._temp = array_ops.identity(temp, name="temp") 53 | check_ops.assert_same_float_dtype([self._low, self._high, self._temp]) 54 | 55 | super(SoftUniform, self).__init__( 56 | dtype=self._low.dtype, 57 | reparameterization_type=ds.FULLY_REPARAMETERIZED, 58 | validate_args=validate_args, 59 | allow_nan_stats=allow_nan_stats, 60 | parameters=parameters, 61 | graph_parents=[self._low, 62 | self._high, 63 | self._temp], 64 | name=name) 65 | 66 | @property 67 | def temp(self): 68 | return self._temp 69 | 70 | def log_prob(self, value, name="log_prob"): 71 | z = self._high - self._low 72 | 73 | log_unnormalized_prob = 0. 74 | 75 | if self._squared: 76 | if self._high.shape.ndims == 0: 77 | log_unnormalized_prob += (self._temp * self._activation(value - self._high)) ** 2 78 | log_unnormalized_prob += (self._temp * self._activation(-value + self._low)) ** 2 79 | else: 80 | z = tf.reduce_prod(z) 81 | log_unnormalized_prob += tf.reduce_sum((self._temp * self._activation(-value + self._low[None])) ** 2, -1) 82 | log_unnormalized_prob += tf.reduce_sum((self._temp * self._activation(value - self._high[None])) ** 2, -1) 83 | else: 84 | if self._high.shape.ndims == 0: 85 | log_unnormalized_prob += self._temp * self._activation(value - self._high) 86 | log_unnormalized_prob += self._temp * self._activation(-value + self._low) 87 | else: 88 | z = tf.reduce_prod(z) 89 | log_unnormalized_prob += tf.reduce_sum(self._temp * self._activation(value - self._high[None]), -1) 90 | log_unnormalized_prob += tf.reduce_sum(self._temp * self._activation(-value + self._low[None]), -1) 91 | 92 | return -tf.log(z) - log_unnormalized_prob 93 | 94 | class SoftUniformNormalCdf(SoftUniform): 95 | """ 96 | SoftUniform with normal log cdf 97 | """ 98 | def __init__(self, 99 | low=0., 100 | high=1., 101 | std=0.1, 102 | temp=1., 103 | reduce_axis=None, 104 | validate_args=False, 105 | allow_nan_stats=True, 106 | name="SoftUniformNormalCdf"): 107 | """ 108 | 109 | :param low: 110 | :param high: 111 | :param std: 112 | :param temp: 113 | :param reduce_axis: if not None, on which axis to reduce log prob, for independant constraints 114 | :param validate_args: 115 | :param allow_nan_stats: 116 | :param name: 117 | """ 118 | 119 | 120 | self._use_low = not low is None 121 | self._use_high = not high is None 122 | 123 | if low is None: 124 | low = 0. 125 | 126 | if high is None: 127 | high = 0. 128 | 129 | with ops.name_scope(name, values=[low, high]) as name: 130 | with ops.control_dependencies([ 131 | check_ops.assert_less( 132 | low, high, 133 | message="uniform not defined when low >= high.") 134 | ] if validate_args else []): 135 | self._std = array_ops.identity(std, name="std") 136 | 137 | self._reduce_axis = reduce_axis 138 | 139 | super(SoftUniformNormalCdf, self).__init__( 140 | low=low, 141 | high=high, 142 | temp=temp, 143 | validate_args=validate_args, 144 | allow_nan_stats=allow_nan_stats, 145 | name=name) 146 | 147 | def log_prob(self, value, name="log_prob"): 148 | log_probs = 0 149 | if self._use_high: 150 | log_probs += self._temp * ds.Normal(-self._high, self._std).log_cdf(-value) 151 | 152 | if self._use_low: 153 | log_probs += self._temp * ds.Normal(self._low, self._std).log_cdf(value) 154 | 155 | 156 | if self._reduce_axis is not None: 157 | return tf.reduce_sum(log_probs, axis=self._reduce_axis) 158 | else: 159 | return log_probs 160 | -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/__init__.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from .frame import Frame, Twist 21 | from .joint import Joint, JointType 22 | from .segment import Segment 23 | from .chain import Chain, FkLayout, ChainDict 24 | from . import rotation 25 | from .urdf_utils import * 26 | from . import utils -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/frame.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | import tensorflow.compat.v1 as tf1 22 | from .utils.tf_utils import * 23 | from .utils import FkLayout 24 | from tensorflow.python.framework.ops import EagerTensor 25 | 26 | class Rotation(tf.Tensor): 27 | pass 28 | 29 | @property 30 | def is_batch(self): 31 | return self.shape.ndims == 3 32 | 33 | 34 | class Twist(object): 35 | def __init__(self, dx=tf.zeros(6)): 36 | self.dx = dx 37 | 38 | @property 39 | def is_batch(self): 40 | return self.dx.shape.ndims == 2 41 | 42 | def dx_mat(self, m, layout=FkLayout.xm): 43 | """ 44 | https://en.wikipedia.org/wiki/Angular_velocity 45 | :param m: 46 | :return: 47 | """ 48 | w = angular_vel_tensor(self.rot) 49 | dm_dphi = matmatmul(w, m) 50 | 51 | if layout is FkLayout.xmv: 52 | if self.is_batch: 53 | return tf.concat( 54 | [self.vel, tf.reshape(tf.transpose(dm_dphi, perm=(0, 2, 1)), (-1, 9))], axis=1 55 | ) 56 | else: 57 | return tf.concat( 58 | [self.vel, tf.reshape(tf.transpose(dm_dphi, perm=(1, 0)), (9, ))], axis=0 59 | ) 60 | 61 | else: 62 | if self.is_batch: 63 | return tf.concat( 64 | [self.vel, tf.reshape(dm_dphi, (-1, 9))], axis=1 65 | ) 66 | else: 67 | return tf.concat( 68 | [self.vel, tf.reshape(dm_dphi, (9, ))], axis=0 69 | ) 70 | @property 71 | def vel(self): 72 | if self.is_batch: 73 | return self.dx[:, :3] 74 | else: 75 | return self.dx[:3] 76 | 77 | @property 78 | def rot(self): 79 | if self.is_batch: 80 | return self.dx[:, 3:] 81 | else: 82 | return self.dx[3:] 83 | 84 | def ref_point(self, v): 85 | if v.shape.ndims > self.rot.shape.ndims: 86 | rot = self.rot[None] * (tf.zeros_like(v) + 1.) 87 | vel = self.vel + tf1.cross(rot, v) 88 | elif v.shape.ndims < self.rot.shape.ndims: 89 | n = self.rot.shape[0].value 90 | vel = self.vel + tf1.cross(self.rot, v[None] * tf.ones((n, 1))) 91 | rot = self.rot 92 | else: 93 | vel = self.vel + tf1.cross(self.rot, v) 94 | rot = self.rot 95 | 96 | 97 | if self.is_batch or v.shape.ndims==2: 98 | return Twist(tf.concat([vel, rot], 1)) 99 | else: 100 | return Twist(tf.concat([vel, rot], 0)) 101 | 102 | def __rmul__(self, other): 103 | if isinstance(other, Frame): 104 | raise NotImplementedError 105 | 106 | else: 107 | rot = matvecmul(other, self.rot) 108 | vel = matvecmul(other, self.vel) 109 | if self.is_batch or rot.shape.ndims==2: 110 | return Twist(tf.concat([vel, rot], 1)) 111 | else: 112 | return Twist(tf.concat([vel, rot], 0)) 113 | 114 | 115 | class Frame(object): 116 | def __init__(self, p=None, m=None, batch_shape=None): 117 | """ 118 | :param p: 119 | Translation vector 120 | :param m: 121 | Rotation matrix 122 | :param batch_shape int 123 | """ 124 | 125 | if batch_shape is None: 126 | p = tf.zeros(3) if p is None else p 127 | m = tf.eye(3) if m is None else m 128 | else: 129 | p = tf.zeros((batch_shape, 3)) if p is None else p 130 | m = tf.eye(3, batch_shape=(batch_shape, )) if m is None else m 131 | 132 | if isinstance(m, tf.Variable): _m = tf.identity(m) 133 | else: _m = m 134 | 135 | 136 | self.p = p 137 | self.m = _m 138 | 139 | def fix_it(self): 140 | # return self 141 | sess = tf.compat.v1.get_default_session() 142 | 143 | p, m = sess.run([self.p, self.m]) 144 | 145 | return Frame(tf.convert_to_tensor(p, tf.float32), tf.convert_to_tensor(m, tf.float32)) 146 | 147 | @property 148 | def is_batch(self): 149 | return self.m.shape.ndims == 3 150 | 151 | @property 152 | def xm(self): 153 | """ 154 | Position and vectorized rotation matrix 155 | (order : 'C' - last index changing the first) 156 | :return: 157 | """ 158 | if self.is_batch: 159 | return tf.concat([self.p, tf.reshape(self.m, [-1, 9])], axis=1) 160 | else: 161 | return tf.concat([self.p, tf.reshape(self.m, [9])], axis=0) 162 | 163 | @property 164 | def xmv(self): 165 | """ 166 | Position and vectorized rotation matrix 167 | (order : 'C' - last index changing the first) 168 | :return: 169 | """ 170 | if self.is_batch: 171 | return tf.concat([self.p, tf.reshape(tf.transpose(self.m, perm=(0, 2, 1)), [-1, 9])], axis=1) 172 | else: 173 | return tf.concat([self.p, tf.reshape(tf.transpose(self.m, perm=(1, 0)), [9])], axis=0) 174 | 175 | @property 176 | def xq(self): 177 | """ 178 | Position and Quaternion 179 | :return: 180 | """ 181 | raise NotImplementedError 182 | 183 | 184 | def inv(self): 185 | return Frame(p=-tf.matmul(self.m, tf.expand_dims(self.p, 1), transpose_a=True)[:,0], 186 | m=tf.transpose(self.m)) 187 | 188 | def __mul__(self, other): 189 | if isinstance(other, Twist): 190 | # TODO check 191 | rot = matvecmul(self.m, other.rot) 192 | vel = matvecmul(self.m, other.vel) + tf1.cross(self.p, rot) # should be cross product 193 | return Twist(tf.concat([vel, rot], 0)) 194 | 195 | elif isinstance(other, tf.Tensor) or isinstance(other, tf.Variable): 196 | if other.shape[-1].value == 3: # only position 197 | if self.is_batch: 198 | # return tf.einsum('aij,bj->abi', self.m, other) + self.p[:, None] 199 | return tf.linalg.LinearOperatorFullMatrix(self.m[:, None]).matvec(other[None]) + self.p[None] 200 | else: 201 | return tf.linalg.LinearOperatorFullMatrix(self.m).matvec(other) + self.p[None] 202 | else: 203 | raise NotImplementedError('Only position supported yet') 204 | else: 205 | return Frame(p=matvecmul(self.m, other.p) + self.p, 206 | m=matmatmul(self.m, other.m)) -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/joint.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | import numpy as np 22 | from .frame import Frame, Twist 23 | from .rotation import * 24 | from enum import IntEnum 25 | 26 | class JointType(IntEnum): 27 | RotX = 0 28 | RotY = 1 29 | RotZ = 2 30 | RotAxis = 3 31 | NoneT = 4 32 | 33 | class Joint(object): 34 | def __init__(self, type, origin=None, axis=None, name='', limits=None): 35 | self.type = type 36 | 37 | if limits is not None: 38 | self.limits = {'up': limits.upper, 'low': limits.lower, 39 | 'vel': limits.velocity, 'effort': limits.effort} 40 | else: 41 | self.limits = {} 42 | 43 | self.axis = axis 44 | self.origin = origin 45 | 46 | 47 | self.pose_0 = self.pose(0.).fix_it() 48 | 49 | def pose(self, a): 50 | # TODO implement origin 51 | 52 | if self.type is JointType.RotX: 53 | return Frame(m=rot_x(a)) 54 | elif self.type is JointType.RotY: 55 | return Frame(m=rot_y(a)) 56 | elif self.type is JointType.RotZ: 57 | return Frame(m=rot_z(a)) 58 | elif self.type is JointType.RotAxis: 59 | return Frame(p=self.origin , m=rot_2(self.axis, a)) 60 | elif self.type is JointType.NoneT: 61 | return Frame() 62 | 63 | def twist(self, a): 64 | if self.type is JointType.RotX: 65 | return Twist(twist_x(a)) 66 | elif self.type is JointType.RotY: 67 | return Twist(twist_y(a)) 68 | elif self.type is JointType.RotZ: 69 | return Twist(twist_z(a)) 70 | elif self.type is JointType.RotAxis: 71 | return Twist(twist_2(self.axis, a)) 72 | elif self.type is JointType.NoneT: 73 | return Twist() 74 | 75 | class Mesh(object): 76 | def __init__(self, mesh, dtype=tf.float32): 77 | self._vertices = tf.convert_to_tensor(mesh.vertices, dtype=dtype) 78 | self._nb_vertices = mesh.vertices.shape[0] 79 | 80 | self._area_faces = tf.convert_to_tensor(mesh.area_faces, dtype=dtype) 81 | 82 | triangles = np.copy(mesh.triangles) 83 | # triangles[:, :, 0] = mesh.triangles[:, :, 2] 84 | # triangles[:, :, 1] = mesh.triangles[:, :, 0] 85 | # triangles[:, :, 2] = mesh.triangles[:, :, 1] 86 | 87 | self._triangles = tf.convert_to_tensor(triangles, dtype=dtype) 88 | sample = self.sample(10) 89 | 90 | def sample(self, size=10): 91 | """ 92 | Random sample on surface 93 | :param size: 94 | :return: 95 | """ 96 | # sample triangle proportional to surface 97 | idx = tf.random.categorical(tf.log(self._area_faces)[None], size)[0] 98 | 99 | triangles_samples = tf.gather( 100 | self._triangles, 101 | idx 102 | ) 103 | 104 | # sample on triangle tf.reduce_sum(tf.transpose(vs)[:, :, None] * triangles_samples, axis=1) 105 | r0, r1 = tf.random_uniform((size, ), 0., 1.), tf.random_uniform((size, ), 0., 1.) 106 | 107 | vs = tf.stack([1. - r0 ** 0.5, r0 ** 0.5 * (1. - r1), r1 * r0 ** 0.5]) 108 | return tf.reduce_sum(tf.transpose(vs)[:, :, None] * triangles_samples, axis=1) 109 | 110 | def sample_face(self, size=10): 111 | return tf.gather(self._vertices, 112 | tf.random_uniform((size, ), 0, self._nb_vertices-1, dtype=tf.int64)) 113 | 114 | class Link(object): 115 | def __init__(self, frame, mass=1.0): 116 | self.mass = mass 117 | self.frame = frame 118 | 119 | self._collision_mesh = None 120 | 121 | @property 122 | def collision_mesh(self): 123 | return self._collision_mesh 124 | 125 | @collision_mesh.setter 126 | def collision_mesh(self, value): 127 | self._collision_mesh = Mesh(value) 128 | 129 | def pose(self): 130 | return self.frame -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/rotation.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | 22 | def twist_x(a): 23 | return tf.stack([0., 0., 0., a, 0., 0.]) 24 | 25 | def twist_y(a): 26 | return tf.stack([0., 0., 0., 0., a, 0.]) 27 | 28 | def twist_z(a): 29 | return tf.stack([0., 0., 0., 0., 0., a]) 30 | 31 | def twist_2(axis, a): 32 | if isinstance(a, float) or a.shape.ndims == 0: 33 | return tf.concat([tf.zeros(3), axis*a], axis=0) 34 | else: 35 | return tf.concat([tf.zeros(3)[None] * a, axis[None]*a], axis=1) 36 | 37 | def skew_x(u): 38 | return tf.stack([[0., -u[2] , u[1]], 39 | [u[2], 0., -u[0]], 40 | [-u[1], u[0] , 0.]]) 41 | 42 | 43 | def rot_2(axis, a): 44 | """ 45 | https://en.wikipedia.org/wiki/Rotation_matrix see Rotation matrix from axis and angle 46 | :param axis: 47 | :param a: 48 | :return: 49 | """ 50 | if isinstance(a, float) or a.shape.ndims == 0: 51 | return tf.cos(a) * tf.eye(3) + tf.sin(a) * skew_x(axis) + (1- tf.cos(a)) * tf.einsum('i,j->ij', axis, axis) 52 | else: 53 | return tf.cos(a)[:, None, None] * tf.eye(3)[None] + tf.sin(a)[:, None, None] * skew_x(axis)[None] +\ 54 | (1- tf.cos(a)[:, None, None]) * tf.einsum('i,j->ij', axis, axis)[None] 55 | 56 | def rpy(rpy): 57 | """ 58 | http://planning.cs.uiuc.edu/node102.html 59 | :param rpy: 60 | :return: 61 | """ 62 | m_r = rot_x(-rpy[..., 0]) 63 | m_p = rot_y(-rpy[..., 1]) 64 | m_y = rot_z(-rpy[..., 2]) 65 | 66 | return tf.matmul(m_y, tf.matmul(m_p, m_r)) 67 | 68 | def rot_x(a): 69 | cs = tf.cos(a) 70 | sn = tf.sin(a) 71 | if a.shape.ndims == 1: 72 | _ones, _zeros = tf.ones_like(cs), tf.zeros_like(cs) 73 | return tf.transpose(tf.stack([[_ones, _zeros, _zeros], 74 | [_zeros, cs, sn], 75 | [_zeros, -sn, cs]]), (2, 0, 1)) 76 | else: 77 | 78 | return tf.stack([[1., 0., 0.], 79 | [0., cs, sn], 80 | [0., -sn, cs]]) 81 | 82 | def rot_y(a): 83 | cs = tf.cos(a) 84 | sn = tf.sin(a) 85 | 86 | if a.shape.ndims == 1: 87 | _ones, _zeros = tf.ones_like(cs), tf.zeros_like(cs) 88 | 89 | return tf.transpose(tf.stack([[cs, _zeros, -sn], 90 | [_zeros, _ones, _zeros], 91 | [sn, _zeros, cs]]), (2, 0, 1)) 92 | 93 | else: 94 | 95 | return tf.stack([[cs, 0., -sn], 96 | [0., 1., 0.], 97 | [sn, 0., cs]]) 98 | 99 | def rot_z(a): 100 | cs = tf.cos(a) 101 | sn = tf.sin(a) 102 | if a.shape.ndims == 1: 103 | _ones, _zeros = tf.ones_like(cs), tf.zeros_like(cs) 104 | 105 | return tf.transpose(tf.stack([[cs, sn, _zeros], 106 | [-sn, cs, _zeros], 107 | [_zeros, _zeros, _ones]]), (2, 0, 1)) 108 | 109 | else: 110 | 111 | return tf.stack([[cs, sn, 0.], 112 | [-sn, cs, 0.], 113 | [0., 0., 1.]]) 114 | 115 | 116 | 117 | -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/segment.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from .frame import Frame 21 | from .joint import Joint 22 | from .utils import * 23 | 24 | class Segment(object): 25 | def __init__(self, joint, f_tip, child_name='', link=None, fixed=True): 26 | """ 27 | Segment of a kinematic chain 28 | 29 | :param joint: 30 | :type joint: tk.Joint 31 | :param f_tip: 32 | :type f_tip: tk.Frame 33 | """ 34 | self.joint = joint 35 | self.f_tip = joint.pose(0.).inv() * f_tip 36 | 37 | if fixed: 38 | self.f_tip = self.f_tip.fix_it() 39 | 40 | # print (self.f_tip) 41 | # print (self.f_tip.fix_it()) 42 | self.child_name = child_name 43 | 44 | self.link = link 45 | 46 | self.pose_0 = self.pose(0.).fix_it() 47 | 48 | def pose(self, q): 49 | return self.joint.pose(q) * self.f_tip 50 | 51 | def twist(self, q, qdot=0.): 52 | return self.joint.twist(qdot).ref_point( 53 | matvecmul(self.joint.pose(q).m, self.f_tip.p)) 54 | 55 | -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/urdf_utils.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import numpy as np 21 | import tensorflow as tf 22 | from .utils.import_pykdl import * 23 | from .utils.urdf_parser_py.urdf import URDF 24 | from .joint import JointType, Joint, Link 25 | from .frame import Frame 26 | from .segment import Segment 27 | from .chain import Chain 28 | 29 | def euler_to_quat(r, p, y): 30 | sr, sp, sy = np.sin(r / 2.0), np.sin(p / 2.0), np.sin(y / 2.0) 31 | cr, cp, cy = np.cos(r / 2.0), np.cos(p / 2.0), np.cos(y / 2.0) 32 | return [sr * cp * cy - cr * sp * sy, 33 | cr * sp * cy + sr * cp * sy, 34 | cr * cp * sy - sr * sp * cy, 35 | cr * cp * cy + sr * sp * sy] 36 | 37 | def quat_to_rot(q): 38 | x, y, z, w = q 39 | 40 | r = np.array([ 41 | [1.- 2. * (y ** 2 + z ** 2), 2 * (x * y - z * w), 2 * (x * z + y * w)], 42 | [2 * (x * y + z * w), 1.- 2. * (x ** 2 + z ** 2), 2 * (y * z - x * w)], 43 | [2 * (x * z - y * w), 2 * (y * z + x * w), 1.- 2. * (y ** 2 + x ** 2)] 44 | ]) 45 | return r 46 | 47 | 48 | def urdf_pose_to_tk_frame(pose): 49 | pos = [0., 0., 0.] 50 | rot = [0., 0., 0.] 51 | 52 | if pose is not None: 53 | if pose.position is not None: 54 | pos = pose.position 55 | if pose.rotation is not None: 56 | rot = pose.rotation 57 | 58 | return Frame(p=tf.constant(pos, dtype=tf.float32), 59 | m=tf.constant(quat_to_rot(euler_to_quat(*rot)), dtype=tf.float32)) 60 | 61 | def urdf_joint_to_tk_joint(jnt): 62 | origin_frame = urdf_pose_to_tk_frame(jnt.origin) 63 | 64 | if jnt.joint_type == 'revolute': 65 | axis = tf.constant(jnt.axis, dtype=tf.float32) 66 | return Joint(JointType.RotAxis, origin=origin_frame.p , 67 | axis=tf.matmul(origin_frame.m, tf.expand_dims(axis, 1))[:,0], name=jnt.name, 68 | limits=jnt.limit), origin_frame 69 | 70 | if jnt.joint_type == 'fixed' or jnt.joint_type == 'prismatic': 71 | return Joint(JointType.NoneT, name=jnt.name), origin_frame 72 | 73 | print("Unknown joint type: %s." % jnt.joint_type) 74 | 75 | def urdf_link_to_tk_link(lnk): 76 | if lnk.inertial is not None and lnk.inertial.origin is not None: 77 | return Link(frame=urdf_pose_to_tk_frame(lnk.inertial.origin), mass=lnk.inertial.mass) 78 | else: 79 | return Link(frame=urdf_pose_to_tk_frame(None), mass=1.) 80 | 81 | 82 | 83 | def tk_tree_from_urdf_model(urdf): 84 | raise NotImplementedError 85 | root = urdf.get_root() 86 | tree = kdl.Tree(root) 87 | 88 | def add_children_to_tree(parent): 89 | if parent in urdf.child_map: 90 | for joint, child_name in urdf.child_map[parent]: 91 | for lidx, link in enumerate(urdf.links): 92 | if child_name == link.name: 93 | for jidx, jnt in enumerate(urdf.joints): 94 | if jnt.name == joint: 95 | tk_jnt, tk_origin = urdf_joint_to_tk_joint( 96 | urdf.joints[jidx]) 97 | tk_origin = urdf_pose_to_tk_frame(urdf.joints[jidx].origin) 98 | 99 | tree.segments += [Segment(joint=tk_jnt, f_tip=tk_origin, 100 | child_name=child_name)] 101 | 102 | tree.addSegment(kdl_sgm, parent) 103 | add_children_to_tree(child_name) 104 | 105 | add_children_to_tree(root) 106 | return tree 107 | 108 | def kdl_chain_from_urdf_model(urdf, root=None, tip=None, 109 | load_collision=False, mesh_path=None): 110 | 111 | if mesh_path is not None and mesh_path[-1] != '/': mesh_path += '/' 112 | 113 | root = urdf.get_root() if root is None else root 114 | segments = [] 115 | 116 | chain = None if tip is None else urdf.get_chain(root, tip)[1:] 117 | 118 | def add_children_to_chain(parent, segments, chain=None): 119 | if parent in urdf.child_map: 120 | # print "parent:", parent 121 | # print "childs:", urdf.child_map[parent] 122 | # 123 | if chain is not None: 124 | childs = [child for child in urdf.child_map[parent] if child[1] in chain] 125 | if len(childs): 126 | joint, child_name = childs[0] 127 | else: 128 | return 129 | else: 130 | if not len(urdf.child_map[parent]) < 2: 131 | print("Robot is not a chain, taking first branch") 132 | 133 | joint, child_name = urdf.child_map[parent][0] 134 | 135 | # print "child name:", child_name 136 | # for lidx, link in enumerate(urdf.links): 137 | # if child_name == link.name: 138 | # child = urdf.links[lidx] 139 | 140 | for jidx, jnt in enumerate(urdf.joints): 141 | if jnt.name == joint and jnt.joint_type in ['revolute', 'fixed', 'prismatic']: 142 | tk_jnt, tk_origin = urdf_joint_to_tk_joint(urdf.joints[jidx]) 143 | tk_origin = urdf_pose_to_tk_frame(urdf.joints[jidx].origin) 144 | 145 | tk_lnk = urdf_link_to_tk_link(urdf.link_map[child_name]) 146 | 147 | if load_collision and urdf.link_map[child_name].collision is not None: 148 | from stl import mesh 149 | import trimesh 150 | filename = mesh_path + \ 151 | urdf.link_map[child_name].collision.geometry.filename.split('/')[-1] 152 | 153 | # filename = filename[:-14] + '.STL' 154 | # tk_lnk.collision_mesh = mesh.Mesh.from_file(filename) 155 | tk_lnk.collision_mesh = trimesh.load(filename) 156 | 157 | segments += [Segment( 158 | joint=tk_jnt, f_tip=tk_origin, child_name=child_name, link=tk_lnk 159 | )] 160 | 161 | 162 | add_children_to_chain(child_name, segments, chain) 163 | 164 | add_children_to_chain(root, segments, chain) 165 | return Chain(segments) 166 | 167 | 168 | def urdf_from_file(file): 169 | return URDF.from_xml_file(file) -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from .tf_utils import * 21 | from .pykdl_utils import forward_kinematic, FKSolver 22 | from . plot_utils import plot_robot 23 | from . import import_pykdl 24 | from .layout import FkLayout -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/utils/import_pykdl.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | try: 21 | import PyKDL as kdl 22 | import rospy 23 | except ImportError as e: 24 | pass 25 | -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/utils/layout.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from enum import IntEnum 21 | 22 | class FkLayout(IntEnum): 23 | x = 0 # only position 24 | xq = 1 # position - quaternion 25 | xm = 2 # position - vectorized rotation matrix - order 'C' 26 | xmv = 4 # position - vectorized rotation matrix - order 'F' - [x_1, x_2, x_3, y_1, ..., z_3] 27 | f = 3 # frame. 28 | -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/utils/plot_utils.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import matplotlib.pyplot as plt 21 | import numpy as np 22 | 23 | def axis_equal_3d(ax): 24 | extents = np.array([getattr(ax, 'get_{}lim'.format(dim))() for dim in 'xyz']) 25 | sz = extents[:,1] - extents[:,0] 26 | centers = np.mean(extents, axis=1) 27 | maxsize = max(abs(sz)) 28 | r = maxsize/2 29 | for ctr, dim in zip(centers, 'xyz'): 30 | getattr(ax, 'set_{}lim'.format(dim))(ctr - r, ctr + r) 31 | 32 | def plot_robot(xs, color='k', xlim=None,ylim=None, **kwargs): 33 | dims = kwargs.pop('dims', [0, 1]) 34 | 35 | l = plt.plot(xs[:, dims[0]], xs[:,dims[1]], marker='o', color=color, lw=10, mfc='w', solid_capstyle='round', 36 | **kwargs) 37 | 38 | plt.axes().set_aspect('equal') 39 | 40 | if xlim is not None: plt.xlim(xlim) 41 | if ylim is not None: plt.ylim(ylim) 42 | 43 | return l -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/utils/pykdl_utils.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from .import_pykdl import * 21 | import numpy as np 22 | 23 | # @profile 24 | def frame_to_np(frame, layout=None, vec=False): 25 | if vec: 26 | return np.array([ 27 | frame.p[0], frame.p[1], frame.p[2], 28 | frame.M[0, 0], frame.M[1, 0], frame.M[2, 0], 29 | frame.M[0, 1], frame.M[1, 1], frame.M[2, 1], 30 | frame.M[0, 2], frame.M[1, 2], frame.M[2, 2], 31 | ]) 32 | else: 33 | return np.array([ 34 | frame.p[0], frame.p[1], frame.p[2], 35 | frame.M[0, 0], frame.M[0, 1], frame.M[0, 2], 36 | frame.M[1, 0], frame.M[1, 1], frame.M[1, 2], 37 | frame.M[2, 0], frame.M[2, 1], frame.M[2, 2], 38 | ]) 39 | 40 | def forward_kinematic(q, chain): 41 | nb_jnt = len(q) if isinstance(q, list) else q.shape[0] 42 | kdl_array = kdl.JntArray(nb_jnt) 43 | 44 | for j in range(nb_jnt): 45 | kdl_array[j] = q[j] 46 | 47 | end_frame = kdl.Frame() 48 | solver = kdl.ChainFkSolverPos_recursive(chain) 49 | 50 | solver.JntToCart(kdl_array, end_frame) 51 | 52 | return frame_to_np(end_frame) 53 | 54 | class FKSolver(object): 55 | def __init__(self, chain, nb_jnt): 56 | """ 57 | 58 | :param chain: 59 | :param nb_jnt: Number of joints 60 | :type nb_jnt: 61 | """ 62 | self.nb_jnt = nb_jnt 63 | self.chain = chain 64 | self.kdl_array = kdl.JntArray(nb_jnt) 65 | 66 | self.end_frame = kdl.Frame() 67 | self.solver = kdl.ChainFkSolverPos_recursive(chain) 68 | 69 | # @profile 70 | def solve(self, q, vec=False): 71 | for j in range(self.nb_jnt): 72 | self.kdl_array[j] = q[j] 73 | 74 | self.solver.JntToCart(self.kdl_array, self.end_frame) 75 | 76 | return frame_to_np(self.end_frame, vec=vec) 77 | 78 | 79 | class JacSolver(object): 80 | def __init__(self, chain, nb_jnt): 81 | self.nb_jnt = nb_jnt 82 | self.chain = chain 83 | self.kdl_array = PyKDL.JntArray(nb_jnt) 84 | 85 | self.end_jacobian = PyKDL.Jacobian() 86 | self.solver = PyKDL.ChainJntToJacSolver(chain) 87 | 88 | # @profile 89 | def solve(self, q): 90 | for j in range(self.nb_jnt): 91 | self.kdl_array[j] = q[j] 92 | 93 | self.solver.JntToJac(self.kdl_array, self.end_frame) 94 | 95 | return frame_to_np(self.end_frame) 96 | 97 | 98 | -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/utils/tf_utils.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | 22 | def matvecmul(mat, vec): 23 | """ 24 | Matrix-vector multiplication 25 | :param mat: 26 | :param vec: 27 | :return: 28 | """ 29 | return tf.linalg.LinearOperatorFullMatrix(mat).matvec(vec) 30 | # 31 | # if mat.shape.ndims == 2 and vec.shape.ndims == 1: 32 | # return tf.einsum('ij,j->i', mat, vec) 33 | # elif mat.shape.ndims == 2 and vec.shape.ndims == 2: 34 | # return tf.einsum('ij,aj->ai', mat, vec) 35 | # elif mat.shape.ndims == 3 and vec.shape.ndims == 1: 36 | # return tf.einsum('aij,j->ai', mat, vec) 37 | # elif mat.shape.ndims == 3 and vec.shape.ndims == 2: 38 | # return tf.einsum('aij,aj->ai', mat, vec) 39 | # else: 40 | # raise NotImplementedError 41 | 42 | def matmatmul(mat1, mat2): 43 | """ 44 | Matrix-matrix multiplication 45 | :param mat1: 46 | :param mat2: 47 | :return: 48 | """ 49 | return tf.linalg.LinearOperatorFullMatrix(mat1).matmul(mat2) 50 | 51 | # if mat1.shape.ndims == 2 and mat2.shape.ndims == 2: 52 | # return tf.matmul(mat1, mat2) 53 | # elif mat1.shape.ndims == 3 and mat2.shape.ndims == 2: 54 | # return tf.einsum('aij,jk->aik', mat1, mat2) 55 | # elif mat1.shape.ndims == 2 and mat2.shape.ndims == 3: 56 | # return tf.einsum('ij,ajk->aik', mat1, mat2) 57 | # elif mat1.shape.ndims == 3 and mat2.shape.ndims == 3: 58 | # return tf.einsum('aij,ajk->aik', mat1, mat2) 59 | # else: 60 | # raise NotImplementedError 61 | 62 | def angular_vel_tensor(w): 63 | if w.shape.ndims == 1: 64 | return tf.stack([[0., -w[2], w[1]], 65 | [w[2], 0. , -w[0]], 66 | [-w[1], w[0], 0.]]) 67 | else: 68 | di = tf.zeros_like(w[:, 0]) 69 | return tf.transpose( 70 | tf.stack([[di, -w[:, 2], w[:, 1]], 71 | [w[:, 2], di , -w[:, 0]], 72 | [-w[:, 1], w[:, 0], di]]), 73 | perm=(2, 0, 1) 74 | ) 75 | 76 | def drotmat_to_w_jac(r): 77 | """ 78 | 79 | :param r: rotation matrix [:, 3, 3] 80 | :return: 81 | """ 82 | return tf.concat([ 83 | tf.compat.v1.matrix_transpose(angular_vel_tensor(r[:, :, i])) for i in range(3)], axis=1) 84 | 85 | # return tf.concat([angular_vel_tensor(tf.matrix_transpose(r)[:, :, i]) for i in range(3)], axis=1) 86 | 87 | def rot_matrix_gains(r, k): 88 | """ 89 | 90 | r: rotation matrix [:, 3, 3] 91 | k: rotation matrix [:, 6, 6] 92 | :param r: 93 | :return: 94 | """ 95 | 96 | if k.shape.ndims == 2: 97 | k = k[None] * tf.ones_like(r[:, :1, :1]) 98 | 99 | k_mat = tf.concat([ 100 | tf.concat([tf.eye(3, batch_shape=(1,)) + tf.zeros_like(k[:, :3, :3]) , 101 | tf.zeros_like(k[:, :3, :3])], axis=2), 102 | tf.concat([tf.zeros_like(k[:, :1, :3]) * tf.ones((1, 9, 1)) , 103 | drotmat_to_w_jac(r)], axis=2)], 1) 104 | 105 | return tf.linalg.matmul(k, k_mat, transpose_b=True) -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/utils/urdf_parser_py/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/emmanuelpignat/tf_robot_learning/0e89791cfc76fc8481b5a592a74a1a24f33a96b0/tf_robot_learning/kinematic/utils/urdf_parser_py/__init__.py -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/utils/urdf_parser_py/sdf.py: -------------------------------------------------------------------------------- 1 | from . xml_reflection.basics import * 2 | from . import xml_reflection as xmlr 3 | 4 | # What is the scope of plugins? Model, World, Sensor? 5 | 6 | xmlr.start_namespace('sdf') 7 | 8 | class Pose(xmlr.Object): 9 | def __init__(self, vec=None, extra=None): 10 | self.xyz = None 11 | self.rpy = None 12 | if vec is not None: 13 | assert isinstance(vec, list) 14 | count = len(vec) 15 | if len == 3: 16 | xyz = vec 17 | else: 18 | self.from_vec(vec) 19 | elif extra is not None: 20 | assert xyz is None, "Cannot specify 6-length vector and 3-length vector" 21 | assert len(extra) == 3, "Invalid length" 22 | self.rpy = extra 23 | 24 | def from_vec(self, vec): 25 | assert len(vec) == 6, "Invalid length" 26 | self.xyz = vec[:3] 27 | self.rpy = vec[3:6] 28 | 29 | def as_vec(self): 30 | xyz = self.xyz if self.xyz else [0, 0, 0] 31 | rpy = self.rpy if self.rpy else [0, 0, 0] 32 | return xyz + rpy 33 | 34 | def read_xml(self, node): 35 | # Better way to do this? Define type? 36 | vec = get_type('vector6').read_xml(node) 37 | self.load_vec(vec) 38 | 39 | def write_xml(self, node): 40 | vec = self.as_vec() 41 | get_type('vector6').write_xml(node, vec) 42 | 43 | def check_valid(self): 44 | assert self.xyz is not None or self.rpy is not None 45 | 46 | name_attribute = xmlr.Attribute('name', str) 47 | pose_element = xmlr.Element('pose', Pose, False) 48 | 49 | class Entity(xmlr.Object): 50 | def __init__(self, name = None, pose = None): 51 | self.name = name 52 | self.pose = pose 53 | 54 | xmlr.reflect(Entity, params = [ 55 | name_attribute, 56 | pose_element 57 | ]) 58 | 59 | 60 | class Inertia(xmlr.Object): 61 | KEYS = ['ixx', 'ixy', 'ixz', 'iyy', 'iyz', 'izz'] 62 | 63 | def __init__(self, ixx=0.0, ixy=0.0, ixz=0.0, iyy=0.0, iyz=0.0, izz=0.0): 64 | self.ixx = ixx 65 | self.ixy = ixy 66 | self.ixz = ixz 67 | self.iyy = iyy 68 | self.iyz = iyz 69 | self.izz = izz 70 | 71 | def to_matrix(self): 72 | return [ 73 | [self.ixx, self.ixy, self.ixz], 74 | [self.ixy, self.iyy, self.iyz], 75 | [self.ixz, self.iyz, self.izz]] 76 | 77 | xmlr.reflect(Inertia, params = [xmlr.Element(key, float) for key in Inertia.KEYS]) 78 | 79 | # Pretty much copy-paste... Better method? 80 | # Use multiple inheritance to separate the objects out so they are unique? 81 | class Inertial(xmlr.Object): 82 | def __init__(self, mass = 0.0, inertia = None, pose=None): 83 | self.mass = mass 84 | self.inertia = inertia 85 | self.pose = pose 86 | 87 | xmlr.reflect(Inertial, params = [ 88 | xmlr.Element('mass', float), 89 | xmlr.Element('inertia', Inertia), 90 | pose_element 91 | ]) 92 | 93 | 94 | class Link(Entity): 95 | def __init__(self, name = None, pose = None, inertial = None, kinematic = False): 96 | Entity.__init__(self, name, pose) 97 | self.inertial = inertial 98 | self.kinematic = kinematic 99 | 100 | xmlr.reflect(Link, parent_cls = Entity, params = [ 101 | xmlr.Element('inertial', Inertial), 102 | xmlr.Attribute('kinematic', bool, False), 103 | xmlr.AggregateElement('visual', Visual, var = 'visuals'), 104 | xmlr.AggregateElement('collision', Collision, var = 'collisions') 105 | ]) 106 | 107 | 108 | class Model(Entity): 109 | def __init__(self, name = None, pose=None): 110 | Entity.__init__(self, name, pose) 111 | self.links = [] 112 | self.joints = [] 113 | self.plugins = [] 114 | 115 | xmlr.reflect(Model, parent_cls = Entity, params = [ 116 | xmlr.AggregateElement('link', Link, var = 'links'), 117 | xmlr.AggregateElement('joint', Joint, var = 'joints'), 118 | xmlr.AggregateElement('plugin', Plugin, var = 'plugins') 119 | ]) 120 | 121 | xmlr.end_namespace('sdf') -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/utils/urdf_parser_py/xml_reflection/__init__.py: -------------------------------------------------------------------------------- 1 | from . core import * 2 | -------------------------------------------------------------------------------- /tf_robot_learning/kinematic/utils/urdf_parser_py/xml_reflection/basics.py: -------------------------------------------------------------------------------- 1 | import string 2 | import yaml 3 | import collections 4 | from lxml import etree 5 | 6 | # Different implementations mix well it seems 7 | # @todo Do not use this? 8 | from xml.etree.ElementTree import ElementTree 9 | 10 | def xml_string(rootXml, addHeader = True): 11 | # Meh 12 | xmlString = etree.tostring(rootXml, pretty_print = True) 13 | if addHeader: 14 | xmlString = '\n' + xmlString 15 | return xmlString 16 | 17 | def dict_sub(obj, keys): 18 | return dict((key, obj[key]) for key in keys) 19 | 20 | def node_add(doc, sub): 21 | if sub is None: 22 | return None 23 | if type(sub) == str: 24 | return etree.SubElement(doc, sub) 25 | elif isinstance(sub, etree._Element): 26 | doc.append(sub) # This screws up the rest of the tree for prettyprint... 27 | return sub 28 | else: 29 | raise Exception('Invalid sub value') 30 | 31 | def pfloat(x): 32 | return str(x).rstrip('.') 33 | 34 | def xml_children(node): 35 | children = node.getchildren() 36 | def predicate(node): 37 | return not isinstance(node, etree._Comment) 38 | return list(filter(predicate, children)) 39 | 40 | def isstring(obj): 41 | try: 42 | return isinstance(obj, basestring) 43 | except NameError: 44 | return isinstance(obj, str) 45 | 46 | def to_yaml(obj): 47 | """ Simplify yaml representation for pretty printing """ 48 | # Is there a better way to do this by adding a representation with yaml.Dumper? 49 | # Ordered dict: http://pyyaml.org/ticket/29#comment:11 50 | if obj is None or isstring(obj): 51 | out = str(obj) 52 | elif type(obj) in [int, float, bool]: 53 | return obj 54 | elif hasattr(obj, 'to_yaml'): 55 | out = obj.to_yaml() 56 | elif isinstance(obj, etree._Element): 57 | out = etree.tostring(obj, pretty_print = True) 58 | elif type(obj) == dict: 59 | out = {} 60 | for (var, value) in obj.items(): 61 | out[str(var)] = to_yaml(value) 62 | elif hasattr(obj, 'tolist'): 63 | # For numpy objects 64 | out = to_yaml(obj.tolist()) 65 | elif isinstance(obj, collections.Iterable): 66 | out = [to_yaml(item) for item in obj] 67 | else: 68 | out = str(obj) 69 | return out 70 | 71 | class SelectiveReflection(object): 72 | def get_refl_vars(self): 73 | return list(vars(self).keys()) 74 | 75 | class YamlReflection(SelectiveReflection): 76 | def to_yaml(self): 77 | raw = dict((var, getattr(self, var)) for var in self.get_refl_vars()) 78 | return to_yaml(raw) 79 | 80 | def __str__(self): 81 | return yaml.dump(self.to_yaml()).rstrip() # Good idea? Will it remove other important things? 82 | -------------------------------------------------------------------------------- /tf_robot_learning/nn/__init__.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from . mlp import MLP 21 | from . density_mlp import DensityMLP 22 | from . inv_net import real_nvp -------------------------------------------------------------------------------- /tf_robot_learning/nn/density_mlp.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | from tensorflow_probability import distributions as ds 22 | from ..distributions import * 23 | from .mlp import MLP 24 | import numpy as np 25 | 26 | 27 | 28 | param_size = { 29 | MultivariateNormalTriL : lambda x : [x, x ** 2], 30 | MultivariateNormalFullCovariance: lambda x : [x, x ** 2], 31 | MultivariateNormalDiag : lambda x : [x, x], 32 | MultivariateNormalIso : lambda x : [x, 1], 33 | } 34 | 35 | param_shape = { 36 | MultivariateNormalTriL : lambda x : [(x, ), (x, x)], 37 | MultivariateNormalFullCovariance : lambda x : [(x, ), (x, x)], 38 | MultivariateNormalDiag : lambda x : [(x, ), (x, )], 39 | MultivariateNormalIso : lambda x : [(x, ), (1, )] 40 | } 41 | 42 | reparam = { 43 | MultivariateNormalTriL : 44 | [ lambda x : x, lambda x : x], 45 | MultivariateNormalFullCovariance: 46 | [ lambda x : x, lambda x : tf.linalg.expm(x + tf.matrix_transpose(x))], 47 | MultivariateNormalDiag : 48 | [ lambda x : x, lambda x : tf.exp(x)], 49 | MultivariateNormalIso : 50 | [ lambda x : x, lambda x : tf.exp(x)] 51 | } 52 | 53 | 54 | class DensityMLP(MLP): 55 | def __init__( 56 | self, n_input, n_output, n_hidden=[12, 12], batch_size_svi=1, 57 | act_fct=tf.nn.relu, density=MultivariateNormalTriL): 58 | 59 | 60 | self._density = density 61 | 62 | self._n_output_event = n_output 63 | 64 | # compute the shape of the parameters of the distribution 65 | assert density in param_size, "Parameters shape not defined for specified distribution" 66 | 67 | n_output_param = sum(param_size[density](n_output)) 68 | 69 | MLP.__init__( 70 | self, n_input, n_output_param, n_hidden=n_hidden, 71 | batch_size_svi=batch_size_svi, act_fct=act_fct 72 | ) 73 | 74 | def density(self, x, vec_weights=None): 75 | # get parameters of conditional distribution 76 | 77 | vec_params = self.pred(x, vec_weights=vec_weights, avoid_concat=True) 78 | 79 | # vec_params is of shape (batch_size_svi, batch_size_data, n_output_param) 80 | _param_size = param_size[self._density](self._n_output_event) 81 | _param_shape = param_shape[self._density](self._n_output_event) 82 | _reparam = reparam[self._density] 83 | # get slices for each parameters 84 | slices = [slice(s, e) for s, e in zip( 85 | np.cumsum([0] + _param_size), 86 | np.cumsum(_param_size) 87 | )] 88 | 89 | params = [ 90 | r(tf.reshape(vec_params[:, :, sl], (self._batch_size_svi, -1,) + s )) 91 | for sl, s, r in zip(slices, _param_shape, _reparam) 92 | ] 93 | 94 | if self._batch_size_svi == 1: 95 | params = [p[0] for p in params] 96 | 97 | 98 | return self._density(*params) 99 | 100 | 101 | 102 | 103 | -------------------------------------------------------------------------------- /tf_robot_learning/nn/inv_net.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import numpy as np 21 | import tensorflow.compat.v1 as tf 22 | import os 23 | import matplotlib.pyplot as plt 24 | from tensorflow.python.ops import init_ops 25 | 26 | class GlorotSmall(init_ops.VarianceScaling): 27 | """The Glorot uniform initializer, also called Xavier uniform initializer. 28 | 29 | It draws samples from a uniform distribution within [-limit, limit] 30 | where `limit` is `sqrt(6 / (fan_in + fan_out))` 31 | where `fan_in` is the number of input units in the weight tensor 32 | and `fan_out` is the number of output units in the weight tensor. 33 | 34 | Args: 35 | seed: A Python integer. Used to create random seeds. See 36 | `tf.set_random_seed` 37 | for behavior. 38 | dtype: Default data type, used if no `dtype` argument is provided when 39 | calling the initializer. Only floating point types are supported. 40 | 41 | References: 42 | [Glorot et al., 2010](http://proceedings.mlr.press/v9/glorot10a.html) 43 | ([pdf](http://jmlr.org/proceedings/papers/v9/glorot10a/glorot10a.pdf)) 44 | """ 45 | 46 | def __init__(self, scale=1., seed=None, dtype=tf.float32): 47 | super(GlorotSmall, self).__init__( 48 | scale=scale, 49 | mode="fan_avg", 50 | distribution="uniform", 51 | seed=seed, 52 | dtype=dtype) 53 | 54 | def get_config(self): 55 | return {"seed": self.seed, "dtype": self.dtype.name} 56 | 57 | 58 | def dense_resnet(inputs, mid_channels, output_channels, num_blocks, activation=None, scale=1.): 59 | assert len(inputs.shape) == 2 60 | initializer = GlorotSmall(scale=scale) 61 | 62 | def _blocks(_x, name): 63 | shortcut = _x 64 | _x = tf.layers.dense(_x, mid_channels, activation=activation, name=name+'1', kernel_initializer=initializer) 65 | _x = tf.layers.dense(_x, mid_channels, activation=activation, name=name+'2', kernel_initializer=initializer) 66 | return _x + shortcut 67 | 68 | inputs = tf.layers.dense(inputs, mid_channels, activation=activation , name='initial', kernel_initializer=initializer) 69 | 70 | for i in range(num_blocks): 71 | inputs = _blocks(inputs, '{}'.format(i)) 72 | inputs = tf.layers.dense(inputs, mid_channels, activation=None, name='final') 73 | return inputs 74 | 75 | #Define mask 76 | def get_mask(inputs, reverse_mask, data_format='NHWC', dtype=tf.float32): 77 | shape = inputs.get_shape().as_list() 78 | if len(shape) == 2: 79 | N = shape[-1] 80 | range_n = tf.range(N) 81 | odd_ind = tf.mod(range_n, 2) 82 | 83 | odd_ind = tf.reshape(odd_ind, [-1, N]) 84 | checker = odd_ind 85 | 86 | 87 | elif len(shape) == 4: 88 | H = shape[2] if data_format == 'NCHW' else shape[1] 89 | W = shape[3] if data_format == 'NCHW' else shape[2] 90 | 91 | range_h = tf.range(H) 92 | range_w = tf.range(W) 93 | 94 | odd_ind_h = tf.cast(tf.mod(range_h, 2), dtype=tf.bool) 95 | odd_ind_w = tf.cast(tf.mod(range_w, 2), dtype=tf.bool) 96 | 97 | odd_h = tf.tile(tf.expand_dims(odd_ind_h, -1), [1, W]) 98 | odd_w = tf.tile(tf.expand_dims(odd_ind_w, 0), [H, 1]) 99 | 100 | checker = tf.logical_xor(odd_h, odd_w) 101 | 102 | reshape = [-1, 1, H, W] if data_format == 'NCHW' else [-1, H, W, 1] 103 | checker = tf.reshape(checker, reshape) 104 | 105 | 106 | else: 107 | raise ValueError('Invalid tensor shape. Dimension of the tensor shape must be ' 108 | '2 (NxD) or 4 (NxCxHxW or NxHxWxC), got {}.'.format(inputs.get_shape().as_list())) 109 | 110 | 111 | if checker.dtype != dtype: 112 | checker = tf.cast(checker, dtype) 113 | 114 | if reverse_mask: 115 | checker = 1. - checker 116 | 117 | return checker 118 | 119 | # Define coupling layer 120 | def coupling_layer(inputs, mid_channels, num_blocks, reverse_mask, activation=None, 121 | name='coupling_layer', backward=False, reuse=None, scale=1.): 122 | mask = get_mask(inputs, reverse_mask) 123 | with tf.variable_scope(name) as scope: 124 | if reuse: 125 | scope.reuse_variables() 126 | 127 | if backward: 128 | v1 = inputs * mask 129 | v2 = inputs * (1-mask) 130 | with tf.variable_scope('st1'): 131 | st1 = dense_resnet( 132 | inputs=v1, mid_channels=mid_channels, 133 | output_channels=inputs.get_shape().as_list()[1]*2, num_blocks=3, 134 | activation=activation, scale=scale 135 | ) 136 | s1 = st1[:, 0:tf.shape(inputs)[1]] 137 | rescale1 = tf.get_variable('rescale_s', shape=[inputs.get_shape().as_list()[1]], dtype=tf.float32, initializer=tf.constant_initializer(1.)) 138 | s1 = rescale1 * tf.nn.tanh(s1) 139 | t1 = st1[:, tf.shape(inputs)[1]:tf.shape(inputs)[1]*2] 140 | 141 | u2 = (1-mask)*(v2 - t1)*tf.exp(-s1) 142 | 143 | with tf.variable_scope('st2'): 144 | st2 = dense_resnet( 145 | inputs=u2, mid_channels=mid_channels, 146 | output_channels=inputs.get_shape().as_list()[1]*2, num_blocks=3, 147 | activation=activation, scale=scale 148 | ) 149 | s2 = st2[:, 0:tf.shape(inputs)[1]] 150 | rescale2 = tf.get_variable('rescale_s', shape=[inputs.get_shape().as_list()[1]], dtype=tf.float32, initializer=tf.constant_initializer(1.)) 151 | s2 = rescale2 * tf.nn.tanh(s2) 152 | t2 = st2[:, tf.shape(inputs)[1]:tf.shape(inputs)[1]*2] 153 | 154 | u1 = mask * (v1 - t2)*tf.exp(-s2) 155 | inputs = u1 + u2 156 | 157 | else: 158 | u1 = inputs * mask 159 | u2 = inputs * (1-mask) 160 | 161 | with tf.variable_scope('st2'): 162 | st2 = dense_resnet( 163 | inputs=u2, mid_channels=mid_channels, 164 | output_channels=inputs.get_shape().as_list()[1]*2, num_blocks=3, 165 | activation=activation, scale=scale 166 | ) 167 | s2 = st2[:, 0:tf.shape(inputs)[1]] 168 | rescale2 = tf.get_variable('rescale_s', shape=[inputs.get_shape().as_list()[1]], dtype=tf.float32, initializer=tf.constant_initializer(1.)) 169 | s2 = rescale2 * tf.nn.tanh(s2) 170 | t2 = st2[:, tf.shape(inputs)[1]:tf.shape(inputs)[1]*2] 171 | 172 | v1 = mask * (u1 * tf.exp(s2) + t2) 173 | 174 | with tf.variable_scope('st1'): 175 | st1 = dense_resnet( 176 | inputs=v1, mid_channels=mid_channels, 177 | output_channels=inputs.get_shape().as_list()[1]*2, num_blocks=3, 178 | activation=activation,scale=scale 179 | ) 180 | s1 = st1[:, 0:tf.shape(inputs)[1]] 181 | rescale1 = tf.get_variable('rescale_s', shape=[inputs.get_shape().as_list()[1]], dtype=tf.float32, initializer=tf.constant_initializer(1.)) 182 | s1 = rescale1 * tf.nn.tanh(s1) 183 | t1 = st1[:, tf.shape(inputs)[1]:tf.shape(inputs)[1]*2] 184 | 185 | 186 | v2 = (1-mask) * (u2 * tf.exp(s1) + t1) 187 | inputs = v1 + v2 188 | 189 | return inputs 190 | 191 | # Code from https://github.com/chrischute/real-nvp 192 | def preprocess(x): 193 | data_constraint = 0.9 194 | y = (x*255. + tf.random.uniform(tf.shape(x), 0, 1))/256. 195 | y = (2 * y - 1) * data_constraint 196 | y = (y + 1) / 2 197 | y = tf.log(y) - tf.log(1-y) 198 | 199 | ldj = tf.nn.softplus(y) + tf.nn.softplus(-y) - tf.nn.softplus(tf.log(1-data_constraint) - tf.log(data_constraint)) 200 | sldj = tf.reduce_sum(tf.reshape(ldj, [tf.shape(ldj)[0], -1]), axis=-1) 201 | return y, sldj 202 | 203 | 204 | def get_nvp_trainable_variables(): 205 | variables = [] 206 | 207 | for i in range(1, 5): 208 | variables += tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='c%d' % i) 209 | 210 | return variables 211 | 212 | def real_nvp(inputs, mid_channels, backward=False, reuse=False, 213 | activation=tf.nn.tanh, scale_kernel=1., name=''): 214 | # 215 | x = inputs 216 | if backward: 217 | 218 | x = coupling_layer(x, mid_channels, 4, activation=activation, 219 | reverse_mask=True, name=name + 'c4', backward=backward, reuse=reuse, 220 | scale=scale_kernel) 221 | x = coupling_layer(x, mid_channels, 4, activation=activation, 222 | reverse_mask=False, name= name + 'c3', backward=backward, reuse=reuse, 223 | scale=scale_kernel) 224 | x = coupling_layer(x, mid_channels, 4, activation=activation, 225 | reverse_mask=True, name=name + 'c2', backward=backward, reuse=reuse, 226 | scale=scale_kernel) 227 | x = coupling_layer(x, mid_channels, 4, activation=activation, 228 | reverse_mask=False, name=name +'c1', backward=backward, reuse=reuse, 229 | scale=scale_kernel) 230 | else: 231 | 232 | # x, sldj = preprocess(inputs) 233 | x = coupling_layer(x, mid_channels, 4, activation=activation, 234 | reverse_mask=False, name=name + 'c1', backward=backward, reuse=reuse, 235 | scale=scale_kernel) 236 | x = coupling_layer(x, mid_channels, 4, activation=activation, 237 | reverse_mask=True, name=name + 'c2', backward=backward, reuse=reuse, 238 | scale=scale_kernel) 239 | x = coupling_layer(x, mid_channels, 4, activation=activation, 240 | reverse_mask=False, name=name + 'c3', backward=backward, reuse=reuse, 241 | scale=scale_kernel) 242 | x = coupling_layer(x, mid_channels, 4, activation=activation, 243 | reverse_mask=True, name=name + 'c4', backward=backward, reuse=reuse, 244 | scale=scale_kernel) 245 | # x = tf.nn.sigmoid(x) 246 | return x -------------------------------------------------------------------------------- /tf_robot_learning/nn/mlp.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | from tensorflow_probability import distributions as ds 22 | from tensorflow.python.framework import constant_op 23 | from tensorflow.python.framework import dtypes 24 | from tensorflow.python.ops import array_ops 25 | from tensorflow.python.ops import random_ops 26 | import numpy as np 27 | 28 | class MLP(object): 29 | def __init__(self, n_input, n_output, n_hidden=[12, 12], batch_size_svi=1, 30 | act_fct=tf.nn.relu, last_linear=True, out_max=1., *args, **kwargs): 31 | """ 32 | 33 | :param n_input: 34 | :param n_output: 35 | :param n_hidden: 36 | :param batch_size_svi: 37 | :param act_fct: 38 | :param last_linear: make last layer linear 39 | :param out_max: maximum value of the output in case of sigmoid of tanh [-out_max, out_max] 40 | :param args: 41 | :param kwargs: 42 | """ 43 | self._act_fct = act_fct 44 | self._batch_size_svi = batch_size_svi 45 | self._weights_size = 0 46 | self._layer_size = len(n_hidden) + 1 47 | 48 | self._last_linear = last_linear 49 | self._out_max = out_max 50 | 51 | shapes = np.array([n_input] + n_hidden + [n_output]) 52 | 53 | self._n_input = n_input 54 | self._n_hidden = n_hidden 55 | self._n_output = n_output 56 | self._shapes = shapes 57 | 58 | # add weights n_input * n_hidden_1 + n_hidden_1 * n_hidden_2 + n_hidden_2 * n_output 59 | self._weights_size += np.sum(shapes[:-1] * shapes[1:]) 60 | 61 | # add biased 62 | self._weights_size += np.sum(shapes[1:]) 63 | 64 | self._vec_weights = None 65 | 66 | def load(self, path, cst=True): 67 | if cst: 68 | self._vec_weights = tf.constant(np.load(path + '.npy'), dtype=tf.float32) 69 | else: 70 | self._vec_weights = tf.Variable(np.load(path + '.npy'), dtype=tf.float32) 71 | @property 72 | def vec_weights(self): 73 | if self._vec_weights is None: 74 | self._vec_weights = tf.Variable(self.glorot_init(self._batch_size_svi)) 75 | 76 | return self._vec_weights 77 | 78 | def save(self, path, sess=None): 79 | sess = tf.compat.v1.get_default_session() 80 | 81 | _vec_weights = sess.run(self._vec_weights) 82 | 83 | np.save(path, _vec_weights) 84 | 85 | def glorot_init(self, batch_size=None): 86 | if batch_size is None: 87 | batch_size = self._batch_size_svi 88 | 89 | init = [] 90 | idx_s = 0 91 | idx_layer = 0 92 | 93 | 94 | for s0, s1 in zip(self._shapes[:-1], self._shapes[1:]): 95 | idx_e = idx_s + s0 * s1 96 | 97 | init += [ 98 | tf.random.normal( 99 | shape=[batch_size ,idx_e-idx_s], 100 | stddev=1. / np.sqrt(idx_e-idx_s / 2.), dtype=tf.float32)] 101 | 102 | idx_layer += 1 103 | idx_s = idx_e 104 | 105 | 106 | idx_layer = 0 107 | 108 | biases = {} 109 | 110 | for s0 in self._shapes[1:]: 111 | idx_e = idx_s + s0 112 | 113 | init += [ 114 | tf.random.normal( 115 | shape=[batch_size, idx_e - idx_s], 116 | stddev=1. / np.sqrt(idx_e - idx_s / 2.), dtype=tf.float32)] 117 | 118 | # biases['b%d' % idx_layer] = vec_weights[:, idx_s:idx_e] 119 | 120 | idx_layer += 1 121 | idx_s = idx_e 122 | 123 | return tf.concat(init, axis=1) 124 | 125 | 126 | @property 127 | def weights_size(self): 128 | return self._weights_size 129 | 130 | @property 131 | def weights_shape(self): 132 | return [self._batch_size_svi, self._weights_size] 133 | 134 | def unpack_weights(self, vec_weights): 135 | idx_s = 0 136 | idx_layer = 0 137 | 138 | weights = {} 139 | 140 | for s0, s1 in zip(self._shapes[:-1], self._shapes[1:]): 141 | idx_e = idx_s + s0 * s1 142 | 143 | weights['h%d' % idx_layer] = tf.reshape(vec_weights[:, idx_s:idx_e], (-1, s0, s1)) 144 | 145 | idx_layer += 1 146 | idx_s = idx_e 147 | 148 | 149 | idx_layer = 0 150 | 151 | biases = {} 152 | 153 | for s0 in self._shapes[1:]: 154 | idx_e = idx_s + s0 155 | 156 | biases['b%d' % idx_layer] = vec_weights[:, idx_s:idx_e] 157 | 158 | idx_layer += 1 159 | idx_s = idx_e 160 | 161 | return weights, biases 162 | 163 | def pred(self, x, vec_weights=None, avoid_concat=False): 164 | """ 165 | :param x: [batch_size, n_input] 166 | :return [batch_size_svi, batch_size, n_output] 167 | """ 168 | 169 | if vec_weights is None: 170 | vec_weights = self.vec_weights 171 | 172 | weights, biases = self.unpack_weights(vec_weights) 173 | if x.shape.ndims <= 2: 174 | batch_size = tf.shape(x)[0] 175 | layer = tf.tile(tf.expand_dims(tf.identity(x), 0), [self._batch_size_svi, 1, 1]) 176 | else: 177 | batch_size = tf.shape(x)[-2] 178 | layer = tf.identity(x) 179 | 180 | for i in range(self._layer_size): 181 | bias = tf.tile(tf.expand_dims(biases['b%d' % i], 1), [1, batch_size, 1]) 182 | if i == self._layer_size - 1 and self._last_linear: 183 | layer = tf.add(tf.matmul(layer, weights['h%d' % i]), bias) 184 | elif i == self._layer_size -1 : 185 | layer = self._act_fct( 186 | tf.add(tf.matmul(layer, weights['h%d' % i]), bias)) 187 | 188 | if self._act_fct is tf.nn.sigmoid: 189 | layer = (layer - 0.5) * 2. * self._out_max 190 | elif self._act_fct is tf.nn.tanh: 191 | layer = layer * self._out_max 192 | 193 | else: 194 | layer = self._act_fct( 195 | tf.add(tf.matmul(layer, weights['h%d' % i]), bias)) 196 | 197 | if self._batch_size_svi == 1 and not avoid_concat: 198 | return layer[0] 199 | else: 200 | return layer 201 | 202 | def RNN(self, x, vec_weights): 203 | 204 | weights, biases = self.unpack_weights(vec_weights) 205 | 206 | # Prepare data shape to match `rnn` function requirements 207 | # Current data input shape: (batch_size, timesteps, n_input) 208 | # Required shape: 'timesteps' tensors list of shape (batch_size, n_input) 209 | 210 | batch_size = tf.shape(x)[0] 211 | timesteps = tf.shape(x)[1] 212 | # Unstack to get a list of 'timesteps' tensors of shape (batch_size, n_input) 213 | x = tf.unstack(x, axis=1) 214 | i = 1 215 | bias = tf.tile(tf.expand_dims(biases['b%d' % i], 1), [1, batch_size, 1]) 216 | bias = tf.tile(tf.expand_dims(bias, 1), [1, timesteps, 1, 1]) 217 | 218 | # Define a lstm cell with tensorflow 219 | lstm_cell = tf.contrib.rnn.LSTMBlockCell(self._n_hidden[0], forget_bias=1.0) 220 | # Get lstm cell output 221 | outputs, states = tf.nn.static_rnn(lstm_cell, x, dtype=tf.float32) 222 | outputs = tf.stack(outputs) 223 | # Linear activation, using rnn inner loop last output 224 | layer = tf.add(tf.einsum('tbh, ahj-> atbj', outputs, weights['h%d' % i]), bias) 225 | 226 | return layer 227 | 228 | -------------------------------------------------------------------------------- /tf_robot_learning/planar_robots/__init__.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from .robot import Robot 21 | from .two_joint import TwoJointRobot 22 | from .three_joint import ThreeJointRobot 23 | from .n_joint import NJointRobot 24 | from .bimanual_robot import BimanualThreeJointRobot -------------------------------------------------------------------------------- /tf_robot_learning/planar_robots/bimanual_robot.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | from .robot import Robot 22 | from .three_joint import ThreeJointRobot 23 | from tensorflow_probability import distributions as ds 24 | 25 | pi = 3.14159 26 | 27 | class BimanualThreeJointRobot(Robot): 28 | def __init__(self, q=None, dq=None, ddq=None, ls=None, session=None): 29 | Robot.__init__(self) 30 | 31 | self._ls = tf.constant([0.25, 0.25, 0.25, 0.25, 0.25]) if ls is None else ls 32 | 33 | 34 | margin = 0.02 35 | self._joint_limits = tf.constant([[0. + margin, pi - margin], 36 | [-pi + margin, pi - margin], 37 | [-pi + margin, pi - margin], 38 | [-pi + margin, pi - margin], 39 | [-pi + margin, pi - margin]], 40 | dtype=tf.float32) 41 | self._arms = [ 42 | ThreeJointRobot(ls=tf.gather(self._ls, [0, 1, 2])), 43 | ThreeJointRobot(ls=tf.gather(self._ls, [0, 1, 2])), 44 | ] 45 | 46 | self._dof = 5 47 | 48 | def joint_limit_cost(self, q, std=0.1): 49 | qs = [q[:, 0:3], tf.concat([q[:, 0][:, None], q[:, 3:5]], axis=1)] 50 | return self._arms[0].joint_limit_cost(qs[0], std=std) + self._arms[0].joint_limit_cost(qs[1], std=std) 51 | 52 | def xs(self, q, concat=True): 53 | if q.shape.ndims == 1: 54 | qs = [q[0:3], tf.concat([q[0][None], q[3:5]], axis=0)] 55 | return tf.concat([self._arms[0].xs(qs[0]), self._arms[1].xs(qs[0])], axis=0) 56 | 57 | else: 58 | qs = [q[:, 0:3], tf.concat([q[:, 0][:, None], q[:, 3:5]], axis=1)] 59 | 60 | fks = [ 61 | self._arms[0].xs(qs[0]), 62 | self._arms[1].xs(qs[1]) 63 | ] 64 | if concat: 65 | return tf.concat(fks, axis=1) 66 | else: 67 | return fks 68 | 69 | -------------------------------------------------------------------------------- /tf_robot_learning/planar_robots/n_joint.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | from .robot import Robot 22 | pi = 3.14159 23 | 24 | class NJointRobot(Robot): 25 | def __init__(self, n=3, session=None): 26 | Robot.__init__(self) 27 | 28 | self._ls = tf.constant(1/float(n) * tf.ones(n)) 29 | 30 | margin = 0.02 31 | 32 | self._joint_limits = tf.constant([ 33 | [0. + margin, pi - margin]] + 34 | [[-pi + margin, pi - margin]] * (n - 1), dtype=tf.float32) 35 | 36 | self._dof = n 37 | 38 | def xs(self, q): 39 | if q.shape.ndims == 1: 40 | q_currents = tf.cumsum(q) 41 | 42 | x = tf.cumsum(self._ls * tf.cos(q_currents)) 43 | x = tf.concat([tf.zeros(1), x], 0) 44 | y = tf.cumsum(self._ls * tf.sin(q_currents)) 45 | y = tf.concat([tf.zeros(1), y], 0) 46 | 47 | return tf.transpose(tf.stack([x, y])) 48 | 49 | 50 | else: 51 | q_currents = tf.cumsum(q, axis=1) 52 | x = tf.cumsum(self._ls[None] * tf.cos(q_currents), axis=1) 53 | x = tf.concat([tf.zeros_like(x[:, 0][:, None]), x], axis=1) 54 | y = tf.cumsum(self._ls[None] * tf.sin(q_currents), axis=1) 55 | y = tf.concat([tf.zeros_like(y[:, 0][:, None]), y], axis=1) 56 | 57 | return tf.concat([x[..., None], y[..., None]], axis=-1) 58 | 59 | def J(self, q): 60 | if q.shape.ndims == 1: 61 | X = self.xs(q)[-1] 62 | dX = tf.reshape(tf.convert_to_tensor([tf.gradients(Xi, q) for Xi in tf.unstack(X)]), (2, self.dof)) 63 | return dX 64 | else: 65 | list_dX = [] 66 | for q_single in tf.unstack(q, axis=0): 67 | X = self.xs(q_single)[-1] 68 | dX = tf.reshape(tf.convert_to_tensor([tf.gradients(Xi, q_single) 69 | for Xi in tf.unstack(X)]), (2, self.dof)) 70 | list_dX.append(dX) 71 | dX = tf.stack(list_dX) 72 | return dX 73 | 74 | 75 | -------------------------------------------------------------------------------- /tf_robot_learning/planar_robots/robot.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | 22 | from tensorflow_probability import distributions as ds 23 | 24 | pi = 3.14159 25 | 26 | class Robot(object): 27 | def __init__(self): 28 | # robot description 29 | self._ls, self._ms, self._ins = [], [], [] 30 | 31 | self._xs, self._J, self._Mn = None, None, None 32 | self._M_chol = None 33 | 34 | self._Mq = None 35 | 36 | self._mass = None 37 | self._joint_limits = None 38 | self._dof = None 39 | 40 | self._q_ev = None 41 | self._ev = {} 42 | 43 | @property 44 | def q_ev(self): 45 | if self._q_ev is None: 46 | self._q_ev = tf.compat.v1.placeholder(tf.float32, (None, self.dof)) 47 | 48 | return self._q_ev 49 | 50 | def __getitem__(self, item): 51 | """ 52 | Get numpy version of the tensorflow function 53 | :param item: ['xs', 'J', ...] 54 | :return: 55 | """ 56 | assert hasattr(self, item), 'Robot does not have this attribute' 57 | 58 | if not item in self._ev: 59 | self._ev[item] = getattr(self, item)(self.q_ev) 60 | 61 | sess = tf.compat.v1.get_default_session() 62 | 63 | return lambda x: sess.run(self._ev[item], {self._q_ev: x}) 64 | 65 | @property 66 | def dof(self): 67 | return self._dof 68 | 69 | def gq(self, q, g): 70 | """ 71 | Gravity 72 | :param q: 73 | :param g: 74 | :return: 75 | """ 76 | raise NotImplementedError 77 | 78 | def Cq(self, q, dq): 79 | """ 80 | Coriolis 81 | :param q: 82 | :param dq: 83 | :return: 84 | """ 85 | raise NotImplementedError 86 | 87 | def Mq(self, q): 88 | """ 89 | Joint Inertia matrix 90 | :param q: 91 | :return: 92 | """ 93 | raise NotImplementedError 94 | 95 | def Mq_inv(self, q): 96 | """ 97 | Inverse of Joint Inertia matrix 98 | :param q: 99 | :return: 100 | """ 101 | raise NotImplementedError 102 | 103 | 104 | def Js_com(self, q): 105 | """ 106 | Jacobian of center of masses 107 | :param q: 108 | :return: 109 | """ 110 | raise NotImplementedError 111 | 112 | def f(self, x_t, u_t): 113 | """ 114 | Dynamic equation 115 | :param x_t: [q, dq] in the shape [bs, 4] 116 | :param u_t: tau 117 | (x_t [bs, xi_dim], u_t [bs, u_dim]) -> (x_t+1 [bs, xi_dim]) 118 | :return: 119 | """ 120 | raise NotImplementedError 121 | 122 | def joint_limit_cost(self, q, std=0.1): 123 | return -ds.Normal(self._joint_limits[:, 0], std).log_cdf(q) - \ 124 | ds.Normal(-self._joint_limits[:, 1], std).log_cdf(-q) 125 | 126 | @property 127 | def mass(self): 128 | """ 129 | Mass matrix of each segment 130 | :return: 131 | """ 132 | raise NotImplementedError 133 | 134 | def J(self, q): 135 | raise NotImplementedError 136 | 137 | def Mn(self, q): 138 | """ 139 | Manipulabililty 140 | :param q: 141 | :return: 142 | """ 143 | return tf.matmul(self.J(q), self.J(q), transpose_b=True) 144 | 145 | 146 | def xs(self, q): 147 | raise NotImplementedError 148 | 149 | @property 150 | def ms(self): 151 | """ 152 | Mass of each segment 153 | :return: 154 | """ 155 | return self._ms 156 | 157 | @property 158 | def ins(self): 159 | """ 160 | Inertia moment of each segment 161 | """ 162 | return self._ins 163 | 164 | @property 165 | def ls(self): 166 | return self._ls 167 | 168 | def segment_samples(self, q, nsamples_segment=10, noise_scale=None): 169 | 170 | if q.shape.ndims == 2: 171 | segments = self.xs(q) # batch_shape, n_points, 2 172 | 173 | n_segments = segments.shape[1].value - 1 174 | samples = [] 175 | 176 | for i in range(n_segments): 177 | u = tf.random_uniform((nsamples_segment,)) 178 | 179 | # linear interpolations between end-points of segments 180 | # batch_shape, nsamples_segment, 2 181 | samples += [u[None, :, None] * segments[:, i][:, None] 182 | + (1. - u[None, :, None]) * segments[:, i + 1][:, None]] 183 | 184 | if noise_scale is not None: 185 | samples[-1] += tf.random_normal(samples[-1].shape, 0., noise_scale) 186 | 187 | return tf.concat(samples, axis=1) 188 | 189 | else: 190 | raise NotImplementedError 191 | 192 | def min_sq_dist_from_point(self, q, x, **kwargs): 193 | """ 194 | 195 | :param q: [batch_shape, 2] 196 | :param x: [2, ] 197 | :return: 198 | """ 199 | if q.shape.ndims == 2: 200 | samples = self.segment_samples(q, **kwargs) # batch_shape, nsamples, 2 201 | 202 | dist = tf.reduce_sum((samples - x[None, None]) ** 2, axis=2) 203 | 204 | return tf.reduce_min(dist, axis=1) 205 | 206 | def plot(self, qs, ax=None, color='k', xlim=None, ylim=None, feed_dict={}, 207 | ee_kwargs=None, 208 | dx=0.02, dy=0.02, fontsize=10, cmap=None, text=True, bicolor=None, 209 | x_base=None, **kwargs): 210 | import numpy as np 211 | import matplotlib.pyplot as plt 212 | 213 | qs = np.array(qs).astype(np.float32) 214 | 215 | if qs.ndim == 1: 216 | qs = qs[None] 217 | 218 | if x_base is None: 219 | xs = self['xs'](qs) 220 | else: 221 | raise NotImplementedError 222 | 223 | if text: 224 | for i in range(xs.shape[1] - 1): 225 | plt.annotate(r'$q_%d$' % i, 226 | (xs[0, i, 0], xs[0, i, 1]), 227 | (xs[0, i, 0] + dx, xs[0, i, 1] + dy), fontsize=fontsize) 228 | 229 | if hasattr(self, '_arms'): 230 | n_joint_arms = xs.shape[1] / 2 231 | for x in xs[:, :2]: 232 | plot = ax if ax is not None else plt 233 | plot.plot(x[:, 0], x[:, 1], marker='o', color='k', lw=10, mfc='w', 234 | solid_capstyle='round', 235 | **kwargs) 236 | 237 | for x in xs[:, 1:n_joint_arms]: 238 | plot = ax if ax is not None else plt 239 | plot.plot(x[:, 0], x[:, 1], marker='o', color=color, lw=10, mfc='w', 240 | solid_capstyle='round', 241 | **kwargs) 242 | 243 | for x in xs[:, n_joint_arms + 1:]: 244 | plot = ax if ax is not None else plt 245 | bicolor = color if bicolor is None else bicolor 246 | plot.plot(x[:, 0], x[:, 1], marker='o', color=bicolor, lw=10, mfc='w', 247 | solid_capstyle='round', 248 | **kwargs) 249 | 250 | else: 251 | alpha = kwargs.pop('alpha', 1.) 252 | plot = ax if ax is not None else plt 253 | plot.plot([], [], marker='o', color=color, lw=10, mfc='w', 254 | solid_capstyle='round', alpha=1., 255 | **kwargs) 256 | 257 | kwargs.pop('label', None) 258 | kwargs['alpha'] = alpha 259 | 260 | for i, x in enumerate(xs): 261 | c = color if cmap is None else cmap(float(i) / (len(xs))) 262 | 263 | plot = ax if ax is not None else plt 264 | plot.plot(x[:, 0], x[:, 1], marker='o', color=c, lw=10, mfc='w', 265 | solid_capstyle='round', 266 | **kwargs) 267 | 268 | kwargs.pop('label', None) 269 | 270 | if ee_kwargs is not None: 271 | ee_kwargs['marker'] = ee_kwargs.pop('marker', 'x') 272 | ee_kwargs['ls'] = ' ' 273 | plot = ax if ax is not None else plt 274 | plot.plot(xs[:, -1, 0], xs[:, -1, 1], **ee_kwargs) 275 | 276 | if ax is None: 277 | plt.axes().set_aspect('equal') 278 | else: 279 | ax.set_aspect('equal') 280 | 281 | if xlim is not None: plt.xlim(xlim) 282 | if ylim is not None: plt.ylim(ylim) 283 | -------------------------------------------------------------------------------- /tf_robot_learning/planar_robots/three_joint.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | from .robot import Robot 22 | from tensorflow_probability import distributions as ds 23 | 24 | pi = 3.14159 25 | 26 | class ThreeJointRobot(Robot): 27 | def __init__(self, q=None, dq=None, ddq=None, ls=None, session=None): 28 | Robot.__init__(self) 29 | 30 | self._ls = tf.constant([0.24, 0.25, 0.25]) if ls is None else ls 31 | 32 | margin = 0.02 33 | self._joint_limits = tf.constant([[0. + margin, pi - margin], 34 | [-pi + margin, pi - margin], 35 | [-pi + margin, pi - margin]], 36 | dtype=tf.float32) 37 | 38 | self._base_limits = tf.constant([[-1., 1.], 39 | [-1., 1.]], 40 | dtype=tf.float32) 41 | 42 | self._dof = 3 43 | 44 | def xs(self, q, x_base=None, angle=False): 45 | """ 46 | 47 | :param q: 48 | :param x_base: [2] or [batch_size, 2] 49 | :param angle: 50 | :return: 51 | """ 52 | if q.shape.ndims == 1: 53 | x = tf.cumsum([0, 54 | self.ls[0] * tf.cos(q[0]), 55 | self.ls[1] * tf.cos(q[0] + q[1]), 56 | self.ls[2] * tf.cos(q[0] + q[1] + q[2])]) 57 | y = tf.cumsum([0, 58 | self.ls[0] * tf.sin(q[0]), 59 | self.ls[1] * tf.sin(q[0] + q[1]), 60 | self.ls[2] * tf.sin(q[0] + q[1] + q[2])]) 61 | 62 | if x_base is not None: 63 | x += x_base[0] 64 | y += x_base[1] 65 | 66 | if angle: 67 | return tf.transpose(tf.stack([x, y, tf.cumsum(tf.concat([q, tf.zeros_like(q[0][None])], axis=0))])) 68 | else: 69 | return tf.transpose(tf.stack([x, y])) 70 | else: 71 | x = tf.cumsum([tf.zeros_like(q[:, 0]), 72 | self.ls[None, 0] * tf.cos(q[:, 0]), 73 | self.ls[None, 1] * tf.cos(q[:, 0] + q[:, 1]), 74 | self.ls[None, 2] * tf.cos(q[:, 0] + q[:, 1] + q[:, 2])]) 75 | 76 | y = tf.cumsum([tf.zeros_like(q[:, 0]), 77 | self.ls[None, 0] * tf.sin(q[:, 0]), 78 | self.ls[None, 1] * tf.sin(q[:, 0] + q[:, 1]), 79 | self.ls[None, 2] * tf.sin(q[:, 0] + q[:, 1] + q[:, 2])]) 80 | 81 | if x_base is not None: 82 | x += x_base[:, 0] 83 | y += x_base[:, 1] 84 | 85 | if angle: 86 | return tf.transpose(tf.stack([x, y, tf.cumsum(tf.concat([tf.transpose(q), tf.zeros_like(q[:, 0][None])], axis=0))]), (2, 1, 0)) 87 | else: 88 | return tf.transpose(tf.stack([x, y]), (2, 1, 0)) 89 | 90 | 91 | 92 | def base_limit_cost(self, x, std=0.1, base_limit=1.): 93 | return -ds.Normal(base_limit * self._base_limits[:, 0], std).log_cdf(x) - ds.Normal(-base_limit * self._base_limits[:, 1], std).log_cdf(-x) 94 | 95 | 96 | def J(self, q): 97 | q0 = q[:, 0] 98 | q01 = q[:, 0] + q[:, 1] 99 | q012 = q[:, 0] + q[:, 1] + q[:, 2] 100 | 101 | J = [[0. for i in range(3)] for j in range(2)] 102 | 103 | J[0][2] = self.ls[2] * -tf.sin(q012) 104 | J[1][2] = self.ls[2] * tf.cos(q012) 105 | J[0][1] = self.ls[1] * -tf.sin(q01) + J[0][2] 106 | J[1][1] = self.ls[1] * tf.cos(q01) + J[1][2] 107 | J[0][0] = self.ls[0] * -tf.sin(q0) + J[0][1] 108 | J[1][0] = self.ls[0] * tf.cos(q0) + J[1][1] 109 | 110 | arr = tf.stack(J) 111 | return tf.transpose(arr, (2, 0, 1)) 112 | -------------------------------------------------------------------------------- /tf_robot_learning/planar_robots/two_joint.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | from .robot import Robot 22 | 23 | from tensorflow_probability import distributions as ds 24 | from tensorflow_probability import math 25 | 26 | pi = 3.14159 27 | 28 | class TwoJointRobot(Robot): 29 | def __init__(self, ms=None, ls=None, dt=0.01, g=9.81): 30 | 31 | Robot.__init__(self) 32 | 33 | self._Js_com = None 34 | 35 | m1 = 5. 36 | m2 = 5. 37 | self._ms = tf.constant([m1,m2]) if ms is None else ms 38 | 39 | # I = tf.constant([ 40 | # [1e-10, 1e-19, 1/12.], 41 | # [1e-10, 1e-10, 1/12.], 42 | # ]) 43 | 44 | I = tf.constant([ 45 | [1e-2, 1e-2, 1/12.], 46 | [1e-2, 1e-2, 1/12.], 47 | ]) 48 | 49 | self._mass = tf.stack([ 50 | tf.compat.v1.matrix_diag( 51 | tf.concat([ 52 | tf.ones(3) * self._ms[0], 53 | self._ms[0] * I[0] 54 | ], axis=0)), 55 | 56 | tf.compat.v1.matrix_diag( 57 | tf.concat([ 58 | tf.ones(3) * self._ms[1], 59 | self._ms[1] * I[1] 60 | ], axis=0)) 61 | 62 | ]) 63 | 64 | 65 | L = 1. 66 | 67 | self._ls = tf.constant([L, L]) if ls is None else ls 68 | 69 | margin = 0.02 70 | 71 | self._joint_limits = tf.constant([ 72 | [0. + margin, pi - margin], 73 | [-pi + margin, pi - margin], 74 | ], dtype=tf.float32) 75 | 76 | self._dof = 2 77 | 78 | self._dt = dt 79 | self._g = g 80 | 81 | 82 | @property 83 | def mass(self): 84 | return self._mass 85 | 86 | def xs(self, q): 87 | if q.shape.ndims == 1: 88 | x = tf.cumsum([0, 89 | self.ls[0] * tf.cos(q[0]), 90 | self.ls[1] * tf.cos(q[0] + q[1])]) 91 | y = tf.cumsum([0, 92 | self.ls[0] * tf.sin(q[0]), 93 | self.ls[1] * tf.sin(q[0] + q[1])]) 94 | 95 | return tf.transpose(tf.stack([x, y])) 96 | 97 | else: 98 | x = tf.cumsum([tf.zeros_like(q[:, 0]), 99 | self.ls[None, 0] * tf.cos(q[:, 0]), 100 | self.ls[None, 1] * tf.cos(q[:, 0] + q[:, 1])]) 101 | y = tf.cumsum([tf.zeros_like(q[:, 0]), 102 | self.ls[None, 0] * tf.sin(q[:, 0]), 103 | self.ls[None, 1] * tf.sin(q[:, 0] + q[: ,1])]) 104 | 105 | return tf.transpose(tf.stack([x, y]), (2, 1, 0)) 106 | 107 | def J(self, q): 108 | if q.shape.ndims == 1: 109 | J = [[0. for i in range(2)] for j in range(2)] 110 | J[0][1] = self.ls[1] * -tf.sin(q[0] + q[1]) 111 | J[1][1] = self.ls[1] * tf.cos(q[0] + q[1]) 112 | J[0][0] = self.ls[0] * -tf.sin(q[0]) + J[0][1] 113 | J[1][0] = self.ls[0] * tf.cos(q[0]) + J[1][1] 114 | 115 | arr = tf.stack(J) 116 | return tf.reshape(arr, (2, 2)) 117 | else: 118 | J = [[0. for i in range(2)] for j in range(2)] 119 | J[0][1] = self.ls[1] * -tf.sin(q[:, 0] + q[:, 1]) 120 | J[1][1] = self.ls[1] * tf.cos(q[:, 0] + q[:, 1]) 121 | J[0][0] = self.ls[0] * -tf.sin(q[:, 0]) + J[0][1] 122 | J[1][0] = self.ls[0] * tf.cos(q[:, 0]) + J[1][1] 123 | 124 | arr = tf.stack(J) 125 | return tf.transpose(arr, (2, 0, 1)) 126 | 127 | def f(self, x_t, u_t): 128 | """ 129 | To be used with tf_oc library 130 | State x_t : [q, dq] in the shape [bs, 4] 131 | (x_t [bs, xi_dim], u_t [bs, u_dim]) -> (x_t+1 [bs, xi_dim]) 132 | ddq = Mq_inv * (u - C*dq - g ), then integrate 133 | """ 134 | 135 | _q = x_t[:,:2] 136 | _dq = x_t[:,2:4] 137 | 138 | 139 | _Mq_inv = self.Mq_inv(_q) 140 | _Cq = self.Cq(_q,_dq) 141 | _gq = self.gq(_q, self._g) 142 | _D = self.D() 143 | 144 | 145 | M_ddq = -tf.einsum('aij,aj->ai',_Cq, _dq) - _gq + u_t - tf.einsum('ij,aj->ai',_D, _dq) 146 | ddq = tf.einsum('aij,aj->ai', _Mq_inv, M_ddq) 147 | dq = _dq + self._dt*ddq 148 | q = _q + self._dt*dq + 0.5 * self._dt ** 2 * ddq 149 | 150 | return tf.concat([q, dq], axis=1) 151 | 152 | def Js_com(self, q): 153 | """ 154 | Recode Js_com with q [batch_size, 2] 155 | Can be found on 156 | http://www-lar.deis.unibo.it/people/cmelchiorri/Files_Robotica/FIR_05_Dynamics.pdf 157 | """ 158 | 159 | if q.shape.ndims == 1: 160 | _Js_com = [[[0. for i in range(2)] for j in range(6)] for k in range(2)] 161 | _Js_com[0][0][0] = self.ls[0] / 2. * -tf.sin(q[0]) 162 | _Js_com[0][1][0] = self.ls[0] / 2. * tf.cos(q[0]) 163 | _Js_com[0][5][0] = 1.0 164 | 165 | _Js_com[1][0][1] = self.ls[1] / 2. * -tf.sin(q[0] + q[1]) 166 | _Js_com[1][1][1] = self.ls[1]/ 2. * tf.cos(q[0] + q[1]) 167 | _Js_com[1][5][1] = 1.0 168 | _Js_com[1][0][0] = self.ls[0] * -tf.sin(q[0]) + _Js_com[1][0][1] 169 | _Js_com[1][1][0] = self.ls[0] * tf.cos(q[0]) + _Js_com[1][1][1] 170 | _Js_com[1][5][0] = 1.0 171 | 172 | return tf.stack(_Js_com) 173 | else: 174 | _Js_com = [[[tf.zeros_like(q[..., 0]) for i in range(2)] for j in range(6)] for k in range(2)] 175 | _Js_com[0][0][0] = self.ls[0] / 2. * -tf.sin(q[..., 0]) 176 | _Js_com[0][1][0] = self.ls[0] / 2. * tf.cos(q[..., 0]) 177 | _Js_com[0][5][0] = tf.ones_like(q[..., 0]) 178 | 179 | _Js_com[1][0][1] = self.ls[1] / 2. * -tf.sin(q[..., 0] + q[..., 1]) 180 | _Js_com[1][1][1] = self.ls[1]/ 2. * tf.cos(q[..., 0] + q[..., 1]) 181 | _Js_com[1][5][1] = tf.ones_like(q[..., 0]) 182 | _Js_com[1][0][0] = self.ls[0] * -tf.sin(q[..., 0]) + _Js_com[1][0][1] 183 | _Js_com[1][1][0] = self.ls[0] * tf.cos(q[..., 0]) + _Js_com[1][1][1] 184 | _Js_com[1][5][0] = tf.ones_like(q[..., 0]) 185 | 186 | return tf.transpose(tf.stack(_Js_com), perm=(3, 0, 1, 2)) 187 | 188 | def Mq(self, q): 189 | Js_com = self.Js_com(q) 190 | 191 | if q.shape.ndims == 1: 192 | _Mq = tf.matmul(Js_com[0], tf.matmul(self.mass[0], Js_com[0]), transpose_a=True) + \ 193 | tf.matmul(Js_com[1], tf.matmul(self.mass[1], Js_com[1]), transpose_a=True) 194 | else: 195 | _Mq = tf.einsum('aji,ajk->aik', Js_com[:, 0], tf.einsum('ij,ajk->aik',self.mass[0], Js_com[:, 0])) + \ 196 | tf.einsum('aji,ajk->aik', Js_com[:, 1], tf.einsum('ij,ajk->aik', self.mass[1], Js_com[:, 1])) 197 | return _Mq 198 | 199 | def Mq_inv(self,q): 200 | _Mq = self.Mq(q) 201 | return tf.linalg.inv(_Mq) 202 | 203 | def Cq(self,q,dq): 204 | if q.shape.ndims == 1: 205 | h = -self._ms[1]*self._ls[0]*(self._ls[1]*0.5)*tf.sin(q[1]) 206 | _Cq = [[0. for i in range(2)] for j in range(2)] 207 | _Cq[0][0] = h*dq[1] 208 | _Cq[0][1] = h*(tf.reduce_sum(dq)) 209 | _Cq[1][0] = -h*dq[0] 210 | return tf.stack(_Cq) 211 | 212 | else: 213 | h = -self._ms[1]*self._ls[0]*(self._ls[1]*0.5)*tf.sin(q[:, 1]) 214 | _Cq = [[tf.zeros_like(q[:, 0]) for i in range(2)] for j in range(2)] 215 | _Cq[0][0] = h*dq[:, 1] 216 | _Cq[0][1] = h*(tf.reduce_sum(dq, axis=1)) 217 | _Cq[1][0] = -h*dq[:, 0] 218 | 219 | return tf.transpose(tf.stack(_Cq), perm=(2, 0, 1)) 220 | 221 | def gq(self, q, g): 222 | if q.shape.ndims==1: 223 | _gq = [0. for i in range(2)] 224 | 225 | a = (self._ms[0]*self._ls[0]/2 + self._ms[1]*self._ls[0])*g 226 | b = self._ms[1]*g*self._ls[1]*0.5 227 | 228 | _gq[0] = a*tf.cos(q[0]) + b*tf.cos(q[0] + q[1]) 229 | _gq[1] = b*tf.cos(q[0] + q[1]) 230 | return tf.transpose(tf.stack(_gq)[None]) 231 | else: 232 | _gq = [tf.zeros_like(q[..., 0]) for i in range(2)] 233 | 234 | a = (self._ms[0]*self._ls[0]/2 + self._ms[1]*self._ls[0])*g 235 | b = self._ms[1]*g*self._ls[1]*0.5 236 | 237 | _gq[0] = a*tf.cos(q[..., 0]) + b*tf.cos(q[..., 0] + q[..., 1]) 238 | _gq[1] = b*tf.cos(q[..., 0] + q[..., 1]) 239 | 240 | return tf.transpose(tf.stack(_gq), perm=(1, 0)) 241 | 242 | def D(self): 243 | return tf.eye(2) * .4 244 | 245 | 246 | 247 | 248 | 249 | -------------------------------------------------------------------------------- /tf_robot_learning/policy/__init__.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from .policy import NNStochasticFeedbackPolicy, NNStochasticPolicy, Policy 21 | from .poe_policy import PoEPolicy, AccLQRPoEPolicy, ForcePoEPolicy, VelocityPoEPolicy, AccPoEPolicy -------------------------------------------------------------------------------- /tf_robot_learning/policy/poe_policy.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | import tensorflow.compat.v1 as tf1 22 | from ..utils.tf_utils import batch_jacobians, matvec, matquad, damped_pinv_right 23 | from ..utils.param_utils import make_loc, make_cov 24 | from tensorflow_probability import distributions as ds 25 | 26 | 27 | class PoEPolicy(object): 28 | def __init__(self, product_size, experts_size, fs, js=None, *args, **kwargs): 29 | """ 30 | 31 | :param product_size int 32 | :param experts_size list 33 | :param fs: list of policy 34 | :param fs: list of transformations 35 | :param js: list of jacobians, optional but will be slower if don't given 36 | """ 37 | 38 | assert isinstance(experts_size, list) 39 | 40 | self._product_size = product_size 41 | self._experts_size = experts_size 42 | 43 | 44 | self._fs = fs 45 | self._js = js 46 | 47 | self._n_experts = len(self._fs) 48 | 49 | if self._js is None: 50 | self._js = [None for i in range(self.n_experts)] 51 | 52 | for i in range(self.n_experts): 53 | if self._js[i] is None: 54 | self._js[i] = lambda x: tf.linalg.LinearOperatorFullMatrix(batch_jacobians(self._fs[i](x), x)) 55 | 56 | @property 57 | def product_size(self): 58 | return self._product_size 59 | 60 | @property 61 | def experts_size(self): 62 | return self._experts_size 63 | 64 | @property 65 | def js(self): 66 | return self._js 67 | 68 | @property 69 | def fs(self): 70 | return self._fs 71 | 72 | @property 73 | def n_experts(self): 74 | return self._n_experts 75 | 76 | def jacobians(self, x): 77 | return [j(x) for j in self._js] 78 | 79 | class ForcePoEPolicy(PoEPolicy): 80 | def __init__(self, pis, u_dim, *args, **kwargs): 81 | """ 82 | 83 | :param pis: list of policies [y, dy, t] -> u, cov_u 84 | :param u_dim: 85 | :param args: 86 | :param kwargs: 87 | """ 88 | self._pis = pis 89 | self._reg = 0.1 90 | 91 | self._u_dim = u_dim 92 | 93 | super(ForcePoEPolicy, self).__init__(*args, **kwargs) 94 | 95 | @property 96 | def reg(self): 97 | return self._reg 98 | 99 | @reg.setter 100 | def reg(self, value): 101 | self._reg = value 102 | 103 | @property 104 | def pis(self): 105 | return self._pis 106 | 107 | # def density(self, xi, t=0): 108 | # x, dx = xi[:, :self._u_dim], xi[:, self._u_dim:] 109 | # 110 | # ys = [f(x) for f in self.fs] # transform state 111 | # js = [j(x) for j in self.js] # get jacobians 112 | # dys = [j.matvec(dx) for j in js] # get velocities in transformed space 113 | # 114 | # # "forces" in transformed space from the different policies 115 | # fys_locs_covs = [self.pis[i](ys[i], dys[i], t) for i in range(self.n_experts)] 116 | # 117 | # # separate locs and covs 118 | # fys_locs = [_y[0] for _y in fys_locs_covs] 119 | # fys_covs = [_y[1] for _y in fys_locs_covs] 120 | # 121 | # # "forces" in original space 122 | # fxs = [js[i].matvec(fys_locs[i], adjoint=True) for i in range(self.n_experts)] 123 | # 124 | # # covariances "forces" in original space 125 | # fxs_covs = [matquad(js[i], fys_covs[i] , adjoint=False) for i in 126 | # range(self.n_experts)] 127 | # 128 | # # precisions with regularization 129 | # fxs_precs = [tf.linalg.inv(cov + self._reg ** 2 * tf.eye(self.product_size)) for i, cov in 130 | # enumerate(fxs_covs)] 131 | # 132 | # # compute product of Gaussian policies 133 | # precs = tf.reduce_sum(fxs_precs, axis=0) 134 | # covs = tf.linalg.inv(precs) 135 | # locs = [tf.linalg.LinearOperatorFullMatrix(fxs_precs[i]).matvec(fxs[i]) for i in 136 | # range(self.n_experts)] 137 | # locs = tf.linalg.LinearOperatorFullMatrix(covs).matvec(tf.reduce_sum(locs, axis=0)) 138 | # 139 | # return ds.MultivariateNormalTriL(locs, tf.linalg.cholesky(covs)) 140 | # 141 | def sample(self, xi, t=0): 142 | return self.density(xi, t).sample() 143 | 144 | def density(self, xi, t=0): 145 | x, dx = xi[:, :self._u_dim], xi[:, self._u_dim:] 146 | 147 | 148 | ys = [f(x) for f in self.fs] # transform state 149 | js = [j(x) for j in self.js] # get jacobians 150 | 151 | pinv_js_trsp = [] 152 | for j in js: 153 | if isinstance(j, tf.linalg.LinearOperatorIdentity): 154 | pinv_js_trsp += [j] 155 | else: 156 | pinv_js_trsp += [tf.linalg.LinearOperatorFullMatrix(damped_pinv_right( 157 | tf1.matrix_transpose(j._matrix)), 1e-5)] 158 | 159 | dys = [j.matvec(dx) for j in js] # get velocities in transformed space 160 | 161 | # "velocities" in transformed space from the different policies 162 | fys_locs_covs = [self.pis[i](ys[i], dys[i], t) for i in range(self.n_experts)] 163 | 164 | # separate locs and covs 165 | fys_locs = [_y[0] for _y in fys_locs_covs] 166 | fys_covs = [_y[1] for _y in fys_locs_covs] 167 | 168 | # precisions with regularization J^T Lambda 169 | fys_precs = [ 170 | tf.linalg.inv(fys_covs[i]) 171 | for i in range(self.n_experts)] 172 | 173 | # fxs_eta = [tf.linalg.LinearOperatorFullMatrix(pinv_js_trsp[i].matmul( 174 | # fys_precs[i], adjoint=True)).matvec( 175 | # fys_locs[i]) for i in range(self.n_experts)] 176 | 177 | fxs_mus = [js[i].matvec( 178 | fys_locs[i], adjoint=True) for i in range(self.n_experts)] 179 | 180 | # fxs_precs = [ 181 | # tf.linalg.inv( 182 | # matquad(js[i], fys_covs[i]) + self._reg ** 2 * tf.eye(self.product_size)) 183 | # for i in range(self.n_experts) 184 | # ] 185 | 186 | fxs_precs = [ 187 | matquad(pinv_js_trsp[i], fys_precs[i]) for i in range(self.n_experts) 188 | ] 189 | 190 | # fxs_precs = tf.Print(fxs_precs, [fxs_precs]) 191 | fxs_eta = [tf.linalg.matvec(fxs_precs[i], fxs_mus[i]) for i in range(self.n_experts)] 192 | # compute product of Gaussian policies 193 | precs = tf.reduce_sum(fxs_precs, axis=0) 194 | 195 | covs = tf.linalg.inv(precs) 196 | 197 | etas = tf.reduce_sum(fxs_eta, axis=0) 198 | locs = tf.linalg.LinearOperatorFullMatrix(covs).matvec(etas) 199 | 200 | return ds.MultivariateNormalTriL(locs, tf.linalg.cholesky(covs)) 201 | 202 | 203 | 204 | class VelocityPoEPolicy(ForcePoEPolicy): 205 | def density(self, xi, t=0): 206 | ys = [f(xi) for f in self.fs] # transform state 207 | js = [j(xi) for j in self.js] # get jacobians 208 | 209 | # "velocities" in transformed space from the different policies 210 | fys_locs_covs = [self.pis[i](ys[i], t) for i in range(self.n_experts)] 211 | 212 | # separate locs and covs 213 | fys_locs = [_y[0] for _y in fys_locs_covs] 214 | fys_covs = [_y[1] for _y in fys_locs_covs] 215 | 216 | # precisions with regularization J^T Lambda 217 | fys_precs = [tf.linalg.inv(fys_covs[i] + self._reg ** 2 * tf.eye(self.experts_size[i])) 218 | for i in range(self.n_experts)] 219 | 220 | fxs_eta = [tf.linalg.LinearOperatorFullMatrix(js[i].matmul( 221 | fys_precs[i], adjoint=True)).matvec( 222 | fys_locs[i]) for i in range(self.n_experts)] 223 | 224 | fxs_precs = [ 225 | matquad(js[i], fys_precs[i]) for i in range(self.n_experts) 226 | ] 227 | 228 | # compute product of Gaussian policies 229 | precs = tf.reduce_sum(fxs_precs, axis=0) 230 | 231 | covs = tf.linalg.inv(precs) 232 | 233 | etas = tf.reduce_sum(fxs_eta, axis=0) 234 | locs = tf.linalg.LinearOperatorFullMatrix(covs).matvec(etas) 235 | 236 | return ds.MultivariateNormalTriL(locs, tf.linalg.cholesky(covs)) 237 | 238 | 239 | class AccPoEPolicy(ForcePoEPolicy): 240 | def density(self, xi, t=0): 241 | x, dx = xi[:, :self._u_dim], xi[:, self._u_dim:] 242 | 243 | ys = [f(x) for f in self.fs] # transform state 244 | js = [j(x) for j in self.js] # get jacobians 245 | dys = [j.matvec(dx) for j in js] # get velocities in transformed space 246 | 247 | # "velocities" in transformed space from the different policies 248 | fys_locs_covs = [self.pis[i](ys[i], dys[i], t) for i in range(self.n_experts)] 249 | 250 | # separate locs and covs 251 | fys_locs = [_y[0] for _y in fys_locs_covs] 252 | fys_covs = [_y[1] for _y in fys_locs_covs] 253 | 254 | # precisions with regularization J^T Lambda 255 | fys_precs = [tf.linalg.inv(fys_covs[i] + self._reg ** 2 * tf.eye(self.experts_size[i])) 256 | for i in range(self.n_experts)] 257 | 258 | fxs_eta = [tf.linalg.LinearOperatorFullMatrix(js[i].matmul( 259 | fys_precs[i], adjoint=True)).matvec( 260 | fys_locs[i]) for i in range(self.n_experts)] 261 | 262 | fxs_precs = [ 263 | matquad(js[i], fys_precs[i]) for i in range(self.n_experts) 264 | ] 265 | 266 | # compute product of Gaussian policies 267 | precs = tf.reduce_sum(fxs_precs, axis=0) 268 | 269 | covs = tf.linalg.inv(precs) 270 | 271 | etas = tf.reduce_sum(fxs_eta, axis=0) 272 | locs = tf.linalg.LinearOperatorFullMatrix(covs).matvec(etas) 273 | 274 | return ds.MultivariateNormalTriL(locs, tf.linalg.cholesky(covs)) 275 | 276 | class AccLQRPoEPolicy(PoEPolicy): 277 | def __init__(self, dt=0.01, 278 | # S_inv_scale=None, S_param=None, 279 | # Sv_inv_scale=1., Sv_param='iso', 280 | # R_inv_scale=20., R_param='iso', 281 | *args, **kwargs): 282 | """ 283 | $K_k = (R + B^T S_{K+1} B)^{-1} B^T S_{k+1} A$ 284 | 285 | :param dt: 286 | :param S_inv_scale: [list of float] 287 | :param S_param: [list of string] 288 | :param Sv_inv_scale: [float] 289 | value function for velocity 290 | :param Sv_param: [string] 291 | :param R_inv_scale: 292 | :param R_param: 293 | :param args: 294 | :param kwargs: 295 | """ 296 | self._dt = dt 297 | super(AccLQRPoEPolicy, self).__init__(*args, **kwargs) 298 | # 299 | # self._S = [] 300 | # for i in range(self.n_experts): 301 | # self._S += [make_cov( 302 | # self.experts_size[i], 15., param='iso', is_prec=False, 303 | # batch_shape=(k_basis_cov,))] 304 | 305 | # S_param is None: 306 | 307 | def Ks(self, x, S, Sv, R): 308 | pass 309 | 310 | def As(self, x): 311 | As = [None for i in range(self.n_experts)] 312 | 313 | for i in range(self.n_experts): 314 | As[i] = tf.linalg.LinearOperatorFullMatrix(tf.concat([ 315 | tf.concat([ 316 | tf.eye(self.experts_size[i], batch_shape=(x.shape[0], )), 317 | self._js[i](x) * self._dt 318 | ], axis=2), 319 | tf.concat([ 320 | tf.zeros((x.shape[0], self.product_size, self.experts_size[i])), 321 | tf.eye(self.product_size, batch_shape=(x.shape[0], )) 322 | ], axis=2), 323 | ], axis=1)) 324 | 325 | return As 326 | 327 | def Bs(self, x): 328 | Bs = [None for i in range(self.n_experts)] 329 | 330 | # for i in range(self.n_experts): 331 | # Bs[i] = tf.linalg.LinearOperatorFullMatrix(tf.concat([ 332 | # 0.5 * self._js[i](x) * self._dt ** 2, 333 | # tf.eye(self.product_size, batch_shape=(x.shape[0],)) * self._dt 334 | # ], axis=1)) 335 | 336 | for i in range(self.n_experts): 337 | Bs[i] = tf.linalg.LinearOperatorFullMatrix( 338 | 0.5 * self._js[i](x) * self._dt ** 2) 339 | 340 | return Bs 341 | -------------------------------------------------------------------------------- /tf_robot_learning/policy/policy.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | from ..nn import MLP 22 | from ..utils.param_utils import make_cov 23 | 24 | class Policy(object): 25 | def __init__(self, xi_dim, u_dim, *args, **kwargs): 26 | self._xi_dim = xi_dim 27 | self._u_dim = u_dim 28 | 29 | 30 | def pi(self, xi, t=0): 31 | """ 32 | 33 | :param xi: [batch_size, xi_dim] 34 | :return: u [batch_size, u_dim] 35 | """ 36 | 37 | raise NotImplementedError 38 | 39 | @property 40 | def u_dim(self): 41 | return self._u_dim 42 | 43 | @property 44 | def xi_dim(self): 45 | return self._xi_dim 46 | 47 | class NNStochasticPolicy(Policy): 48 | def __init__( 49 | self, xi_dim, u_dim, noise_dim, n_hidden=[50, 50,], 50 | act_fct=tf.nn.tanh, noise_scale=1.): 51 | """ 52 | u = f(x) + noise 53 | 54 | :param xi_dim: 55 | :param u_dim: 56 | :param noise_dim: 57 | :param n_hidden: 58 | :param act_fct: 59 | """ 60 | 61 | Policy.__init__(self, xi_dim, u_dim) 62 | 63 | self._nn = MLP( 64 | n_input=xi_dim + noise_dim, 65 | n_output=u_dim, 66 | n_hidden=n_hidden, 67 | batch_size_svi=1, 68 | act_fct=act_fct 69 | ) 70 | 71 | self._noise_dim = noise_dim 72 | self._noise_scale = noise_scale 73 | 74 | @property 75 | def params(self): 76 | return self._nn.vec_weights 77 | 78 | def pi(self, xi, t=0): 79 | return self._nn.pred( 80 | tf.concat([xi, tf.random.normal((xi.shape[0].value, self._noise_dim), 0., self._noise_scale) ], axis=1)) 81 | 82 | class NNStochasticFeedbackPolicy(Policy): 83 | def __init__(self, xi_dim, u_dim, noise_dim, A, B , n_hidden=[50, 50,], diag=False, 84 | act_fct=tf.nn.tanh, R_init=10, S_init=[0.02, 0.02, 0.5, 0.5], noise_scale=1.): 85 | """ 86 | 87 | u = K(x)(x_d(x) - x) 88 | 89 | :param xi_dim: 90 | :param u_dim: 91 | :param noise_dim: 92 | :param n_hidden: 93 | :param act_fct: 94 | """ 95 | Policy.__init__(self, xi_dim, u_dim) 96 | 97 | _n_output = xi_dim * xi_dim + xi_dim + 1 if not diag else xi_dim + xi_dim + 1 98 | 99 | self._nn = MLP( 100 | n_input=xi_dim + noise_dim, 101 | n_output=_n_output, 102 | n_hidden=n_hidden, 103 | batch_size_svi=1, 104 | act_fct=act_fct 105 | ) 106 | 107 | self._diag = diag 108 | 109 | self._noise_scale = noise_scale 110 | 111 | self._S_init = S_init 112 | self._R_init = R_init 113 | 114 | self._A = A 115 | self._B = B 116 | 117 | self._noise_dim = noise_dim 118 | 119 | def pi(self, xi, t=0): 120 | pi_params = self._nn.pred( 121 | tf.concat([xi, tf.random.normal((xi.shape[0].value, self._noise_dim), 0., self._noise_scale)], 122 | axis=1)) 123 | 124 | xi_d, r, s = pi_params[:, :self.xi_dim], pi_params[:, self.xi_dim:self.xi_dim+1], \ 125 | pi_params[:, self.xi_dim+1:] 126 | 127 | R = tf.eye(self.u_dim, batch_shape=(xi.shape[0].value, )) * tf.math.exp(r)[:, None] 128 | 129 | if self._diag: 130 | S = tf.compat.v1.matrix_diag(tf.math.exp(s)) 131 | else: 132 | s = tf.reshape(s, (-1, self.xi_dim, self.xi_dim)) 133 | S = tf.linalg.expm(tf.compat.v1.matrix_transpose(s) + s) 134 | 135 | 136 | K = tf.matmul( 137 | tf.linalg.inv(R + self._B.matmul( 138 | self._B.matmul(S, adjoint=True), adjoint_arg=True, adjoint=True)), 139 | tf.linalg.LinearOperatorFullMatrix(self._B.matmul(S, adjoint=True)).matmul(self._A._matrix) 140 | ) 141 | 142 | 143 | return tf.linalg.LinearOperatorFullMatrix(K).matvec(xi_d - xi) 144 | 145 | @property 146 | def params(self): 147 | return self._nn.vec_weights -------------------------------------------------------------------------------- /tf_robot_learning/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | from . import tf_utils as tf 21 | from . import plot_utils as plot 22 | from . import basis_utils as basis 23 | from . import param_utils as param 24 | from . import data_utils as data 25 | from . import img_utils as img -------------------------------------------------------------------------------- /tf_robot_learning/utils/basis_utils.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | 22 | import numpy as np 23 | from .tf_utils import block_diagonal 24 | 25 | 26 | def build_fixed_psi(n_step, n_dim, n_state, scale=None, dtype=tf.float32): 27 | _hs, _psis, _ = build_psis([n_step], n_state, n_dim, scale=scale) 28 | 29 | h = tf.constant(_hs[0].eval()) 30 | psi = tf.constant(_psis[0].eval()) 31 | 32 | return psi, h 33 | 34 | 35 | def build_psi_partial(h, n_dim, dim_obs, n_step=None, n_state=None, dtype=tf.float32): 36 | """ 37 | For conditioning, build psi with partial observation 38 | 39 | :param dim_obs: 40 | :type dim_obs: list or slice 41 | :param h: [n_timestep, n_state] 42 | :return: 43 | """ 44 | if isinstance(dim_obs, slice): 45 | dim_obs = [i for i in range(dim_obs.start, dim_obs.stop)] 46 | 47 | n_dim_in = len(dim_obs) 48 | 49 | n_dim_traj = n_dim_in * tf.shape(h)[0] if n_step is None else n_dim_in * n_step 50 | n_state = tf.shape(h)[1] if n_state is None else n_state 51 | 52 | obs_matrix = np.zeros((n_dim_in, n_dim)) 53 | 54 | for j, i in enumerate(dim_obs): 55 | obs_matrix[j, i] = 1. 56 | 57 | obs_matrix = tf.convert_to_tensor(obs_matrix, dtype=dtype) 58 | 59 | return tf.transpose(tf.reshape( 60 | tf.transpose(h[:, :, None, None] * obs_matrix[None, None], 61 | perm=(0, 2, 1, 3)), 62 | shape=(n_dim_traj, n_dim * n_state))) 63 | 64 | def build_psi(h, n_dim, n_step=None, n_state=None, dtype=tf.float32): 65 | """ 66 | 67 | :param h: [n_timestep, n_state] 68 | :return: 69 | """ 70 | n_dim_traj = n_dim * tf.shape(h)[0] if n_step is None else n_dim * n_step 71 | n_state = tf.shape(h)[1] if n_state is None else n_state 72 | 73 | return tf.transpose(tf.reshape( 74 | tf.transpose(h[:, :, None, None] * tf.eye(n_dim, dtype=dtype)[None, None], 75 | perm=(0, 2, 1, 3)), 76 | shape=(n_dim_traj, n_dim * n_state))) 77 | 78 | def build_psis(n_steps, n_state, n_dim, scale=None, fixed_handles=True, dtype=tf.float32, 79 | dim_obs=None): 80 | """ 81 | 82 | :param n_steps: Length of each trajectory 83 | :type n_steps: list of int 84 | :param n_state: Number of basis function 85 | :type n_state: int 86 | :type scale: float 0. - 1. 87 | :return: psis, hs, handles 88 | else: 89 | :return: psis, hs 90 | 91 | """ 92 | from tensorflow_probability import distributions as ds 93 | 94 | # create handles for each demos 95 | if fixed_handles: 96 | handles = tf.stack([tf.linspace(tf.cast(0., dtype=dtype), tf.cast(n, dtype=dtype), n_state) 97 | for n in n_steps]) 98 | else: 99 | handles = tf.Variable([tf.linspace(tf.cast(0., dtype=dtype), tf.cast(n, dtype=dtype), n_state) 100 | for n in n_steps]) 101 | 102 | n_traj = len(n_steps) 103 | 104 | if scale is None: 105 | scale = 1./ n_state 106 | 107 | # create mixtures whose p_z will be the activations 108 | 109 | 110 | 111 | h_mixture = [ 112 | ds.MixtureSameFamily( 113 | mixture_distribution=ds.Categorical(logits=tf.ones(n_state)), 114 | components_distribution=ds.MultivariateNormalDiag( 115 | loc=handles[j][:, None], 116 | scale_diag=tf.cast(scale, dtype)[None, None] * n_steps[j] 117 | ) 118 | ) 119 | for j in range(n_traj)] 120 | 121 | 122 | # create evaluation points of the mixture for each demo 123 | idx = [tf.range(n, dtype=dtype) for n in n_steps] 124 | from .tf_utils import log_normalize 125 | j = 0 126 | # create activations 127 | # print tf.transpose(h_mixture[0].components_log_prob(idx[0][:, None])) 128 | hs = [tf.exp(log_normalize( 129 | h_mixture[j].components_distribution.log_prob(idx[j][:, None, None]), axis=1)) 130 | for j in range(n_traj)] 131 | # hs = [tf for h in hs] 132 | if dim_obs is None: 133 | psis = [build_psi(hs[i], n_dim, n, n_state, dtype=dtype) 134 | for i, n in enumerate(n_steps)] 135 | else: 136 | psis = [build_psi_partial(hs[i], n_dim, dim_obs, n, n_state, dtype=dtype) 137 | for i, n in enumerate(n_steps)] 138 | 139 | return hs, psis, handles 140 | 141 | 142 | def build_obs(y): 143 | """ 144 | 145 | :param y: 146 | :type y: tf.Tensor [..., n_timestep, n_dim] 147 | :return: 148 | """ 149 | if y.shape.ndims == 3: 150 | return tf.reshape(y, (tf.shape(y)[0], -1)) 151 | elif y.shape.ndims == 2: 152 | return tf.reshape(y, (-1,)) 153 | 154 | 155 | def build_promp_transform(A, b): 156 | """ 157 | Build A, b such to move from the object coordinate system to the parent 158 | 159 | :param A: Rotation matrix of the object 160 | :type A: tf.Tensor or np.ndarray [..., n_dim, n_dim] 161 | :param b: Position of the object 162 | :type b: tf.Tensor or np.ndarray [..., n_dim] 163 | :return: 164 | """ 165 | 166 | b = tf.reshape(tf.transpose(b, (1, 0, 2)), (b.shape[1].value, -1)) 167 | 168 | A = block_diagonal([A[i] for i in range(A.shape[0].value)]) 169 | 170 | return A, b 171 | 172 | 173 | def build_dct_matrix(m, n_comp=None, 174 | dtype=tf.float32, name='dct'): 175 | """ 176 | 177 | :param m: 178 | :param dtype: 179 | :param name: 180 | :return: 181 | """ 182 | 183 | if n_comp is None: 184 | n_comp = m 185 | 186 | dctm = np.zeros((n_comp, m)) 187 | sqrt_m = np.sqrt(m) 188 | 189 | for p in range(n_comp): 190 | for q in range(m): 191 | if p == 0: 192 | dctm[p, q] = 1/sqrt_m 193 | elif p > 0: 194 | dctm[p, q] = 2**0.5/sqrt_m * np.cos(np.pi * (2 * q + 1) * p/2./m ) 195 | 196 | return tf.convert_to_tensor(dctm, dtype=dtype, name=name) 197 | 198 | 199 | def seq_dct(samples, dct_m, inverse=False): 200 | """ 201 | samples : [batch_size, timesteps, xi_dim] 202 | """ 203 | return tf.transpose(dct_m.matvec(tf.transpose(samples, (0, 2, 1)), adjoint=inverse), (0, 2, 1)) 204 | 205 | def seq_lp(samples): 206 | return tf.log(samples ** 2) 207 | 208 | def correlated_normal_noise(shape, dct_components, mean=0., scale=1.): 209 | """ 210 | shape, (batch_size, horizon, xi_dim) 211 | """ 212 | k_basis = dct_components.shape[0] if isinstance(dct_components.shape[0], int) else \ 213 | dct_components.shape[0].value 214 | 215 | _, h = build_fixed_psi( 216 | n_step=shape[1], n_dim=shape[-1], n_state=k_basis, scale=0.5 / k_basis) 217 | 218 | dct_m = tf.linalg.LinearOperatorFullMatrix(build_dct_matrix(shape[1])) 219 | 220 | noise = tf.random.normal(shape, mean, scale) 221 | dct_components_t = tf.linalg.matvec(h, dct_components) 222 | dct_norm_t = dct_components_t / tf.linalg.norm(dct_components_t) * shape[1] ** 0.5 223 | return seq_dct(seq_dct(noise, dct_m) * dct_norm_t[None, :, None], dct_m, inverse=True) -------------------------------------------------------------------------------- /tf_robot_learning/utils/data_utils.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import numpy as np 21 | import os, sys 22 | 23 | def load_letter(letter='P', get_x=True, get_dx=True, get_ddx=False, concat=True, 24 | fill_zeros=None): 25 | path = os.path.abspath(__file__) 26 | 27 | datapath = path[:-38] 28 | if datapath[-1] != '/': 29 | datapath += '/' 30 | 31 | datapath += 'data/2Dletters/' 32 | 33 | try: 34 | data = np.load(datapath + '%s.npy' % letter, allow_pickle=True)[()] 35 | except ValueError: 36 | data = np.load(datapath + '%s_2.npy' % letter, allow_pickle=True)[()] 37 | except: 38 | print("Unexpected error:", sys.exc_info()[0]) 39 | raise 40 | 41 | demos = [] 42 | 43 | if get_x: 44 | if fill_zeros is None: 45 | demos += [data['x']] 46 | else: 47 | demos += [[np.concatenate( 48 | [d, d[[-1]] * np.ones((fill_zeros, data['x'][0].shape[-1]))], axis=0) 49 | for d in data['x']]] 50 | 51 | if get_dx: 52 | if fill_zeros is None: 53 | demos += [data['dx']] 54 | else: 55 | demos += [[np.concatenate( 56 | [d, np.zeros((fill_zeros, data['dx'][0].shape[-1]))], axis=0) 57 | for d in data['dx']]] 58 | 59 | if get_ddx: 60 | if fill_zeros is None: 61 | demos += [data['ddx']] 62 | else: 63 | demos += [[np.concatenate( 64 | [d, np.zeros((fill_zeros, data['ddx'][0].shape[-1]))], axis=0) 65 | for d in data['ddx']]] 66 | 67 | if concat: 68 | return np.concatenate(demos, axis=2) 69 | else: 70 | return demos 71 | 72 | def load_letter_bimodal(letter='P', get_x=True, get_dx=True, get_ddx=False, concat=True, 73 | fill_zeros=None, center=[0.2, 0.], vec=[1, 0.], dist=0.1, displ=0.2): 74 | 75 | demos = load_letter(letter=letter, get_x=get_x, get_dx=get_dx, get_ddx=get_ddx, 76 | concat=False, fill_zeros=fill_zeros) 77 | 78 | direction = 2. * (np.random.randint(0, 2, len(demos[0])) - 0.5) 79 | 80 | act = np.exp(-np.sum(((np.array(demos[0]) - center) / dist) ** 2, axis=-1)) 81 | 82 | d_x = direction[:, None, None] * displ * np.array(vec)[None, None, :] * act[:, :, None] 83 | demos[0] = np.array(demos[0]) + d_x 84 | 85 | if get_dx: 86 | demos[1] = np.array(demos[1]) + np.gradient(d_x, axis=1) / 0.01 87 | if get_ddx: 88 | demos[2] = np.array(demos[2]) + np.gradient(np.gradient(d_x, axis=1), axis=1) / 0.01 ** 2 89 | 90 | if concat: 91 | return np.concatenate(demos, axis=2) 92 | else: 93 | return demos -------------------------------------------------------------------------------- /tf_robot_learning/utils/img_utils.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | import numpy as np 22 | from .tf_utils import log_normalize 23 | 24 | def spatial_soft_argmax(img_array, temp=1.): 25 | """ 26 | Compute spatial soft argmax 27 | 28 | :param img_array: [nb_samples, height, width] 29 | :type img_array: tf.Tensor 30 | :return: 31 | """ 32 | 33 | if len(img_array.get_shape()) == 3: 34 | height = img_array.get_shape()[-2].value 35 | width = img_array.get_shape()[-1].value 36 | else: 37 | height = img_array.get_shape()[-3].value 38 | width = img_array.get_shape()[-2].value 39 | n_channel = img_array.get_shape()[-1].value 40 | 41 | img_coords = tf.constant( 42 | np.array([np.resize(np.arange(height), (height, width)), 43 | np.resize(np.arange(width), (width, height)).T]).astype( 44 | np.float32)) 45 | 46 | # softmax = tf.nn.relu(vec_img)/tf.reduce_sum(tf.nn.relu(vec_img)) 47 | if len(img_array.get_shape()) == 3: 48 | # tf.reduce_sum( 49 | # img_coords[None, :, :, :, None] * tf.exp( 50 | # log_normalize(temp * img_array, axis=(0, 1)))[:, None], axis=(2, 3)) 51 | # 52 | vec_img = tf.reshape(img_array, [-1, height * width]) 53 | softmax = tf.nn.softmax(vec_img) 54 | softmax = tf.reshape(softmax, [-1, height, width]) 55 | return tf.einsum('aij,bij->ab', softmax, img_coords) 56 | else: 57 | softmax = tf.exp(log_normalize(temp * img_array, axis=(1, 2))) 58 | return tf.reduce_sum( 59 | img_coords[None, :, :, :, None] *softmax[:, None], axis=(2, 3)), softmax 60 | 61 | 62 | def spatial_soft_argmax_temporal(img_array, prev_argmax, temp=1., dist=0.2): 63 | """ 64 | Compute spatial soft argmax 65 | 66 | :param img_array: [nb_samples, height, width] 67 | :type img_array: tf.Tensor 68 | :return: 69 | """ 70 | 71 | height = img_array.get_shape()[-3].value 72 | width = img_array.get_shape()[-2].value 73 | 74 | img_coords = tf.constant( 75 | np.array([np.resize(np.arange(height), (height, width)), 76 | np.resize(np.arange(width), (width, height)).T]).astype( 77 | np.float32)) 78 | 79 | temp_log_prob = tf.reduce_sum( 80 | ((prev_argmax[:, :, None, None] - img_coords[None, :, :, :, None]) ** 2)/(height * width * dist), axis=(1)) 81 | 82 | softmax = tf.exp(log_normalize(temp * img_array - temp_log_prob, axis=(1, 2))) 83 | return tf.reduce_sum( 84 | img_coords[None, :, :, :, None] *softmax[:, None], axis=(2, 3)), softmax -------------------------------------------------------------------------------- /tf_robot_learning/utils/param_utils.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import tensorflow as tf 21 | from .tf_utils import log_normalize 22 | 23 | def has_parent_variable(tensor, variable): 24 | """ 25 | Check if tensor depends on variable 26 | :param tensor: 27 | :param variable: 28 | :return: 29 | """ 30 | var_list = len([var for var in tensor.op.values() if var == variable._variable]) 31 | 32 | if var_list: 33 | return True 34 | 35 | for i in tensor.op.inputs: 36 | if has_parent_variable(i, variable): 37 | return True 38 | 39 | return False 40 | 41 | 42 | def get_parent_variables(tensor, sess=None): 43 | """ 44 | Get all variables in the graph on which the tensor depends 45 | :param tensor: 46 | :param sess: 47 | :return: 48 | """ 49 | if sess is None: 50 | sess = tf.compat.v1.get_default_session() 51 | 52 | if tensor.__class__ == tf.Variable: 53 | return [tensor] 54 | 55 | return [v for v in sess.graph._collections['variables'] if 56 | has_parent_variable(tensor, v)] 57 | # 58 | 59 | class Param(tf.Tensor): 60 | def set_inverse_transform(self, f): 61 | self._inverse_transform = f 62 | 63 | def assign_op(self, value): 64 | return self._parent.assign(self._inverse_transform(value)) 65 | 66 | def assign(self, value, sess=None): 67 | if sess is None: 68 | sess = tf.compat.v1.get_default_session() 69 | 70 | return sess.run(self.assign_op(value)) 71 | 72 | @property 73 | def variable(self): 74 | return self._parent 75 | 76 | @variable.setter 77 | def variable(self, value): 78 | self._parent = value 79 | 80 | 81 | def make_logits_from_value(priors): 82 | import numpy as np 83 | init_log = np.log(priors) 84 | 85 | v = tf.Variable(init_log, dtype=tf.float32) 86 | m = log_normalize(v, axis=-1) 87 | 88 | m.__class__ = Param 89 | m.variable = v 90 | 91 | return m 92 | 93 | def make_loc_from_value(loc): 94 | v = tf.Variable(loc, dtype=tf.float32) 95 | m = tf.identity(v) 96 | 97 | m.__class__ = Param 98 | m.variable = v 99 | 100 | return m 101 | 102 | def make_cov_from_value(cov, param='expm'): 103 | import numpy as np 104 | 105 | if param == 'expm': 106 | init_cov = tf.cast( 107 | tf.linalg.logm(cov.astype(np.complex64)), 108 | tf.float32) 109 | 110 | v = tf.Variable(init_cov) 111 | m = tf.linalg.expm(0.5 * (v + tf.compat.v1.matrix_transpose(v))) 112 | else: 113 | raise NotImplementedError 114 | 115 | m.__class__ = Param 116 | m.variable = v 117 | 118 | return m 119 | 120 | def make_loc(shape, mean=0., scale=1.): 121 | v = tf.Variable(tf.random.normal(shape, mean, scale)) 122 | m = tf.identity(v) 123 | 124 | m.__class__ = Param 125 | m.variable = v 126 | 127 | def inverse_transform(x): 128 | return tf.cast(x, dtype=tf.float32) 129 | 130 | m.set_inverse_transform(inverse_transform) 131 | return m 132 | 133 | def make_rp(shape, mean=0.): 134 | v = tf.Variable(tf.zeros(shape) + tf.log(mean)) 135 | m = tf.math.exp(v) 136 | 137 | m.__class__ = Param 138 | m.variable = v 139 | 140 | return m 141 | 142 | 143 | def make_cov(k, scale=1., param='expm', batch_shape=(), is_prec=False, var=None): 144 | """ 145 | 146 | :param k: event_dim spd will have shape (..., k, k) 147 | :type k : int 148 | :param scale: scale of the diagonal (std if covariance) 149 | :type scale: float, list of float, or tf.Tensor 150 | :param param: type of parametrization 151 | :type param: str in ['expm', 'tril', 'iso', 'diag'] 152 | :param batch_shape: 153 | :param is_prec: if True return a precision matrix whose std is scale 154 | :return: 155 | """ 156 | if isinstance(scale, float) and param is not 'iso': 157 | scale = scale * tf.ones(k) 158 | elif isinstance(scale, list): 159 | scale = tf.convert_to_tensor(scale) 160 | 161 | 162 | if param == 'expm': 163 | p = 2. if not is_prec else -2. 164 | v = tf.Variable(tf.eye(k, batch_shape=batch_shape) * tf.math.log(scale ** p)) if var is None else var 165 | m = tf.linalg.expm(0.5 * (v + tf.compat.v1.matrix_transpose(v))) 166 | 167 | m.__class__ = Param 168 | m.variable = v 169 | 170 | def inverse_transform(cov): 171 | import numpy as np 172 | return tf.cast( 173 | tf.linalg.logm(cov.astype(np.complex64)), 174 | tf.float32) 175 | 176 | m.set_inverse_transform( 177 | inverse_transform 178 | ) 179 | 180 | elif param == 'tril': 181 | v = tf.Variable(tf.eye(k, batch_shape=batch_shape)) if var is None else var 182 | p = 1. if not is_prec else -1. 183 | m = v * (tf.linalg.diag(scale ** p) + tf.ones((k, k)) - tf.eye(k)) 184 | m = tf.matmul(m, m, transpose_b=True) 185 | 186 | m.__class__ = Param 187 | m.variable = v 188 | 189 | elif param == 'iso': 190 | p = 2. if not is_prec else -2. 191 | v = tf.Variable(tf.ones(batch_shape + (1, 1))) 192 | m = v * tf.math.log(scale) 193 | m = tf.math.exp(m) ** p * tf.eye(k, batch_shape=batch_shape) if var is None else var 194 | 195 | m.__class__ = Param 196 | m.variable = v 197 | 198 | elif param == 'diag': 199 | p = 2. if not is_prec else -2. 200 | v = tf.Variable(tf.ones(batch_shape + (k,)) * tf.math.log(scale)) if var is None else var 201 | m = tf.compat.v1.matrix_diag(tf.math.exp(v) ** p) 202 | 203 | m.__class__ = Param 204 | m.variable = v 205 | 206 | def inverse_transform(cov): 207 | import numpy as np 208 | return 1./p * tf.log(tf.cast(tf.compat.v1.matrix_diag_part(cov), dtype=tf.float32)) 209 | 210 | m.set_inverse_transform( 211 | inverse_transform 212 | ) 213 | 214 | else: 215 | raise ValueError('param should be in [expm, tril, iso, diag]') 216 | 217 | 218 | 219 | return m 220 | -------------------------------------------------------------------------------- /tf_robot_learning/utils/plot_utils.py: -------------------------------------------------------------------------------- 1 | # tf_robot_learning, a all-around tensorflow library for robotics. 2 | # 3 | # Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/ 4 | # Written by Emmanuel Pignat , 5 | # 6 | # This file is part of tf_robot_learning. 7 | # 8 | # tf_robot_learning is free software: you can redistribute it and/or modify 9 | # it under the terms of the GNU General Public License version 3 as 10 | # published by the Free Software Foundation. 11 | # 12 | # tf_robot_learning is distributed in the hope that it will be useful, 13 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | # GNU General Public License for more details. 16 | # 17 | # You should have received a copy of the GNU General Public License 18 | # along with tf_robot_learning. If not, see . 19 | 20 | import matplotlib.pyplot as plt 21 | import tensorflow as tf 22 | import numpy as np 23 | plt.style.use('ggplot') 24 | 25 | 26 | class DensityPlotter(object): 27 | def __init__(self, f, nb_sub=20, np=False, vectorized=True): 28 | if not np: 29 | self.x_batch = tf.compat.v1.placeholder(tf.float32, (nb_sub ** 2, 2)) 30 | self.fx_batch = f(self.x_batch) 31 | 32 | self.f = f 33 | self.nb_sub = nb_sub 34 | self.vmin = 0. 35 | self.vmax = 0. 36 | 37 | self.vectorized = vectorized 38 | 39 | self.np = np 40 | 41 | def f_np(self, x, feed_dict={}): 42 | if self.np: 43 | return self.f(x) 44 | else: 45 | fd = {self.x_batch: x} 46 | fd.update(feed_dict) 47 | return self.fx_batch.eval(fd) 48 | 49 | def plot(self, ax0=None, xlim=[-1, 1], ylim=[-1, 1], cmap='viridis', 50 | lines=True, heightmap=True, inv=False, vmin=None, vmax=None, use_tf=False, 51 | ax=None, feed_dict={}, exp=True, img=False, act_fct=None, kwargs_lines={}): 52 | 53 | z = plot_density(lambda x: self.f_np(x, feed_dict=feed_dict), 54 | nb_sub=self.nb_sub, ax0=ax0, xlim=xlim, ylim=ylim, cmap=cmap, 55 | lines=lines, heightmap=heightmap, inv=inv, vmin=vmin, vmax=vmax, use_tf=False, 56 | ax=ax, feed_dict=feed_dict, exp=exp, img=img, kwargs_lines=kwargs_lines, 57 | vectorized=self.vectorized, act_fct=act_fct) 58 | 59 | if exp: 60 | self.vmax, self.vmin = np.max(np.exp(z)), np.min(np.exp(z)) 61 | else: 62 | self.vmax, self.vmin = np.max(z), np.min(z) 63 | 64 | 65 | return z 66 | 67 | 68 | def plot_density(f, nb_sub=10, ax0=None, xlim=[-1, 1], ylim=[-1, 1], cmap='viridis', 69 | lines=True, heightmap=True, inv=False, vmin=None, vmax=None, use_tf=False, 70 | ax=None, feed_dict={}, exp=True, img=False, kwargs_lines={}, vectorized=True, 71 | act_fct=None): 72 | 73 | if use_tf: 74 | x_batch = tf.placeholder(tf.float32, (nb_sub ** 2, 2)) 75 | fx_batch = f(x_batch) 76 | def f_np(x): 77 | fd = {x_batch: x} 78 | fd.update(**feed_dict) 79 | return fx_batch.eval(fd) 80 | 81 | f = f_np 82 | 83 | 84 | x = np.linspace(*xlim, num=nb_sub) 85 | y = np.linspace(*ylim, num=nb_sub) 86 | 87 | if img: 88 | Y, X = np.meshgrid(x, y) 89 | else: 90 | X, Y = np.meshgrid(x, y) 91 | 92 | if vectorized: 93 | zs = f(np.concatenate([np.atleast_2d(X.ravel()), np.atleast_2d(Y.ravel())]).T) 94 | 95 | else: 96 | _xys = np.concatenate([np.atleast_2d(X.ravel()), np.atleast_2d(Y.ravel())]).T 97 | zs = [] 98 | for _xy in _xys: 99 | try: 100 | zs += [f(_xy)] 101 | except: 102 | zs += [0.] 103 | 104 | zs = np.array(zs) 105 | 106 | if act_fct is not None: 107 | zs = act_fct(zs) 108 | 109 | Z = zs.reshape(X.shape) 110 | 111 | 112 | if lines: 113 | if ax is None: 114 | plt.contour(X, Y, Z, **kwargs_lines) 115 | else: 116 | ax.contour(X, Y, Z, **kwargs_lines) 117 | 118 | 119 | if heightmap: 120 | if exp: 121 | _img = np.exp(Z) if not inv else -np.exp(Z) 122 | else: 123 | _img = Z if not inv else -Z 124 | 125 | kwargs = {'origin': 'lower'} 126 | 127 | if ax is None: 128 | plt.imshow(_img, interpolation='bilinear', extent=xlim + ylim, 129 | alpha=0.5, cmap=cmap, vmax=vmax, vmin=vmin, **kwargs) 130 | else: 131 | ax.imshow(_img, interpolation='bilinear', extent=xlim + ylim, 132 | alpha=0.5, cmap=cmap, vmax=vmax, vmin=vmin, **kwargs) 133 | 134 | return Z 135 | 136 | 137 | class PolicyPlot(object): 138 | def __init__(self, pi, nb_sub=20): 139 | 140 | self._x = tf.compat.v1.placeholder(tf.float32, (nb_sub**2, 2)) 141 | self._u = tf.transpose(pi(self._x)) 142 | self._nb_sub = nb_sub 143 | def plot(self, ax=None, xlim=[-1, 1], ylim=[-1, 1], scale=0.01, 144 | name=None, equal=False, feed_dict={}, sess=None, **kwargs): 145 | """ 146 | Plot a dynamical system dx = f(x) 147 | :param f: a function that takes as input x as [N,2] and return dx [N, 2] 148 | :param nb_sub: 149 | :param ax0: 150 | :param xlim: 151 | :param ylim: 152 | :param scale: 153 | :param kwargs: 154 | :return: 155 | """ 156 | 157 | if sess is None: 158 | sess = tf.compat.v1.get_default_session() 159 | 160 | Y, X = np.mgrid[ 161 | ylim[0]:ylim[1]:complex(self._nb_sub), 162 | xlim[0]:xlim[1]:complex(self._nb_sub)] 163 | mesh_data = np.vstack([X.ravel(), Y.ravel()]) 164 | 165 | feed_dict[self._x] = mesh_data.T 166 | field = sess.run(self._u, feed_dict) 167 | 168 | U = field[0] 169 | V = field[1] 170 | U = U.reshape(self._nb_sub, self._nb_sub) 171 | V = V.reshape(self._nb_sub, self._nb_sub) 172 | speed = np.sqrt(U * U + V * V) 173 | 174 | if name is not None: 175 | plt.suptitle(name) 176 | 177 | if ax is not None: 178 | strm = ax.streamplot(X, Y, U, V, linewidth=scale * speed, **kwargs) 179 | ax.set_xlim(xlim) 180 | ax.set_ylim(ylim) 181 | 182 | if equal: 183 | ax.set_aspect('equal') 184 | 185 | else: 186 | strm = plt.streamplot(X, Y, U, V, linewidth=scale * speed, **kwargs) 187 | plt.xlim(xlim) 188 | plt.ylim(ylim) 189 | 190 | if equal: 191 | plt.axes().set_aspect('equal') 192 | 193 | return [strm] 194 | 195 | 196 | 197 | class MVNPlot(object): 198 | def __init__(self, ds, nb_segments=20): 199 | 200 | from ..distributions import GaussianMixtureModelML, GaussianMixtureModelFromSK 201 | self._ds = ds 202 | if isinstance(ds, GaussianMixtureModelML) or isinstance(ds, GaussianMixtureModelFromSK): 203 | self._loc = ds.components_distribution.loc 204 | self._cov = ds.components_distribution.covariance() 205 | else: 206 | self._cov = ds.covariance() 207 | self._loc = ds.loc 208 | 209 | self._t = np.linspace(-np.pi, np.pi, nb_segments) 210 | 211 | def plot(self, *args, **kwargs): 212 | return self.plot_gmm(*args, **kwargs) 213 | 214 | 215 | def plot_gmm(self, dim=None, color=[1, 0, 0], alpha=0.5, linewidth=1, 216 | markersize=6, batch_idx=0, 217 | ax=None, empty=False, edgecolor=None, edgealpha=None, 218 | border=False, nb=1, center=True, zorder=20, equal=True, sess=None, 219 | feed_dict={}, axis=0): 220 | 221 | if sess is None: 222 | sess = tf.compat.v1.get_default_session() 223 | 224 | loc, cov = sess.run([self._loc, self._cov], feed_dict) 225 | 226 | 227 | if loc.ndim == 3: 228 | loc = loc[batch_idx] if axis==0 else loc[:, batch_idx] 229 | if cov.ndim == 4: 230 | cov = cov[batch_idx] if axis==0 else cov[:, batch_idx] 231 | 232 | if loc.ndim == 1: 233 | loc = loc[None] 234 | if cov.ndim == 2: 235 | cov = cov[None] 236 | 237 | nb_states = loc.shape[0] 238 | 239 | if dim: 240 | loc = loc[:, dim] 241 | cov = cov[np.ix_(range(cov.shape[0]), dim, dim)] if isinstance(dim, 242 | list) else cov[ 243 | :, 244 | dim, 245 | dim] 246 | if not isinstance(color, list) and not isinstance(color, np.ndarray): 247 | color = [color] * nb_states 248 | elif not isinstance(color[0], str) and not isinstance(color, np.ndarray): 249 | color = [color] * nb_states 250 | 251 | if not isinstance(alpha, np.ndarray): 252 | alpha = [alpha] * nb_states 253 | else: 254 | alpha = np.clip(alpha, 0.1, 0.9) 255 | 256 | rs = tf.linalg.sqrtm(cov).eval() 257 | 258 | pointss = np.einsum('aij,js->ais', rs, np.array([np.cos(self._t), np.sin(self._t)])) 259 | pointss += loc[:, :, None] 260 | 261 | for i, c, a in zip(range(0, nb_states, nb), color, alpha): 262 | points = pointss[i] 263 | 264 | if edgecolor is None: 265 | edgecolor = c 266 | 267 | polygon = plt.Polygon(points.transpose().tolist(), facecolor=c, alpha=a, 268 | linewidth=linewidth, zorder=zorder, 269 | edgecolor=edgecolor) 270 | 271 | if edgealpha is not None: 272 | plt.plot(points[0, :], points[1, :], color=edgecolor) 273 | 274 | if ax: 275 | ax.add_patch(polygon) # Patch 276 | 277 | l = None 278 | if center: 279 | a = alpha[i] 280 | else: 281 | a = 0. 282 | 283 | ax.plot(loc[i, 0], loc[i, 1], '.', color=c, alpha=a) # Mean 284 | 285 | if border: 286 | ax.plot(points[0, :], points[1, :], color=c, linewidth=linewidth, 287 | markersize=markersize) # Contour 288 | if equal: 289 | ax.set_aspect('equal') 290 | else: 291 | if empty: 292 | plt.gca().grid('off') 293 | # ax[-1].set_xlabel('x position [m]') 294 | plt.gca().set_axis_bgcolor('w') 295 | plt.axis('off') 296 | 297 | plt.gca().add_patch(polygon) # Patch 298 | l = None 299 | 300 | if center: 301 | a = alpha[i] 302 | else: 303 | a = 0.0 304 | 305 | l, = plt.plot(loc[i, 0], loc[i, 1], '.', color=c, alpha=a) # Mean 306 | # plt.plot(points[0,:], points[1,:], color=c, linewidth=linewidth , markersize=markersize) # Contour 307 | if equal: 308 | plt.gca().set_aspect('equal') 309 | return l 310 | 311 | def plot_coordinate_system(A, b, scale=1., text=None, equal=True, text_kwargs={}, 312 | ax=None, **kwargs): 313 | """ 314 | 315 | :param A: nb_dim x nb_dim 316 | Rotation matrix 317 | :param b: nb_dim 318 | Translation 319 | :param scale: float 320 | Scaling of the axis 321 | :param equal: bool 322 | Set matplotlib axis to equal 323 | :param ax: plt.axes() 324 | :param kwargs: 325 | :return: 326 | """ 327 | a0 = np.vstack([b, b + scale * A[:,0]]) 328 | a1 = np.vstack([b, b + scale * A[:,1]]) 329 | 330 | if ax is None: 331 | p, a = (plt, plt.gca()) 332 | else: 333 | p, a = (ax, ax) 334 | 335 | if equal and a is not None: 336 | a.set_aspect('equal') 337 | 338 | p.plot(a0[:, 0], a0[:, 1], 'r', **kwargs) 339 | p.plot(a1[:, 0], a1[:, 1], 'b', **kwargs) 340 | 341 | if not text is None: 342 | p.text(b[0]-0.1 * scale, b[1]- 0.15 * scale, text, **text_kwargs) 343 | 344 | def plot_coordinate_system_3d( 345 | A, b, scale=1., equal=True, dim=None, ax=None, text=None, dx_text=[0., 0.], 346 | text_kwargs={}, **kwargs): 347 | """ 348 | 349 | :param A: nb_dim x nb_dim 350 | Rotation matrix 351 | :param b: nb_dim 352 | Translation 353 | :param scale: float 354 | Scaling of the axis 355 | :param equal: bool 356 | Set matplotlib axis to equal 357 | :param ax: plt.axes() 358 | :param kwargs: 359 | :return: 360 | """ 361 | 362 | if dim is None: 363 | dim = [0, 1] 364 | 365 | a0 = np.vstack([b[dim], b[dim] + scale * A[dim, 0]]) 366 | a1 = np.vstack([b[dim], b[dim] + scale * A[dim, 1]]) 367 | a2 = np.vstack([b[dim], b[dim] + scale * A[dim, 2]]) 368 | 369 | if ax is None: 370 | p, a = (plt, plt.gca()) 371 | else: 372 | p, a = (ax, ax) 373 | 374 | 375 | if equal and a is not None: 376 | a.set_aspect('equal') 377 | 378 | 379 | if text is not None: 380 | a.text(b[dim[0]] + dx_text[0], b[dim[1]] + dx_text[1], text, **text_kwargs) 381 | 382 | label = kwargs.pop('label', None) 383 | color = kwargs.get('color', None) 384 | if label is not None: kwargs['label'] = label + ' x' 385 | if color is not None: kwargs['label'] = label 386 | 387 | p.plot(a0[:, 0], a0[:, 1], 'r', **kwargs) 388 | 389 | if color is not None: 390 | label = kwargs.pop('label', None) 391 | 392 | if label is not None and color is None: kwargs['label'] = label + ' y' 393 | p.plot(a1[:, 0], a1[:, 1], 'g', **kwargs) 394 | if label is not None and color is None: kwargs['label'] = label + ' z' 395 | p.plot(a2[:, 0], a2[:, 1], 'b', **kwargs) 396 | --------------------------------------------------------------------------------