class TimeSteppingMultibodyPlant(): """ """ def __init__(self, file=None, terrain=FlatTerrain(), dlevel=1): """ Initialize TimeSteppingMultibodyPlant with a model from a file and an arbitrary terrain geometry. Initialization also welds the first frame in the MultibodyPlant to the world frame """ self.builder = DiagramBuilder() self.multibody, self.scene_graph = AddMultibodyPlantSceneGraph( self.builder, 0.001) # Store the terrain self.terrain = terrain self._dlevel = dlevel # Build the MultibodyPlant from the file, if one exists self.model_index = [] if file is not None: # Parse the file self.model_index = Parser(self.multibody).AddModelFromFile( FindResource(file)) # Initialize the collision data self.collision_frames = [] self.collision_poses = [] self.collision_radius = [] def Finalize(self): """ Cements the topology of the MultibodyPlant and identifies all available collision geometries. """ # Finalize the underlying plant model self.multibody.Finalize() # Idenify and store collision geometries self.__store_collision_geometries() def num_contacts(self): """ returns the number of contact points""" return len(self.collision_frames) def num_friction(self): """ returns the number of friction components""" return 4 * self.dlevel * self.num_contacts() def GetNormalDistances(self, context): """ Returns an array of signed distances between the contact geometries and the terrain, given the current system context Arguments: context: a pyDrake MultibodyPlant context Return values: distances: a numpy array of signed distance values """ qtype = self.multibody.GetPositions(context).dtype nCollisions = len(self.collision_frames) distances = np.zeros((nCollisions, ), dtype=qtype) for n in range(0, nCollisions): # Transform collision frames to world coordinates collision_pt = self.multibody.CalcPointsPositions( context, self.collision_frames[n], self.collision_poses[n].translation(), self.multibody.world_frame()) # Squeeze collision point (necessary for AutoDiff plants) collision_pt = np.squeeze(collision_pt) # Calc nearest point on terrain in world coordinates terrain_pt = self.terrain.nearest_point(collision_pt) # Calc normal distance to terrain terrain_frame = self.terrain.local_frame(terrain_pt) normal = terrain_frame[0, :] distances[n] = normal.dot(collision_pt - terrain_pt) - self.collision_radius[n] # Return the distances as a single array return distances def GetContactJacobians(self, context): """ Returns a tuple of numpy arrays representing the normal and tangential contact Jacobians evaluated at each contact point Arguments: context: a pyDrake MultibodyPlant context Return Values (Jn, Jt): the tuple of contact Jacobians. Jn represents the normal components and Jt the tangential components """ qtype = self.multibody.GetPositions(context).dtype numN = self.num_contacts() numT = int(self.num_friction() / numN) D = self.friction_discretization_matrix().transpose() Jn = np.zeros((numN, self.multibody.num_velocities()), dtype=qtype) Jt = np.zeros((numN * numT, self.multibody.num_velocities()), dtype=qtype) for n in range(0, numN): # Transform collision frames to world coordinates collision_pt = self.multibody.CalcPointsPositions( context, self.collision_frames[n], self.collision_poses[n].translation(), self.multibody.world_frame()) # Calc nearest point on terrain in world coordinates terrain_pt = self.terrain.nearest_point(collision_pt) # Calc normal distance to terrain terrain_frame = self.terrain.local_frame(terrain_pt) normal, tangent = np.split(terrain_frame, [1], axis=0) # Discretize to the chosen level Dtangent = D.dot(tangent) # Get the contact point Jacobian J = self.multibody.CalcJacobianTranslationalVelocity( context, JacobianWrtVariable.kV, self.collision_frames[n], self.collision_poses[n].translation(), self.multibody.world_frame(), self.multibody.world_frame()) # Calc contact Jacobians Jn[n, :] = normal.dot(J) Jt[n * numT:(n + 1) * numT, :] = Dtangent.dot(J) # Return the Jacobians as a tuple of np arrays return (Jn, Jt) def GetFrictionCoefficients(self, context): """ Return friction coefficients for nearest point on terrain Arguments: context: the current MultibodyPlant context Return Values: friction_coeff: list of friction coefficients """ friction_coeff = [] for frame, pose in zip(self.collision_frames, self.collision_poses): # Transform collision frames to world coordinates collision_pt = self.multibody.CalcPointsPositions( context, frame, pose.translation(), self.multibody.world_frame()) # Calc nearest point on terrain in world coordiantes terrain_pt = self.terrain.nearest_point(collision_pt) friction_coeff.append(self.terrain.get_friction(terrain_pt)) # Return list of friction coefficients return friction_coeff def getTerrainPointsAndFrames(self, context): """ Return the nearest points on the terrain and the local coordinate frame Arguments: context: current MultibodyPlant context Return Values: terrain_pts: a 3xN array of points on the terrain terrain_frames: a 3x3xN array, specifying the local frame of the terrain """ terrain_pts = [] terrain_frames = [] for frame, pose in zip(self.collision_frames, self.collision_poses): # Calc collision point collision_pt = self.multibody.CalcPointsPositions( context, frame, pose.translation(), self.multibody.world_frame()) # Calc nearest point on terrain in world coordinates terrain_pt = self.terrain.nearest_point(collision_pt) terrain_pts.append(terrain_pt) # Calc local coordinate frame terrain_frames.append(self.terrain.local_frame(terrain_pt)) return (terrain_pts, terrain_frames) def toAutoDiffXd(self): """Covert the MultibodyPlant to use AutoDiffXd instead of Float""" # Create a new TimeSteppingMultibodyPlant model copy_ad = TimeSteppingMultibodyPlant(file=None, terrain=self.terrain, dlevel=self._dlevel) # Instantiate the plant as the Autodiff version copy_ad.multibody = self.multibody.ToAutoDiffXd() copy_ad.scene_graph = self.scene_graph.ToAutoDiffXd() copy_ad.model_index = self.model_index # Store the collision frames to finalize the model copy_ad.__store_collision_geometries() return copy_ad def set_discretization_level(self, dlevel=0): """Set the friction discretization level. The default is 0""" self._dlevel = dlevel def __store_collision_geometries(self): """Identifies the collision geometries in the model and stores their parent frame and pose in parent frame in lists""" # Create a diagram and a scene graph inspector = self.scene_graph.model_inspector() # Locate collision geometries and contact points body_inds = self.multibody.GetBodyIndices(self.model_index) # Get the collision frames for each body in the model for body_ind in body_inds: body = self.multibody.get_body(body_ind) collision_ids = self.multibody.GetCollisionGeometriesForBody(body) for id in collision_ids: # get and store the collision geometry frames frame_name = inspector.GetName( inspector.GetFrameId(id)).split("::") self.collision_frames.append( self.multibody.GetFrameByName(frame_name[-1])) self.collision_poses.append(inspector.GetPoseInFrame(id)) # Check for a spherical geometry geoms = inspector.GetGeometries(inspector.GetFrameId(id), Role.kProximity) shape = inspector.GetShape(geoms[0]) if type(shape) is Sphere: self.collision_radius.append(shape.radius()) else: self.collision_radius.append(0.) def __discretize_friction(self, normal, tangent): """ Rotates the terrain tangent vectors to discretize the friction cone Arguments: normal: The terrain normal direction, (1x3) numpy array tangent: The terrain tangent directions, (2x3) numpy array Return Values: all_tangents: The discretized friction vectors, (2nx3) numpy array This method is now deprecated """ # Reflect the current friction basis tangent = np.concatenate((tangent, -tangent), axis=0) all_tangents = np.zeros((4 * (self._dlevel), tangent.shape[1])) all_tangents[0:4, :] = tangent # Rotate the tangent basis around the normal vector for n in range(1, self._dlevel): # Create an angle-axis representation of rotation R = RotationMatrix(theta_lambda=AngleAxis( angle=n * pi / (2 * (self._dlevel)), axis=normal)) # Apply the rotation matrix all_tangents[n * 4:(n + 1) * 4, :] = R.multiply( tangent.transpose()).transpose() return all_tangents def simulate(self, h, x0, u=None, N=1): # Initialize arrays nx = x0.shape[0] x = np.zeros(shape=(nx, N)) x[:, 0] = x0 t = np.zeros(shape=(N, )) nf = 1 if u is None: B = self.multibody.MakeAcutatorMatrix() u = np.zeros(shape=(B.shape[1], N)) context = self.multibody.CreateDefaultContext() Jn, Jt = self.GetContactJacobians(context) f = np.zeros(shape=(Jn.shape[0] + Jt.shape[0], N)) # Integration loop for n in range(0, N - 1): f[:, n] = self.contact_impulse(h, x[:, n], u[:, n]) x[:, n + 1] = self.integrate(h, x[:, n], u[:, n], f[:, n]) t[n + 1] = t[n] + h f[:, n] = f[:, n] / h return (t, x, f) def integrate(self, h, x, u, f): # Get the configuration and the velocity q, dq = np.split(x, 2) # Estimate the next configuration, assuming constant velocity qhat = q + h * dq # Set the context context = self.multibody.CreateDefaultContext() self.multibody.SetPositions(context, qhat) v = self.multibody.MapQDotToVelocity(context, dq) self.multibody.SetVelocities(context, v) # Get the current system properties M = self.multibody.CalcMassMatrixViaInverseDynamics(context) C = self.multibody.CalcBiasTerm(context) G = self.multibody.CalcGravityGeneralizedForces(context) B = self.multibody.MakeActuationMatrix() Jn, Jt = self.GetContactJacobians(context) J = np.vstack((Jn, Jt)) # Calculate the next state b = h * (B.dot(u) - C.dot(dq) + G) + J.transpose().dot(f) v = np.linalg.solve(M, b) dq += v q += h * dq # Collect the configuration and velocity into a state vector return np.concatenate((q, dq), axis=0) def contact_impulse(self, h, x, u): # Get the configuration and generalized velocity q, dq = np.split(x, 2) # Estimate the configuration at the next time step qhat = q + h * dq # Get the system parameters context = self.multibody.CreateDefaultContext() self.multibody.SetPositions(context, q) v = self.multibody.MapQDotToVelocity(context, dq) self.multibody.SetVelocities(context, v) M = self.multibody.CalcMassMatrixViaInverseDynamics(context) C = self.multibody.CalcBiasTerm(context) G = self.multibody.CalcGravityGeneralizedForces(context) B = self.multibody.MakeActuationMatrix() tau = B.dot(u) - C + G Jn, Jt = self.GetContactJacobians(context) phi = self.GetNormalDistances(context) alpha = Jn.dot(qhat) - phi # Calculate the force size from the contact Jacobian numT = Jt.shape[0] numN = Jn.shape[0] S = numT + 2 * numN # Initialize LCP parameters P = np.zeros(shape=(S, S), dtype=float) w = np.zeros(shape=(numN, S)) numF = int(numT / numN) for n in range(0, numN): w[n, n * numF + numN:numN + (n + 1) * numF] = 1 # Construct LCP matrix J = np.vstack((Jn, Jt)) JM = J.dot(np.linalg.inv(M)) P[0:numN + numT, 0:numN + numT] = JM.dot(J.transpose()) P[:, numN + numT:] = w.transpose() P[numN + numT:, :] = -w P[numN + numT:, 0:numN] = np.diag(self.GetFrictionCoefficients(context)) #Construct LCP bias vector z = np.zeros(shape=(S, ), dtype=float) z[0:numN + numT] = h * JM.dot(tau) + J.dot(dq) #z[0:numN] += (Jn.dot(q) - alpha)/h z[0:numN] += phi / h # Solve the LCP for the reaction impluses f, status = solve_lcp(P, z) if f is None: return np.zeros(shape=(numN + numT, )) else: # Strip the slack variables from the LCP solution return f[0:numN + numT] def get_multibody(self): return self.multibody def joint_limit_jacobian(self): """ Returns a matrix for mapping the positive-only joint limit torques to """ # Create the Jacobian as if all joints were limited nV = self.multibody.num_velocities() I = np.eye(nV) # Check for infinite limits qhigh = self.multibody.GetPositionUpperLimits() qlow = self.multibody.GetPositionLowerLimits() # Assert two sided limits low_inf = np.isinf(qlow) assert all(low_inf == np.isinf(qhigh)) # Make the joint limit Jacobian if not all(low_inf): # Remove floating base limits floating = self.multibody.GetFloatingBaseBodies() floating_pos = [] while len(floating) > 0: body = self.multibody.get_body(floating.pop()) if body.has_quaternion_dofs(): floating_pos.append(body.floating_positions_start()) # Remove entries corresponding to the extra dof from quaternions low_inf = np.delete(low_inf, floating_pos) low_inf = np.squeeze(low_inf) # Zero-out the entries that don't correspond to joint limits I[low_inf, low_inf] = 0 # Remove the columns that don't correspond to joint limits I = np.delete(I, low_inf, axis=1) return np.concatenate([I, -I], axis=1) else: return None def get_contact_points(self, context): """ Returns a list of positions of contact points expressed in world coordinates given the system context. """ contact_pts = [] for frame, pose in zip(self.collision_frames, self.collision_poses): contact_pts.append( self.multibody.CalcPointsPositions( context, frame, pose.translation(), self.multibody.world_frame())) return contact_pts def resolve_contact_forces_in_world(self, context, forces): """ Transform non-negative discretized force components used in complementarity constraints into 3-vectors in world coordinates Returns a list of (3,) numpy arrays """ # First remove the discretization forces = self.resolve_forces(forces) # Reorganize forces from (Normal, Tangential) to a list force_list = [] for n in range(0, self.num_contacts()): force_list.append(forces[[ n, self.num_contacts() + 2 * n, self.num_contacts() + 2 * n + 1 ]]) # Transform the forces into world coordinates using the terrain frames _, frames = self.getTerrainPointsAndFrames(context) world_forces = [] for force, frame in zip(force_list, frames): world_forces.append(frame.dot(force)) # Return a list of forces in world coordinates return world_forces def resolve_limit_forces(self, jl_forces): """ Combine positive and negative components of joint limit torques Arguments: jl_forces: (2n, m) numpy array Return values: (n, m) numpy array """ JL = self.joint_limit_jacobian() has_limits = np.sum(abs(JL), axis=1) > 0 f_jl = JL.dot(jl_forces) if f_jl.ndim > 1: return f_jl[has_limits, :] else: return f_jl[has_limits] def duplicator_matrix(self): """Returns a matrix of 1s and 0s for combining friction forces or duplicating sliding velocities""" numN = self.num_contacts() numT = self.num_friction() w = np.zeros((numN, numT)) nD = int(numT / numN) for n in range(numN): w[n, n * nD:(n + 1) * nD] = 1 return w def friction_discretization_matrix(self): """ Make a matrix for converting discretized friction into a single vector""" n = 4 * self.dlevel theta = np.linspace(0, n - 1, n) * 2 * np.pi / n D = np.vstack((np.cos(theta), np.sin(theta))) # Threshold out small values D[np.abs(D) < 1e-10] = 0 return D def resolve_forces(self, forces): """ Convert discretized friction & normal forces into a non-discretized 3-vector""" numN = self.num_contacts() n = 4 * self.dlevel fN = forces[0:numN, :] fT = forces[numN:numN * (n + 1), :] D_ = self.friction_discretization_matrix() D = np.zeros((2 * numN, n * numN)) for k in range(numN): D[2 * k:2 * k + 2, k * n:(k + 1) * n] = D_ ff = D.dot(fT) return np.concatenate((fN, ff), axis=0) def static_controller(self, qref, verbose=False): """ Generates a controller to maintain a static pose Arguments: qref: (N,) numpy array, the static pose to be maintained Return Values: u: (M,) numpy array, actuations to best achieve static pose f: (C,) numpy array, associated normal reaction forces static_controller generates the actuations and reaction forces assuming the velocity and accelerations are zero. Thus, the equation to solve is: N(q) = B*u + J*f where N is a vector of gravitational and conservative generalized forces, B is the actuation selection matrix, and J is the contact-Jacobian transpose. Currently, static_controller only considers the effects of the normal forces. Frictional forces are not yet supported """ #TODO: Solve for friction forces as well # Check inputs if qref.ndim > 1 or qref.shape[0] != self.multibody.num_positions(): raise ValueError( f"Reference position mut be ({self.multibody.num_positions(),},) array" ) # Set the context context = self.multibody.CreateDefaultContext() self.multibody.SetPositions(context, qref) # Get the necessary properties G = self.multibody.CalcGravityGeneralizedForces(context) Jn, _ = self.GetContactJacobians(context) phi = self.GetNormalDistances(context) B = self.multibody.MakeActuationMatrix() #Collect terms A = np.concatenate([B, Jn.transpose()], axis=1) # Create the mathematical program prog = MathematicalProgram() l_var = prog.NewContinuousVariables(self.num_contacts(), name="forces") u_var = prog.NewContinuousVariables(self.multibody.num_actuators(), name="controls") # Ensure dynamics approximately satisfied prog.AddL2NormCost(A=A, b=-G, vars=np.concatenate([u_var, l_var], axis=0)) # Enforce normal complementarity prog.AddBoundingBoxConstraint(np.zeros(l_var.shape), np.full(l_var.shape, np.inf), l_var) prog.AddConstraint(phi.dot(l_var) == 0) # Solve result = Solve(prog) # Check for a solution if result.is_success(): u = result.GetSolution(u_var) f = result.GetSolution(l_var) if verbose: printProgramReport(result, prog) return (u, f) else: print(f"Optimization failed. Returning zeros") if verbose: printProgramReport(result, prog) return (np.zeros(u_var.shape), np.zeros(l_var.shape)) @property def dlevel(self): return self._dlevel @dlevel.setter def dlevel(self, val): """Check that the value of dlevel is a positive integer""" if type(val) is int and val > 0: self._dlevel = val else: raise ValueError("dlevel must be a positive integer")
body = plant.get_body(ind) collisionIds.append(plant.GetCollisionGeometriesForBody(body)) # Flatten the list collisionIds = [id for id_list in collisionIds for id in id_list if id_list] # Move the box vertically upward to make the example slightly more interesting plant.SetPositions(context, [0, 0, 1, 0, 0, 0]) # Loop over the collision identifiers for id in collisionIds: # Get the frame in which the collision geometry resides frameId = inspector.GetFrameId(id) frame_name = inspector.GetName(frameId) # Get the pose of the collision geometry in the frame R = inspector.GetPoseInFrame(id) # Print the pose in homogeneous coordinates print( f"Collision geometry in frame {frame_name} with pose \n{R.GetAsMatrix4()}" ) # Note that the frameId is NOT the same as the FrameIndex, and cannot be used to the get the frame from MultibodyPlant. We can strip the frame name from the ID and use that instead name = frame_name.split("::") body_frame = plant.GetFrameByName(name[-1]) # Then we can calculate the position in world coordinates world_pt = plant.CalcPointsPositions(context, body_frame, R.translation(), world_frame) print(f"In world coordinates, the collision point is\n {world_pt}") # Finally, we can calculate the Jacobian for the contact point Jc = plant.CalcJacobianTranslationalVelocity(context, JacobianWrtVariable.kQDot, body_frame, R.translation(), world_frame, world_frame) print(f"The Jacobian for the contact point in world coordinates is \n{Jc}")
class iiwa_sys(): def __init__(self, builder, dt=5e-4, N=150, params=None, trj_decay=0.7, x_w_cov=1e-5, door_angle_ref=1.0, visualize=False): self.plant_derivs = MultibodyPlant(time_step=dt) parser = Parser(self.plant_derivs) self.derivs_iiwa, _, _ = self.add_models(self.plant_derivs, parser, params=params) self.plant_derivs.Finalize() self.plant_derivs_context = self.plant_derivs.CreateDefaultContext() self.plant_derivs.get_actuation_input_port().FixValue( self.plant_derivs_context, [0., 0., 0., 0., 0., 0., 0.]) null_force = BasicVector([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) self.plant_derivs.GetInputPort("applied_generalized_force").FixValue( self.plant_derivs_context, null_force) self.plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=dt) parser = Parser(self.plant, scene_graph) self.iiwa, self.hinge, self.bushing = self.add_models(self.plant, parser, params=params) self.plant.Finalize( ) # Finalize will assign ports for compatibility w/ the scene_graph; could be cause of the issue w/ first order taylor. self.meshcat = ConnectMeshcatVisualizer(builder, scene_graph, zmq_url=zmq_url) self.sym_derivs = False # If system should use symbolic derivatives; if false, autodiff self.custom_sim = False # If rollouts should be gathered with sys.dyn() calls nq = self.plant.num_positions() nv = self.plant.num_velocities() self.n_x = nq + nv self.n_u = self.plant.num_actuators() self.n_y = self.plant.get_state_output_port(self.iiwa).size() self.N = N self.dt = dt self.decay = trj_decay self.V = 1e-2 * np.ones(self.n_y) self.W = np.concatenate((1e-7 * np.ones(nq), 1e-4 * np.ones(nv))) self.W0 = np.concatenate((1e-9 * np.ones(nq), 1e-6 * np.ones(nv))) self.x_w_cov = x_w_cov self.door_angle_ref = door_angle_ref self.q0 = np.array( [-3.12, -0.17, 0.52, -3.11, 1.22, -0.75, -1.56, 0.55]) #self.q0 = np.array([-3.12, -0.27, 0.52, -3.11, 1.22, -0.75, -1.56, 0.55]) self.x0 = np.concatenate((self.q0, np.zeros(nv))) self.door_index = None self.phi = {} def ports_init(self, context): self.plant_context = self.plant.GetMyMutableContextFromRoot(context) self.plant.SetPositionsAndVelocities(self.plant_context, self.x0) #door_angle = self.plant.GetPositionsFromArray(self.hinge, self.x0[:8]) self.door_index = self.plant.GetJointByName( 'right_door_hinge').position_start() null_force = BasicVector([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) self.plant.GetInputPort("applied_generalized_force").FixValue( self.plant_context, null_force) self.plant.get_actuation_input_port().FixValue( self.plant_context, [0., 0., 0., 0., 0., 0., 0.]) self.W0[self.door_index] = self.x_w_cov def add_models(self, plant, parser, params=None): iiwa = parser.AddModelFromFile( "/home/hanikevi/drake/manipulation/models/iiwa_description/sdf/iiwa14_no_collision_no_grav.sdf" ) plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("iiwa_link_0")) #box = Box(10., 10., 10.) #X_WBox = RigidTransform([0, 0, -5]) #mu = 0.6 #plant.RegisterCollisionGeometry(plant.world_body(), X_WBox, box, "ground", CoulombFriction(mu, mu)) #plant.RegisterVisualGeometry(plant.world_body(), X_WBox, box, "ground", [.9, .9, .9, 1.0]) #planar_joint_frame = plant.AddFrame(FixedOffsetFrame("planar_joint_frame", plant.world_frame(), RigidTransform(RotationMatrix.MakeXRotation(np.pi/2)))) X_WCylinder = RigidTransform([-0.75, 0, 0.5]) hinge = parser.AddModelFromFile( "/home/hanikevi/drake/examples/manipulation_station/models/simple_hinge.sdf" ) plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("base"), X_WCylinder) #cupboard_door_spring = plant.AddForceElement(RevoluteSpring_[float](plant.GetJointByName("right_door_hinge"), nominal_angle = -0.4, stiffness = 10)) if params is None: bushing = LinearBushingRollPitchYaw_[float]( plant.GetFrameByName("iiwa_link_7"), plant.GetFrameByName("handle"), [50, 50, 50], # Torque stiffness [2., 2., 2.], # Torque damping [5e4, 5e4, 5e4], # Linear stiffness [80, 80, 80], # Linear damping ) else: print('setting custom stiffnesses') bushing = LinearBushingRollPitchYaw_[float]( plant.GetFrameByName("iiwa_link_7"), plant.GetFrameByName("handle"), [params['k4'], params['k5'], params['k6']], # Torque stiffness [2, 2, 2], # Torque damping [params['k1'], params['k2'], params['k3']], # Linear stiffness [100, 100, 100], # Linear damping ) bushing_element = plant.AddForceElement(bushing) return iiwa, hinge, bushing def cost_stage(self, x, u): ctrl = 1e-5 * np.sum(u**2) pos = 15.0 * (x[self.door_index] - self.door_angle_ref)**2 vel = 1e-5 * np.sum(x[8:]**2) return pos + ctrl + vel def cost_final(self, x): return 50 * (1.0 * (x[self.door_index] - self.door_angle_ref)**2 + np.sum(2.5e-4 * x[8:]**2)) def get_deriv(self, x, u): self.plant_derivs.SetPositionsAndVelocities(self.plant_derivs_context, x) lin = FirstOrderTaylorApproximation( self.plant_derivs, self.plant_derivs_context, self.plant.get_actuation_input_port().get_index(), self.plant.get_state_output_port(self.iiwa).get_index()) return lin.A(), lin.B(), lin.C() def get_param_deriv(self, x, u): # Using a closed-form solution; as currently DRAKE doesn't do support autodiff on parameters for LinearBusing. # Note dC is 0 - stiffness does not affect measurements of q, v W = self.plant.world_frame() I = self.plant.GetFrameByName("iiwa_link_7") H = self.plant.GetFrameByName("handle") self.plant.SetPositionsAndVelocities(self.plant_context, x) M = self.plant.CalcMassMatrixViaInverseDynamics(self.plant_context) Jac_I = self.plant.CalcJacobianSpatialVelocity( self.plant_context, JacobianWrtVariable.kQDot, I, [0, 0, 0], W, W) Jac_H = self.plant.CalcJacobianSpatialVelocity( self.plant_context, JacobianWrtVariable.kQDot, H, [0, 0, 0], W, W) dA = np.zeros((self.n_x**2, 3)) dA_sq = np.zeros((self.n_x, self.n_x)) for param_ind in range(3): JH = np.outer(Jac_H[param_ind, :], Jac_H[param_ind, :]) JI = np.outer(Jac_I[param_ind, :], Jac_I[param_ind, :]) dA_sq[8:, :8] = self.dt * np.linalg.inv(M).dot(JH + JI) dA[:, param_ind] = deepcopy(dA_sq.ravel()) #print(np.sum(np.abs(dA), axis=0)) return dA def reset(self): x_traj_new = np.zeros((self.N + 1, self.n_x)) x_traj_new[0, :] = self.x0 + np.multiply(np.sqrt(self.W0), np.random.randn(self.n_x)) u_traj_new = np.zeros((self.N, self.n_u)) #self.plant_context.SetDiscreteState(x_traj_new[0,:]) self.plant.SetPositionsAndVelocities(self.plant_context, x_traj_new[0, :]) return x_traj_new, u_traj_new def rollout(self): self.u_trj = np.random.randn(self.N, self.n_u) * 0.001 self.x_trj, _ = self.reset() for i in range(self.N): self.x_trj[i + 1, :] = self.dyn(self.x_trj[i, :], self.u_trj[i], noise=True) def dyn(self, x, u, phi=None, noise=False): x_next = self.plant.AllocateDiscreteVariables() self.plant.SetPositionsAndVelocities(self.plant_context, x) self.plant.get_actuation_input_port(self.iiwa).FixValue( self.plant_context, u) self.plant.CalcDiscreteVariableUpdates(self.plant_context, x_next) #print(x_next.get_mutable_vector().get_value()) x_new = x_next.get_mutable_vector().get_value() if noise: noise = np.multiply(np.sqrt(self.W), np.random.randn(self.n_x)) x_new += noise return x_new def obs(self, x, mode=None, noise=False, phi=None): y = self.plant.get_state_output_port(self.iiwa).Eval( self.plant_context) #print(self.bushing.CalcBushingSpatialForceOnFrameA(self.plant_context).translational()) if noise: y += np.multiply(np.sqrt(self.V), np.random.randn(self.n_y)) return y def cost(self, x_trj=None, u_trj=None): cost_trj = 0.0 if x_trj is None: for i in range(self.N): cost_trj += self.cost_stage(self.x_trj[i, :], self.u_trj[i, :]) cost_trj += self.cost_final(self.sys.plant, self.x_trj[-1, :]) else: for i in range(self.N): cost_trj += self.cost_stage(x_trj[i, :], u_trj[i, :]) cost_trj += self.cost_final(x_trj[-1, :]) return cost_trj
# plant.get_actuation_input_port().FixValue(context, np.zeros(7)) # simulator = Simulator(plant, context) # simulator.AdvanceTo(5.0) builder = DiagramBuilder() # Adds both MultibodyPlant and SceneGraph and wires them together plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=1e-4) # Note that we parse into both the plant and the scene_graph here Parser(plant, scene_graph).AddModelFromFile( FindResourceOrThrow( "drake/manipulation/models/iiwa_description/sdf/iiwa14_no_collision.sdf" )) plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("iiwa_link_0")) # Adds the MeshcatVisualizer and wires it to the Scene Graph meshcat = ConnectMeshcatVisualizer(builder, scene_graph, open_browser=True) plant.Finalize() diagram = builder.Build() context = diagram.CreateDefaultContext() meshcat.load() diagram.Publish(context) plant_context = plant.GetMyMutableContextFromRoot(context) plant.SetPositions(plant_context, [-1.57, 0.1, 0, -1.2, 0, 1.6, 0]) plant.get_actuation_input_port().FixValue(plant_context, np.zeros(7))