{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Intro: Configurations\n",
    "\n",
    "A configuration is essentially a tree of (coordinate) frames, where each frame can have an associated shape, joint, inertia, etc. This tutorial introduces first steps to creating & loading configurations in rai, setting/getting the joint and frame state, computing features, and handling the view window.\n",
    "\n",
    "Please see the [Robot Kinematics and Dynamics Essentials](https://www.user.tu-berlin.de/mtoussai/notes/robotKin.html) lecture notes for some background."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a configuration from scratch\n",
    "\n",
    "One typically loads robot configurations from a file. But to learn the concepts, let's first build a configuration from scratch by adding frames and setting properties such as pose, shape and parentship of frames.\n",
    "\n",
    "The following creates an empty configuration, and opens a view window for it.\n",
    "\n",
    "**Tip:** Make the view window appear \"Always On Top\" (right click on the window bar)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import robotic as ry\n",
    "import numpy as np\n",
    "import time\n",
    "print('robotic package version:', ry.__version__, ry.compiled())\n",
    "\n",
    "C = ry.Config()\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `C.view()` command opened a window that shows an empty configuration. It is still responsive (as the underlying glfw event loop is running in a different thread).\n",
    "\n",
    "Next we declare it a *marker* frame, which means it has no real shape, but is displayed using its base vectors in the view:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C.clear()\n",
    "f = C.addFrame(name='first')\n",
    "f.setShape(type=ry.ST.marker, size=[.3])\n",
    "f.setPosition([0.,0.,.5])\n",
    "f.setQuaternion([1., .3, .0, .0]) # more about quaternions below\n",
    "print('frame name:', f.name, 'pos:', f.getPosition(), 'quat:', f.getQuaternion())\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's add a second frame, but a bit larget, with the first as parent, and with a hinge joint:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = C.addFrame(name='second', parent='first')\n",
    "f.setJoint(ry.JT.hingeX)\n",
    "f.setShape(type=ry.ST.marker, size=[.4])\n",
    "print('frame name:', f.name, 'pos:', f.getPosition(), 'quat:', f.getQuaternion())\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the `getPosition` and `getQuaternion` methods understood this is a child frame of the above (i.e., computed forward kinematics). Since we now have a configuration with a joint, we can articulate it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q = C.getJointState()\n",
    "q[0] = q[0] + .2\n",
    "C.setJointState(q)\n",
    "print('joint state:', q)\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For illustration, let's add various other child frames with shapes to the 'second' frame and animate it with a trivial loop:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add new frames: a box, a ball and a capsule\n",
    "C.addFrame('ball', 'second') .setShape(ry.ST.sphere, [.1]) .setColor([.7,.0,.0]) .setRelativePosition([-.3,.0,.2])\n",
    "C.addFrame('box', 'second') .setShape(ry.ST.ssBox, [.3,.2,.1,.02]) .setColor([.0,.7,.0]) .setRelativePosition([.0,.0,.2])\n",
    "C.addFrame('capsule', 'second') .setShape(ry.ST.capsule, [.3, .05]) .setColor([.0,.0,.7]) .setRelativePosition([.3,.0,.2])\n",
    "\n",
    "# Articulate the new configuration by moving the hinge joint\n",
    "for t in range(50):\n",
    "    C.setJointState([np.cos(.1*t)])\n",
    "    C.view()\n",
    "    time.sleep(.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following lists all predefined shape types:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ry.ST.__members__.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the primitive shapes, the `size` specifies them (box: 3 widths, sphere: radius, capsule: [length radius], cylinder: [length radius], marker: [axes lengths]). The `ssBox` is a sphere-swept-Box with a 4D size: [3 widths, corner radius], which is used a lot. Also mesh shapes can be set manually from python (`setMesh`) by directly giving vertex and triangle arrays, or loaded from a file (`setMeshFile`). `setConvexMesh` means sphere-swept convex mesh, which is specified by a mesh and 1D size (radius), and also a very useful shape type. Other methods to set shapes include: `setImplicitSurface` (for a given 3D tensor, using Marching Cube to find the iso-surface mesh), `setTensorShape` (a 3D tensor with densities, displayed as transparent density), `setPointCloud` (optionally with normals), and `setLines`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also briefly list the possible joint types:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ry.JT.__members__.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please see the [Robot Kinematics and Dynamics Essentials](https://www.user.tu-berlin.de/mtoussai/notes/robotKin.html) for background on such joints (=parameterized transformations). The `quatBall` perhaps deserves special mentioning: it realizes a ball joint by introducing the 4D quaternion as joint state - this sometimes requires special attention. But the library implements correct Jacobians (accounting also for implicit normalization) w.r.t. the quaternion parameterization. The `free` joint is a 7D joint with 3D position and `quatBall`.\n",
    "\n",
    "There is also a `rigid` joint, which simply means no DOF and seems fully redundant: having no joint or a rigid joint is the same. However, there is a semantic difference. A child frame without joint is really considered to be *a part* of the parent: The parent could be a *link* with multiple shapes (e.g. convex parts) or masses associated, but it is semantically one big link. When a frame has a joint (including a `rigid` joint), it is not considered part of its parent. That semantics makes a difference, e.g. in picking or walking scenarios, where objects may *kinematically switch* their parent."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specifying transformations: Quaternions & turtle strings\n",
    "\n",
    "See the [Quaternions](https://www.user.tu-berlin.de/mtoussai/notes/quaternions.html) lecture note for background. Note that `ry.Quaternion` implements exactly the methods described in that note for converting between rotation representations. We use the (w,x,y,z) convention for quaternions. Whenever you set a quaternion manually, it is internally normalized. So setting a quaterniong (1,1,0,0) is totally fine and means 90-degrees around x:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C.clear()\n",
    "C.addFrame('A') .setShape(ry.ST.marker, [.3]) .setPosition([0.,0.,.5]) .setQuaternion([1., 1., .0, .0])\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But specifying transformations using translation and quaternion is not always intuitive as a human.\n",
    "Therefore, we also allow for a *turtle string* notation, namely a string that \n",
    "chains little interpretable 'translate' and 'rotate' commands to define a \n",
    "transformation, as in the old turtle language. Here is an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C.clear()\n",
    "f = C.addFrame('A') .setShape(ry.ST.marker, [.3])\n",
    "f.setPoseByText('t(0 0 .5) d(30 1 0 0) t(0 0 .5) d(30 1 0 0)') #t=translate, #d=rotate-by-degrees around axis\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a specification of the possible commands:\n",
    "<pre>\n",
    "  t(x y z)       # translation by (x,y,z)\n",
    "  q(q0 q1 q2 q3) # rotation by a quaternion\n",
    "  r(r x y z)     # rotation by `r` _radians_ around the axis (x,y,z)\n",
    "  d(d x y z)     # rotation by `d` _degrees_ around the axis (x,y,z)\n",
    "  E(r p y)       # rotation by roll-pitch-yaw Euler angles\n",
    "</pre>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading existing configurations\n",
    "\n",
    "You will usually load predefined configurations from file, and then perhaps edit and combine them manually. The rai package has several robot models predefined and installed in the `raiPath`. These are yaml files listing the frames (for historical reasons they have the `.g`-file extension (see [Notes:Yaml-Graph Files](https://marctoussaint.github.io/robotic/notes.html#yaml-graph-files)). See the [Import/Edit tutorial](config_3_import_edit.html), e.g. for loading URDF or Mujoco."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C.clear()\n",
    "C.addFile(ry.raiPath('panda/panda.g'))\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's add a second panda, but to avoid frame name conflicts we prefix all frame names. That's often important when importing multiple models into a single configuration. We can move the second robot simply by redefining the position of its base frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C.addFile(ry.raiPath('panda/panda.g'), 'r_')\n",
    "base_r = C.getFrame('r_panda_base')\n",
    "base_r.setPosition([.0, .5, .0])\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These models have several joints. We can get the joint state of the full configuration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(C.getJointState())\n",
    "print('joints:', C.getJointNames())\n",
    "print('frames:', C.getFrameNames())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's animate - without respect for limits or collisions!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q0 = C.getJointState()\n",
    "for t in range(20):\n",
    "    q = q0 + .1*np.random.randn(q0.shape[0])\n",
    "    C.setJointState(q)\n",
    "    C.view()\n",
    "    time.sleep(.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Looping through frames and retrieving frame information\n",
    "\n",
    "Sometimes we want to loop through all frames (e.g. to export a configuration). In python, we can get the specs of a frame as a dict, and set/add attributes of a frame. And we can write all frames in a fresh .yml file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C = ry.Config()\n",
    "C.addFile(ry.raiPath('panda/panda.g'))\n",
    "q0 = C.getJointState()\n",
    "C.view(False)\n",
    "\n",
    "for f in C.getFrames():\n",
    "    print(f.name, f.asDict()) #info returns all attributes, similar to what is defined in the .g-files\n",
    "    #see also all the f.get... methods\n",
    "\n",
    "# setting attributes:\n",
    "f = C.getFrames()[0]\n",
    "f.setAttributes({'myvalue': 12.345, 'myname': 'Hari Seldon'})\n",
    "print(f.asDict())\n",
    "\n",
    "# writing a configuration to a file:\n",
    "with open(\"z.yml\", \"w\") as fil:\n",
    "    fil.write(C.write())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Features: computing geometric properties\n",
    "\n",
    "For every frame we can query its pose:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = C.getFrame('gripper')\n",
    "print('gripper pos:', f.getPosition())\n",
    "print('gripper quat:', f.getQuaternion())\n",
    "print('gripper rot:', f.getRotationMatrix())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above provides basic forward kinematics: After `setJointState` you can query the pose of any frame.\n",
    "\n",
    "However, there is a more general way to query *features*, i.e. properties of the configuration in a differentiable manner. Features are defined and explained in more detail in the [Features tutorial](config_2_features.html).\n",
    "\n",
    "Features allow you not only to query the value, but also the Jacobian. In practise, one rarely queries features directly (as below). But they are essential to define constraints and objectives for inverse kinematics and path optimization problems (see [KOMO tutorial](komo_1_intro.html)). Here an example of how to query them directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[y,J] = C.eval(ry.FS.position, ['gripper'])\n",
    "print('position of gripper:', y, '\\nJacobian:', J)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# negative(!) distance between two convex shapes (or origin of marker)\n",
    "C.eval(ry.FS.negDistance, ['panda_coll7', 'panda_coll7'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the x-axis of the given frame in world coordinates\n",
    "C.eval(ry.FS.vectorX, ['gripper'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See the next [Features tutorial](config_2_features.html) on all the things you can compute over configurations, and that can be used to formulate optimization problems."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Joint and Frame State\n",
    "\n",
    "A configuration is a tree of $n$ frames. Every frame has a pose (position & quaternion), which is represented as a 7D vector $(x,y,z, q_w,q_x,q_y,q_z)^T$. The frame state is the $n\\times 7$ matrix, where the $i$-th row is the pose of the $i$-th frame.\n",
    "\n",
    "A configuration also defines joints, which means that the relative transfromation from a parent to a child frame is parameterized by degrees-of-freedom (DOFs). If the configuration has in total $d$ DOFs, the joint state is a $d$-dimensional vector.\n",
    "\n",
    "Setting the joint state implies computing all relative transformations, and then forward chaining all transformations to compute all frame poses. So setting the joint state also sets the frame state.\n",
    "     \n",
    "Setting the frame state allows you to set frame poses that are inconsistent/impossible w.r.t. the joints! Setting the frame state implies computing all relative transformations from the frame poses, and then assigning the joint state to the *projection* onto the actual DOFs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C.setJointState(q0)\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The *frame state* is a $n\\times 7$ matrix, which contains for all of $n$ frames the 7D pose. A pose is stored as [p_x, p_y, p_z, q_w, q_x, q_y, q_z], with position p and quaternion q."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X0 = C.getFrameState()\n",
    "print('frame state: ', X0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's do a very questionable thing: adding .1 to all numbers in the frame matrix!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = X0 + .1\n",
    "C.setFrameState(X)\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That totally broke the original design of the robot! Setting global frame states overwrites the relative transformations between frames! (Also, the rows of X have non-normalized quaternions. But these are normalized when setting the frame state.)\n",
    "\n",
    "Let's reset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C.setFrameState(X0)\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## Selecting joints\n",
    "\n",
    "Often one would like to choose which joints are actually active, that is, which joints are referred to in q. This allows one to sub-select joints and work only with projections of the full dof state. This changes the joint state dimensionality, including ordering of entries in q. The frame state is not affected by such a selection of active joints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C.selectJoints(['panda_joint1', 'panda_joint2'])\n",
    "print('joint state: ', C.getJointState())\n",
    "print('joint names: ', C.getJointNames() )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C.selectJoints([], notThose=True)\n",
    "print('joint state: ', C.getJointState())\n",
    "print('joint names: ', C.getJointNames() )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look closely at the above: We suddenly have 8 joints! (When loading the configuration we only saw 7.) The reason is that the last joint `panda_finger_joint1` is (when loaded) non-selected (inactive).\n",
    "\n",
    "However, even if a joint is non-selected, one can still query or set them directly, by extra arguments of setJointState, as shown below. However, a really important implication of selecting (or de-selecting) joints is w.r.t. the Jacobian: The Jacobians are always returns only w.r.t. the currently selected joints. (And also optimization problems (in KOMO) are only w.r.t. the currently selected joints).\n",
    "\n",
    "Let's demonstrate explicit access:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C.clear()\n",
    "C.addFile(ry.raiPath('panda/panda.g'))\n",
    "q = C.getJointState()\n",
    "print('joint state shape: ', q.shape)\n",
    "\n",
    "C.setJointState([.04], ['panda_finger_joint1']) #closing the fingers to a 2cm gap (panda_finger_joint2 is a mimic joint)\n",
    "C.setJointState([.1,.1], ['panda_joint2','panda_joint4']) #stretching these joints\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## View interaction and releasing objects"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can close and re-open the view window"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C.view_close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# things are still there\n",
    "C.view(pause=False, message='this is a message')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For user interaction it is often useful to wait for a keypress (i.e., by making `view` a blocking call via `pause=True`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pressed_key = C.view(pause=True, message='press some key!')\n",
    "print('pressed key:', pressed_key, chr(pressed_key))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see above, `.view()` also returnes the ID of key that was pressed, which can be helpful for debugging purposes.\n",
    "\n",
    "To visualize your configuration outside of the viewer, you can also get a screenshot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = C.viewer().getRgb()\n",
    "print(type(img), img.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.imshow(img)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cleanup\n",
    "\n",
    "Don't forget to release everything, including closing the view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}