{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4dfcd570-a1eb-4535-8ea7-7c023153aafb",
   "metadata": {},
   "source": [
    "© Copyright, 2025 G. Schaer.\n",
    "\n",
    "SPDX-License-Identifier: GPL-3.0-only"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bd6dd9a-a0f6-4248-89d4-657c318e0d79",
   "metadata": {},
   "source": [
    "# Tutorial 2: Joint Torques"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c18bd9e-7470-4b09-b7a7-fac8bb2dec17",
   "metadata": {},
   "source": [
    "## Tutorial Description"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7e48aef-41be-4597-be37-1bd196bc97a5",
   "metadata": {},
   "source": [
    "This tutorial covers creating a `condynsate` `Project` in which torques are applied to specific joints of an articulated body. We will cover:\n",
    "1. Applying torques to individual continuous joints of a .urdf object.\n",
    "2. Measuring the state of individual joints of a .urdf object."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be64a904-72de-4411-8a77-37a630f0c3a4",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ca5c838-ffae-45eb-9e60-0ad1e0ac8698",
   "metadata": {},
   "source": [
    "To begin, we import the same modules for the same reasons as tutorial 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b2b6dc1e-8420-4242-91a3-e73422ce713e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from condynsate import Project\n",
    "from condynsate import __assets__ as assets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "846403b0-0ac0-4cd8-bc64-0ce4aa5e7d60",
   "metadata": {},
   "source": [
    "## Initializing the Project Class"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e8b4d04-e298-491d-bffe-1e4f87f5dbaf",
   "metadata": {},
   "source": [
    "Here we initialize the `Project` in the same way as Tutorial 0, but instead of loading a cube object, we load a pendulum. Additionally, we leave the ground grid on and do not load a ground plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "54974e36-d9c3-417a-b12a-fb9259c6f76a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You can open the visualizer by visiting the following URL:\n",
      "http://127.0.0.1:7032/static/\n"
     ]
    }
   ],
   "source": [
    "# Create the project\n",
    "proj = Project()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d15cee85-9874-453b-82c3-25726a344d32",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "proj.visualizer.set_axes(False) # Returns 0 on success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8e3b69f-ab5e-4db6-8163-6be3d8371cac",
   "metadata": {},
   "source": [
    "This time, we make the pendulum fixed. This ensures the base of the pendulum has 0 degrees of freedom. Setting the `fixed` flag doe not affect the degrees of freedom of the joints of the pendulum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5431e17b-4926-44c0-b544-f652d7ee5f1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load a pendulum object and set its initial position to resting on the ground\n",
    "pendulum = proj.load_urdf(assets['pendulum.urdf'], \n",
    "                          fixed=True\n",
    "                         )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc265f8f-74ec-4f26-83ce-77ee635a7fdf",
   "metadata": {},
   "source": [
    "Additionally to setting the initial state of the body as a whole, we can also set the initial state of each joint of articulated bodies. Here we will set the initial angle of the pendulum arm to 10 degrees. To access the joints of a body, we use the `joints` attribute. The `joints` attribute is a dictionary that indexes joint names (as defined in the .urdf files from which bodies are loaded) to members of the `condynsate.simulator.objects.Joint` class.\n",
    "\n",
    "The name of the pendulum joint is `base_to_arm`. This name is defined in the pendulum.urdf file.\n",
    "\n",
    "We can then call the `condynsate.simulator.objects.Joint.set_initial_state` function to set the initial state of the arm.\n",
    "\n",
    "    -----------------------------------------------------------------------------\n",
    "    | condynsate.simulator.objects.Joint.set_initial_state                      |\n",
    "    -----------------------------------------------------------------------------\n",
    "    Sets the initial state of the joint. When the simulation is reset\n",
    "    the joint will be reset to this value\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    **kwargs\n",
    "\n",
    "    Keyword Args\n",
    "    ------------\n",
    "    angle : float, optional\n",
    "        The (angle in radians) of the joint about the joint axis.\n",
    "    omega : float, optional\n",
    "        The angular velocity (angle in radians / second) of the joint\n",
    "        about the joint axis.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    ret_code : int\n",
    "        0 if successful, -1 if something went wrong."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "399f40d8-fc5f-4b26-8a07-5c4fdbc5ea7c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set the initial angle of the pendulum arm\n",
    "pendulum.joints['base_to_arm'].set_initial_state(angle=0.1745) # All angles in condynsate are radians"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1f58e78f-780c-4604-878d-5b4e162aa364",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Refresh the visualizer to show changes to the pendulums's position and angle\n",
    "proj.refresh_visualizer() # Returns 0 on success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7c03a28-215d-4b7c-a41c-062592923084",
   "metadata": {},
   "source": [
    "## Running a Simulation Loop"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "516a7db7-7e24-4522-86f1-27698df7ff98",
   "metadata": {},
   "source": [
    "Similary to Tutorial 0, we start the simulation loop by calling `condynsate.Project.reset`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d8de6343-31a0-4525-8791-bb7bdb130e49",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "proj.reset() # Returns 0 on success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f536c1e-bd87-4454-9348-e343c4adb618",
   "metadata": {},
   "source": [
    "In each step of the loop we take 3 steps\n",
    "1. Get the pendulum's joint angle\n",
    "2. Apply a torque to the pendulum joint proportional to the joint angle \n",
    "4. Take a single simulation step"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e127fedf-9fe5-48a6-b7e3-f768c1fd2fe7",
   "metadata": {},
   "source": [
    "To get state information about joints, we access a joint's `state` attribute. The `state` attribute has the following child attributes:\n",
    "\n",
    "    -----------------------------------------------------------------------------\n",
    "    | condynsate.simulator.dataclasses.JointState                               |\n",
    "    -----------------------------------------------------------------------------\n",
    "    Attributes\n",
    "    ----------\n",
    "    angle : float\n",
    "        The angle of the joint about the joint axis.\n",
    "    omega : float\n",
    "        The angular velocity of the joint about the joint axis.\n",
    "\n",
    "To get the pendulum's joint angle, we access the `angle` child attribute."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88aa3ab7-00d3-46f0-aa69-dca04d50e821",
   "metadata": {},
   "source": [
    "To apply a torque to a joint, we call the `condynsate.simulator.objects.Joint.apply_torque` function\n",
    "\n",
    "    -----------------------------------------------------------------------------\n",
    "    | condynsate.simulator.objects.Joint.apply_torque                           |\n",
    "    -----------------------------------------------------------------------------\n",
    "    Applies torque to a joint for a single simulation step.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    torque : float\n",
    "        The torque being applied about the joint's axis..\n",
    "    **kwargs\n",
    "\n",
    "    Keyword Args\n",
    "    ------------\n",
    "    draw_arrow : bool, optional\n",
    "        A Boolean flag that indicates if an arrow should be drawn\n",
    "        to represent the applied torque. The default is False.\n",
    "    arrow_scale : float, optional\n",
    "        The scaling factor, relative to the size of the applied torque,\n",
    "        that is used to size the torque arrow. The default is 1.0.\n",
    "    arrow_offset : float, optional\n",
    "        The amount by which the drawn is offset from the center of the\n",
    "        joint's child link along the joint axis. The default is 0.0.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    ret_code : int\n",
    "        0 if successful, -1 if something went wrong."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ccf9fb0-4c0f-41d8-94e7-77e36a7aa8be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run a 10 second simulation\n",
    "while proj.simtime <= 10.:\n",
    "    \n",
    "    # Get the pendulum's joint angle\n",
    "    joint_state = pendulum.joints['base_to_arm'].state\n",
    "    angle = joint_state.angle\n",
    "\n",
    "    # Apply a torque to the pendulum's joint\n",
    "    torque = -14.0 * angle\n",
    "    pendulum.joints['base_to_arm'].apply_torque(torque, draw_arrow=True, arrow_scale=0.25)\n",
    "\n",
    "    # Take a single simulation step\n",
    "    proj.step(real_time=True, # Run the simulation in real time\n",
    "              stable_step=False # Dynamically adjust the refresh rate for best total run time\n",
    "             )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d4d047e-286a-4fab-a79b-839c5f62bec3",
   "metadata": {},
   "source": [
    "Note that exactly like the previous tutorial, after `condynsate.Project.step` is called, the torques applied in the previous simulation step will not be applied in the next step. Therefore, we must call `apply_torque` function before every time step in which we want to apply joint torques."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfab525b-b921-43ee-91d6-44a8a4edf317",
   "metadata": {},
   "source": [
    "Finally, we ensure all children threads exit gracefully."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a95f4782-c4bd-4ab1-9b4f-bc3310d079f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "proj.terminate() # Returns 0 on success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36d1e2de-e3f4-4dfb-a530-766a4420f876",
   "metadata": {},
   "source": [
    "## Challenge"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ff6e017-b96d-4668-879d-5ebaf083ae88",
   "metadata": {},
   "source": [
    "This tutorial is now complete. For an added challenge, think of how you would modify the simulation loop to implement a PD controller. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "457d9e63-4353-424c-ae6c-fe7571527ef3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.14.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}