example: also account for a cart-pole

examples/simulation-inverted-pendulum.py

@@ -5,23 +5,61 @@ import math
 import time
 import sys
 
-N = 10 # number of pendulums
+# Parse input arguments
+import argparse
+parser = argparse.ArgumentParser()
+parser.add_argument("--with-cart", help="Add a cart at the base of the pendulum to simulate a cart pole system.",
+                    action="store_true")
+parser.add_argument("-N", help="Number of pendulums compositing the dynamical system.",
+                    type=int)
+args = parser.parse_args()
+
+if args.N:
+    N = args.N
+else:
+    N = 1 # number of pendulums
+
 model = pin.Model()
 geom_model = pin.GeometryModel()
 
 parent_id = 0
+
+if args.with_cart:
+    cart_radius = 0.1
+    cart_length = 5 * cart_radius
+    cart_mass = 2.
+    joint_name = "joint_cart"
+
+    geometry_placement = pin.SE3.Identity()
+    geometry_placement.rotation = pin.Quaternion(np.array([0.,0.,1.]),np.array([0.,1.,0.])).toRotationMatrix()
+
+    joint_id = model.addJoint(parent_id, pin.JointModelPY(), pin.SE3.Identity(), joint_name)
+
+    body_inertia = pin.Inertia.FromCylinder(cart_mass,cart_radius,cart_length)
+    body_placement = geometry_placement
+    model.appendBodyToJoint(joint_id,body_inertia,body_placement) # We need to rotate the inertia as it is expressed in the LOCAL frame of the geometry
+
+    shape_cart = fcl.Cylinder(cart_radius, cart_length)
+
+    geom_cart = pin.GeometryObject("shape_cart", joint_id, shape_cart, geometry_placement)
+    geom_cart.meshColor = np.array([1.,0.1,0.1,1.])
+    geom_model.addGeometryObject(geom_cart)
+
+    parent_id = joint_id
+else:
+    base_radius = 0.2
+    shape_base = fcl.Sphere(base_radius)
+    geom_base = pin.GeometryObject("base", 0, shape_base, pin.SE3.Identity())
+    geom_base.meshColor = np.array([1.,0.1,0.1,1.])
+    geom_model.addGeometryObject(geom_base)
+
 joint_placement = pin.SE3.Identity()
 body_mass = 1.
 body_radius = 0.1
 
-shape0 = fcl.Sphere(body_radius)
-geom0_obj = pin.GeometryObject("base", 0, shape0, pin.SE3.Identity())
-geom0_obj.meshColor = np.array([1.,0.1,0.1,1.])
-geom_model.addGeometryObject(geom0_obj)
-
 for k in range(N):
     joint_name = "joint_" + str(k+1)
-    joint_id = model.addJoint(parent_id,pin.JointModelRY(),joint_placement,joint_name)
+    joint_id = model.addJoint(parent_id,pin.JointModelRX(),joint_placement,joint_name)
 
     body_inertia = pin.Inertia.FromSphere(body_mass,body_radius)
     body_placement = joint_placement.copy()
@@ -46,10 +84,10 @@ for k in range(N):
     parent_id = joint_id
     joint_placement = body_placement.copy()
 
-from pinocchio.visualize import GepettoVisualizer
+from pinocchio.visualize import MeshcatVisualizer as Visualizer
 
 visual_model = geom_model
-viz = GepettoVisualizer(model, geom_model, visual_model)
+viz = Visualizer(model, geom_model, visual_model)
 
 # Initialize the viewer.
 try:
@@ -78,19 +116,31 @@ N = math.floor(T/dt)
 
 model.lowerPositionLimit.fill(-math.pi)
 model.upperPositionLimit.fill(+math.pi)
-q = pin.randomConfiguration(model)
-v = np.zeros((model.nv))
 
-t = 0.
+if args.with_cart:
+    model.lowerPositionLimit[0] = model.upperPositionLimit[0] = 0. 
+
 data_sim = model.createData()
+
+t = 0.
+q = pin.randomConfiguration(model)
+v = np.zeros((model.nv))
+tau_control = np.zeros((model.nv))
+damping_value = 0.1
 for k in range(N):
-    tau_control = np.zeros((model.nv))
+
+    tic = time.time()
+    tau_control = -damping_value * v # small damping
     a = pin.aba(model,data_sim,q,v,tau_control) # Forward dynamics
 
+    # Semi-explicit integration
     v += a*dt
-    #q += v*dt
-    q = pin.integrate(model,q,v*dt)
+    q = pin.integrate(model,q,v*dt) # Configuration integration
 
     viz.display(q)
-    time.sleep(dt)
+    toc = time.time()
+    ellapsed = toc - tic
+    
+    dt_sleep = max(0,dt - (ellapsed))
+    time.sleep(dt_sleep)
     t += dt