{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b3344021",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "# BotOp-3: Step Inferace (env-like) example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58cd1685",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "import robotic as ry\n",
    "import numpy as np\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7dc6066",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "A minimal IK method (could be replace by just using the Jacobian to translate the delta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "952ba11e",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def mini_IK(C: ry.Config, pos, qHome):\n",
    "    q = C.getJointState()\n",
    "    komo = ry.KOMO(C, 1, 1, 0, False) #one phase one time slice problem, with 'delta_t=1', order=0\n",
    "    komo.addObjective([], ry.FS.jointState, [], ry.OT.sos, [1e-1], q) #cost: close to 'current state'\n",
    "    komo.addObjective([], ry.FS.jointState, [], ry.OT.sos, [1e-1], qHome) #cost: close to qHome\n",
    "    komo.addObjective([], ry.FS.position, ['l_gripper'], ry.OT.eq, [1e2], pos) #constraint: gripper position\n",
    "    \n",
    "    ret = ry.NLP_Solver(komo.nlp()) .setOptions(verbose=0, stopEvals=20, stopTolerance=1e-3) .solve()\n",
    "\n",
    "    if ret.feasible:\n",
    "        return ret.x - q\n",
    "    else:\n",
    "        return np.zeros((q.size))\n",
    "\n",
    "def mini_JacIK(C: ry.Config, pos, qHome):\n",
    "    q = C.getJointState()\n",
    "    y, J = C.eval(ry.FS.position, ['l_gripper'])\n",
    "    Jinv = J.T @ np.linalg.pinv(J@J.T+1e-3*np.eye(y.size))\n",
    "    dq = Jinv @ (pos-y) + 0.1*(np.eye(q.size) - Jinv@J) @ (qHome-q)\n",
    "    return dq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71dc5fe3",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "basic setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43f2d5b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "C = ry.Config()\n",
    "\n",
    "C.addFile(ry.raiPath(\"../rai-robotModels/scenarios/pandaSingle.g\"))\n",
    "qHome = C.getJointState()\n",
    "\n",
    "target = C.addFrame('target', 'table')\n",
    "target.setShape(ry.ST.marker, [.1])\n",
    "target.setRelativePosition([0., .3, .3])\n",
    "pos = target.getPosition()\n",
    "cen = pos.copy()\n",
    "C.view()\n",
    "\n",
    "bot = ry.BotOp(C, useRealRobot=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37905a58",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "A basic observe-decide-step loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "048b282e",
   "metadata": {},
   "outputs": [],
   "source": [
    "tau_step = .05\n",
    "lmbda = .2\n",
    "for t in range(100):\n",
    "    bot.sync(C, tau_step) #keep the workspace C sync'ed to real/sim, and idle .1 sec\n",
    "    pos = cen + .98 * (pos-cen) + 0.02 * np.random.randn(3)\n",
    "    target.setPosition(pos)\n",
    "    \n",
    "    # observe\n",
    "    obs = bot.stepObservation()\n",
    "    \n",
    "    # decide\n",
    "    # dq = mini_IK(C, pos, qHome)\n",
    "    dq = mini_JacIK(C, pos, qHome)\n",
    "    \n",
    "    # step\n",
    "    # bot.moveTo(ret.x, timeCost=2., overwrite=True)\n",
    "    bot.stepAction(dq, obs, lmbda, maxAccel=2.)        \n",
    "\n",
    "bot.home(C)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}