diff --git a/python/least_squares.ipynb b/python/least_squares.ipynb index e382cdc7..ed3b8063 100644 --- a/python/least_squares.ipynb +++ b/python/least_squares.ipynb @@ -11,27 +11,17 @@ "#

Least Squares

\n", "\n", "This notebook describes a utility function included in the MuJoCo Python library performing box-bounded nonlinear least squares optimization. We provide some theoretical background, describe our implementation and show example usage.\n", - "" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "pOyD_TXFrM_4" - }, - "source": [ - "### Copyright notice" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "zYLuuNFmrEfo" - }, - "source": [ - ">

Copyright 2022 DeepMind Technologies Limited

\n", - ">

Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0.

\n", - ">

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

" + "\n", + "\n", + "" ] }, { @@ -86,6 +76,7 @@ " print('Checking that the installation succeeded:')\n", " import mujoco\n", " from mujoco import minimize\n", + " from mujoco import rollout\n", " mujoco.MjModel.from_xml_string('')\n", "except Exception as e:\n", " raise e from RuntimeError(\n", @@ -107,7 +98,8 @@ "import time\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", - "from matplotlib.patches import Rectangle\n" + "from matplotlib.patches import Rectangle\n", + "from typing import Tuple, Optional, Union\n" ] }, { @@ -262,6 +254,26 @@ "This completes the background section and allows us to accurately describe the function explored in the rest of the notebook. " ] }, + { + "cell_type": "markdown", + "metadata": { + "id": "IaaI2ZIY5FIw" + }, + "source": [ + "## Least Norm Generalization\n", + "\n", + "The Least Squares problem described above can be generalized as follows. Instead of the quadratic norm $\\frac{1}{2} r(x)^T\\cdot r(x)$, we can use some other smooth, convex norm $f(x) = n(r(x))$. Letting $\\nabla n = \\frac{\\partial n}{\\partial r}$ and $\\nabla^2 n = \\frac{\\partial^2 n}{\\partial r^2}$ be respectively the gradient and Hessian of $n$ with respect to $r$, we have\n", + "$$\n", + "\\begin{align}\n", + "g &= J^T\\cdot\\nabla n\\\\\n", + "H_{GN} &= J^T\\cdot \\nabla^2 n \\cdot J\n", + "\\end{align}\n", + "$$\n", + "It is easy to verify that for the quadratic norm, these expression reduce to the ones above as $\\nabla (\\frac{1}{2}r^T\\cdot r) = r$ and $\\nabla^2 (\\frac{1}{2}r^T\\cdot r) = I_m$.\n", + "\n", + "This completes the background section and allows us to accurately describe the function explored in the rest of the notebook. " + ] + }, { "cell_type": "markdown", "metadata": { @@ -293,7 +305,7 @@ "3. Since each evaluation of the residual involves rolling out the physics to obtain simulated sensor values, computing $r(x)$ is the most expensive part of the optimization.\n", "4. Analytic Jacobians $J = \\frac{\\partial r}{\\partial x}$ are usually not available and must be obtained with **finite-differencing**.\n", "5. Due to the sematics of $x$, box-bounds are usually sufficient (for example, masses and friction coefficients cannot be negative, joint angles should not exceed their limits).\n", - "6. The implementation should be efficient yet readable. The [least_squares function](https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/minimize.py) takes less than 250 lines of code.\n", + "6. The implementation should be efficient yet readable. The [least_squares function](https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/minimize.py) takes up ~250 lines of code.\n", "\n", "Let's look at the function's docstring and then discuss some implementation notes." ] @@ -317,11 +329,13 @@ "source": [ "## Implementation notes\n", "\n", + "1. The residual funciton must be vectorized: besides taking a column vector $x$ and returning the residual $r(x)$, it must accept an $n\\times k$ matrix $X$, returning an $m\\times k$ matrix $R$. The vectorized format is used by the internal finite-difference implementation and can be exploited to speed up the minimization by using multi-threading inside the residual function implementation.\n", "1. Bounds must be `None` or fully specified for all dimensions of $x$.\n", - "2. The `jacobian` callback can be supplied by the user and is finite-differenced otherwise. Note that this callback is not made available in the case the user knows the analytic Jacobian (this is very rare in the sysID context), but in case the user wants to implement their own multi-threaded fin-diff callback.\n", - "3. Automatic forward/backward differencing, chosen to avoid crossing the bounds, with optional central differencing. The fin-diff epsilon `eps` is scaled by the size of the bounds, if provided.\n", - "4. The termination criterion is based on small step size $||\\delta x|| < \\textrm{tol}$.\n", - "5. We use the simple yet affective $\\mu$-search strategy described in [Bazaraa et-al.](https://onlinelibrary.wiley.com/doi/book/10.1002/0471787779). Backtracking $\\mu$-increases are *careful*, attempting to find the smallest $\\mu$ where sufficient reduction is found. $\\mu$-decreases are *aggressive*, allowing fast quadratic convergence to a local minimum." + "1. The `jacobian` callback can be supplied by the user and is finite-differenced otherwise. Note that this callback is not made available in the case the user knows the analytic Jacobian (this is very rare in the sysID context), but in case the user wants to implement their own multi-threaded fin-diff callback.\n", + "1. Automatic forward/backward differencing, chosen to avoid crossing the bounds, with optional central differencing. The fin-diff epsilon `eps` is scaled by the size of the bounds, if provided.\n", + "1. The termination criterion is based on small step size $||\\delta x|| < \\textrm{tol}$.\n", + "1. We use the simple yet affective $\\mu$-search strategy described in [Bazaraa et-al.](https://onlinelibrary.wiley.com/doi/book/10.1002/0471787779). Backtracking $\\mu$-increases are *careful*, attempting to find the smallest $\\mu$ where sufficient reduction is found. $\\mu$-decreases are *aggressive*, allowing fast quadratic convergence to a local minimum.\n", + "1. The user may optionally provide a `norm` different than the quadratic norm (the default), this is covered in more detail below." ] }, { @@ -330,7 +344,7 @@ "id": "MvnHwh2ZgGT2" }, "source": [ - "# Toy examples: 2D" + "# Toy examples" ] }, { @@ -429,7 +443,7 @@ "source": [ "# Minimize Rosenbrock function.\n", "def rosenbrock(x):\n", - " return np.array((1-x[0], 10*(x[1]-x[0]**2)))\n", + " return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)])\n", "\n", "x0 = np.array((0.0, 0.0))\n", "x, rb_trace = minimize.least_squares(x0, rosenbrock);" @@ -477,9 +491,9 @@ "source": [ "#@title Minimize, visualize Beale\n", "def beale(x):\n", - " return np.array((1.5-x[0]+x[0]*x[1],\n", - " 2.25-x[0]+x[0]*x[1]*x[1],\n", - " 2.625-x[0]+x[0]*x[1]*x[1]*x[1]))\n", + " return np.stack((1.5-x[0, :]+x[0, :]*x[1, :],\n", + " 2.25-x[0, :]+x[0, :]*x[1, :]*x[1, :],\n", + " 2.625-x[0, :]+x[0, :]*x[1, :]*x[1, :]*x[1, :]))\n", "\n", "x0 = np.array((-3.0, -3.0))\n", "x, bl_trace = minimize.least_squares(x0, beale)\n", @@ -516,7 +530,7 @@ "# Choose bounds.\n", "lower = np.array([-.3, -1.])\n", "upper = np.array([0.9, 1.9])\n", - "bounds = (lower, upper)\n", + "bounds = [lower, upper]\n", "\n", "# Make some initial points, minimize, save taces.\n", "num_points = 4\n", @@ -549,7 +563,7 @@ "# Choose bounds.\n", "lower = np.array([-2, -1.3])\n", "upper = np.array([1.5, 3.])\n", - "bounds = (lower, upper)\n", + "bounds = [lower, upper]\n", "\n", "# Make some initial points, minimize, save taces.\n", "num_points = 5\n", @@ -569,15 +583,6 @@ "plot_2D(beale, 'Beale Function', plot_range, minimum, traces, bounds)" ] }, - { - "cell_type": "markdown", - "metadata": { - "id": "jX81C8JjGOwc" - }, - "source": [ - "# Toy examples: high-D" - ] - }, { "cell_type": "markdown", "metadata": { @@ -604,8 +609,8 @@ "n = 20\n", "\n", "def rosenbrock_n(x):\n", - " res0 = [1-x[i] for i in range(n-1)]\n", - " res1 = [10*(x[i]-x[i+1]**2) for i in range(n-1)]\n", + " res0 = [1 - x[i, :] for i in range(n - 1)]\n", + " res1 = [10 * (x[i, :] - x[i + 1, :] ** 2) for i in range(n - 1)]\n", " return np.asarray(res0 + res1)\n", "\n", "x0 = np.zeros(n)\n", @@ -615,570 +620,6 @@ "assert np.linalg.norm(x-1) < 1e-8" ] }, - { - "cell_type": "markdown", - "metadata": { - "id": "B23ybrGiVjN6" - }, - "source": [ - "## Simple humanoid control\n", - "\n", - "We'll now use `least_squares` to solve a humanoid control problem. While this is not its intended purpose, it demonstrates the power and general usability of the function.\n", - "\n", - "Below, we copy MuJoCo's [standard humanoid model](https://github.com/google-deepmind/mujoco/blob/main/model/humanoid/humanoid.xml), with the following modifications:\n", - "1. Added a \"target\" mocap body with a pink spherical site.\n", - "2. Changed the color of the right hand to pink.\n", - "3. Replaced the torque actuators with position actuators." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "cellView": "form", - "id": "fQlg1Pwl6eRG" - }, - "outputs": [], - "source": [ - "#@title Humanoid XML\n", - "xml = \"\"\"\n", - "\n", - " \n", - "\"\"\"" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "K4enFXBZa9TE" - }, - "source": [ - "Let's load the model and render the initial state for our control problem, chosen to be the \"squat\" [keyframe](https://mujoco.readthedocs.io/en/latest/XMLreference.html#keyframe)." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "Sz8eUNsyYnyr" - }, - "outputs": [], - "source": [ - "# Load model, make data\n", - "model = mujoco.MjModel.from_xml_string(xml)\n", - "data = mujoco.MjData(model)\n", - "\n", - "# Set the state to the \"squat\" keyframe, call mj_forward.\n", - "key = model.key('squat').id\n", - "mujoco.mj_resetDataKeyframe(model, data, key)\n", - "mujoco.mj_forward(model, data)\n", - "\n", - "# If a renderer exists, close it.\n", - "if 'renderer' in locals():\n", - " renderer.close()\n", - "\n", - "# Make a Renderer and a camera.\n", - "renderer = mujoco.Renderer(model, height=480, width=640)\n", - "camera = mujoco.MjvCamera()\n", - "mujoco.mjv_defaultFreeCamera(model, camera)\n", - "camera.distance = 3\n", - "camera.elevation = -10\n", - "\n", - "# Point the camera at the humanoid, render.\n", - "camera.lookat = data.body('torso').subtree_com\n", - "renderer.update_scene(data, camera)\n", - "media.show_image(renderer.render())" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "4xpd_ioic6Wq" - }, - "source": [ - "### Problem definition\n", - "\n", - "We choose the following optimal control problem defintion:\n", - "- A trajectory is rolled out from $t=0\\ldots T$, starting at the \"squat\" keyframe shown above.\n", - "- Controls $u_t$ are applied during the rollout which are a linear interpolation of first control $u_0$ and the last one $u_T$. These two control vectors are our decision variable $x = \\begin{pmatrix} u_0 & u_T \\end{pmatrix}$.\n", - "- The residual is a concatenation, over all time steps, of the vector from the right hand to the target and the torques applied by the actuators, scaled by some factor (the torques are numerically much larger than the hand-target distances).\n", - "\n", - "Let's see what this residual looks like:" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "QrqGsFKMK3of" - }, - "outputs": [], - "source": [ - "def reach(ctrl0T, target, T, torque_scale, traj=None):\n", - " \"\"\"Residual for target-reaching task.\n", - "\n", - " Args:\n", - " ctrl0T: contatenation of the first and last control vectors.\n", - " target: target to which the right hand should reach.\n", - " T: final time for the rollout.\n", - " torque_scale: coefficient by which to scale the torques.\n", - " traj: optional list of positions to be recorded.\n", - "\n", - " Returns:\n", - " The residual of the target-reaching task.\n", - " \"\"\"\n", - "\n", - " # Reset to the \"squat\" keyframe.\n", - " key = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, 'squat')\n", - " mujoco.mj_resetDataKeyframe(model, data, key)\n", - "\n", - " # Move the mocap body to the target (for visualization only)\n", - " mocapid = model.body('target').mocapid\n", - " data.mocap_pos[mocapid] = target\n", - "\n", - " # Extract the first and last ctrl vectors\n", - " ctrl0 = ctrl0T[:model.nu]\n", - " ctrlT = ctrl0T[model.nu:]\n", - "\n", - " # Roll out the trajectory, accumulate the residual.\n", - " res = []\n", - " while data.time < T:\n", - " # Interpolate ctrl from ctrl0 and ctrlT.\n", - " f0 = (T - data.time) / T\n", - " f1 = 1 - f0\n", - " data.ctrl = f0*ctrl0 + f1*ctrlT\n", - "\n", - " # Step.\n", - " mujoco.mj_step(model, data)\n", - "\n", - " # Append the task residual: hand to target.\n", - " res.append(data.geom('hand_right').xpos - np.array(target))\n", - "\n", - " # Append the energy residual: actuator torques.\n", - " res.append(torque_scale * data.actuator_force.flatten())\n", - "\n", - " # Save state to traj, if requested.\n", - " if traj is not None:\n", - " traj.append(data.qpos.copy())\n", - "\n", - " # The normalizer keeps objective values the same when changing T or timestep.\n", - " normalizer = 100 * model.opt.timestep / T\n", - " return np.hstack(res).flatten() * normalizer" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "MvKh4XWIK8Bx" - }, - "source": [ - "Now let's give the rest of the problem definition:\n", - "1. The trajectory is integrated for 0.7s.\n", - "2. Torques are scaled by 0.003.\n", - "3. Since our decision variable is two copies of `mjData.ctrl`, the bounds are two concatenated copies of `mjData.actuator_ctrlrange`.\n", - "4. Our initial guess $x_0 = \\begin{pmatrix} u_0 & u_T \\end{pmatrix} = \\begin{pmatrix} q_\\textrm{squat} & q_\\textrm{stand} \\end{pmatrix}$ is the joint angles at the squatting position, followed by the angles at the default (standing position). We can use angles to initialize our controls because the position actuators have angle semantics.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "ZpT5JeyMFNjy" - }, - "outputs": [], - "source": [ - "T = 0.7 # Rollout length (seconds)\n", - "torque_scale = 0.003 # Scaling for the torques\n", - "\n", - "# Bounds are the stacked control bounds.\n", - "lower = model.actuator_ctrlrange[:,0]\n", - "upper = model.actuator_ctrlrange[:,1]\n", - "bounds = [np.hstack((lower, lower)), np.hstack((upper, upper))]\n", - "\n", - "# Initial guess is midpoint of the bounds\n", - "x0 = 0.5 * (bounds[1] + bounds[0])" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "RAD9cZkFKRpp" - }, - "source": [ - "Let's define a utility function for rendering frames and visualize the initial guess:" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "kqpsjWtPKf77" - }, - "outputs": [], - "source": [ - "def render_solution(x, target):\n", - " # Ask reach to save positions to traj.\n", - " traj = []\n", - " reach(x, target, T, torque_scale, traj=traj);\n", - "\n", - " frames = []\n", - " counter = 0\n", - " print('Rendering frames:', flush=True, end='')\n", - " for qpos in traj:\n", - " # Set positions, call mj_forward to update kinematics.\n", - " data.qpos = qpos\n", - " mujoco.mj_forward(model, data)\n", - "\n", - " # Render and save frames.\n", - " camera.lookat = data.body('torso').subtree_com\n", - " renderer.update_scene(data, camera)\n", - " pixels = renderer.render()\n", - " frames.append(pixels)\n", - " counter += 1\n", - " if counter % 10 == 0:\n", - " print(f' {counter}', flush=True, end='')\n", - " return frames\n", - "\n", - "# Visualize the initial guess.\n", - "target = (0., 0., -1.) # Target irrelevant, put it under the floor.\n", - "media.show_video(render_solution(x0, target))" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "ZX4NJFxWLNPr" - }, - "source": [ - "### Solutions to the reach task\n", - "\n", - "Let's solve once for some target and look at the optimization printout:" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "h8jRWeyJ9_A7" - }, - "outputs": [], - "source": [ - "target = (.4, -.3, 1.2)\n", - "\n", - "reach_target = lambda x: reach(x, target, T, torque_scale, traj=None)\n", - "\n", - "r0 = reach_target(x0)\n", - "print(f'The decision variable x has size {x0.size}')\n", - "print(f'The residual r(x) has size {r0.size}\\n')\n", - "\n", - "x, _ = minimize.least_squares(x0, reach_target, bounds);" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "aOpVVTb5L9gZ" - }, - "source": [ - "Let's see what this solution looks like:" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "LV9g5ExyJKrU" - }, - "outputs": [], - "source": [ - "media.show_video(render_solution(x, target))" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "igXc849HOdO6" - }, - "source": [ - "Let's rerun this for several target values and make a video of all of them:" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "id": "kbNubkPyIdJu" - }, - "outputs": [], - "source": [ - "targets = [(0.4, 0., 0.), (0.2, -1., 0.5), (-1., -.3, 1.), (0., -.2, 2.2)]\n", - "\n", - "frames = []\n", - "for target in targets:\n", - " res_target = lambda x: reach(x, target, T, torque_scale)\n", - " print(f'Optimizing for target at {target}', flush=True)\n", - " x, trace = minimize.least_squares(x0, res_target, bounds,\n", - " verbose=minimize.Verbosity.FINAL)\n", - " frames += render_solution(x, target)\n", - " print('\\n')\n", - "\n", - "print('Making video', flush=True)\n", - "media.show_video(frames)" - ] - }, { "cell_type": "markdown", "metadata": { @@ -1612,7 +1053,7 @@ "outputs": [], "source": [ "# Bounds at the joint limits.\n", - "bounds = (model.jnt_range[:, 0], model.jnt_range[:, 1])\n", + "bounds = [model.jnt_range[:, 0], model.jnt_range[:, 1]]\n", "\n", "# Inital guess is the 'home' keyframe.\n", "x0 = model.key('home').qpos" @@ -1654,29 +1095,34 @@ " data.mocap_quat[id] = model.body('target').quat if quat is None else quat\n", "\n", " # Set qpos, compute forward kinematics.\n", - " data.qpos = x\n", - " mujoco.mj_kinematics(model, data)\n", + " res = []\n", + " for i in range(x.shape[1]):\n", + " data.qpos = x[:, i]\n", + " mujoco.mj_kinematics(model, data)\n", "\n", - " # Position residual.\n", - " res_pos = data.site('effector').xpos - data.site('target').xpos\n", + " # Position residual.\n", + " res_pos = data.site('effector').xpos - data.site('target').xpos\n", "\n", - " # Effector quat, use mju_mat2quat.\n", - " effector_quat = np.empty(4)\n", - " mujoco.mju_mat2Quat(effector_quat, data.site('effector').xmat)\n", + " # Effector quat, use mju_mat2quat.\n", + " effector_quat = np.empty(4)\n", + " mujoco.mju_mat2Quat(effector_quat, data.site('effector').xmat)\n", "\n", - " # Target quat, exploit the fact that the site is aligned with the body.\n", - " target_quat = data.body('target').xquat\n", + " # Target quat, exploit the fact that the site is aligned with the body.\n", + " target_quat = data.body('target').xquat\n", "\n", - " # Orientation residual: quaternion difference.\n", - " res_quat = np.empty(3)\n", - " mujoco.mju_subQuat(res_quat, target_quat, effector_quat)\n", - " res_quat *= radius\n", + " # Orientation residual: quaternion difference.\n", + " res_quat = np.empty(3)\n", + " mujoco.mju_subQuat(res_quat, target_quat, effector_quat)\n", + " res_quat *= radius\n", "\n", - " # Regularization residual.\n", - " reg_target = model.key('home').qpos if reg_target is None else reg_target\n", - " res_reg = reg * (x - reg_target)\n", + " # Regularization residual.\n", + " reg_target = model.key('home').qpos if reg_target is None else reg_target\n", + " res_reg = reg * (x[:, i] - reg_target)\n", "\n", - " return np.hstack((res_pos, res_quat, res_reg))" + " res_i = np.hstack((res_pos, res_quat, res_reg))\n", + " res.append(np.atleast_2d(res_i).T)\n", + "\n", + " return np.hstack(res)" ] }, { @@ -1714,10 +1160,8 @@ " # useful, but we don't need it here.\n", " del res\n", "\n", - " # We can assume x has been copied into qpos\n", - " # and that mj_kinematics has been called by ik()\n", - "\n", - " # Call mj_comPos (required for Jacobians).\n", + " # Call mj_kinematics and mj_comPos (required for Jacobians).\n", + " mujoco.mj_kinematics(model, data)\n", " mujoco.mj_comPos(model, data)\n", "\n", " # Get end-effector site Jacobian.\n", @@ -1763,7 +1207,8 @@ "print('Finite-differenced Jacobian:')\n", "x_fd, _ = minimize.least_squares(x0, ik, bounds, verbose=1);\n", "print('Analytic Jacobian:')\n", - "x_analytic, _ = minimize.least_squares(x0, ik, bounds, jacobian=ik_jac, verbose=1);\n", + "x_analytic, _ = minimize.least_squares(x0, ik, bounds, jacobian=ik_jac,\n", + " verbose=1, check_derivatives=True);\n", "\n", "# Assert that we got a nearly identical solution\n", "assert np.linalg.norm(x_fd - x_analytic) < 1e-5" @@ -1775,9 +1220,9 @@ "id": "UP9UamTWzWM4" }, "source": [ - "Nice speed-up! This will become more pronounced the harder the specific IK problem (more dofs, more difficult configuration).\n", + "Nice speed-up! This will become more pronounced the harder the specific IK problem. We'll do a more comprehensive timing comparison a few cells down (this specific configuration happens to be solved rather slowly).\n", "\n", - "We'll do a more comprehensive timing comparison a few cells down." + "Note that we passed `check_derivatives=True` to ask the function to verify that our analytic Jacobian is correct, by making a comparison to the internal finite-difference function at the first timestep." ] }, { @@ -2117,13 +1562,760 @@ "\n", "media.show_video(frames, loop=False)" ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "69c7LaZ07VGY" + }, + "source": [ + "# Non-quadratic norms\n", + "\n", + "We'll now use `least_squares` to solve a trajectory optimization (control) problem. While this is not its intended purpose, it demonstrates the power and general usability of the function.\n", + "\n", + "After using regular Least Squares, we'll define a custom **non-quadratic norm** and solve again.\n", + "\n", + "Below, we copy MuJoCo's [standard humanoid model](https://github.com/google-deepmind/mujoco/blob/main/model/humanoid/humanoid.xml), with the following modifications:\n", + "1. Added a \"target\" mocap body with a pink spherical site.\n", + "2. Changed the color of the right hand to pink.\n", + "3. Replaced the torque actuators with position actuators.\n", + "4. Added sensors corresponding to the residual." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "cellView": "form", + "id": "1Cg4ABJa7VGd" + }, + "outputs": [], + "source": [ + "#@title Humanoid XML\n", + "xml = \"\"\"\n", + "\n", + " \n", + "\"\"\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "goe2DuOv7VGd" + }, + "source": [ + "Let's load the model and render the initial state for our control problem, chosen to be the \"squat\" [keyframe](https://mujoco.readthedocs.io/en/latest/XMLreference.html#keyframe)." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "xhvrJ9bX7VGd" + }, + "outputs": [], + "source": [ + "# Load model, make data\n", + "model = mujoco.MjModel.from_xml_string(xml)\n", + "data = mujoco.MjData(model)\n", + "\n", + "# Set the state to the \"squat\" keyframe, call mj_forward.\n", + "key = model.key('squat').id\n", + "mujoco.mj_resetDataKeyframe(model, data, key)\n", + "mujoco.mj_forward(model, data)\n", + "\n", + "# If a renderer exists, close it.\n", + "if 'renderer' in locals():\n", + " renderer.close()\n", + "\n", + "# Make a Renderer and a camera.\n", + "renderer = mujoco.Renderer(model, height=480, width=640)\n", + "camera = mujoco.MjvCamera()\n", + "mujoco.mjv_defaultFreeCamera(model, camera)\n", + "camera.distance = 3\n", + "camera.elevation = -10\n", + "\n", + "# Point the camera at the humanoid, render.\n", + "camera.lookat = data.body('torso').subtree_com\n", + "renderer.update_scene(data, camera)\n", + "media.show_image(renderer.render())" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "yxIzoZ_O7VGd" + }, + "source": [ + "### Problem definition\n", + "\n", + "We define the following optimal control problem defintion:\n", + "- A trajectory is rolled out from $t=0\\ldots T$, starting at the \"squat\" keyframe shown above.\n", + "- Controls $u_t$ are applied during the rollout which are a linear interpolation of first control $u_0$ and the last one $u_T$. These two vectors are our decision variable $x = \\begin{pmatrix} u_0 & u_T \\end{pmatrix}$.\n", + "- The residual is a concatenation, over all time steps, of:\n", + " - The vector from the right hand to the target.\n", + " - The torques applied by the actuators, scaled by some factor (they are numerically much larger than the hand-target distances).\n", + "\n", + "Our residual uses `mujoco.rollout` to evaluate parallel trajectories:" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "cIX95v757VGd" + }, + "outputs": [], + "source": [ + "def reach(ctrl0T, target, T, torque_scale, traj=None):\n", + " \"\"\"Residual for target-reaching task.\n", + "\n", + " Args:\n", + " ctrl0T: contatenation of the first and last control vectors.\n", + " target: target to which the right hand should reach.\n", + " T: final time for the rollout.\n", + " torque_scale: coefficient by which to scale the torques.\n", + " traj: optional list of positions to be recorded.\n", + "\n", + " Returns:\n", + " The residual of the target-reaching task.\n", + " \"\"\"\n", + " # Extract the initial and final ctrl vectors, transpose to row vectors\n", + " ctrl0 = ctrl0T[:model.nu, :].T\n", + " ctrlT = ctrl0T[model.nu:, :].T\n", + "\n", + " # Move the mocap body to the target\n", + " mocapid = model.body('target').mocapid\n", + " data.mocap_pos[mocapid] = target\n", + "\n", + " # Append the mocap targets to the controls\n", + " nroll = ctrl0.shape[0]\n", + " mocap = np.tile(data.mocap_pos[mocapid], (nroll, 1))\n", + " ctrl0 = np.hstack((ctrl0, mocap))\n", + " ctrlT = np.hstack((ctrlT, mocap))\n", + "\n", + " # Define control spec (ctrl + mocap_pos)\n", + " mjtState = mujoco.mjtState\n", + " control_spec = mjtState.mjSTATE_CTRL | mjtState.mjSTATE_MOCAP_POS\n", + "\n", + " # Interpolate and stack the control sequences\n", + " nstep = int(np.round(T / model.opt.timestep))\n", + " control = np.stack(np.linspace(ctrl0, ctrlT, nstep), axis=1)\n", + "\n", + " # Reset to the \"squat\" keyframe, get the initial state\n", + " key = model.key('squat').id\n", + " mujoco.mj_resetDataKeyframe(model, data, key)\n", + " spec = mjtState.mjSTATE_FULLPHYSICS\n", + " nstate = mujoco.mj_stateSize(model, spec)\n", + " state = np.empty(nstate)\n", + " mujoco.mj_getState(model, data, state, spec)\n", + "\n", + " # Perform rollouts (sensors.shape == nroll, nstep, nsensordata)\n", + " states, sensors = rollout.rollout(model, data, state, control,\n", + " control_spec=control_spec)\n", + "\n", + " # If requested, extract qpos into traj\n", + " if traj is not None:\n", + " assert states.shape[0] == 1\n", + " # Skip the first element in state (mjData.time)\n", + " traj.extend(np.split(states[0, :, 1:model.nq+1], nstep))\n", + "\n", + " # Scale torque sensors\n", + " sensors[:, :, 3:] *= torque_scale\n", + "\n", + " # Reshape to stack the sensor values, transpose to column vectors\n", + " sensors = sensors.reshape((sensors.shape[0], -1)).T\n", + "\n", + " # The normalizer keeps objective values similar when changing T or timestep.\n", + " normalizer = 100 * model.opt.timestep / T\n", + " return normalizer * sensors" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Bgugbs987VGe" + }, + "source": [ + "Now let's give the rest of the problem definition:\n", + "1. The trajectory is integrated for 0.7s.\n", + "2. Torques are scaled by 0.003.\n", + "3. Since our decision variable is two copies of `mjData.ctrl`, the bounds are two concatenated copies of `mjData.actuator_ctrlrange`.\n", + "4. Our initial guess $x_0 = \\begin{pmatrix} u_0 & u_T \\end{pmatrix} = \\begin{pmatrix} q_\\textrm{squat} & q_\\textrm{stand} \\end{pmatrix}$ is the joint angles at the squatting position, followed by the angles at the default (standing position). We can use angles to initialize our controls because the position actuators have angle semantics.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "Q_6PshTn7VGe" + }, + "outputs": [], + "source": [ + "T = 0.7 # Rollout length (seconds)\n", + "torque_scale = 0.003 # Scaling for the torques\n", + "\n", + "# Bounds are the stacked control bounds.\n", + "lower = np.atleast_2d(model.actuator_ctrlrange[:,0]).T\n", + "upper = np.atleast_2d(model.actuator_ctrlrange[:,1]).T\n", + "bounds = [np.vstack((lower, lower)), np.vstack((upper, upper))]\n", + "\n", + "# Initial guess is midpoint of the bounds\n", + "x0 = 0.5 * (bounds[1] + bounds[0])" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "DLkrsjeE7VGe" + }, + "source": [ + "Let's define a utility function for rendering frames and visualize the initial guess:" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "Ghnvrsyt7VGe" + }, + "outputs": [], + "source": [ + "def render_solution(x, target):\n", + " # Ask reach to save positions to traj.\n", + " traj = []\n", + " reach(x, target, T, torque_scale, traj=traj);\n", + "\n", + " frames = []\n", + " counter = 0\n", + " print('Rendering frames:', flush=True, end='')\n", + " for qpos in traj:\n", + " # Set positions, call mj_forward to update kinematics.\n", + " data.qpos = qpos\n", + " mujoco.mj_forward(model, data)\n", + "\n", + " # Render and save frames.\n", + " camera.lookat = data.body('torso').subtree_com\n", + " renderer.update_scene(data, camera)\n", + " pixels = renderer.render()\n", + " frames.append(pixels)\n", + " counter += 1\n", + " if counter % 10 == 0:\n", + " print(f' {counter}', flush=True, end='')\n", + " return frames\n", + "\n", + "# Use default target.\n", + "target = data.mocap_pos[model.body('target').mocapid]\n", + "\n", + "# Visualize the initial guess.\n", + "media.show_video(render_solution(x0, target))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "0L_W3Y5r7VGe" + }, + "source": [ + "### Solutions to the reach task\n", + "\n", + "Let's solve once for some target and look at the optimization printout:" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "RrEuIHJF7VGe" + }, + "outputs": [], + "source": [ + "target = (.4, -.3, 1.2)\n", + "\n", + "reach_target = lambda x: reach(x, target, T, torque_scale, traj=None)\n", + "\n", + "r0 = reach_target(x0)\n", + "print(f'The decision variable x has size {x0.size}')\n", + "print(f'The residual r(x) has size {r0.size}\\n')\n", + "\n", + "x, _ = minimize.least_squares(x0, reach_target, bounds);" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "S3r1KWHF7VGe" + }, + "source": [ + "Let's see what this solution looks like:" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "GcRuhJX-7VGe" + }, + "outputs": [], + "source": [ + "media.show_video(render_solution(x, target))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "zqgKu71p7VGe" + }, + "source": [ + "Let's rerun this for several target values and make a video of all of them:" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "GCsIOtdn7VGe" + }, + "outputs": [], + "source": [ + "targets = [(0.4, 0., 0.), (0.2, -1., 0.5), (-1., -.3, 1.), (0., -.2, 2.2)]\n", + "\n", + "frames = []\n", + "for target in targets:\n", + " res_target = lambda x: reach(x, target, T, torque_scale)\n", + " print(f'Optimizing for target at {target}', flush=True)\n", + " x, trace = minimize.least_squares(x0, res_target, bounds,\n", + " verbose=minimize.Verbosity.FINAL)\n", + " frames += render_solution(x, target)\n", + " print('\\n')\n", + "\n", + "print('Making video', flush=True)\n", + "media.show_video(frames)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d12KzVKU_TFR" + }, + "source": [ + "### Non-quadratic norms\n", + "\n", + "As explained in the background section, Least Squares can be generalized to norms other than the quadratic. We are now in a position to show how to define a non-quadratic norm, which is important in the estimation and system-identification contexts, where long-taled, disturbance-rejecting distributions are proportional to the exponent of a non-quadratic function.\n", + "\n", + "Let's say that we wish the task residual i.e., the vector from the hand to the target, to be evaluated with the \"Smooth L2\" function $c(r)$ which, for a given smoothing radius $d \\gt 0$ is\n", + "$$\n", + "c(r) = \\sqrt{r^T\\cdot r + d^2 } - d\n", + "$$\n", + "This function is quadratic in a $d$-sized neighborhood of the origin, and then grows linearly thereafter, like the L2 norm. The first and second derivatives are\n", + "$$\n", + "\\begin{align}\n", + "s&=\\sqrt{r^T\\cdot r + d^2 }\\\\\n", + "g &= \\tfrac{\\partial c}{\\partial r} = \\frac{r}{s} \\\\\n", + "H &=\\tfrac{\\partial^2 c}{\\partial r^2} = \\frac{I_{n_r} - g\\cdot g^T}{s}\n", + "\\end{align}\n", + "$$\n", + "There is no particularly good reason to use this norm for this optimization task, it is meerly an example.\n", + "\n", + "Let's read the documentation of the `minimize.Norm` class:\n" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "TnRn0Mk7QR6Z" + }, + "outputs": [], + "source": [ + "print(minimize.Norm.__doc__)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "3JGpR47DQ9S2" + }, + "source": [ + "Our sensors are 3 `r_pos` residual values for the hand-to-object vector followed by 21 `r_torque` actuator torques, for a total of `ns = 24` sensors. These are concatented for the entire trajectory, leading to a residual of size `24*N`, where `N` is the number of timesteps in a trajectory. After reshaping and slicing appropriately, the norm implementation looks like" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "f3MDXm4e_Wup" + }, + "outputs": [], + "source": [ + "class SmoothL2(minimize.Norm):\n", + " def __init__(self):\n", + " self.n = model.nsensordata # Equals 24.\n", + " self.d = 0.1 # The smoothing radius (length).\n", + "\n", + " def value(self, r):\n", + " rr = r.reshape((self.n, -1), order='F')\n", + " r_pos = rr[:3, :]\n", + " s = np.sqrt(np.sum(r_pos**2, axis=0) + self.d**2)\n", + " y_pos = (s - self.d).sum()\n", + " r_torque = rr[3:, :]\n", + " y_torque = 0.5 * (r_torque.T**2).sum()\n", + " return y_pos + y_torque\n", + "\n", + " def grad_hess(self, r, proj):\n", + " rr = r.reshape((self.n, -1), order='F')\n", + " r_pos = rr[:3, :]\n", + " s = np.sqrt(np.sum(r_pos**2, axis=0) + self.d**2)\n", + " g_pos = r_pos / s\n", + " g_torque = rr[3:, :]\n", + " g = np.vstack((g_pos, g_torque))\n", + " grad = proj.T @ g.reshape((-1, 1), order='F')\n", + " h_proj = proj.copy() # norm Hessian * projection matrix\n", + " for i in range(g_pos.shape[1]):\n", + " h_i = (np.eye(3) - g_pos[:, i:i+1] @ g_pos[:, i:i+1].T) / s[i]\n", + " j = self.n*i\n", + " h_proj[j:j+3, :] = h_i @ proj[j:j+3, :]\n", + " hess = proj.T @ h_proj\n", + " return grad, hess" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "oJrh-rfL0LQh" + }, + "source": [ + "Before running the optimization, let's ask `least_squares` to check our norm implementation. We'll do this with a short trajectory simulation time `T`, to avoid creating huge matrices." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "MUCDTV6JV0ix" + }, + "outputs": [], + "source": [ + "target = (.4, -.3, 1.2)\n", + "\n", + "T_short = 0.02\n", + "\n", + "reach_target = lambda x: reach(x, target, T=T_short,\n", + " torque_scale=torque_scale, traj=None)\n", + "\n", + "x, _ = minimize.least_squares(x0, reach_target, bounds, norm=SmoothL2(),\n", + " max_iter=1, check_derivatives=True);" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "J0brs9pG0-8Y" + }, + "source": [ + "Now that we are confident of our implemetation, we can see what the solution looks like:" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "AICTHd-9z1rL" + }, + "outputs": [], + "source": [ + "target = (.4, -.3, 1.2)\n", + "\n", + "reach_target = lambda x: reach(x, target, T, torque_scale, traj=None)\n", + "\n", + "x, _ = minimize.least_squares(x0, reach_target, bounds, norm=SmoothL2());\n", + "\n", + "media.show_video(render_solution(x, target))" + ] } ], "metadata": { "accelerator": "GPU", "colab": { "collapsed_sections": [ - "pOyD_TXFrM_4", "zhVv8-0Tvlrl" ], "gpuType": "T4", diff --git a/python/mujoco/minimize.py b/python/mujoco/minimize.py index 857cbec3..377733f6 100644 --- a/python/mujoco/minimize.py +++ b/python/mujoco/minimize.py @@ -14,6 +14,7 @@ # ============================================================================== """Nonlinear Least Squares minimization with box bounds.""" +import abc import dataclasses import enum import time @@ -32,14 +33,14 @@ class Verbosity(enum.Enum): class Status(enum.Enum): FACTORIZATION_FAILED = enum.auto() - NO_IMPORVEMENT = enum.auto() + NO_IMPROVEMENT = enum.auto() MAX_ITER = enum.auto() DX_TOL = enum.auto() _STATUS_MESSAGE = { Status.FACTORIZATION_FAILED: 'factorization failed.', - Status.NO_IMPORVEMENT: 'insufficient reduction.', + Status.NO_IMPROVEMENT: 'insufficient reduction.', Status.MAX_ITER: 'maximum iterations reached.', Status.DX_TOL: 'norm(dx) < tol.', } @@ -68,90 +69,70 @@ class IterLog: step: Optional[np.ndarray] = None -def jacobian_fd( - residual: Callable[[np.ndarray], np.ndarray], - x: np.ndarray, - r: np.ndarray, - eps: np.float64, - central: bool, - n_res: int, - bounds: Optional[List[np.ndarray]] = None, -): - """Finite-difference Jacobian of a residual function. +class Norm(abc.ABC): + """Abstract interface for norm functions, measuring the magnitude of vectors. - Args: - residual: function that returns the residual for a given point. - x: point at which to evaluate the Jacobian. - r: residual at x. - eps: finite-difference step size. - central: whether to use central differences. - n_res: number or residual evaluations so far. - bounds: optional pair of lower and upper bounds. + Key Concepts: - Returns: - jac: Jacobian of the residual at x. - n_res: updated number of residual evaluations. + * Norm Value: The value of the norm for a given input vector. + * Gradient and Hessian: The gradient (first derivative) and Hessian (second + derivative) of the norm function with respect to the input vector. + Subclasses Must Implement: + + * `value(self, r: np.ndarray)`: Computes and returns the norm value for the + input vector `r`. + * `grad_hess(self, r: np.ndarray, proj: np.ndarray)`: Computes and returns + both the gradient and Hessian of the norm at `r`, projected onto `proj`. + The reason we ask the user to perform the projection themselves is that + norm Hessians are often large and sparse, and the "sandwich" projection + operator `proj.T @ hess @ proj` can be computed efficiently by taking the + specific norm structure into account. """ - nx = x.size - nr = r.size - jac = np.zeros((nr, nx)) - xh = x.copy() - if bounds is None: - # No bounds, simple forward or central differencing. - for i in range(nx): - xh[i] = x[i] + eps - rp = residual(xh) - if central: - xh[i] = x[i] - eps - rm = residual(xh) - jac[:, i] = (rp - rm) / (2*eps) - else: - jac[:, i] = (rp - r) / eps - xh[i] = x[i] - n_res += 2*nx if central else nx - else: - lower, upper = bounds - midpoint = 0.5 * (upper - lower) - for i in range(nx): - # Scale eps, don't cross bounds. - eps_i = eps * (upper[i] - lower[i]) - if central: - # Use central differencing if away from bounds. - if x[i] - eps_i < lower[i]: - # Near lower bound, use forward. - xh[i] = x[i] + eps_i - rp = residual(xh) - jac[:, i] = (rp - r) / eps_i - n_res += 1 - elif x[i] + eps_i > upper[i]: - # Near upper bound, use backward. - xh[i] = x[i] - eps_i - rm = residual(xh) - jac[:, i] = (r - rm) / eps_i - n_res += 1 - else: - # Use central. - xh[i] = x[i] + eps_i - rp = residual(xh) - xh[i] = x[i] - eps_i - rm = residual(xh) - jac[:, i] = (rp - rm) / (2*eps_i) - n_res += 2 - else: - # Below midpoint use forward differencing, otherwise backward. - if x[i] < midpoint[i]: - xh[i] = x[i] + eps_i - rp = residual(xh) - jac[:, i] = (rp - r) / eps_i - else: - xh[i] = x[i] - eps_i - rm = residual(xh) - jac[:, i] = (r - rm) / eps_i - n_res += 1 - # Reset. - xh[i] = x[i] - return jac, n_res + + @abc.abstractmethod + def value(self, r: np.ndarray) -> np.float64: + """Returns the value of the norm at the input vector `y = norm(r)`.""" + pass + + @abc.abstractmethod + def grad_hess(self, r: np.ndarray, proj: np.ndarray): + """Computes the projected gradient and Hessian of the norm at `r`. + + Args: + r: A NumPy column vector (nr x 1). + proj: A pre-computed projection matrix (nr x nx). + + Returns: + A tuple containing: + * Projected gradient: proj.T @ (d_norm/d_r). + * Projected Hessian: proj.T @ (d^2_norm/d_r^2) @ proj. + """ + pass + + +class Quadratic(Norm): + """Implementation of the quadratic norm.""" + + def value(self, r: np.ndarray): + """Returns the quadratic norm of `r`.""" + return 0.5 * (r.T @ r).item() + + def grad_hess(self, r: np.ndarray, proj: np.ndarray): + """Computes the projected gradient and Hessian of the quadratic norm at `r`. + + Args: + r: A NumPy column vector (nr x 1). + proj: A pre-computed projection matrix (nr x nx). + + Returns: + A tuple containing: + * Projected gradient: `proj.T @ r`. + * Projected Hessian: `proj.T @ proj`. + """ + grad = proj.T @ r + hess = proj.T @ proj # Notionally proj.T @ np.eye(r.size) @ proj + return grad, hess def least_squares( @@ -159,33 +140,38 @@ def least_squares( residual: Callable[[np.ndarray], np.ndarray], bounds: Optional[List[np.ndarray]] = None, jacobian: Optional[Callable[[np.ndarray, np.ndarray], np.ndarray]] = None, + norm: Norm = Quadratic(), eps: float = 1e-6, - central: bool = False, mu_min: float = 1e-6, mu_max: float = 1e8, mu_factor: float = 10.0**0.1, - tol: float = 1e-7, + tol: float = 1e-6, max_iter: int = 100, verbose: Union[Verbosity, int] = Verbosity.ITER, output: Optional[TextIO] = None, + iter_callback: Optional[Callable[[List[IterLog]], None]] = None, + check_derivatives: bool = False, ) -> Tuple[np.ndarray, List[IterLog]]: """Nonlinear Least Squares minimization with box bounds. Args: - x0: initial guess - residual: function that returns the residual for a given point x. - bounds: optional pair of lower and upper bounds on the solution. - jacobian: optional function that returns Jacobian of the residual at a given + x0: Initial guess + residual: Vectorized function returning the residual for 1 or more points. + bounds: Optional pair of lower and upper bounds on the solution. + jacobian: Optional function that returns Jacobian of the residual at a given point and residual. If not given, `residual` will be finite-differenced. - eps: perurbation used for automatic finite-differencing. - central: whether to use central differences. - mu_min: minimum value of the regularizer. - mu_max: maximum value of the regularizer. - mu_factor: factor increasing or decreasing the regularizer. - tol: termination tolerance on the step size. - max_iter: maximum number of iterations. - verbose: verbosity level. - output: optional file or StringIO to which to print messages. + norm: Norm object returning norm scalar or its projected gradient and + Hessian. See Norm class for detailed documentation. + eps: Perurbation used for automatic finite-differencing. + mu_min: Minimum value of the regularizer. + mu_max: Maximum value of the regularizer. + mu_factor: Factor for increasing or decreasing the regularizer. + tol: Termination tolerance on the step size. + max_iter: Maximum number of iterations. + verbose: Verbosity level. + output: Optional file or StringIO to which to print messages. + iter_callback: Optional iteration callback, takes trace argument. + check_derivatives: Compare user-defined Jacobian and norm against fin-diff. Returns: x: best solution found @@ -202,10 +188,10 @@ def least_squares( # Initialize locals. status = Status.MAX_ITER i = 0 - x = x0.astype(np.float64) - n = x.size - xnew = np.zeros((n,)) - dx = np.zeros((n,)) + n = x0.size + x = x0.astype(np.float64).reshape((n, 1)) + xnew = np.zeros((n, 1)) + dx = np.zeros((n, 1)) scratch = np.zeros((n, n + 7)) eps = np.float64(eps) mu = np.float64(0.0) # Optimistically start with no regularization. @@ -234,6 +220,8 @@ def least_squares( n_reduc = 0 # Reset n_reduc. return mu, n_reduc + # Make local copy of bounds to avoid reshaping user input. + bounds = None if bounds is None else bounds.copy() if bounds is not None: # Checks bounds. if len(bounds) != 2: @@ -244,7 +232,10 @@ def least_squares( raise ValueError('bounds must be finite.') if not np.all(bounds[0] < bounds[1]): raise ValueError('bounds[0] must be smaller than bounds[1].') - # Clip. + + # Reshape and clip. + bounds[0] = bounds[0].reshape(n, 1) + bounds[1] = bounds[1].reshape(n, 1) np.clip(x, bounds[0], bounds[1], out=x) # Check for NaNs. @@ -267,21 +258,28 @@ def least_squares( break # Get objective y. - y = 0.5 * r.dot(r) + y = norm.value(r) # Get Jacobian jac. t_start = time.time() if jacobian is None: - jac, n_res = jacobian_fd(residual, x, r, eps, central, n_res, bounds) + jac, n_res = jacobian_fd(residual, x, r, eps, n_res, bounds) t_res += time.time() - t_start else: jac = jacobian(x, r) t_jac += time.time() - t_start n_jac += 1 + # Check user-provided Jacobian + if i == 0 and check_derivatives: + n_res = check_jacobian(residual, x, r, jac, eps, n_res, bounds, output) + + # Check user-provided norm + if i == 0 and check_derivatives and not isinstance(norm, Quadratic): + check_norm(r, norm, eps, output) + # Get gradient, Gauss-Newton Hessian. - grad = jac.T @ r - hess = jac.T @ jac + grad, hess = norm.grad_hess(r, jac) # Bounds relative to x dlower = None if bounds is None else bounds[0] - x @@ -316,13 +314,13 @@ def least_squares( n_res += 1 # New objective, evaluate reduction. - ynew = 0.5 * rnew.dot(rnew) + ynew = norm.value(rnew) reduction = y - ynew - armijo = reduction + armijo_c1*grad.dot(dx) + armijo = reduction + armijo_c1 * (grad.T @ dx).item() if armijo < 0: if mu >= mu_max: - status = Status.NO_IMPORVEMENT + status = Status.NO_IMPROVEMENT break mu, n_reduc = increase_mu(mu) @@ -330,7 +328,7 @@ def least_squares( break # Compute reduction ratio. - expected_reduction = -(grad.dot(dx) + 0.5 * dx.T @ hess @ dx) + expected_reduction = -(grad.T @ dx + 0.5 * dx.T @ hess @ dx).item() reduction_ratio = 0.0 if expected_reduction <= 0: if verbose > Verbosity.SILENT.value: @@ -352,11 +350,13 @@ def least_squares( ) print(message, file=output) - # Append log to trace. + # Append log to trace, call iter_callback. log = IterLog(candidate=x, objective=y, reduction=reduction, regularizer=mu) if verbose >= Verbosity.FULLITER.value: log = dataclasses.replace(log, residual=r, jacobian=jac, step=dx) trace.append(log) + if iter_callback is not None: + iter_callback(trace) # Check for success. if dx_norm < tol: @@ -373,12 +373,14 @@ def least_squares( x = xnew r = rnew - # Append final log to trace. - # Note: unlike other iter logs, this is at the end point. - yfinal = 0.5 * r.dot(r) - red = np.float64(0.0) + # Append final log to trace, call iter_callback. + # Note: unlike other iter logs, values are computed at the end point. + yfinal = norm.value(r) + red = np.float64(0.0) # No reduction sice we didn't take a step. log = IterLog(candidate=x, objective=yfinal, reduction=red, regularizer=mu) trace.append(log) + if iter_callback is not None: + iter_callback(trace) # Print final diagnostics. if verbose > Verbosity.SILENT.value: @@ -401,4 +403,124 @@ def least_squares( message += f' Jacobian {jac_percent:<.1f}%' print(message, file=output) - return x, trace + return x.reshape(x0.shape), trace + + +def jacobian_fd( + residual: Callable[[np.ndarray], np.ndarray], + x: np.ndarray, + r: np.ndarray, + eps: np.float64, + n_res: int, + bounds: Optional[List[np.ndarray]] = None, +) -> Tuple[np.ndarray, int]: + """Finite-difference Jacobian of a residual function. + + Args: + residual: vectorized function that returns the residual of a vector array. + x: point at which to evaluate the Jacobian. + r: residual at x. + eps: finite-difference step size. + n_res: number or residual evaluations so far. + bounds: optional pair of lower and upper bounds. + + Returns: + jac: Jacobian of the residual at x. + n_res: updated number of residual evaluations (add x.size). + + """ + n = x.size + if bounds is None: + eps_vec = eps * np.ones(n) + else: + mid = 0.5 * (bounds[1] - bounds[0]) + eps_vec = np.where(x > mid, -eps, eps).flatten() + xh = x + np.diag(eps_vec) + rh = residual(xh) + jac = (rh - r) / eps_vec + return jac, n_res+n + + +def check_jacobian( + residual: Callable[[np.ndarray], np.ndarray], + x: np.ndarray, + r: np.ndarray, + jac: np.ndarray, + eps: np.float64, + n_res: int, + bounds: Optional[List[np.ndarray]] = None, + output: Optional[TextIO] = None, + name: Optional[str] = 'Jacobian', +) -> int: + """Check user-provided Jacobian against internal finite-differencing. + + Args: + residual: vectorized function that returns the residual of a vector array. + x: point at which the r and jac were evaluated. + r: residual at x. + jac: Jacobian at x. + eps: finite-difference step size. + n_res: number or residual evaluations so far. + bounds: optional pair of lower and upper bounds. + output: Optional file or StringIO to which to print messages. + name: Optional name of the function being tested. + + Returns: + n_res: updated number of residual evaluations. + + """ + jac_fd, n_res = jacobian_fd(residual, x, r, eps, n_res, bounds) + denom = np.abs(jac).sum() + np.abs(jac_fd).sum() + 1e-8 + rel_diff = np.abs(jac - jac_fd) / denom + if np.any(rel_diff > 1e-5): + raise ValueError(f'User-provided {name} does not match finite-differences ' + 'to a relative tolerance of 1e-5.') + print(f'User-provided {name} matches finite-differences.', file=output) + return n_res + + +def check_norm( + r: np.ndarray, + norm: Norm, + eps: np.float64, + output: Optional[TextIO] = None, +): + """Check user-provided norm against internal finite-differencing. + + Args: + r: residual vector. + norm: Norm function returning either the norm scalar or its gradient + and Gauss-Newton Hessian. + eps: finite-difference step size. + output: Optional file or StringIO to which to print messages. + """ + # Get norm(r) value and 1st, 2nd derivatives. + n = np.atleast_2d(norm.value(r)) # norm value as 1x1 array. + eye = np.eye(r.size) # Identity projection. + n_g, n_h = norm.grad_hess(r, eye) # Gradient and Hessian. + + # Check that Hessian is symmetric. + if not np.allclose(n_h, n_h.T): + raise ValueError('User-provided norm Hessian is not symmetric.') + + # Check that Hessian is positive-definite. + if np.any(np.linalg.eigvals(n_h) < 0): + h_min = np.min(np.linalg.eigvals(n_h)) + raise ValueError('User-provided norm Hessian is not positive definite. ' + f'Minimum eigenvalue is {h_min:<.4g}') + + # Local function returning norm values (vectorized). + def norm_vec(v): + norms = [np.atleast_2d(norm.value(v[:, i:i+1])) for i in range(v.shape[1])] + return np.hstack(norms) + + # Check the norm gradient. + check_jacobian(norm_vec, r, n, n_g.T, eps, 0, None, output, 'norm gradient') + + # Local function returning norm gradients (vectorized). + def grad_vec(v): + gradients = [norm.grad_hess(v[:, i:i+1], eye)[0] for i in range(v.shape[1])] + return np.hstack(gradients) + + # Check the norm Hessian. + check_jacobian(grad_vec, r, n_g, n_h, eps, 0, None, output, 'norm Hessian') diff --git a/python/mujoco/minimize_test.py b/python/mujoco/minimize_test.py index 4de82fab..6d96f52c 100644 --- a/python/mujoco/minimize_test.py +++ b/python/mujoco/minimize_test.py @@ -15,7 +15,6 @@ """Tests for minimize.py.""" import io -from typing import Tuple from absl.testing import absltest from mujoco import minimize @@ -25,20 +24,19 @@ import numpy as np class MinimizeTest(absltest.TestCase): def test_basic(self) -> None: - def residual(x: np.ndarray) -> np.ndarray: - return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)], dtype=np.float64) + def residual(x): + return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)]) - for central in [False, True]: - out = io.StringIO() - x0 = np.array((0.0, 0.0)) - x, _ = minimize.least_squares(x0, residual, output=out, central=central) - expected_x = np.array((1.0, 1.0)) - np.testing.assert_array_almost_equal(x, expected_x) - self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol') + out = io.StringIO() + x0 = np.array((0.0, 0.0)) + x, _ = minimize.least_squares(x0, residual, output=out) + expected_x = np.array((1.0, 1.0)) + np.testing.assert_array_almost_equal(x, expected_x) + self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol') def test_start_at_minimum(self) -> None: - def residual(x: np.ndarray) -> np.ndarray: - return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)]) + def residual(x): + return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)]) out = io.StringIO() x0 = np.array((1.0, 1.0)) @@ -49,33 +47,36 @@ class MinimizeTest(absltest.TestCase): self.assertContainsSubsequence(out.getvalue(), 'exact minimum found') def test_jac_callback(self) -> None: - def residual(x: np.ndarray) -> np.ndarray: - return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)]) + def residual(x): + return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)]) - def jacobian(x: np.ndarray, r: np.ndarray) -> Tuple[float, np.ndarray]: + def jacobian(x, r): del r # Unused. - return np.array([[-1, 0], [-20 * x[0], 10]]) + return np.array([[-1, 0], [-20 * x[0, 0], 10]]) x0 = np.array((0.0, 0.0)) out = io.StringIO() - x, _ = minimize.least_squares(x0, residual, jacobian=jacobian, output=out) + x, _ = minimize.least_squares(x0, residual, jacobian=jacobian, output=out, + check_derivatives=True) expected_x = np.array((1.0, 1.0)) np.testing.assert_array_almost_equal(x, expected_x) self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol') + self.assertContainsSubsequence(out.getvalue(), 'Jacobian matches') - # Try with bad Jacobian, expect no improvement. - def jac_bad1(x: np.ndarray, r: np.ndarray) -> Tuple[float, np.ndarray]: - return -jacobian(x, r) - out1 = io.StringIO() - minimize.least_squares(x0, residual, jacobian=jac_bad1, output=out1) - self.assertContainsSubsequence(out1.getvalue(), 'insufficient reduction') + # Try with bad Jacobian, ask least_squares to check it. + def bad_jacobian(x, r): + del r # Unused. + return np.array([[-1, 0], [-20 * x[0, 0], 15]]) + with self.assertRaisesRegex(ValueError, r'\bJacobian does not match\b'): + minimize.least_squares(x0, residual, jacobian=bad_jacobian, output=out, + check_derivatives=True) def test_max_iter(self) -> None: dim = 20 # High-D Rosenbrock - def residual(x: np.ndarray) -> np.ndarray: - res0 = [1 - x[i] for i in range(dim - 1)] - res1 = [10 * (x[i] - x[i + 1] ** 2) for i in range(dim - 1)] + def residual(x): + res0 = [1 - x[i, :] for i in range(dim - 1)] + res1 = [10 * (x[i, :] - x[i + 1, :] ** 2) for i in range(dim - 1)] return np.asarray(res0 + res1) # Fail to reach minimum after 20 iterations. @@ -90,8 +91,8 @@ class MinimizeTest(absltest.TestCase): np.testing.assert_array_almost_equal(x, expected_x) def test_bounds(self) -> None: - def residual(x: np.ndarray) -> np.ndarray: - return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)]) + def residual(x): + return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)]) out = io.StringIO() x0 = np.array((0.0, 0.0)) @@ -108,32 +109,28 @@ class MinimizeTest(absltest.TestCase): self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol') # Test different bounds conditions. - verbose = minimize.Verbosity.FULLITER - - for central in [False, True]: - for bounds in bounds_types.values(): - out = io.StringIO() - x, trace = minimize.least_squares( - x0, - residual, - bounds=bounds, - output=out, - central=central, - verbose=verbose, - ) - self.assertContainsSubsequence(out.getvalue(), ' < tol') - grad = trace[-2].jacobian.T @ trace[-2].residual - # If x_i is on the boundary, gradient points out, otherwise it is 0. - for i, xi in enumerate(x): - if xi == bounds[0][i]: - self.assertGreater(grad[i], 0) - elif xi == bounds[1][i]: - self.assertLess(grad[i], 0) - else: - self.assertAlmostEqual(grad[i], 0, places=4) + for bounds in bounds_types.values(): + out = io.StringIO() + x, trace = minimize.least_squares( + x0, + residual, + bounds=bounds, + output=out, + verbose=minimize.Verbosity.FULLITER, + ) + self.assertContainsSubsequence(out.getvalue(), ' < tol') + grad = trace[-2].jacobian.T @ trace[-2].residual + # If x_i is on the boundary, gradient points out, otherwise it is 0. + for i, xi in enumerate(x): + if xi == bounds[0][i]: + self.assertGreater(grad[i], 0) + elif xi == bounds[1][i]: + self.assertLess(grad[i], 0) + else: + self.assertAlmostEqual(grad[i].item(), 0, places=4) def test_bad_bounds(self) -> None: - def residual(x: np.ndarray) -> np.ndarray: + def residual(x): return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)]) out = io.StringIO() @@ -150,5 +147,114 @@ class MinimizeTest(absltest.TestCase): with self.assertRaises(ValueError): minimize.least_squares(x0, residual, bounds=bounds, output=out) + def test_iter_callback(self) -> None: + def residual(x): + return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)]) + + out = io.StringIO() + + def iter_callback(trace): + print(f'Hello iteration {len(trace)}!', file=out) + + x0 = np.array((0.0, 0.0)) + x, _ = minimize.least_squares(x0, residual, output=out, + iter_callback=iter_callback) + expected_x = np.array((1.0, 1.0)) + np.testing.assert_array_almost_equal(x, expected_x) + self.assertContainsSubsequence(out.getvalue(), 'Hello iteration 3!') + + def test_norm(self) -> None: + def residual(x): + return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)]) + + p = 0.01 # Smoothing radius for smooth-L2 norm. + + class SmoothL2(minimize.Norm): + def value(self, r): + return np.sqrt((r.T @ r).item() + p*p) - p + + def grad_hess(self, r, proj): + s = np.sqrt((r.T @ r).item() + p*p) + y_r = r / s + grad = proj.T @ y_r + y_rr = (np.eye(r.size) - y_r @ y_r.T) / s + hess = proj.T @ y_rr @ proj + return grad, hess + + out = io.StringIO() + x0 = np.array((0.0, 0.0)) + x, _ = minimize.least_squares(x0, residual, norm=SmoothL2(), output=out, + check_derivatives=True) + expected_x = np.array((1.0, 1.0)) + np.testing.assert_array_almost_equal(x, expected_x) + self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol') + self.assertContainsSubsequence(out.getvalue(), + 'User-provided norm gradient matches') + self.assertContainsSubsequence(out.getvalue(), + 'User-provided norm Hessian matches') + + class SmoothL2BadGrad(minimize.Norm): + def value(self, r): + return np.sqrt((r.T @ r).item() + p*p) - p + + def grad_hess(self, r, proj): + s = np.sqrt((r.T @ r).item() + p*p) + y_r = r / s + grad = proj.T @ (y_r + 0.001) # 0.001 is erronous. + y_rr = (np.eye(r.size) - y_r @ y_r.T) / s + hess = proj.T @ y_rr @ proj + return grad, hess + + with self.assertRaisesRegex(ValueError, r'\bgradient does not match\b'): + minimize.least_squares(x0, residual, norm=SmoothL2BadGrad(), output=out, + check_derivatives=True) + + class SmoothL2BadHess(minimize.Norm): + def value(self, r): + return np.sqrt((r.T @ r).item() + p*p) - p + + def grad_hess(self, r, proj): + s = np.sqrt((r.T @ r).item() + p*p) + y_r = r / s + grad = proj.T @ y_r + y_rr = (1.001 * np.eye(r.size) - y_r @ y_r.T) / s # 1.001 is erronous. + hess = proj.T @ y_rr @ proj + return grad, hess + + with self.assertRaisesRegex(ValueError, r'\bHessian does not match\b'): + minimize.least_squares(x0, residual, norm=SmoothL2BadHess(), output=out, + check_derivatives=True) + + class SmoothL2AsymHess(minimize.Norm): + def value(self, r): + return np.sqrt((r.T @ r).item() + p*p) - p + + def grad_hess(self, r, proj): + s = np.sqrt((r.T @ r).item() + p*p) + y_r = r / s + grad = proj.T @ y_r + y_rr = (np.eye(r.size) - (y_r + 0.0001) @ y_r.T) / s + hess = proj.T @ y_rr @ proj + return grad, hess + + with self.assertRaisesRegex(ValueError, r'\bnot symmetric\b'): + minimize.least_squares(x0, residual, norm=SmoothL2AsymHess(), output=out, + check_derivatives=True) + + class SmoothL2NegHess(minimize.Norm): + def value(self, r): + return np.sqrt((r.T @ r).item() + p*p) - p + + def grad_hess(self, r, proj): + s = np.sqrt((r.T @ r).item() + p*p) + y_r = r / s + grad = proj.T @ y_r + y_rr = -(np.eye(r.size) - y_r @ y_r.T) / s # Negative-definite. + hess = proj.T @ y_rr @ proj + return grad, hess + + with self.assertRaisesRegex(ValueError, r'\bnot positive definite\b'): + minimize.least_squares(x0, residual, norm=SmoothL2NegHess(), output=out, + check_derivatives=True) if __name__ == '__main__': absltest.main()