|
| 1 | +import cv2 # importing cv2 |
| 2 | +from functools import partial |
| 3 | +import jax |
| 4 | + |
| 5 | +jax.config.update("jax_enable_x64", True) # double precision |
| 6 | +from diffrax import diffeqsolve, Euler, ODETerm, SaveAt, Tsit5 |
| 7 | +from jax import Array, vmap |
| 8 | +from jax import numpy as jnp |
| 9 | +import matplotlib.pyplot as plt |
| 10 | +import numpy as onp |
| 11 | +from pathlib import Path |
| 12 | +from typing import Callable, Dict |
| 13 | + |
| 14 | +import jsrm |
| 15 | +from jsrm import ode_factory |
| 16 | +from jsrm.systems import tendon_actuated_planar_pcs as planar_pcs |
| 17 | + |
| 18 | +num_segments = 1 |
| 19 | + |
| 20 | +# filepath to symbolic expressions |
| 21 | +sym_exp_filepath = ( |
| 22 | + Path(jsrm.__file__).parent |
| 23 | + / "symbolic_expressions" |
| 24 | + / f"planar_pcs_ns-{num_segments}.dill" |
| 25 | +) |
| 26 | + |
| 27 | +# set parameters |
| 28 | +rho = 1070 * jnp.ones((num_segments,)) # Volumetric density of Dragon Skin 20 [kg/m^3] |
| 29 | +params = { |
| 30 | + "th0": jnp.array(0.0), # initial orientation angle [rad] |
| 31 | + "l": 1e-1 * jnp.ones((num_segments,)), |
| 32 | + "r": 2e-2 * jnp.ones((num_segments,)), |
| 33 | + "rho": rho, |
| 34 | + "g": jnp.array([0.0, 9.81]), |
| 35 | + "E": 2e3 * jnp.ones((num_segments,)), # Elastic modulus [Pa] |
| 36 | + "G": 1e3 * jnp.ones((num_segments,)), # Shear modulus [Pa] |
| 37 | + "d": 2e-2 * jnp.array([[1.0, -1.0]]).repeat(num_segments, axis=0), # distance of tendons from the central axis [m] |
| 38 | +} |
| 39 | +params["D"] = 1e-3 * jnp.diag( |
| 40 | + (jnp.repeat( |
| 41 | + jnp.array([[1e0, 1e3, 1e3]]), num_segments, axis=0 |
| 42 | + ) * params["l"][:, None]).flatten() |
| 43 | +) |
| 44 | + |
| 45 | +# activate all strains (i.e. bending, shear, and axial) |
| 46 | +strain_selector = jnp.ones((3 * num_segments,), dtype=bool) |
| 47 | +# actuation selector for the segments |
| 48 | +segment_actuation_selector = jnp.ones((num_segments,), dtype=bool) |
| 49 | +# segment_actuation_selector = jnp.array([False, True]) # only the last segment is actuated |
| 50 | + |
| 51 | +# define initial configuration |
| 52 | +q0 = jnp.repeat(jnp.array([5.0 * jnp.pi, 0.1, 0.2])[None, :], num_segments, axis=0).flatten() |
| 53 | +# number of generalized coordinates |
| 54 | +n_q = q0.shape[0] |
| 55 | + |
| 56 | +# set simulation parameters |
| 57 | +dt = 1e-4 # time step |
| 58 | +ts = jnp.arange(0.0, 10.0, dt) # time steps |
| 59 | +skip_step = 10 # how many time steps to skip in between video frames |
| 60 | +video_ts = ts[::skip_step] # time steps for video |
| 61 | + |
| 62 | +# video settings |
| 63 | +video_width, video_height = 700, 700 # img height and width |
| 64 | +video_path = Path(__file__).parent / "videos" / f"planar_pcs_ns-{num_segments}.mp4" |
| 65 | + |
| 66 | + |
| 67 | +def draw_robot( |
| 68 | + batched_forward_kinematics_fn: Callable, |
| 69 | + params: Dict[str, Array], |
| 70 | + q: Array, |
| 71 | + width: int, |
| 72 | + height: int, |
| 73 | + num_points: int = 50, |
| 74 | +) -> onp.ndarray: |
| 75 | + # plotting in OpenCV |
| 76 | + h, w = height, width # img height and width |
| 77 | + ppm = h / (2.0 * jnp.sum(params["l"])) # pixel per meter |
| 78 | + base_color = (0, 0, 0) # black robot_color in BGR |
| 79 | + robot_color = (255, 0, 0) # black robot_color in BGR |
| 80 | + |
| 81 | + # we use for plotting N points along the length of the robot |
| 82 | + s_ps = jnp.linspace(0, jnp.sum(params["l"]), num_points) |
| 83 | + |
| 84 | + # poses along the robot of shape (3, N) |
| 85 | + chi_ps = batched_forward_kinematics_fn(params, q, s_ps) |
| 86 | + |
| 87 | + img = 255 * onp.ones((w, h, 3), dtype=jnp.uint8) # initialize background to white |
| 88 | + curve_origin = onp.array( |
| 89 | + [w // 2, 0.1 * h], dtype=onp.int32 |
| 90 | + ) # in x-y pixel coordinates |
| 91 | + # draw base |
| 92 | + cv2.rectangle(img, (0, h - curve_origin[1]), (w, h), color=base_color, thickness=-1) |
| 93 | + # transform robot poses to pixel coordinates |
| 94 | + # should be of shape (N, 2) |
| 95 | + curve = onp.array((curve_origin + chi_ps[:2, :].T * ppm), dtype=onp.int32) |
| 96 | + # invert the v pixel coordinate |
| 97 | + curve[:, 1] = h - curve[:, 1] |
| 98 | + cv2.polylines(img, [curve], isClosed=False, color=robot_color, thickness=10) |
| 99 | + |
| 100 | + return img |
| 101 | + |
| 102 | + |
| 103 | +if __name__ == "__main__": |
| 104 | + strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = ( |
| 105 | + planar_pcs.factory(num_segments, sym_exp_filepath, strain_selector, segment_actuation_selector=segment_actuation_selector) |
| 106 | + ) |
| 107 | + actuation_mapping_fn = auxiliary_fns["actuation_mapping_fn"] |
| 108 | + # jit the functions |
| 109 | + dynamical_matrices_fn = jax.jit(partial(dynamical_matrices_fn)) |
| 110 | + batched_forward_kinematics = vmap( |
| 111 | + forward_kinematics_fn, in_axes=(None, None, 0), out_axes=-1 |
| 112 | + ) |
| 113 | + |
| 114 | + # test the actuation mapping function |
| 115 | + xi_eq = jnp.array([0.0, 0.0, 1.0])[None].repeat(num_segments, axis=0).flatten() |
| 116 | + B_xi = strain_basis |
| 117 | + # call the actuation mapping function |
| 118 | + A = actuation_mapping_fn( |
| 119 | + forward_kinematics_fn, |
| 120 | + auxiliary_fns["jacobian_fn"], |
| 121 | + params, |
| 122 | + B_xi, |
| 123 | + xi_eq, |
| 124 | + jnp.zeros_like(q0), |
| 125 | + ) |
| 126 | + print("A =\n", A) |
| 127 | + |
| 128 | + x0 = jnp.concatenate([q0, jnp.zeros_like(q0)]) # initial condition |
| 129 | + u = jnp.array([1.0, 1.0])[None].repeat(num_segments, axis=0).flatten() # tendon tensions |
| 130 | + # u = 2e-1 * jnp.array([2.0, 0.0, 0.0, 1.0]) |
| 131 | + print("u =\n", u) |
| 132 | + |
| 133 | + ode_fn = ode_factory(dynamical_matrices_fn, params, u) |
| 134 | + term = ODETerm(ode_fn) |
| 135 | + |
| 136 | + sol = diffeqsolve( |
| 137 | + term, |
| 138 | + solver=Tsit5(), |
| 139 | + t0=ts[0], |
| 140 | + t1=ts[-1], |
| 141 | + dt0=dt, |
| 142 | + y0=x0, |
| 143 | + max_steps=None, |
| 144 | + saveat=SaveAt(ts=video_ts), |
| 145 | + ) |
| 146 | + |
| 147 | + print("sol.ys =\n", sol.ys) |
| 148 | + # the evolution of the generalized coordinates |
| 149 | + q_ts = sol.ys[:, :n_q] |
| 150 | + # the evolution of the generalized velocities |
| 151 | + q_d_ts = sol.ys[:, n_q:] |
| 152 | + |
| 153 | + # evaluate the forward kinematics along the trajectory |
| 154 | + chi_ee_ts = vmap(forward_kinematics_fn, in_axes=(None, 0, None))( |
| 155 | + params, q_ts, jnp.array([jnp.sum(params["l"])]) |
| 156 | + ) |
| 157 | + # plot the configuration vs time |
| 158 | + plt.figure() |
| 159 | + for segment_idx in range(num_segments): |
| 160 | + plt.plot( |
| 161 | + video_ts, q_ts[:, 3 * segment_idx + 0], |
| 162 | + label=r"$\kappa_\mathrm{be," + str(segment_idx + 1) + "}$ [rad/m]" |
| 163 | + ) |
| 164 | + plt.plot( |
| 165 | + video_ts, q_ts[:, 3 * segment_idx + 1], |
| 166 | + label=r"$\sigma_\mathrm{sh," + str(segment_idx + 1) + "}$ [-]" |
| 167 | + ) |
| 168 | + plt.plot( |
| 169 | + video_ts, q_ts[:, 3 * segment_idx + 2], |
| 170 | + label=r"$\sigma_\mathrm{ax," + str(segment_idx + 1) + "}$ [-]" |
| 171 | + ) |
| 172 | + plt.xlabel("Time [s]") |
| 173 | + plt.ylabel("Configuration") |
| 174 | + plt.legend() |
| 175 | + plt.grid(True) |
| 176 | + plt.tight_layout() |
| 177 | + plt.show() |
| 178 | + # plot end-effector position vs time |
| 179 | + plt.figure() |
| 180 | + plt.plot(video_ts, chi_ee_ts[:, 0], label="x") |
| 181 | + plt.plot(video_ts, chi_ee_ts[:, 1], label="y") |
| 182 | + plt.xlabel("Time [s]") |
| 183 | + plt.ylabel("End-effector Position [m]") |
| 184 | + plt.legend() |
| 185 | + plt.grid(True) |
| 186 | + plt.box(True) |
| 187 | + plt.tight_layout() |
| 188 | + plt.show() |
| 189 | + # plot the end-effector position in the x-y plane as a scatter plot with the time as the color |
| 190 | + plt.figure() |
| 191 | + plt.scatter(chi_ee_ts[:, 0], chi_ee_ts[:, 1], c=video_ts, cmap="viridis") |
| 192 | + plt.axis("equal") |
| 193 | + plt.grid(True) |
| 194 | + plt.xlabel("End-effector x [m]") |
| 195 | + plt.ylabel("End-effector y [m]") |
| 196 | + plt.colorbar(label="Time [s]") |
| 197 | + plt.tight_layout() |
| 198 | + plt.show() |
| 199 | + # plt.figure() |
| 200 | + # plt.plot(chi_ee_ts[:, 0], chi_ee_ts[:, 1]) |
| 201 | + # plt.axis("equal") |
| 202 | + # plt.grid(True) |
| 203 | + # plt.xlabel("End-effector x [m]") |
| 204 | + # plt.ylabel("End-effector y [m]") |
| 205 | + # plt.tight_layout() |
| 206 | + # plt.show() |
| 207 | + |
| 208 | + # plot the energy along the trajectory |
| 209 | + kinetic_energy_fn_vmapped = vmap( |
| 210 | + partial(auxiliary_fns["kinetic_energy_fn"], params) |
| 211 | + ) |
| 212 | + potential_energy_fn_vmapped = vmap( |
| 213 | + partial(auxiliary_fns["potential_energy_fn"], params) |
| 214 | + ) |
| 215 | + U_ts = potential_energy_fn_vmapped(q_ts) |
| 216 | + T_ts = kinetic_energy_fn_vmapped(q_ts, q_d_ts) |
| 217 | + plt.figure() |
| 218 | + plt.plot(video_ts, U_ts, label="Potential energy") |
| 219 | + plt.plot(video_ts, T_ts, label="Kinetic energy") |
| 220 | + plt.xlabel("Time [s]") |
| 221 | + plt.ylabel("Energy [J]") |
| 222 | + plt.legend() |
| 223 | + plt.grid(True) |
| 224 | + plt.box(True) |
| 225 | + plt.tight_layout() |
| 226 | + plt.show() |
| 227 | + |
| 228 | + # create video |
| 229 | + fourcc = cv2.VideoWriter_fourcc(*"mp4v") |
| 230 | + video_path.parent.mkdir(parents=True, exist_ok=True) |
| 231 | + video = cv2.VideoWriter( |
| 232 | + str(video_path), |
| 233 | + fourcc, |
| 234 | + 1 / (skip_step * dt), # fps |
| 235 | + (video_width, video_height), |
| 236 | + ) |
| 237 | + |
| 238 | + for time_idx, t in enumerate(video_ts): |
| 239 | + x = sol.ys[time_idx] |
| 240 | + img = draw_robot( |
| 241 | + batched_forward_kinematics, |
| 242 | + params, |
| 243 | + x[: (x.shape[0] // 2)], |
| 244 | + video_width, |
| 245 | + video_height, |
| 246 | + ) |
| 247 | + video.write(img) |
| 248 | + |
| 249 | + video.release() |
| 250 | + print(f"Video saved at {video_path}") |
0 commit comments