diff --git a/demos/control_demo.py b/demos/control_demo.py new file mode 100644 index 0000000..f8646d9 --- /dev/null +++ b/demos/control_demo.py @@ -0,0 +1,86 @@ +import sys +import torch +import pygame + +from lcp_physics.physics.bodies import Circle, Rect, Hull +from lcp_physics.physics.constraints import Joint, YConstraint, XConstraint, RotConstraint, TotalConstraint +from lcp_physics.physics.constraints import FixedJoint +from lcp_physics.physics.forces import ExternalForce, Gravity, vert_impulse, hor_impulse, MDP +from lcp_physics.physics.utils import Defaults, Recorder +from lcp_physics.physics.world import World, run_world_traj, Trajectory +import numpy as np +from numpy import loadtxt + +vel = np.array([[1.0,1.0,1.0,1.0,1.0],[-0.0,-0.0,-0.0,-0.0,-0.0]])/10 +traj1 = vel.copy() +traj1[0,:] = 2500*vel[0,:] +traj1[1,:] = -2500*vel[1,:] + +traj2 = vel.copy() +traj2[0,:] = 2500*vel[0,:] +traj2[1,:] = 2500*vel[1,:] + +pygame.init() +screen = pygame.display.set_mode((1000, 1000), pygame.DOUBLEBUF) +screen.set_alpha(None) + +polygon = np.array((loadtxt("../DynamicAffordances/data/polygons_1_2f_sq.csv", delimiter=','))) + +bodies = [] +joints = [] +restitution = 0.00 # no impacts in quasi-dynamics +fric_coeff = 0.01 +n_pol = int(polygon[0,0]) + +xr = 500 +yr = 500 + +# adds body based on triangulation +r0 = Hull([xr, yr], [[1, 1], [-1, 1], [-1, -1], [1, -1]], + restitution=0.00, fric_coeff=0.00, mass = 0.01, name="obj") +bodies.append(r0) + +for i in range(n_pol): + x2 = [polygon[0,1+8*i], -polygon[0,2+8*i]] + x1 = [polygon[0,3+8*i], -polygon[0,4+8*i]] + x0 = [polygon[0,5+8*i], -polygon[0,6+8*i]] + verts = 250*np.array([x0, x1, x2]) + print(verts) + p0 = np.array([xr + polygon[0,7+8*i], yr - polygon[0,8+8*i]]) + r1 = Hull(p0, verts, restitution=restitution, mass = 0.0001, fric_coeff=1, name="obj_"+str(i)) + print('disp1') + r1.add_force(MDP(g=100)) + # r1.add_force(Gravity(g=100)) + bodies.append(r1) + joints += [FixedJoint(r1, r0)] + r1.add_no_contact(r0) + r0 = r1 + +# Manipulators +c = Circle([200, 550], 5, mass = 100000000, vel=(0, 0, 0), restitution=restitution, + fric_coeff=1, name = "f1") +bodies.append(c) + +c = Circle([200, 450], 5, mass = 100000000, vel=(0, 0, 0), restitution=restitution, + fric_coeff=1, name = "f2") +bodies.append(c) + +vel = torch.tensor(traj1) +traj_f = [] +traj_f.append(Trajectory(vel = vel, name = "f1")) + +vel = torch.tensor(traj2) +traj_f.append(Trajectory(vel = vel, name = "f2")) + +# Environment +# r = Rect([0, 500, 505], [900, 10], +# restitution=restitution, fric_coeff=1) +# bodies.append(r) +# joints.append(TotalConstraint(r)) + +recorder = None +world = World(bodies, joints, dt=0.1) +run_world_traj(world, run_time=5, screen=screen, recorder=recorder, traj=traj_f) +print('\n') +print(world.states) +print('\n') \ No newline at end of file diff --git a/demos/demo.py b/demos/demo.py index 6a10f26..29f6c36 100644 --- a/demos/demo.py +++ b/demos/demo.py @@ -1,3 +1,4 @@ + import math import sys @@ -165,7 +166,7 @@ def timed_force(t): screen.set_alpha(None) pygame.display.set_caption('2D Engine') - slide_demo(screen) - fric_demo(screen) - chain_demo(screen) - debug_demo(screen) + # slide_demo(screen) + # fric_demo(screen) + # chain_demo(screen) + debug_demo(screen) \ No newline at end of file diff --git a/lcp_physics/lcp/lcp.py b/lcp_physics/lcp/lcp.py index 06c6dcb..b2faee7 100644 --- a/lcp_physics/lcp/lcp.py +++ b/lcp_physics/lcp/lcp.py @@ -1,24 +1,19 @@ import torch -from torch.autograd import Function - from .solvers import pdipm +import numpy as np from .util import bger, extract_batch_size - -class LCPFunction(Function): +class LCPFunction(torch.autograd.Function): """A differentiable LCP solver, uses the primal dual interior point method implemented in pdipm. """ - def __init__(self, eps=1e-12, verbose=-1, not_improved_lim=3, - max_iter=10): + # @staticmethod + def __init__(self): super().__init__() - self.eps = eps - self.verbose = verbose - self.not_improved_lim = not_improved_lim - self.max_iter = max_iter self.Q_LU = self.S_LU = self.R = None # @profile + @staticmethod def forward(self, Q, p, G, h, A, b, F): _, nineq, nz = G.size() neq = A.size(1) if A.ndimension() > 1 else 0 @@ -28,20 +23,23 @@ def forward(self, Q, p, G, h, A, b, F): self.Q_LU, self.S_LU, self.R = pdipm.pre_factor_kkt(Q, G, F, A) zhats, self.nus, self.lams, self.slacks = pdipm.forward( Q, p, G, h, A, b, F, self.Q_LU, self.S_LU, self.R, - eps=self.eps, max_iter=self.max_iter, verbose=self.verbose, - not_improved_lim=self.not_improved_lim) + eps=1e-18, max_iter=10, verbose=-1, + not_improved_lim=3) self.save_for_backward(zhats, Q, p, G, h, A, b, F) return zhats + @staticmethod def backward(self, dl_dzhat): + + # print(dl_dzhat) zhats, Q, p, G, h, A, b, F = self.saved_tensors batch_size = extract_batch_size(Q, p, G, h, A, b) neq, nineq, nz = self.neq, self.nineq, self.nz # D = torch.diag((self.lams / self.slacks).squeeze(0)).unsqueeze(0) - d = self.lams / self.slacks + d = self.lams / (self.slacks + 1e-30) pdipm.factor_kkt(self.S_LU, self.R, d) dx, _, dlam, dnu = pdipm.solve_kkt(self.Q_LU, d, G, A, self.S_LU, @@ -61,12 +59,26 @@ def backward(self, dl_dzhat): dQs = 0.5 * (bger(dx, zhats) + bger(zhats, dx)) grads = (dQs, dps, dGs, dhs, dAs, dbs, dFs) + + if np.isnan(grads[0].sum()).item() > 0: + import pdb + pdb.set_trace() + + # print('\n\n\na\n\n') + # print(grads) + # print('\n\n\nb\n\n') + return grads + @staticmethod def numerical_backward(self, dl_dzhat): # XXX experimental # adapted from pytorch's grad check # from torch.autograd.gradcheck import get_numerical_jacobian + + + print('dl_dzhat') + from torch.autograd import Variable from collections import Iterable diff --git a/lcp_physics/lcp/lemkelcp.py b/lcp_physics/lcp/lemkelcp.py new file mode 100755 index 0000000..be1b5dc --- /dev/null +++ b/lcp_physics/lcp/lemkelcp.py @@ -0,0 +1,217 @@ +import torch +import numpy as np + +EPSILON = 1e-6 + + +class lemketableau: + def __init__(self, M, q, maxIter=1e4): + n = len(q) + self.T = torch.cat((torch.eye(n), -M, -torch.ones((n, 1)), q.reshape((n, 1))), dim=1) + self.n = n + self.wPos = list(range(n)) + self.zPos = list(range(n, 2 * n)) + self.W = 0 + self.Z = 1 + self.Y = 2 + self.Q = 3 + TbInd = torch.stack((self.W * torch.ones(n, dtype=int), + torch.arange(n, dtype=int)), dim=0) + TnbInd = torch.stack((self.Z * torch.ones(n, dtype=int), + torch.arange(n, dtype=int)), dim=0) + DriveInd = torch.tensor([[self.Y], [0]]) + QInd = torch.tensor([[self.Q], [0]]) + self.Tind = torch.cat((TbInd, TnbInd, DriveInd, QInd), dim=1) + self.maxIter = maxIter + + def lemkeAlgorithm(self): + initVal = self.initialize() + if not initVal: + return torch.zeros(self.n), 0, 'Solution Found' + + for k in range(self.maxIter): + stepVal = self.step() + if self.Tind[0, -2] == self.Y: + # Solution Found + z = self.extractSolution() + return z, 0, 'Solution Found' + elif not stepVal: + # pass + return None, 1, 'Secondary ray found' + + return None, 2, 'Max Iterations Exceeded' + # z = self.extractSolution() + # print('Returning approximate solution.') + # return z, 0, 'Max Iterations Exceeded' + + def initialize(self): + q = self.T[:, -1] + minQ = torch.min(q) + if minQ < 0: + ind = torch.argmin(q) + self.clearDriverColumn(ind) + self.pivot(ind) + + return True + else: + return False + + def step(self): + q = self.T[:, -1] + a = self.T[:, -2] + ind = torch.tensor(float('nan')) + minRatio = torch.tensor(float('inf')) + for i in range(self.n): + if a[i] > EPSILON: + newRatio = q[i] / a[i] + if newRatio < minRatio - EPSILON: + ind = i + minRatio = newRatio + + if minRatio < torch.tensor(float('inf')): + self.clearDriverColumn(ind) + self.pivot(ind) + return True + else: + return False + + def extractSolution(self): + z = torch.zeros(self.n) + q = self.T[:, -1] + for i in range(self.n): + if self.Tind[0, i] == self.Z: + z[self.Tind[1, i]] = q[i] + return z + + def partnerPos(self, pos): + v, ind = self.Tind[:, pos] + if v == self.W: + ppos = self.zPos[ind] + elif v == self.Z: + ppos = self.wPos[ind] + else: + ppos = None + return ppos + + def pivot(self, pos): + ppos = self.partnerPos(pos) + if ppos is not None: + self.swapColumns(pos, ppos) + self.swapColumns(pos, -2) + return True + else: + self.swapColumns(pos, -2) + return False + + def swapMatColumns(self, M, i, j): + Mi = M[:, i].clone() + Mj = M[:, j].clone() + M[:, i] = Mj + M[:, j] = Mi + return M + + def swapPos(self, v, ind, newPos): + if v == self.W: + self.wPos[ind] = newPos % (2 * self.n + 2) + elif v == self.Z: + self.zPos[ind] = newPos % (2 * self.n + 2) + + def swapColumns(self, i, j): + iInd = self.Tind[:, i] + jInd = self.Tind[:, j] + + v, ind = iInd + self.swapPos(v, ind, j) + v, ind = jInd + self.swapPos(v, ind, i) + + self.Tind = self.swapMatColumns(self.Tind, i, j) + self.T = self.swapMatColumns(self.T, i, j) + + def clearDriverColumn(self, ind): + a = self.T[ind, -2] + self.T[ind] /= a.clone() + for i in range(self.n): + if i != ind: + b = self.T[i, -2] + self.T[i] -= b * self.T[ind] + + def ind2str(self, indvec): + v, pos = indvec + if v == self.W: + s = 'w%d' % pos + elif v == self.Z: + s = 'z%d' % pos + elif v == self.Y: + s = 'y' + else: + s = 'q' + return s + + def indexStringArray(self): + indstr = np.array([self.ind2str(indvec) for indvec in self.Tind.T], dtype=object) + return indstr + + def indexedTableau(self): + indstr = self.indexStringArray() + return torch.cat((indstr, self.T), dim=0) + + def __repr__(self): + IT = self.indexedTableau() + return IT.__repr__() + + def __str__(self): + IT = self.indexedTableau() + return IT.__str__() + + +def lemkelcp(M, q, maxIter=100): + """ + sol = lemkelcp(M,q,maxIter) + Uses Lemke's algorithm to copute a solution to the + linear complementarity problem: + Mz + q >= 0 + z >= 0 + z'(Mz+q) = 0 + The inputs are given by: + M - an nxn numpy array + q - a length n numpy array + maxIter - an optional number of pivot iterations. Set to 100 by default + The solution is a tuple of the form: + z,exit_code,exit_string = sol + The entries are summaries in the table below: + + |z | exit_code | exit_string | + ----------------------------------------------------------- + | solution to LCP | 0 | 'Solution Found' | + | None | 1 | 'Secondary ray found' | + | None | 2 | 'Max Iterations Exceeded' | + """ + + tableau = lemketableau(M, q, maxIter) + return tableau.lemkeAlgorithm() + + +class LemkeLCP(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + self.prev_sol = None + + def forward(self, P, Q, R, S, u, v): + """From Cline 2.7.2""" + Pinv = torch.inverse(P) + M = S - ((R @ Pinv) @ Q) + q = (v - ((R @ Pinv) @ u)).flatten() + + z, code, msg = lemkelcp(M, q, maxIter=100) + if code != 0: + if self.prev_sol is not None: + print(f'Could not solve LCP ({msg}), returning previous solution.') + return self.prev_sol + raise Exception('Could not solve LCP: ' + msg) + z = z.type_as(M) + w = (M @ z) + q + x = Pinv @ (-u -(Q @ z)) + self.prev_sol = (z, w, x) + return z, w, x + diff --git a/lcp_physics/lcp/solvers/dev_pdipm.py b/lcp_physics/lcp/solvers/dev_pdipm.py index 4cb2c1c..11a44ba 100644 --- a/lcp_physics/lcp/solvers/dev_pdipm.py +++ b/lcp_physics/lcp/solvers/dev_pdipm.py @@ -33,14 +33,11 @@ def btrifact_hack(x): INACC_ERR = """ -------- lcp warning: Returning an inaccurate and potentially incorrect solution. - Some residual is large. Your problem may be infeasible or difficult. - You can try using the CVXPY solver to see if your problem is feasible and you can use the verbose option to check the convergence status of our solver while increasing the number of iterations. - Advanced users: You can also try to enable iterative refinement in the solver: https://github.com/locuslab/lcp/issues/6 @@ -431,8 +428,8 @@ def factor_solve_kkt_reg(Q_tilde, D, G, A, C_tilde, rx, rs, rz, ry, eps): H_LU = btrifact_hack(H_) - invH_A_ = A_.transpose(1, 2).btrisolve(*H_LU) # H-1 AT - invH_g_ = g_.btrisolve(*H_LU) # H-1 g + invH_A_ = A_.transpose(1, 2).lu_solve(*H_LU) # H-1 AT + invH_g_ = g_.lu_solve(*H_LU) # H-1 g S_ = torch.bmm(A_, invH_A_) # A H-1 AT # A H-1 AT + C_tilde @@ -441,11 +438,11 @@ def factor_solve_kkt_reg(Q_tilde, D, G, A, C_tilde, rx, rs, rz, ry, eps): # [(H-1 g)T AT]T - h = A H-1 g - h t_ = torch.bmm(invH_g_.unsqueeze(1), A_.transpose(1, 2)).squeeze(1) - h_ # w = (A H-1 AT + C_tilde)-1 (A H-1 g - h) <= Av - eps I w = h - w_ = -t_.btrisolve(*S_LU) + w_ = -t_.lu_solve(*S_LU) # XXX Shouldn't it be just g (no minus)? # (Doesn't seem to make a difference, though...) t_ = -g_ - w_.unsqueeze(1).bmm(A_).squeeze() # -g - AT w - v_ = t_.btrisolve(*H_LU) # v = H-1 (-g - AT w) + v_ = t_.lu_solve(*H_LU) # v = H-1 (-g - AT w) dx = v_[:, :nz] ds = v_[:, nz:] @@ -542,12 +539,12 @@ def sparse_factor_solve_kkt_reg(spH_, A_, spA_, spC_tilde, rx, rs, rz, ry, neq, # [(H-1 g)T AT]T - h = A H-1 g - h t_ = spA_.dot(invH_g_) - h_ # w = (A H-1 AT + C_tilde)-1 (A H-1 g - h) <= Av - eps I w = h - # w_ = -t_.btrisolve(*S_LU) + # w_ = -t_.lu_solve(*S_LU) w_ = -lu_solve(S_LU, t_) # XXX Shouldn't it be just g (no minus)? # (Doesn't seem to make a difference, though...) t_ = -g_ - spA_.transpose().dot(w_) # -g - AT w - # v_ = t_.btrisolve(*H_LU) # v = H-1 (-g - AT w) + # v_ = t_.lu_solve(*H_LU) # v = H-1 (-g - AT w) v_ = H_LU.solve(t_) dx = v_[:nz] @@ -694,15 +691,15 @@ def factor_solve_kkt(Q_tilde, D_tilde, A_, C_tilde, rx, rs, rz, ry, ns): H_LU = btrifact_hack(H_) - invH_A_ = A_.transpose(1, 2).btrisolve(*H_LU) - invH_g_ = g_.btrisolve(*H_LU) + invH_A_ = A_.transpose(1, 2).lu_solve(*H_LU) + invH_g_ = g_.lu_solve(*H_LU) S_ = torch.bmm(A_, invH_A_) + C_tilde S_LU = btrifact_hack(S_) t_ = torch.bmm(invH_g_.unsqueeze(1), A_.transpose(1, 2)).squeeze(1) - h_ - w_ = -t_.btrisolve(*S_LU) + w_ = -t_.lu_solve(*S_LU) t_ = -g_ - w_.unsqueeze(1).bmm(A_).squeeze() - v_ = t_.btrisolve(*H_LU) + v_ = t_.lu_solve(*H_LU) dx = v_[:, :nz] ds = v_[:, nz:] @@ -761,7 +758,7 @@ def solve_kkt(Q_LU, d, G, A, S_LU, rx, rs, rz, ry): nineq, nz, neq, nBatch = get_sizes(G, A) - invQ_rx = rx.btrisolve(*Q_LU) # Q-1 rx + invQ_rx = rx.lu_solve(*Q_LU) # Q-1 rx if neq > 0: # A Q-1 rx - ry # G Q-1 rx + rs / d - rz @@ -770,14 +767,14 @@ def solve_kkt(Q_LU, d, G, A, S_LU, rx, rs, rz, ry): else: h = invQ_rx.unsqueeze(1).bmm(G.transpose(1, 2)).squeeze(1) + rs / d - rz - w = -(h.btrisolve(*S_LU)) # S-1 h = + w = -(h.lu_solve(*S_LU)) # S-1 h = g1 = -rx - w[:, neq:].unsqueeze(1).bmm(G).squeeze(1) # -rx - GT w = -rx -GT S-1 h if neq > 0: g1 -= w[:, :neq].unsqueeze(1).bmm(A).squeeze(1) # - AT w = -AT S-1 h g2 = -rs - w[:, neq:] - dx = g1.btrisolve(*Q_LU) # Q-1 g1 = - Q-1 AT S-1 h + dx = g1.lu_solve(*Q_LU) # Q-1 g1 = - Q-1 AT S-1 h ds = g2 / d # g2 / d = (-rs - w) / d dz = w[:, neq:] dy = w[:, :neq] if neq > 0 else None @@ -806,12 +803,12 @@ def pre_factor_kkt(Q, G, F, A): # See the 'Block LU factorization' part of our website # for more details. - G_invQ_GT = torch.bmm(G, G.transpose(1, 2).btrisolve(*Q_LU)) + F + G_invQ_GT = torch.bmm(G, G.transpose(1, 2).lu_solve(*Q_LU)) + F R = G_invQ_GT.clone() S_LU_pivots = torch.IntTensor(range(1, 1 + neq + nineq)).unsqueeze(0) \ .repeat(nBatch, 1).type_as(Q).int() if neq > 0: - invQ_AT = A.transpose(1, 2).btrisolve(*Q_LU) + invQ_AT = A.transpose(1, 2).lu_solve(*Q_LU) A_invQ_AT = torch.bmm(A, invQ_AT) G_invQ_AT = torch.bmm(G, invQ_AT) @@ -821,9 +818,9 @@ def pre_factor_kkt(Q, G, F, A): S_LU_11 = LU_A_invQ_AT[0] U_A_invQ_AT_inv = (P_A_invQ_AT.bmm(L_A_invQ_AT) - ).btrisolve(*LU_A_invQ_AT) + ).lu_solve(*LU_A_invQ_AT) S_LU_21 = G_invQ_AT.bmm(U_A_invQ_AT_inv) - T = G_invQ_AT.transpose(1, 2).btrisolve(*LU_A_invQ_AT) + T = G_invQ_AT.transpose(1, 2).lu_solve(*LU_A_invQ_AT) S_LU_12 = U_A_invQ_AT.bmm(T) S_LU_22 = torch.zeros(nBatch, nineq, nineq).type_as(Q) S_LU_data = torch.cat((torch.cat((S_LU_11, S_LU_12), 2), diff --git a/lcp_physics/lcp/solvers/pdipm.py b/lcp_physics/lcp/solvers/pdipm.py index c808378..a5aff05 100644 --- a/lcp_physics/lcp/solvers/pdipm.py +++ b/lcp_physics/lcp/solvers/pdipm.py @@ -2,6 +2,7 @@ """ import torch +import numpy as np from enum import Enum from ..util import efficient_btriunpack @@ -15,7 +16,7 @@ def btrifact_hack(x): global shown_btrifact_warning try: - return x.btrifact(pivot=not x.is_cuda) + return x.lu(pivot=not x.is_cuda) except TypeError: if not shown_btrifact_warning: print("""---------- @@ -31,13 +32,10 @@ def btrifact_hack(x): INACC_ERR = """ -------- lcp warning: Returning an inaccurate and potentially incorrect solution. - Some residual is large. Your problem may be infeasible or difficult. - You can try using the verbose option to check the convergence status of our solver while increasing the number of iterations. - Advanced users: You can also try to enable iterative refinement in the solver: https://github.com/locuslab/qpth/issues/6 @@ -65,13 +63,13 @@ def forward(Q, p, G, h, A, b, F, Q_LU, S_LU, R, # Make all of the slack variables >= 1. M = torch.min(s, 1)[0] M = M.view(M.size(0), 1).repeat(1, nineq) - I = M <= 0 + I = (M <= 0).bool() s[I] -= M[I] - 1 # Make all of the inequality dual variables >= 1. M = torch.min(z, 1)[0] M = M.view(M.size(0), 1).repeat(1, nineq) - I = M <= 0 + I = (M <= 0).bool() z[I] -= M[I] - 1 best = {'resids': None, 'x': None, 'z': None, 's': None, 'y': None} @@ -88,14 +86,14 @@ def forward(Q, p, G, h, A, b, F, Q_LU, S_LU, R, - torch.bmm(z.unsqueeze(1), F.transpose(1, 2)).squeeze(1) ry = torch.bmm(x.unsqueeze(1), A.transpose( 1, 2)).squeeze(1) - b if neq > 0 else 0.0 - mu = torch.abs((s * z).sum(1).squeeze() / nineq) + mu = torch.abs((s * z).sum(1).squeeze() / (nineq + 1e-30)) z_resid = torch.norm(rz, 2, 1).squeeze() y_resid = torch.norm(ry, 2, 1).squeeze() if neq > 0 else 0 pri_resid = y_resid + z_resid dual_resid = torch.norm(rx, 2, 1).squeeze() resids = pri_resid + dual_resid + nineq * mu - d = z / s + d = z / (s + 1e-30) try: factor_kkt(S_LU, R, d) except: @@ -135,6 +133,11 @@ def forward(Q, p, G, h, A, b, F, Q_LU, S_LU, R, print(INACC_ERR) return best['x'], best['y'], best['z'], best['s'] + if np.isnan(rs.sum()).item() > 0: + # print(rs) + rs = rs*0 + # import pdb + # pdb.set_trace() dx_aff, ds_aff, dz_aff, dy_aff = solve_kkt( Q_LU, d, G, A, S_LU, rx, rs, rz, ry) @@ -142,18 +145,19 @@ def forward(Q, p, G, h, A, b, F, Q_LU, S_LU, R, alpha = torch.min(torch.min(get_step(z, dz_aff), get_step(s, ds_aff)), torch.ones(batch_size).type_as(Q)) + alpha_nineq = alpha.repeat(nineq, 1).t() - t1 = s + alpha_nineq * ds_aff - t2 = z + alpha_nineq * dz_aff + t1 = s + alpha_nineq * (ds_aff + 1e-30) + t2 = z + alpha_nineq * (dz_aff + 1e-30) t3 = torch.sum(t1 * t2, 1).squeeze() t4 = torch.sum(s * z, 1).squeeze() - sig = (t3 / t4)**3 + sig = (t3 / (t4 + 1e-30))**3 rx = Q.new_zeros(batch_size, nz) - rs = ((-mu * sig).repeat(nineq, 1).t() + ds_aff * dz_aff) / s + rs = ((-mu * sig).repeat(nineq, 1).t() + ds_aff * dz_aff) / (s + 1e-30) rz = Q.new_zeros(batch_size, nineq) ry = Q.new_zeros(batch_size, neq) - + dx_cor, ds_cor, dz_cor, dy_cor = solve_kkt( Q_LU, d, G, A, S_LU, rx, rs, rz, ry) @@ -180,11 +184,19 @@ def forward(Q, p, G, h, A, b, F, Q_LU, S_LU, R, def get_step(v, dv): - a = -v / dv + a = -v / (dv + 1e-12) a[dv > 0] = max(1.0, a.max()) step = a.min(1)[0].squeeze() return step +# def get_step(v, dv): + # I = dv < 1e-12 + # if torch.sum(I) > 0: # TODO: Use something like torch.any(dv < 0) + # a = -v / (dv + 1e-30) + # return torch.min(a[I]) + # else: + # return 1 + def unpack_kkt(v, nz, nineq, neq): i = 0 @@ -213,115 +225,6 @@ def kkt_resid_reg(Q_tilde, D_tilde, G, A, F_tilde, eps, dx, ds, dz, dy, rx, rs, return resx, ress, resz, resy -def solve_kkt_ir(Q, D, G, A, F, rx, rs, rz, ry, niter=1): - """Inefficient iterative refinement.""" - nineq, nz, neq, nBatch = get_sizes(G, A) - - eps = 1e-7 - Q_tilde = Q + eps * torch.eye(nz).type_as(Q).repeat(nBatch, 1, 1) - D_tilde = D + eps * torch.eye(nineq).type_as(Q).repeat(nBatch, 1, 1) - - # XXX Might not workd for batch size > 1 - C_tilde = -eps * torch.eye(neq + nineq).type_as(Q_tilde).repeat(nBatch, 1, 1) - if F is not None: # XXX inverted sign for F below - C_tilde[:, :nineq, :nineq] -= F - F_tilde = C_tilde[:, :nineq, :nineq] - - dx, ds, dz, dy = factor_solve_kkt_reg( - Q_tilde, D_tilde, G, A, C_tilde, rx, rs, rz, ry, eps) - resx, ress, resz, resy = kkt_resid_reg(Q, D, G, A, F_tilde, eps, - dx, ds, dz, dy, rx, rs, rz, ry) - for k in range(niter): - ddx, dds, ddz, ddy = factor_solve_kkt_reg(Q_tilde, D_tilde, G, A, C_tilde, - -resx, -ress, -resz, - -resy if resy is not None else None, - eps) - dx, ds, dz, dy = [v + dv if v is not None else None - for v, dv in zip((dx, ds, dz, dy), (ddx, dds, ddz, ddy))] - resx, ress, resz, resy = kkt_resid_reg(Q, D, G, A, F_tilde, eps, - dx, ds, dz, dy, rx, rs, rz, ry) - - return dx, ds, dz, dy - - -def factor_solve_kkt_reg(Q_tilde, D, G, A, C_tilde, rx, rs, rz, ry, eps): - nineq, nz, neq, nBatch = get_sizes(G, A) - - H_ = torch.zeros(nBatch, nz + nineq, nz + nineq).type_as(Q_tilde) - H_[:, :nz, :nz] = Q_tilde - H_[:, -nineq:, -nineq:] = D - if neq > 0: - # H_ = torch.cat([torch.cat([Q, torch.zeros(nz,nineq).type_as(Q)], 1), - # torch.cat([torch.zeros(nineq, nz).type_as(Q), D], 1)], 0) - A_ = torch.cat([torch.cat([G, torch.eye(nineq).type_as(Q_tilde).repeat(nBatch, 1, 1)], 2), - torch.cat([A, torch.zeros(nBatch, neq, nineq).type_as(Q_tilde)], 2)], 1) - g_ = torch.cat([rx, rs], 1) - h_ = torch.cat([rz, ry], 1) - else: - A_ = torch.cat( - [G, torch.eye(nineq).type_as(Q_tilde).repeat(nBatch, 1, 1)], 2) - g_ = torch.cat([rx, rs], 1) - h_ = rz - - H_LU = btrifact_hack(H_) - - invH_A_ = A_.transpose(1, 2).btrisolve(*H_LU) # H-1 AT - invH_g_ = g_.btrisolve(*H_LU) # H-1 g - - S_ = torch.bmm(A_, invH_A_) # A H-1 AT - # A H-1 AT + C_tilde - S_ -= C_tilde - S_LU = btrifact_hack(S_) - # [(H-1 g)T AT]T - h = A H-1 g - h - t_ = torch.bmm(invH_g_.unsqueeze(1), A_.transpose(1, 2)).squeeze(1) - h_ - # w = (A H-1 AT + C_tilde)-1 (A H-1 g - h) <= Av - eps I w = h - w_ = -t_.btrisolve(*S_LU) - # Shouldn't it be just g (no minus)? - # (Doesn't seem to make a difference, though...) - t_ = -g_ - w_.unsqueeze(1).bmm(A_).squeeze() # -g - AT w - v_ = t_.btrisolve(*H_LU) # v = H-1 (-g - AT w) - - dx = v_[:, :nz] - ds = v_[:, nz:] - dz = w_[:, :nineq] - dy = w_[:, nineq:] if neq > 0 else None - - return dx, ds, dz, dy - - -def factor_solve_kkt(Q_tilde, D_tilde, A_, C_tilde, rx, rs, rz, ry, ns): - nineq, nz, neq, nBatch = ns - - H_ = torch.zeros(nBatch, nz + nineq, nz + nineq).type_as(Q_tilde) - H_[:, :nz, :nz] = Q_tilde - H_[:, -nineq:, -nineq:] = D_tilde - if neq > 0: - g_ = torch.cat([rx, rs], 1) - h_ = torch.cat([rz, ry], 1) - else: - g_ = torch.cat([rx, rs], 1) - h_ = rz - - H_LU = btrifact_hack(H_) - - invH_A_ = A_.transpose(1, 2).btrisolve(*H_LU) - invH_g_ = g_.btrisolve(*H_LU) - - S_ = torch.bmm(A_, invH_A_) + C_tilde - S_LU = btrifact_hack(S_) - t_ = torch.bmm(invH_g_.unsqueeze(1), A_.transpose(1, 2)).squeeze(1) - h_ - w_ = -t_.btrisolve(*S_LU) - t_ = -g_ - w_.unsqueeze(1).bmm(A_).squeeze() - v_ = t_.btrisolve(*H_LU) - - dx = v_[:, :nz] - ds = v_[:, nz:] - dz = w_[:, :nineq] - dy = w_[:, nineq:] if neq > 0 else None - - return dx, ds, dz, dy - - def solve_kkt(Q_LU, d, G, A, S_LU, rx, rs, rz, ry): """ Solve KKT equations for the affine step""" @@ -330,30 +233,41 @@ def solve_kkt(Q_LU, d, G, A, S_LU, rx, rs, rz, ry): nineq, nz, neq, nBatch = get_sizes(G, A) - invQ_rx = rx.btrisolve(*Q_LU) # Q-1 rx + invQ_rx = rx.T.lu_solve(*Q_LU).view(1, -1) # Q-1 rx if neq > 0: # A Q-1 rx - ry # G Q-1 rx + rs / d - rz h = torch.cat([invQ_rx.unsqueeze(1).bmm(A.transpose(1, 2)).squeeze(1) - ry, - invQ_rx.unsqueeze(1).bmm(G.transpose(1, 2)).squeeze(1) + rs / d - rz], 1) + invQ_rx.unsqueeze(1).bmm(G.transpose(1, 2)).squeeze(1) + rs / (d + 1e-30) - rz], 1) else: - h = invQ_rx.unsqueeze(1).bmm(G.transpose(1, 2)).squeeze(1) + rs / d - rz + h = invQ_rx.unsqueeze(1).bmm(G.transpose(1, 2)).squeeze(1) + rs / (d + 1e-30) - rz - w = -(h.btrisolve(*S_LU)) # S-1 h = + S_LU[0] = S_LU[0] + 1e-30 + Q_LU[0][0] = Q_LU[0][0] + 1e-30 + + w = -(h.T.lu_solve(*S_LU).view(1, -1)) # S-1 h = + + if np.isnan(w.sum()).item() > 0 or np.isnan(h.sum()).item() > 0: + # import pdb + # print(h) + # print(w) + # pdb.set_trace() + + h = h*0 + w = w*0 g1 = -rx - w[:, neq:].unsqueeze(1).bmm(G).squeeze(1) # -rx - GT w = -rx -GT S-1 h if neq > 0: g1 -= w[:, :neq].unsqueeze(1).bmm(A).squeeze(1) # - AT w = -AT S-1 h g2 = -rs - w[:, neq:] - dx = g1.btrisolve(*Q_LU) # Q-1 g1 = - Q-1 AT S-1 h - ds = g2 / d # g2 / d = (-rs - w) / d + dx = g1.T.lu_solve(*Q_LU).view(1, -1) # Q-1 g1 = - Q-1 AT S-1 h + ds = g2 / ( d + 1e-30) # g2 / d = (-rs - w) / d dz = w[:, neq:] dy = w[:, :neq] if neq > 0 else None return dx, ds, dz, dy - def pre_factor_kkt(Q, G, F, A): """ Perform all one-time factorizations and cache relevant matrix products""" nineq, nz, neq, nBatch = get_sizes(G, A) @@ -375,12 +289,12 @@ def pre_factor_kkt(Q, G, F, A): # See the 'Block LU factorization' part of our website # for more details. - G_invQ_GT = torch.bmm(G, G.transpose(1, 2).btrisolve(*Q_LU)) + F + G_invQ_GT = torch.bmm(G, G.transpose(1, 2).lu_solve(*Q_LU)) + F R = G_invQ_GT.clone() S_LU_pivots = torch.IntTensor(range(1, 1 + neq + nineq)).unsqueeze(0) \ .repeat(nBatch, 1).type_as(Q).int() if neq > 0: - invQ_AT = A.transpose(1, 2).btrisolve(*Q_LU) + invQ_AT = A.transpose(1, 2).lu_solve(*Q_LU) A_invQ_AT = torch.bmm(A, invQ_AT) G_invQ_AT = torch.bmm(G, invQ_AT) @@ -390,9 +304,9 @@ def pre_factor_kkt(Q, G, F, A): S_LU_11 = LU_A_invQ_AT[0] U_A_invQ_AT_inv = (P_A_invQ_AT.bmm(L_A_invQ_AT) - ).btrisolve(*LU_A_invQ_AT) + ).lu_solve(*LU_A_invQ_AT) S_LU_21 = G_invQ_AT.bmm(U_A_invQ_AT_inv) - T = G_invQ_AT.transpose(1, 2).btrisolve(*LU_A_invQ_AT) + T = G_invQ_AT.transpose(1, 2).lu_solve(*LU_A_invQ_AT) S_LU_12 = U_A_invQ_AT.bmm(T) S_LU_22 = Q.new_zeros(nBatch, nineq, nineq) S_LU_data = torch.cat((torch.cat((S_LU_11, S_LU_12), 2), @@ -412,6 +326,7 @@ def pre_factor_kkt(Q, G, F, A): # @profile def factor_kkt(S_LU, R, d): + """ Factor the U22 block that we can only do after we know D. """ nBatch, nineq = d.size() neq = S_LU[1].size(1) - nineq @@ -425,7 +340,7 @@ def factor_kkt(S_LU, R, d): # T[factor_kkt_eye] += (1. / d).view(-1) # more efficient version of these two lines in pytorch versions > 0.3.1 T = torch.zeros_like(R) - T.masked_scatter_(factor_kkt_eye, (1. / d).view(-1)) + T.masked_scatter_(factor_kkt_eye.bool(), (1. / (d + 1e-30)).view(-1)) T += R.clone() T_LU = btrifact_hack(T) @@ -451,4 +366,4 @@ def factor_kkt(S_LU, R, d): S_LU[1][:, -nineq:] = newPivotsPacked + neq # Add the new S_LU_22 block. - S_LU[0][:, -nineq:, -nineq:] = T_LU[0] + S_LU[0][:, -nineq:, -nineq:] = T_LU[0] \ No newline at end of file diff --git a/lcp_physics/physics/%%%%%%%%%%%% b/lcp_physics/physics/%%%%%%%%%%%% new file mode 100644 index 0000000..c2ce875 --- /dev/null +++ b/lcp_physics/physics/%%%%%%%%%%%% @@ -0,0 +1,30 @@ +%%%%%%%%%%%% +% delta % +%%%%%%%%%%%% + +ReducedBaseOffline +mu = [0.4, 0.6, 0.8, 1.2, 0.1]; +[u, T0] = ReducedBaseOnline(mu, 10, ANq, FN); + +T0 + +ReducedBaseOffline +mu = [1.8, 4.2, 5.7, 2.9, 0.3]; +[u, T1] = ReducedBaseOnline(mu, 10, ANq, FN); + +T1 + +%%%%%%%%%%%% +% epsilon % +%%%%%%%%%%%% + +Bi_rnge = 0.1:10; +C = []; + +for bi = 0.1:10 + mu = [0.4, 0.6, 0.8, 1.2, bi]; + [u, Ti] = ReducedBaseOnline(mu, 10, ANq, FN); + C = [C, 0.2*bi + Ti]; +end + +plot(Bi_rnge, C) \ No newline at end of file diff --git a/lcp_physics/physics/bodies.py b/lcp_physics/physics/bodies.py index 411f2c9..556fda5 100644 --- a/lcp_physics/physics/bodies.py +++ b/lcp_physics/physics/bodies.py @@ -4,6 +4,7 @@ import pygame import torch +import numpy as np from .utils import Indices, Defaults, get_tensor, cross_2d, rotation_matrix @@ -20,7 +21,6 @@ def __init__(self, pos, vel=(0, 0, 0), mass=1, restitution=Defaults.RESTITUTION, col=(255, 0, 0), thickness=1): # get base tensor to define dtype, device and layout for others self._set_base_tensor(locals().values()) - self.eps = get_tensor(eps, base_tensor=self._base_tensor) # rotation & position vectors pos = get_tensor(pos, base_tensor=self._base_tensor) @@ -50,10 +50,14 @@ def __init__(self, pos, vel=(0, 0, 0), mass=1, restitution=Defaults.RESTITUTION, self.restitution = get_tensor(restitution, base_tensor=self._base_tensor) self.forces = [] + self.idx = 0 + self.col = col self.thickness = thickness + self.name = 'NameNo' self._create_geom() + self.no_contact = set() def _set_base_tensor(self, args): """Check if any tensor provided and if so set as base tensor to @@ -79,6 +83,7 @@ def _get_ang_inertia(self, mass): def move(self, dt, update_geom_rotation=True): new_p = self.p + self.v * dt + # print(self.v) self.set_p(new_p, update_geom_rotation) def set_p(self, new_p, update_geom_rotation=True): @@ -87,23 +92,30 @@ def set_p(self, new_p, update_geom_rotation=True): self.rot = self.p[0:1] self.pos = self.p[1:] - self.geom.setPosition([self.pos[0], self.pos[1], 0.0]) - if update_geom_rotation: - # XXX sign correction - s = math.sin(-self.rot.item() / 2) - c = math.cos(-self.rot.item() / 2) - quat = [s, 0, 0, c] # Eq 2.3 - self.geom.setQuaternion(quat) + # self.geom.setPosition([self.pos[0], self.pos[1], 0.0]) + # if update_geom_rotation: + # # XXX sign correction + # s = math.sin(-self.rot.item() / 2) + # c = math.cos(-self.rot.item() / 2) + # quat = [s, 0, 0, c] # Eq 2.3 + # self.geom.setQuaternion(quat) def apply_forces(self, t): if len(self.forces) == 0: return self.v.new_zeros(len(self.v)) else: - return sum([f.force(t) for f in self.forces]) + if np.isnan(sum([f.force(t) for f in self.forces])).sum().item() > 0: + f = sum([f.force(t) for f in self.forces]) + f[:] = 0 + return f + else: + return sum([f.force(t) for f in self.forces]) def add_no_contact(self, other): - self.geom.no_contact.add(other.geom) - other.geom.no_contact.add(self.geom) + # self.geom.no_contact.add(other.geom) + # other.geom.no_contact.add(self.geom) + self.no_contact.add(other) + other.no_contact.add(self) def add_force(self, f): self.forces.append(f) @@ -116,7 +128,8 @@ def draw(self, screen, pixels_per_meter=1): class Circle(Body): def __init__(self, pos, rad, vel=(0, 0, 0), mass=1, restitution=Defaults.RESTITUTION, fric_coeff=Defaults.FRIC_COEFF, eps=Defaults.EPSILON, - col=(255, 0, 0), thickness=1): + col=(255, 0, 0), thickness=1, name = None, traj = None): + self.name = name self._set_base_tensor(locals().values()) self.rad = get_tensor(rad, base_tensor=self._base_tensor) super().__init__(pos, vel=vel, mass=mass, restitution=restitution, @@ -126,10 +139,11 @@ def _get_ang_inertia(self, mass): return mass * self.rad * self.rad / 2 def _create_geom(self): - self.geom = ode.GeomSphere(None, self.rad.item() + self.eps.item()) - self.geom.setPosition(torch.cat([self.pos, - self.pos.new_zeros(1)])) - self.geom.no_contact = set() + # self.geom = ode.GeomSphere(None, self.rad.item() + self.eps.item()) + # self.geom.setPosition(torch.cat([self.pos, + # self.pos.new_zeros(1)])) + # self.geom.no_contact = set() + pass def move(self, dt, update_geom_rotation=False): super().move(dt, update_geom_rotation=update_geom_rotation) @@ -139,7 +153,7 @@ def set_p(self, new_p, update_geom_rotation=False): def draw(self, screen, pixels_per_meter=1): center = (self.pos.detach().numpy() * pixels_per_meter).astype(int) - rad = int(self.rad.item() * pixels_per_meter) + rad = int(self.rad * pixels_per_meter) # draw radius to visualize orientation r = pygame.draw.line(screen, (0, 0, 255), center, center + [math.cos(self.rot.item()) * rad, @@ -147,7 +161,7 @@ def draw(self, screen, pixels_per_meter=1): self.thickness) # draw circle c = pygame.draw.circle(screen, self.col, center, - rad, self.thickness) + 10, self.thickness) return [c, r] @@ -161,7 +175,8 @@ class Hull(Body): """ def __init__(self, ref_point, vertices, vel=(0, 0, 0), mass=1, restitution=Defaults.RESTITUTION, fric_coeff=Defaults.FRIC_COEFF, eps=Defaults.EPSILON, - col=(255, 0, 0), thickness=1): + col=(255, 0, 0), thickness=1, name = None, traj = None): + self.name = name self._set_base_tensor(locals().values()) ref_point = get_tensor(ref_point, base_tensor=self._base_tensor) # center vertices around centroid @@ -189,15 +204,16 @@ def _get_ang_inertia(self, mass): return 1 / 6 * mass * numerator / denominator def _create_geom(self): - # find vertex furthest from centroid - max_rad = max([v.dot(v).item() for v in self.verts]) - max_rad = math.sqrt(max_rad) + # # find vertex furthest from centroid + # max_rad = max([v.dot(v).item() for v in self.verts]) + # max_rad = math.sqrt(max_rad) - # XXX Using sphere with largest vertex ray for broadphase for now - self.geom = ode.GeomSphere(None, max_rad + self.eps.item()) - self.geom.setPosition(torch.cat([self.pos, - self.pos.new_zeros(1)])) - self.geom.no_contact = set() + # # XXX Using sphere with largest vertex ray for broadphase for now + # self.geom = ode.GeomSphere(None, max_rad + self.eps.item()) + # self.geom.setPosition(torch.cat([self.pos, + # self.pos.new_zeros(1)])) + # self.geom.no_contact = set() + pass def set_p(self, new_p, update_geom_rotation=False): rot = new_p[0] - self.p[0] @@ -253,7 +269,8 @@ def draw(self, screen, draw_center=True, pixels_per_meter=1): class Rect(Hull): def __init__(self, pos, dims, vel=(0, 0, 0), mass=1, restitution=Defaults.RESTITUTION, fric_coeff=Defaults.FRIC_COEFF, eps=Defaults.EPSILON, - col=(255, 0, 0), thickness=1): + col=(255, 0, 0), thickness=1, name="env", traj = None): + self.name = name self._set_base_tensor(locals().values()) self.dims = get_tensor(dims, base_tensor=self._base_tensor) pos = get_tensor(pos, base_tensor=self._base_tensor) @@ -270,10 +287,11 @@ def _get_ang_inertia(self, mass): return mass * torch.sum(self.dims ** 2) / 12 def _create_geom(self): - self.geom = ode.GeomBox(None, torch.cat([self.dims + 2 * self.eps.item(), - self.dims.new_ones(1)])) - self.geom.setPosition(torch.cat([self.pos, self.pos.new_zeros(1)])) - self.geom.no_contact = set() + # self.geom = ode.GeomBox(None, torch.cat([self.dims + 2 * self.eps.item(), + # self.dims.new_ones(1)])) + # self.geom.setPosition(torch.cat([self.pos, self.pos.new_zeros(1)])) + # self.geom.no_contact = set() + pass def rotate_verts(self, rot): rot_mat = rotation_matrix(rot) @@ -299,3 +317,105 @@ def draw(self, screen, pixels_per_meter=1): p = super().draw(screen, pixels_per_meter=pixels_per_meter, draw_center=False) return [l1, l2] + p + + +class NonConvex(Body): + """Body's position will not necessarily match reference point. + Reference point is used as a world frame reference for setting the position + of vertices, which maintain the world frame position that was passed in. + After vertices are positioned in world frame using reference point, centroid + of hull is calculated and the vertices' representation is adjusted to the + centroid's frame. Object position is set to centroid. + """ + def __init__(self, ref_point, vertices, vel=(0, 0, 0), mass=1, restitution=Defaults.RESTITUTION, + fric_coeff=Defaults.FRIC_COEFF, eps=Defaults.EPSILON, + col=(255, 0, 0), thickness=1, name = None, traj = None): + self.name = name + self._set_base_tensor(locals().values()) + ref_point = get_tensor(ref_point, base_tensor=self._base_tensor) + # center vertices around centroid + verts = [get_tensor(v, base_tensor=self._base_tensor) for v in vertices] + assert len(verts) > 2 and self._is_clockwise(verts) + centroid = self._get_centroid(verts) + self.verts = [v for v in verts] + # center position at centroid + pos = ref_point + # store last separating edge for SAT + self.last_sat_idx = 0 + super().__init__(pos, vel=vel, mass=mass, restitution=restitution, + fric_coeff=fric_coeff, eps=eps, col=col, thickness=thickness) + + def _get_ang_inertia(self, mass): + numerator = 0 + denominator = 0 + for i in range(len(self.verts)): + v1 = self.verts[i] + v2 = self.verts[(i+1) % len(self.verts)] + norm_cross = torch.norm(cross_2d(v2, v1)) + numerator = numerator + norm_cross * \ + (torch.dot(v1, v1) + torch.dot(v1, v2) + torch.dot(v2, v2)) + denominator = denominator + norm_cross + return 1 / 6 * mass * numerator / denominator + + def _create_geom(self): + # # find vertex furthest from centroid + # max_rad = max([v.dot(v).item() for v in self.verts]) + # max_rad = math.sqrt(max_rad) + + # # XXX Using sphere with largest vertex ray for broadphase for now + # self.geom = ode.GeomSphere(None, max_rad + self.eps.item()) + # self.geom.setPosition(torch.cat([self.pos, + # self.pos.new_zeros(1)])) + # self.geom.no_contact = set() + pass + + def set_p(self, new_p, update_geom_rotation=False): + rot = new_p[0] - self.p[0] + if rot.item() != 0: + self.rotate_verts(rot) + super().set_p(new_p, update_geom_rotation=update_geom_rotation) + + def move(self, dt, update_geom_rotation=False): + super().move(dt, update_geom_rotation=update_geom_rotation) + + def rotate_verts(self, rot): + rot_mat = rotation_matrix(rot) + for i in range(len(self.verts)): + self.verts[i] = rot_mat.matmul(self.verts[i]) + + @staticmethod + def _get_centroid(verts): + numerator = 0 + denominator = 0 + for i in range(len(verts)): + v1 = verts[i] + v2 = verts[(i + 1) % len(verts)] + cross = cross_2d(v2, v1) + numerator = numerator + cross * (v1 + v2) + denominator = denominator + cross / 2 + return 1 / 6 * numerator / denominator + + @staticmethod + def _is_clockwise(verts): + total = 0 + for i in range(len(verts)): + v1 = verts[i] + v2 = verts[(i+1) % len(verts)] + total = total + ((v2[X] - v1[X]) * (v2[Y] + v1[Y])).item() + return total < 0 + + def draw(self, screen, draw_center=True, pixels_per_meter=1): + # vertices in global frame + pts = [(v + self.pos).detach().cpu().numpy() * pixels_per_meter + for v in self.verts] + + # draw hull + p = pygame.draw.polygon(screen, self.col, pts, self.thickness) + # draw center + if draw_center: + c_pos = (self.pos.data.numpy() * pixels_per_meter).astype(int) + c = pygame.draw.circle(screen, (0, 0, 255), c_pos, 2) + return [p, c] + else: + return [p] + diff --git a/lcp_physics/physics/constraints.py b/lcp_physics/physics/constraints.py index 3f6e1c2..3b6bf28 100644 --- a/lcp_physics/physics/constraints.py +++ b/lcp_physics/physics/constraints.py @@ -65,7 +65,7 @@ def __init__(self, body1, body2): self.rot1 = self.pos.new_tensor(0) self.rot2 = None self.pos2 = self.pos - self.body2.pos - self.rot2 = self.body2.p[0] - self.body1.p[0] # inverted sign? + self.rot2 = self.body1.p[0] - self.body2.p[0] # inverted sign? def J(self): J1 = torch.cat([torch.cat([-self.pos1[Y:Y+1], self.pos1[X:X+1]]).unsqueeze(1), @@ -79,6 +79,12 @@ def J(self): def move(self, dt): self.update_pos() + def stabilize(self): + pass + # self.body1.set_p(torch.cat([self.body1.p[0:1],self.pos.view(2)])) + # if self.body2 is not None: + # self.body2.set_p(torch.cat([self.body2.p[0:1],self.pos2.view(2)+self.body2.pos])) + def update_pos(self): self.pos = self.body1.pos self.pos1 = self.pos - self.body1.pos @@ -199,6 +205,9 @@ def update_pos(self): self.pos1 = polar_to_cart(self.r1, self.rot1) self.pos = self.body1.pos + self.pos1 + def stabilize(self): + pass + def draw(self, screen, pixels_per_meter=1): pos = (self.pos.detach().numpy() * pixels_per_meter).astype(int) return [pygame.draw.circle(screen, (0, 255, 0), pos + 1, 5, 1), diff --git a/lcp_physics/physics/contacts.py b/lcp_physics/physics/contacts.py index 99f3c93..c5b271e 100644 --- a/lcp_physics/physics/contacts.py +++ b/lcp_physics/physics/contacts.py @@ -4,9 +4,11 @@ import torch -from .bodies import Circle +from .bodies import Circle, Hull, NonConvex from .utils import Indices, Defaults, left_orthogonal - +import pdb +import numpy as np +import math X = Indices.X Y = Indices.Y @@ -20,15 +22,15 @@ def __init__(self): def __call__(self, *args, **kwargs): raise NotImplementedError - class OdeContactHandler(ContactHandler): - def __call__(self, args, geom1, geom2): - if geom1 in geom2.no_contact: + def __call__(self, args, geom1, geom2, gamma = 1): + # raise NotImplementedError + if geom1.body_ref in geom2.no_contact: return world = args[0] base_tensor = world.bodies[0].p - - contacts = ode.collide(geom1, geom2) + # pdb.set_trace() + contacts = ode.collide(geom1.body_ref, geom2.body_ref) for c in contacts: point, normal, penetration, geom1, geom2 = c.getContactGeomParams() # XXX Simple disambiguation of 3D repetition of contacts @@ -44,7 +46,451 @@ def __call__(self, args, geom1, geom2): p2 = point - base_tensor.new_tensor(geom2.getPosition()) world.contacts.append(((normal, p1[:DIM], p2[:DIM], penetration), geom1.body, geom2.body)) - # world.contacts_debug = world.contacts # XXX + + world.contacts_debug = world.contacts # XXX + + +class NonConvexContactHandler(ContactHandler): + """Differentiable contact handler, operations to calculate contact manifold + are done in autograd. + """ + def __init__(self): + self.debug_callback = OdeContactHandler() + + def __call__(self, args, geom1, geom2, gamma = 1): + # self.debug_callback(args, geom1, geom2) + + if geom1.body_ref in geom2.no_contact: + return + world = args[0] + + b1 = world.bodies[geom1.body] + b2 = world.bodies[geom2.body] + is_circle_g1 = isinstance(b1, Circle) + is_circle_g2 = isinstance(b2, Circle) + is_hull_g1 = isinstance(b1, Hull) + is_hull_g2 = isinstance(b2, Hull) + is_ncvx_g1 = isinstance(b1, NonConvex) + is_ncvx_g2 = isinstance(b2, NonConvex) + + if is_circle_g1 and is_circle_g2: + # Simple circle vs circle + r = b1.rad + b2.rad + normal = b1.pos - b2.pos + dist = normal.norm() + penetration = r - dist + if penetration.item() < -world.eps: + return + normal = normal / dist + + # contact points on surface of object if not interpenetrating, + # otherwise its the point midway between two objects inside of them + p1 = -normal * b1.rad + p2 = normal * b2.rad + if penetration > 0: + p1 = p1 + normal * penetration / 2 # p1 = -normal * (b1.rad - penetration / 2) + p2 = p2 - normal * penetration / 2 # p2 = normal * (b2.rad - penetration / 2) + + pts = [[normal, p1, p2, penetration]] + elif is_circle_g1 or is_circle_g2: + if is_circle_g2: + # set circle to b1 + b1, b2 = b2, b1 + + if is_ncvx_g1 or is_ncvx_g2: + # finger in objec frame + best_dist, best_pt1, best_pt2, best_normal = self.resolve_ncvx(world.facets, b1, b2) + if is_circle_g2: + # flip back values for circle as g2 + best_normal = -best_normal + best_pt1, best_pt2 = best_pt2, best_pt1 + pts = [[best_normal.view(2).double(), best_pt1.view(2).double(), best_pt2.view(2).double(), -torch.tensor(best_dist)]] + else: + # Shallow penetration with GJK + test_point = b1.pos - b2.pos + simplex = [random.choice(b2.verts)] + while True: + closest, ids_used = self.get_closest(test_point, simplex) + if len(ids_used) == 3: + break + if len(ids_used) == 2: + # use orthogonal when closest is in segment + search_dir = left_orthogonal(simplex[ids_used[0]] - simplex[ids_used[1]]) + if search_dir.dot(test_point - simplex[ids_used[0]]).item() < 0: + search_dir = -search_dir + else: + search_dir = test_point - closest + if search_dir[0].item() == 0 and search_dir[1].item() == 0: + break + support, _ = self.get_support(b2.verts, search_dir) + if support in set(simplex): + break + simplex = [simplex[idx] for idx in ids_used] # remove unused points + simplex.append(support) + if len(ids_used) < 3: + best_pt2 = closest + closest = closest + b2.pos + best_pt1 = closest - b1.pos + best_dist = torch.norm(closest - b1.pos) - b1.rad + if best_dist.item() > world.eps: + print('this should not be happening look at contacts.py') + return + # normal points from closest point to circle center + best_normal = -best_pt1 / torch.norm(best_pt1) + else: + # SAT for circle vs hull if deep penetration + best_dist = torch.tensor(-1e5) + num_verts = len(b2.verts) + start_edge = b2.last_sat_idx + for i in range(start_edge, num_verts + start_edge): + idx = i % num_verts + edge = b2.verts[(idx+1) % num_verts] - b2.verts[idx] + edge_norm = edge.norm() + normal = left_orthogonal(edge) / edge_norm + # adjust to hull1's frame + center = b1.pos - b2.pos + # get distance from circle point to edge + dist = normal.dot(center - b2.verts[idx]) - b1.rad + + if dist.item() > best_dist.item(): + b2.last_sat_idx = idx + if dist.item() > world.eps: + # exit early if separating axis found + return + best_dist = dist + best_normal = normal + best_pt2 = center + normal * -(dist + b1.rad) + best_pt1 = best_pt2 + b2.pos - b1.pos + + if is_circle_g2: + # flip back values for circle as g2 + best_normal = -best_normal + best_pt1, best_pt2 = best_pt2, best_pt1 + pts = [[best_normal, best_pt1, best_pt2, -best_dist]] + else: + # SAT for hull x hull contact + # TODO Optimize for rectangle vs rectangle? + contact1 = self.test_separations(b1, b2, eps=0.1) + b1.last_sat_idx = contact1[6] + if contact1[0].item() > 0.1: + return + contact2 = self.test_separations(b2, b1, eps=0.1) + b2.last_sat_idx = contact2[6] + if contact2[0].item() > 0.1: + return + if contact2[0].item() > contact1[0].item(): + normal = -contact2[3] + half_edge_norm = contact2[5] / 2 + ref_edge_idx = contact2[6] + incident_vertex_idx = contact2[4] + incident_edge_idx = self.get_incident_edge(normal, b1, incident_vertex_idx) + incident_verts = [b1.verts[incident_edge_idx], + b1.verts[(incident_edge_idx + 1) % len(b1.verts)]] + incident_verts = [v + b1.pos - b2.pos for v in incident_verts] + clip_plane = left_orthogonal(normal) + clipped_verts = self.clip_segment_to_line(incident_verts, clip_plane, + half_edge_norm) + if len(clipped_verts) < 2: + return + clipped_verts = self.clip_segment_to_line(clipped_verts, -clip_plane, + half_edge_norm) + pts = [] + for v in clipped_verts: + dist = normal.dot(v - b2.verts[ref_edge_idx]) + if dist.item() <= 0.1: + pt1 = v + normal * -dist + pt2 = pt1 + b2.pos - b1.pos + pts.append([normal, pt2, pt1, -dist]) + else: + normal = -contact1[3] + half_edge_norm = contact1[5] / 2 + ref_edge_idx = contact1[6] + incident_vertex_idx = contact1[4] + incident_edge_idx = self.get_incident_edge(normal, b2, incident_vertex_idx) + incident_verts = [b2.verts[incident_edge_idx], + b2.verts[(incident_edge_idx+1) % len(b2.verts)]] + incident_verts = [v + b2.pos - b1.pos for v in incident_verts] + clip_plane = left_orthogonal(normal) + clipped_verts = self.clip_segment_to_line(incident_verts, clip_plane, + half_edge_norm) + if len(clipped_verts) < 2: + return + clipped_verts = self.clip_segment_to_line(clipped_verts, -clip_plane, + half_edge_norm) + pts = [] + for v in clipped_verts: + dist = normal.dot(v - b1.verts[ref_edge_idx]) + # import pdb + # pdb.set_trace() + if dist.item() <= 0.1: + pt1 = v + normal * -dist + pt2 = pt1 + b1.pos - b2.pos + pts.append([-normal, pt1, pt2, -dist]) + + for p in pts: + world.contacts.append([p, geom1.body, geom2.body]) + + # smooth contact hack + for i, contact in enumerate(world.contacts): + # at 0 penetration (objects exact contact) we want p percent of contact normal. + # compute adjustment with inverse of sigmoid + p = torch.tensor(0.97) + delta = torch.log(p / (1 - p)) + + # contact[0] = (normal, pt1, pt2, penetration_dist) + # print('MESSAGE !!! ') + gamma = world.gamma + contact[0][0] = contact[0][0] * torch.sigmoid(gamma*contact[0][3] + delta) + if np.isnan(sum(contact[0][0])).sum().item() > 0: + contact[0][0][:][:] = 0 + + world.contacts_debug = world.contacts # XXX + + @staticmethod + def get_support(points, direction): + best_point = None + best_norm = -1. + + found = True + for i, p in enumerate(points): + cur_norm = p.dot(direction).item() + if (cur_norm >= best_norm) or found: + best_point = p + best_idx = i + best_norm = cur_norm + found = False + + return best_point, best_idx + + + @staticmethod + def test_separations(hull1, hull2, eps=0): + verts1, verts2 = hull1.verts, hull2.verts + num_verts = len(verts1) + best_dist = torch.tensor(-1e10) + best_normal = None + best_vertex = -1 + start_edge = hull1.last_sat_idx + for i in range(start_edge, num_verts + start_edge): + idx = i % num_verts + edge = verts1[(idx+1) % num_verts] - verts1[idx] + edge_norm = edge.norm() + normal = left_orthogonal(edge) / edge_norm + support_point, support_idx = DiffContactHandler.get_support(verts2, -normal) + # adjust to hull1's frame + support_point = support_point + hull2.pos - hull1.pos + # get distance from support point to edge + dist = normal.dot(support_point - verts1[idx]) + + if dist.item() > best_dist.item(): + if dist.item() > 0.1: + # exit early if separating axis found + return dist, None, None, None, None, None, idx + best_dist = dist + best_normal = normal + best_pt1 = support_point + normal * -dist + best_pt2 = best_pt1 + hull1.pos - hull2.pos + best_vertex = support_idx + best_edge_norm = edge_norm + best_edge = idx + return best_dist, best_pt1, best_pt2, -best_normal, \ + best_vertex, best_edge_norm, best_edge + + @staticmethod + def test_separations_all(hull1, hull2, eps=0): + verts1, verts2 = hull1.verts, hull2.verts + num_verts = len(verts1) + + # saves a list + best_dist = [] + best_normal = [] + best_vertex = [] + best_edge_norm = [] + best_edge = [] + best_pt1 = [] + best_pt2 = [] + + start_edge = hull1.last_sat_idx + for i in range(start_edge, num_verts + start_edge): + idx = i % num_verts + edge = verts1[(idx+1) % num_verts] - verts1[idx] + edge_norm = edge.norm() + normal = left_orthogonal(edge) / edge_norm + support_point, support_idx = DiffContactHandler.get_support(verts2, -normal) + # adjust to hull1's frame + support_point = support_point + hull2.pos - hull1.pos + # get distance from support point to edge + dist = normal.dot(support_point - verts1[idx]) + + # if dist.item() > best_dist.item(): + # if dist.item() > eps: + # # exit early if separating axis found + # return dist, None, None, None, None, None, idx + best_dist.append(dist) + best_normal.append(normal) + best_pt1.append(support_point + normal * -dist) + best_pt2.append(best_pt1 + hull1.pos - hull2.pos) + best_vertex.append(support_idx) + best_edge_norm.append(edge_norm) + best_edge.append(idx) + return best_dist, best_pt1, best_pt2, -best_normal, \ + best_vertex, best_edge_norm, best_edge + + @staticmethod + def get_incident_edge(ref_normal, inc_hull, inc_vertex): + inc_verts = inc_hull.verts + # two possible incident edges (pointing to and from incident vertex) + edges = [(inc_vertex-1) % len(inc_verts), inc_vertex] + min_dot = 1e10 + best_edge = -1 + for i in edges: + edge = inc_verts[(i+1) % len(inc_verts)] - inc_verts[i] + edge_norm = edge.norm() + inc_normal = left_orthogonal(edge) / edge_norm + dot = ref_normal.dot(inc_normal).item() + if dot < min_dot: + min_dot = dot + best_edge = i + return best_edge + + @staticmethod + def clip_segment_to_line(verts, normal, offset): + clipped_verts = [] + + # Calculate the distance of end points to the line + distance0 = normal.dot(verts[0]) + offset + distance1 = normal.dot(verts[1]) + offset + + # If the points are behind the plane + if distance0.item() >= 0.0: + clipped_verts.append(verts[0]) + if distance1.item() >= 0.0: + clipped_verts.append(verts[1]) + + # If the points are on different sides of the plane + if distance0.item() * distance1.item() < 0.0 or len(clipped_verts) < 2: + # Find intersection point of edge and plane + interp = distance0 / (distance0 - distance1) + + # Vertex is hitting edge. + cv = verts[0] + interp * (verts[1] - verts[0]) + clipped_verts.append(cv) + + return clipped_verts + + @staticmethod + def get_closest(point, simplex): + if len(simplex) == 1: + return simplex[0], [0] + elif len(simplex) == 2: + u, v = DiffContactHandler.get_barycentric_coords(point, simplex) + if u.item() <= 0: + return simplex[1], [1] + elif v.item() <= 0: + return simplex[0], [0] + else: + return u * simplex[0] + v * simplex[1], [0, 1] + elif len(simplex) == 3: + uAB, vAB = DiffContactHandler.get_barycentric_coords(point, simplex[0:2]) + uBC, vBC = DiffContactHandler.get_barycentric_coords(point, simplex[1:]) + uCA, vCA = DiffContactHandler.get_barycentric_coords(point, [simplex[2], simplex[0]]) + uABC, vABC, wABC = DiffContactHandler.get_barycentric_coords(point, simplex) + + if vAB.item() <= 0 and uCA.item() <= 0: + return simplex[0], [0] + elif vBC.item() <= 0 and uAB.item() <= 0: + return simplex[1], [1] + elif vCA.item() <= 0 and uBC.item() <= 0: + return simplex[2], [2] + elif uAB.item() > 0 and vAB.item() > 0 and wABC.item() <= 0: + return uAB * simplex[0] + vAB * simplex[1], [0, 1] + elif uBC.item() > 0 and vBC.item() > 0 and uABC.item() <= 0: + return uBC * simplex[1] + vBC * simplex[2], [1, 2] + elif uCA.item() > 0 and vCA.item() > 0 and vABC.item() <= 0: + return uCA * simplex[2] + vCA * simplex[0], [2, 0] + elif uABC.item() > 0 and vABC.item() > 0 and wABC.item() > 0: + return point, [0, 1, 2] + else: + print(uAB, vAB, uBC, vBC, uCA, vCA, uABC, vABC, wABC) + raise ValueError('Point does not satisfy any condition in get_closest()') + else: + raise ValueError('Simplex should not have more than 3 points in GJK.') + + @staticmethod + def resolve_ncvx(facets, b1, b2): + inside = False + irot = torch.tensor([[torch.cos(b2.p[0]),torch.sin(b2.p[0])],[-torch.sin(b2.p[0]),torch.cos(b2.p[0])]]).view(2,2).float() + b1_wrt_b2 = torch.matmul(irot,b1.pos.float()-b2.pos.float()).view(2,1) + + # pdb.set_trace() + for f in facets: + # checks if the ray from b1 in y+ intercepts the facet (Jordan Polygon Theorem) + if (b1_wrt_b2[0] <= max(f[0][0],f[1][0])) and (b1_wrt_b2[0] >= min(f[0][0],f[1][0])): + if f[0][0] == f[1][0]: + if (b1_wrt_b2[1] <= f[0][1]): + inside = True - inside + if (b1_wrt_b2[1] <= f[1][1]): + inside = True - inside + else: + if (b1_wrt_b2[1] <= ((b1_wrt_b2[0]-f[1][0])*(f[0][1]-f[1][1])/(f[0][0]-f[1][0]) + f[1][1])): + inside = True - inside + # pdb.set_trace() + # finds closest facet + best_dist = 1e6 + for f in facets: + v1 = torch.tensor(f[1].copy()).view(2,1)-torch.tensor(f[0]).view(2,1) + v2 = b1_wrt_b2 - torch.tensor(f[0].copy()).view(2,1) + + d1 = torch.matmul(v1.view(1,2),v2)/torch.norm(v1) + if d1.item() <= 0: + p1 = torch.tensor(f[0].copy()).view(2,1) + elif d1.item() >= torch.norm(torch.tensor(f[1]).view(2,1) - torch.tensor(f[0]).view(2,1)).item(): + p1 = torch.tensor(f[1].copy()).view(2,1) + else: + p1 = torch.tensor(f[0].copy()).view(2,1) + d1.view(1,1)*v1/torch.norm(v1) + + if torch.norm(p1 - b1_wrt_b2).item() < best_dist: + best_dist = torch.norm(p1 - b1_wrt_b2).item() + correction = (p1 - b1_wrt_b2).view(2,1) + bp1 = p1 + normal = torch.tensor([0,0]).float().view(2,1) + normal[0] = f[0][1].copy()-f[1][1].copy() + normal[1] = f[1][0].copy()-f[0][0].copy() + + if inside: + best_dist = -best_dist + + pos_f0 = torch.matmul(irot.transpose(0,1),torch.tensor(facets[0][0].copy()).float().view(2,1)).view(2,1) + b2.pos.float().view(2,1) + best_pt1 = 0*b1_wrt_b2 + best_pt2 = torch.matmul(irot.transpose(0,1),bp1.float()).view(2,1) + + best_normal = torch.matmul(irot.transpose(0,1),normal.float().view(2,1)).view(2,1)/(1e-6 + torch.norm(normal)) + + return best_dist, best_pt1, best_pt2, best_normal + + @staticmethod + def get_barycentric_coords(point, verts): + if len(verts) == 2: + diff = verts[1] - verts[0] + diff_norm = torch.norm(diff) + normalized_diff = diff / diff_norm + u = torch.dot(verts[1] - point, normalized_diff) / diff_norm + v = torch.dot(point - verts[0], normalized_diff) / diff_norm + return u, v + elif len(verts) == 3: + # TODO Area method instead of LinAlg + M = torch.cat([ + torch.cat([verts[0], verts[0].new_ones(1)]).unsqueeze(1), + torch.cat([verts[1], verts[1].new_ones(1)]).unsqueeze(1), + torch.cat([verts[2], verts[2].new_ones(1)]).unsqueeze(1), + ], dim=1) + invM = torch.inverse(M) + uvw = torch.matmul(invM, torch.cat([point, point.new_ones(1)]).unsqueeze(1)) + return uvw + else: + raise ValueError('Barycentric coords only works for 2 or 3 points') + + class DiffContactHandler(ContactHandler): @@ -54,10 +500,10 @@ class DiffContactHandler(ContactHandler): def __init__(self): self.debug_callback = OdeContactHandler() - def __call__(self, args, geom1, geom2): + def __call__(self, args, geom1, geom2, gamma = 1): # self.debug_callback(args, geom1, geom2) - if geom1 in geom2.no_contact: + if geom1.body_ref in geom2.no_contact: return world = args[0] @@ -74,9 +520,16 @@ def __call__(self, args, geom1, geom2): if penetration.item() < -world.eps: return normal = normal / dist - p1 = -normal * (b1.rad - penetration / 2) - p2 = normal * (b2.rad - penetration / 2) - pts = [(normal, p1, p2, penetration)] + + # contact points on surface of object if not interpenetrating, + # otherwise its the point midway between two objects inside of them + p1 = -normal * b1.rad + p2 = normal * b2.rad + if penetration > 0: + p1 = p1 + normal * penetration / 2 # p1 = -normal * (b1.rad - penetration / 2) + p2 = p2 - normal * penetration / 2 # p2 = normal * (b2.rad - penetration / 2) + + pts = [[normal, p1, p2, penetration]] elif is_circle_g1 or is_circle_g2: if is_circle_g2: # set circle to b1 @@ -109,12 +562,13 @@ def __call__(self, args, geom1, geom2): best_pt1 = closest - b1.pos best_dist = torch.norm(closest - b1.pos) - b1.rad if best_dist.item() > world.eps: + print('this should not be happening look at contacts.py') return # normal points from closest point to circle center best_normal = -best_pt1 / torch.norm(best_pt1) else: # SAT for circle vs hull if deep penetration - best_dist = torch.tensor(-1e10) + best_dist = torch.tensor(-1e5) num_verts = len(b2.verts) start_edge = b2.last_sat_idx for i in range(start_edge, num_verts + start_edge): @@ -141,17 +595,17 @@ def __call__(self, args, geom1, geom2): # flip back values for circle as g2 best_normal = -best_normal best_pt1, best_pt2 = best_pt2, best_pt1 - pts = [(best_normal, best_pt1, best_pt2, -best_dist)] + pts = [[best_normal, best_pt1, best_pt2, -best_dist]] else: # SAT for hull x hull contact # TODO Optimize for rectangle vs rectangle? - contact1 = self.test_separations(b1, b2, eps=world.eps) + contact1 = self.test_separations(b1, b2, eps=0.1) b1.last_sat_idx = contact1[6] - if contact1[0].item() > world.eps: + if contact1[0].item() > 0.1: return - contact2 = self.test_separations(b2, b1, eps=world.eps) + contact2 = self.test_separations(b2, b1, eps=0.1) b2.last_sat_idx = contact2[6] - if contact2[0].item() > world.eps: + if contact2[0].item() > 0.1: return if contact2[0].item() > contact1[0].item(): normal = -contact2[3] @@ -172,10 +626,10 @@ def __call__(self, args, geom1, geom2): pts = [] for v in clipped_verts: dist = normal.dot(v - b2.verts[ref_edge_idx]) - if dist.item() <= world.eps: + if dist.item() <= 0.1: pt1 = v + normal * -dist pt2 = pt1 + b2.pos - b1.pos - pts.append((normal, pt2, pt1, -dist)) + pts.append([normal, pt2, pt1, -dist]) else: normal = -contact1[3] half_edge_norm = contact1[5] / 2 @@ -195,27 +649,51 @@ def __call__(self, args, geom1, geom2): pts = [] for v in clipped_verts: dist = normal.dot(v - b1.verts[ref_edge_idx]) - if dist.item() <= world.eps: + # import pdb + # pdb.set_trace() + if dist.item() <= 0.1: pt1 = v + normal * -dist pt2 = pt1 + b1.pos - b2.pos - pts.append((-normal, pt1, pt2, -dist)) + pts.append([-normal, pt1, pt2, -dist]) for p in pts: - world.contacts.append((p, geom1.body, geom2.body)) - # world.contacts_debug = world.contacts # XXX + world.contacts.append([p, geom1.body, geom2.body]) + + # smooth contact hack + for i, contact in enumerate(world.contacts): + # at 0 penetration (objects exact contact) we want p percent of contact normal. + # compute adjustment with inverse of sigmoid + p = torch.tensor(0.97) + delta = torch.log(p / (1 - p)) + + # contact[0] = (normal, pt1, pt2, penetration_dist) + # print('MESSAGE !!! ') + gamma = world.gamma + contact[0][0] = contact[0][0] * torch.sigmoid(gamma*contact[0][3] + delta) + if np.isnan(sum(contact[0][0])).sum().item() > 0: + contact[0][0][:][:] = 0 + if sum(contact[0][0]).sum().item() > 100: + contact[0][0][:][:] = 0 + + world.contacts_debug = world.contacts # XXX @staticmethod def get_support(points, direction): best_point = None best_norm = -1. + + found = True for i, p in enumerate(points): cur_norm = p.dot(direction).item() - if cur_norm >= best_norm: + if (cur_norm >= best_norm) or found: best_point = p best_idx = i best_norm = cur_norm + found = False + return best_point, best_idx + @staticmethod def test_separations(hull1, hull2, eps=0): verts1, verts2 = hull1.verts, hull2.verts @@ -236,7 +714,7 @@ def test_separations(hull1, hull2, eps=0): dist = normal.dot(support_point - verts1[idx]) if dist.item() > best_dist.item(): - if dist.item() > eps: + if dist.item() > 0.1: # exit early if separating axis found return dist, None, None, None, None, None, idx best_dist = dist @@ -249,6 +727,46 @@ def test_separations(hull1, hull2, eps=0): return best_dist, best_pt1, best_pt2, -best_normal, \ best_vertex, best_edge_norm, best_edge + @staticmethod + def test_separations_all(hull1, hull2, eps=0): + verts1, verts2 = hull1.verts, hull2.verts + num_verts = len(verts1) + + # saves a list + best_dist = [] + best_normal = [] + best_vertex = [] + best_edge_norm = [] + best_edge = [] + best_pt1 = [] + best_pt2 = [] + + start_edge = hull1.last_sat_idx + for i in range(start_edge, num_verts + start_edge): + idx = i % num_verts + edge = verts1[(idx+1) % num_verts] - verts1[idx] + edge_norm = edge.norm() + normal = left_orthogonal(edge) / edge_norm + support_point, support_idx = DiffContactHandler.get_support(verts2, -normal) + # adjust to hull1's frame + support_point = support_point + hull2.pos - hull1.pos + # get distance from support point to edge + dist = normal.dot(support_point - verts1[idx]) + + # if dist.item() > best_dist.item(): + # if dist.item() > eps: + # # exit early if separating axis found + # return dist, None, None, None, None, None, idx + best_dist.append(dist) + best_normal.append(normal) + best_pt1.append(support_point + normal * -dist) + best_pt2.append(best_pt1 + hull1.pos - hull2.pos) + best_vertex.append(support_idx) + best_edge_norm.append(edge_norm) + best_edge.append(idx) + return best_dist, best_pt1, best_pt2, -best_normal, \ + best_vertex, best_edge_norm, best_edge + @staticmethod def get_incident_edge(ref_normal, inc_hull, inc_vertex): inc_verts = inc_hull.verts diff --git a/lcp_physics/physics/engines.py b/lcp_physics/physics/engines.py index b720e47..3f74b89 100644 --- a/lcp_physics/physics/engines.py +++ b/lcp_physics/physics/engines.py @@ -1,11 +1,7 @@ -""" -Author: Filipe de Avila Belbute Peres -Based on: M. B. Cline, Rigid body simulation with contact and constraints, 2002 -""" - import torch from lcp_physics.lcp.lcp import LCPFunction +from lcp_physics.lcp.lemkelcp import LemkeLCP class Engine: @@ -50,7 +46,7 @@ def solve_dynamics(self, world, dt): else: # Solve Mixed LCP (Kline 2.7.2) Jc = world.Jc() - v = torch.matmul(Jc, world.get_v()) * world.restitutions() + v = torch.matmul(Jc, world.get_v()) * 0*world.restitutions() M = world.M().unsqueeze(0) if neq > 0: b = Je.new_zeros(Je.size(0)).unsqueeze(0) @@ -72,8 +68,8 @@ def solve_dynamics(self, world, dt): F[:, -mu.size(1):, mu.size(2):mu.size(2) + E.size(1)] = \ -E.transpose(1, 2) h = torch.cat([v, v.new_zeros(v.size(0), Jf.size(1) + mu.size(1))], 1) - - x = -self.lcp_solver(max_iter=self.max_iter, verbose=-1)(M, u, G, h, Je, b, F) + solver = self.lcp_solver() + x = -solver.apply(M, u, G, h, Je, b, F) new_v = x[:world.vec_len * len(world.bodies)].squeeze(0) return new_v @@ -87,7 +83,7 @@ def post_stabilization(self, world): ge = torch.matmul(Je, v) gc = None if Jc is not None: - gc = torch.matmul(Jc, v) + torch.matmul(Jc, v) * -world.restitutions() + gc = torch.matmul(Jc, v) + torch.matmul(Jc, v) * -0*world.restitutions() u = torch.cat([Je.new_zeros(Je.size(1)), ge]) if Jc is None: @@ -111,6 +107,117 @@ def post_stabilization(self, world): M = M.unsqueeze(0) v = v.unsqueeze(0) F = Jc.new_zeros(Jc.size(1), Jc.size(1)).unsqueeze(0) - x = self.lcp_solver()(M, h, Jc, v, Je, b, F) + + solver = self.lcp_solver() + x = solver.apply(M, h, Jc, v, Je, b, F) dp = -x[:M.size(0)] return dp + + +class LemkeEngine(Engine): + """Engine that uses the primal dual interior point method LCP solver. + """ + def __init__(self, max_iter=10): + # self.lcp_solver = LCPFunction + self.lcp_solver = LemkeLCP() + self.cached_inverse = None + self.max_iter = max_iter + + # @profile + def solve_dynamics(self, world, dt): + t = world.t + Je = world.Je() + neq = Je.size(0) if Je.ndimension() > 0 else 0 + + f = world.apply_forces(t) + u = torch.matmul(world.M(), world.get_v()) + dt * f + if neq > 0: + u = torch.cat([u, u.new_zeros(neq)]) + if not world.contacts: + # No contact constraints, no complementarity conditions + if neq > 0: + P = torch.cat([torch.cat([world.M(), -Je.t()], dim=1), + torch.cat([Je, Je.new_zeros(neq, neq)], + dim=1)]) + else: + P = world.M() + if self.cached_inverse is None: + inv = torch.inverse(P) + if world.static_inverse: + self.cached_inverse = inv + else: + inv = self.cached_inverse + x = torch.matmul(inv, u) # Kline Eq. 2.41 + else: + M = world.M() + Jc = -world.Jc() + Jf = world.Jf() + Je = world.Je() + E = world.E() + mu = world.mu() + + u = torch.cat(((M.float() @ world.get_v().float()) + dt * f.float(), torch.zeros(Je.shape[0]))) + P = torch.cat([torch.cat([M.float(), -Je.t().float()], dim=1), + torch.cat([Je.float(), torch.zeros((Je.shape[0], Je.shape[0]))], dim=1)]) + + Q = torch.cat((-Jc.t().float(), -Jf.t().float(), torch.zeros((Jc.shape[1], E.shape[1]))), dim=1) + Q = torch.cat((Q, torch.zeros((Je.shape[0], Q.shape[1])))) + + R = torch.cat((Jc.float(), Jf.float(), torch.zeros((mu.shape[0], Jc.shape[1])))) + R = torch.cat((R, torch.zeros((R.shape[0], Je.shape[0]))), dim=1) + + S = torch.zeros((Jc.shape[0], Q.shape[1])) + S = torch.cat( + (S, torch.cat((torch.zeros((Jf.shape[0], mu.shape[1])), torch.zeros((Jf.shape[0], Jf.shape[0])), E.float()), dim=1))) + S = torch.cat((S, torch.cat((mu.float(), -E.t().float(), torch.zeros((mu.shape[0], E.shape[1]))), dim=1))) + + v = torch.cat(((Jc.float() @ world.get_v().float()) * -world.restitutions().float(), + torch.zeros(Jf.shape[0] + mu.shape[0]))) + + _, _, x = self.lcp_solver(P, Q, R, S, u, v) + x = -x.double() + + new_v = x[:world.vec_len * len(world.bodies)].squeeze(0) + return new_v + + def post_stabilization(self, world): + M = world.M() + Je = world.Je() + Jc = None + if world.contacts: + Jc = world.Jc() + ge = torch.matmul(Je, world.get_v()) + gc = None + if Jc is not None: + gc = torch.matmul(Jc, world.get_v()) + torch.matmul(Jc, world.get_v()) * -world.restitutions() + + u = torch.cat([Je.new_zeros(Je.size(1)), ge]) + if Jc is None: + neq = Je.size(0) if Je.ndimension() > 0 else 0 + if neq > 0: + P = torch.cat([torch.cat([M, -Je.t()], dim=1), + torch.cat([Je, Je.new_zeros(neq, neq)], dim=1)]) + else: + P = M + if self.cached_inverse is None: + inv = torch.inverse(P) + else: + inv = self.cached_inverse + x = torch.matmul(inv, u) + else: + neq = Je.size(0) if Je.ndimension() > 0 else 0 + if neq > 0: + P = torch.cat([torch.cat([M, -Je.t()], dim=1), + torch.cat([Je, Je.new_zeros(neq, neq)], dim=1)]) + else: + P = M + + Q = torch.cat((-Jc.t(), torch.zeros((Je.shape[0], Jc.shape[0])))) + R = torch.cat((Jc, torch.zeros((Jc.shape[0], Je.shape[0]))), dim=1) + S = torch.zeros((R.shape[0], Q.shape[1])) + v = gc + _, _, x = self.lcp_solver(P, Q, R, S, u, v) + x = -x.double() + + dp = -x[:M.size(0)] + return dp \ No newline at end of file diff --git a/lcp_physics/physics/forces.py b/lcp_physics/physics/forces.py index 3ea7097..4721eaf 100644 --- a/lcp_physics/physics/forces.py +++ b/lcp_physics/physics/forces.py @@ -1,5 +1,5 @@ from .utils import get_tensor - +import torch def down_force(t): return ExternalForce.DOWN @@ -26,6 +26,27 @@ def rot_impulse(t): return ExternalForce.ZEROS +class ExternalForce: + """Generic external force to be added to objects. + Takes in a force_function which returns a force vector as a function of time, + and a multiplier that multiplies such vector. + """ + # Pre-store basic forces + DOWN = get_tensor([0, 0, 1]) + RIGHT = get_tensor([0, 1, 0]) + ROT = get_tensor([1, 0, 0]) + ZEROS = get_tensor([0, 0, 0]) + + def __init__(self, force_func=down_force, multiplier=100.): + self.multiplier = multiplier + self.force = lambda t: force_func(t) * self.multiplier + self.body = None + + def set_body(self, body): + self.body = body + # match body's tensor type and device + self.multiplier = get_tensor(self.multiplier, base_tensor=body._base_tensor) + class ExternalForce: """Generic external force to be added to objects. Takes in a force_function which returns a force vector as a function of time, @@ -65,3 +86,54 @@ def set_body(self, body): super().set_body(body) down_tensor = ExternalForce.DOWN.type_as(body._base_tensor).to(body._base_tensor) self.cached_force = down_tensor * self.body.mass * self.multiplier + + +class MDP(ExternalForce): + """Gravity force object, constantly returns a downwards pointing force of + magnitude body.mass * g. + """ + + def __init__(self, g=10.0): + self.multiplier = g + self.body = None + self.cached_force = None + + def force(self, t): + agains_vel_tensor = 2*torch.nn.Sigmoid()(self.body.v) - 1 + self.cached_force = -agains_vel_tensor * self.body.mass * self.multiplier + self.cached_force[0] = self.cached_force[0]*1000 + return self.cached_force + + def set_body(self, body): + super().set_body(body) + vel = body.v + agains_vel_tensor = 2*torch.nn.Sigmoid()(vel) - 1 + self.cached_force = -agains_vel_tensor * self.body.mass * self.multiplier + +class FingerTrajectory(ExternalForce): + """Object trajectory ass force: body.mass * ddp. + """ + + def __init__(self, traj): + self.Traj = traj + self.multiplier = 1 + self.body = None + self.t_prev = 0 + self.idx = 0 + self.cached_force = None + print(self.Traj) + + def force(self, t): + # if t - self.t_prev > 0.1: + import math + self.idx = int(math.floor(t*10)) + if self.idx > 4: + self.idx = 4 + self.cached_force = self.Traj[:,self.idx] * self.body.mass * 3 + # print(self.cached_force) + return self.cached_force + + def set_body(self, body): + super().set_body(body) + vel = body.v + self.cached_force = self.Traj[:,0] \ No newline at end of file diff --git a/lcp_physics/physics/untitled.py b/lcp_physics/physics/untitled.py new file mode 100644 index 0000000..328d497 --- /dev/null +++ b/lcp_physics/physics/untitled.py @@ -0,0 +1,59 @@ +import os +import sys +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +from src.model import * +from src.trainers import * +from src.datatools import * + +import torch +import pdb +import numpy as np +import torch.optim as optim +import matplotlib.pyplot as plt + +import time + +print("loading training data...") +# loads the training data +data, vids, pols = load_dataset(0,0) +N_data = np.shape(data)[0] +print("parsing training data...") +inputs_1, inputs_2, inputs_img, _, labels = parse_dataVids(data) +print(np.shape(vids)) + +# define network +print("Setting up network...") +use_cuda = torch.cuda.is_available() # check if GPU exists +device = torch.device("cuda" if use_cuda else "cpu") # use CPU or GPU +net = ContactNet(N_data).to(device) +net.addFrameVAELayers() +net.addVideoLayers() + +net.load() +net.eval() + +print("training video cod. autoencoders") +TrainVideoCondDecoders(net, vids, inputs_1, inputs_img, epochs = 5) +VizVideoCondDecoders(net, vids, inputs_1, inputs_img) +# net.save() +# TrainVideoParams(net, vids, inputs_2, epochs = 500) +# net.save() + +criterion = torch.nn.MSELoss(reduction='mean') +optimizer = optim.Adam(net.parameters(), lr=0.0001) +for epoch in range(200): # loop over the dataset multiple times + loss_t = 0 + optimizer.zero_grad() + + outputs = net.forwardVideo(torch.tensor(vids).float()) + loss = criterion(10*outputs, 10*labels.float()) + + loss_t = loss.item() + loss.backward() + optimizer.step() + + print("Train loss at epoch ",epoch," = ",loss_t) + +net.save() +net.gen_resVid(vids,'trainVid_57') \ No newline at end of file diff --git a/lcp_physics/physics/utils.py b/lcp_physics/physics/utils.py index 3f6c19b..7d37f12 100644 --- a/lcp_physics/physics/utils.py +++ b/lcp_physics/physics/utils.py @@ -13,13 +13,13 @@ class Defaults: DIM = 2 # Contact detectopm parameter - EPSILON = 0.1 + EPSILON = float('inf') # Penetration tolerance parameter TOL = 1e-6 # Default simulation parameters - RESTITUTION = 0.5 + RESTITUTION = 0 FRIC_COEFF = 0.9 FRIC_DIRS = 2 @@ -28,7 +28,7 @@ class Defaults: DT = 1.0 / FPS ENGINE = 'PdipmEngine' - CONTACT = 'DiffContactHandler' + CONTACT = 'NonConvexContactHandler' # Tensor defaults DTYPE = torch.double diff --git a/lcp_physics/physics/world.py b/lcp_physics/physics/world.py index 39da5ad..e12017b 100644 --- a/lcp_physics/physics/world.py +++ b/lcp_physics/physics/world.py @@ -1,320 +1,514 @@ import time -from functools import lru_cache +# from functools import lru_cache +from argparse import Namespace import ode import torch +import pdb +import math from . import engines as engines_module from . import contacts as contacts_module from .utils import Indices, Defaults, cross_2d, get_instance, left_orthogonal - +import numpy as np X, Y = Indices.X, Indices.Y DIM = Defaults.DIM +class Trajectory(object): + """Fingers velocity trajectory""" + def __init__(self, vel=np.zeros((2,5)), name='TrajNo'): + # super(Trajectory, self).__init__() + self.vel = vel + self.name = name + +class Reference(object): + """Fingers pose trajectory""" + def __init__(self, pos=np.zeros((3,5)), name='RefNo'): + # super(Trajectory, self).__init__() + self.ref = pos + self.name = name + class World: - """A physics simulation world, with bodies and constraints. - """ - def __init__(self, bodies, constraints=[], dt=Defaults.DT, engine=Defaults.ENGINE, - contact_callback=Defaults.CONTACT, eps=Defaults.EPSILON, - tol=Defaults.TOL, fric_dirs=Defaults.FRIC_DIRS, - post_stab=Defaults.POST_STABILIZATION, strict_no_penetration=True): - # self.contacts_debug = None # XXX - - # Load classes from string name defined in utils - self.engine = get_instance(engines_module, engine) - self.contact_callback = get_instance(contacts_module, contact_callback) - - self.t = 0 - self.dt = dt - self.eps = eps - self.tol = tol - self.fric_dirs = fric_dirs - self.post_stab = post_stab - - self.bodies = bodies - self.vec_len = len(self.bodies[0].v) - - # XXX Using ODE for broadphase for now - self.space = ode.HashSpace() - for i, b in enumerate(bodies): - b.geom.body = i - self.space.add(b.geom) - - self.static_inverse = True - self.num_constraints = 0 - self.joints = [] - for j in constraints: - b1, b2 = j.body1, j.body2 - i1 = bodies.index(b1) - i2 = bodies.index(b2) if b2 else None - self.joints.append((j, i1, i2)) - self.num_constraints += j.num_constraints - if not j.static: - self.static_inverse = False - - M_size = bodies[0].M.size(0) - self._M = bodies[0].M.new_zeros(M_size * len(bodies), M_size * len(bodies)) - # XXX Better way for diagonal block matrix? - for i, b in enumerate(bodies): - self._M[i * M_size:(i + 1) * M_size, i * M_size:(i + 1) * M_size] = b.M - - self.set_v(torch.cat([b.v for b in bodies])) - - self.contacts = None - self.find_contacts() - self.strict_no_pen = strict_no_penetration - if self.strict_no_pen: - assert all([c[0][3].item() <= self.tol for c in self.contacts]), \ - 'Interpenetration at start:\n{}'.format(self.contacts) - - def step(self, fixed_dt=False): - dt = self.dt - if fixed_dt: - end_t = self.t + self.dt - while self.t < end_t: - dt = end_t - self.t - self.step_dt(dt) - else: - self.step_dt(dt) - - # @profile - def step_dt(self, dt): - start_p = torch.cat([b.p for b in self.bodies]) - start_rot_joints = [(j[0].rot1, j[0].rot2) for j in self.joints] - new_v = self.engine.solve_dynamics(self, dt) - self.set_v(new_v) - while True: - # try step with current dt - for body in self.bodies: - body.move(dt) - for joint in self.joints: - joint[0].move(dt) - self.find_contacts() - if all([c[0][3].item() <= self.tol for c in self.contacts]): - break - else: - if not self.strict_no_pen and dt < self.dt / 4: - # if step becomes too small, just continue - break - dt /= 2 - # reset positions to beginning of step - # XXX Clones necessary? - self.set_p(start_p.clone()) - for j, c in zip(self.joints, start_rot_joints): - j[0].rot1 = c[0].clone() - j[0].update_pos() - - if self.post_stab: - tmp_v = self.v - dp = self.engine.post_stabilization(self).squeeze(0) - dp /= 2 # XXX Why 1/2 factor? - # XXX Clean up / Simplify this update? - self.set_v(dp) - for body in self.bodies: - body.move(dt) - for joint in self.joints: - joint[0].move(dt) - self.set_v(tmp_v) - - self.find_contacts() # XXX Necessary to recheck contacts? - self.t += dt - - def get_v(self): - return self.v - - def set_v(self, new_v): - self.v = new_v - for i, b in enumerate(self.bodies): - b.v = self.v[i * len(b.v):(i + 1) * len(b.v)] - - def set_p(self, new_p): - for i, b in enumerate(self.bodies): - b.set_p(new_p[i * self.vec_len:(i + 1) * self.vec_len]) - - def apply_forces(self, t): - return torch.cat([b.apply_forces(t) for b in self.bodies]) - - def find_contacts(self): - self.contacts = [] - # ODE contact detection - self.space.collide([self], self.contact_callback) - - def restitutions(self): - restitutions = self._M.new_empty(len(self.contacts)) - for i, c in enumerate(self.contacts): - r1 = self.bodies[c[1]].restitution - r2 = self.bodies[c[2]].restitution - restitutions[i] = (r1 + r2) / 2 - # restitutions[i] = math.sqrt(r1 * r2) - return restitutions - - def M(self): - return self._M - - def Je(self): - Je = self._M.new_zeros(self.num_constraints, - self.vec_len * len(self.bodies)) - row = 0 - for joint in self.joints: - J1, J2 = joint[0].J() - i1 = joint[1] - i2 = joint[2] - Je[row:row + J1.size(0), - i1 * self.vec_len:(i1 + 1) * self.vec_len] = J1 - if J2 is not None: - Je[row:row + J2.size(0), - i2 * self.vec_len:(i2 + 1) * self.vec_len] = J2 - row += J1.size(0) - return Je - - def Jc(self): - Jc = self._M.new_zeros(len(self.contacts), self.vec_len * len(self.bodies)) - for i, contact in enumerate(self.contacts): - c = contact[0] # c = (normal, contact_pt_1, contact_pt_2) - i1 = contact[1] - i2 = contact[2] - J1 = torch.cat([cross_2d(c[1], c[0]).reshape(1, 1), - c[0].unsqueeze(0)], dim=1) - J2 = -torch.cat([cross_2d(c[2], c[0]).reshape(1, 1), - c[0].unsqueeze(0)], dim=1) - Jc[i, i1 * self.vec_len:(i1 + 1) * self.vec_len] = J1 - Jc[i, i2 * self.vec_len:(i2 + 1) * self.vec_len] = J2 - return Jc - - def Jf(self): - Jf = self._M.new_zeros(len(self.contacts) * self.fric_dirs, - self.vec_len * len(self.bodies)) - for i, contact in enumerate(self.contacts): - c = contact[0] # c = (normal, contact_pt_1, contact_pt_2) - dir1 = left_orthogonal(c[0]) - dir2 = -dir1 - i1 = contact[1] # body 1 index - i2 = contact[2] # body 2 index - J1 = torch.cat([ - torch.cat([cross_2d(c[1], dir1).reshape(1, 1), - dir1.unsqueeze(0)], dim=1), - torch.cat([cross_2d(c[1], dir2).reshape(1, 1), - dir2.unsqueeze(0)], dim=1), - ], dim=0) - J2 = torch.cat([ - torch.cat([cross_2d(c[2], dir1).reshape(1, 1), - dir1.unsqueeze(0)], dim=1), - torch.cat([cross_2d(c[2], dir2).reshape(1, 1), - dir2.unsqueeze(0)], dim=1), - ], dim=0) - Jf[i * self.fric_dirs:(i + 1) * self.fric_dirs, - i1 * self.vec_len:(i1 + 1) * self.vec_len] = J1 - Jf[i * self.fric_dirs:(i + 1) * self.fric_dirs, - i2 * self.vec_len:(i2 + 1) * self.vec_len] = -J2 - return Jf - - def mu(self): - return self._memoized_mu(*[(c[1], c[2]) for c in self.contacts]) - - def _memoized_mu(self, *contacts): - # contacts is argument so that cacheing can be implemented at some point - mu = self._M.new_zeros(len(self.contacts)) - for i, contacts in enumerate(self.contacts): - i1 = contacts[1] - i2 = contacts[2] - # mu[i] = torch.sqrt(self.bodies[i1].fric_coeff * self.bodies[i2].fric_coeff) - mu[i] = 0.5 * (self.bodies[i1].fric_coeff + self.bodies[i2].fric_coeff) - return torch.diag(mu) - - def E(self): - return self._memoized_E(len(self.contacts)) - - def _memoized_E(self, num_contacts): - n = self.fric_dirs * num_contacts - E = self._M.new_zeros(n, num_contacts) - for i in range(num_contacts): - E[i * self.fric_dirs: (i + 1) * self.fric_dirs, i] += 1 - return E - - def save_state(self): - raise NotImplementedError - - def load_state(self, state_dict): - raise NotImplementedError - - def reset_engine(self): - raise NotImplementedError + """A physics simulation world, with bodies and constraints. + """ + def __init__(self, bodies, constraints=[], dt=Defaults.DT, engine=Defaults.ENGINE, + contact_callback=Defaults.CONTACT, eps=Defaults.EPSILON, + tol=Defaults.TOL, fric_dirs=Defaults.FRIC_DIRS, + post_stab=Defaults.POST_STABILIZATION, strict_no_penetration=True, facets = []): + self.contacts_debug = False # XXX + + # Load classes from string name defined in utils + self.engine = get_instance(engines_module, engine) + self.contact_callback = get_instance(contacts_module, contact_callback) + self.states = [] + self.fingers1 = [] + self.fingers2 = [] + self.times = [] + self.t = 0 + self.t_prev = -1 + self.idx = 0 + self.dt = dt + self.traj = [] + self.ref = [] + self.eps = eps + self.tol = tol + self.fric_dirs = fric_dirs + self.post_stab = post_stab + self.gamma = 0.01 + self.applied = True + self.facets = facets + + self.bodies = bodies + self.vec_len = len(self.bodies[0].v) + + # XXX Using ODE for broadphase for now + # self.space = ode.HashSpace() + # for i, b in enumerate(bodies): + # b.geom.body = i + # self.space.add(b.geom) + + self.static_inverse = True + self.num_constraints = 0 + self.joints = [] + for j in constraints: + b1, b2 = j.body1, j.body2 + i1 = bodies.index(b1) + i2 = bodies.index(b2) if b2 else None + self.joints.append((j, i1, i2)) + self.num_constraints += j.num_constraints + if not j.static: + self.static_inverse = False + + M_size = bodies[0].M.size(0) + self._M = bodies[0].M.new_zeros(M_size * len(bodies), M_size * len(bodies)) + # XXX Better way for diagonal block matrix? + for i, b in enumerate(bodies): + self._M[i * M_size:(i + 1) * M_size, i * M_size:(i + 1) * M_size] = b.M + + self.set_v(torch.cat([b.v for b in bodies])) + + self.contacts = None + self.find_contacts() + self.strict_no_pen = strict_no_penetration + # if self.strict_no_pen: + # for b in self.bodies: + # print(f'{b.__class__}: {vars(b)}\n') + # assert all([c[0][3].item() <= self.tol for c in self.contacts]),'Interpenetration at start:\n{}'.format(self.contacts) + + def step(self, fixed_dt=True): + dt = self.dt + if fixed_dt: + end_t = self.t + self.dt + while self.t < end_t: + dt = end_t - self.t + self.step_dt(dt) + else: + self.step_dt(dt) + + # @profile + def step_dt(self, dt): + + # PI contrrol weights + w1 = 1 # P-term + w2 = 1 - w1 # I-term + + # gets velocities + if self.idx >= 0 and self.applied: + for body in self.bodies: + for tr in self.traj: + for ref in self.ref: + if body.name == tr.name and body.name == ref.name: + if self.idx < np.shape(tr.vel)[1]: + vel = w1*tr.vel[:,self.idx] + vel = torch.cat([vel.new_zeros(1), vel]) + dt_ = math.ceil((self.t+1e-6)*10)/10 - self.t + if self.idx == np.shape(tr.vel)[1]-1: + vel += w2*(ref.ref[:,self.idx] - body.p)/dt_ + else: + vel += w2*(ref.ref[:,self.idx+1] - body.p)/dt_ + body.v = vel + + # # updates velocities + self.set_v(torch.cat([b.v for b in self.bodies])) + self.applied = True + + + start_p = torch.cat([b.p for b in self.bodies]) + start_rot_joints = [(j[0].rot1, j[0].rot2) for j in self.joints] + new_v = self.engine.solve_dynamics(self, dt) + self.set_v(new_v) + # print('orig') + # print(new_v) + + if self.idx >= 0 and self.applied: + for body in self.bodies: + for tr in self.traj: + for ref in self.ref: + if body.name == tr.name and body.name == ref.name: + if self.idx < np.shape(tr.vel)[1]: + vel = w1*tr.vel[:,self.idx] + vel = torch.cat([vel.new_zeros(1), vel]) + dt_ = math.ceil((self.t+1e-6)*10)/10 - self.t + if self.idx == np.shape(tr.vel)[1]-1: + vel += w2*(ref.ref[:,self.idx] - body.p)/dt_ + else: + vel += w2*(ref.ref[:,self.idx+1] - body.p)/dt_ + body.v = vel + + # # updates velocities + self.set_v(torch.cat([b.v for b in self.bodies])) + self.applied = True + + while True: + # try step with current dt + for body in self.bodies: + body.move(dt) + for joint in self.joints: + joint[0].move(dt) + joint[0].stabilize() + self.find_contacts() + + if all([c[0][3].item() <= self.tol for c in self.contacts]): + break + else: + break + # print('refining') + if not self.strict_no_pen and dt < self.dt / 8: + # if step becomes too small, just continue + break + dt /= 2 + # reset positions to beginning of step + # XXX Clones necessary? + self.set_p(start_p.clone()) + for j, c in zip(self.joints, start_rot_joints): + j[0].rot1 = c[0].clone() + j[0].update_pos() + + self.correct_penetrations() + + if self.post_stab: + tmp_v = self.v + dp = self.engine.post_stabilization(self).squeeze(0) + dp /= 2 # XXX Why 1/2 factor? + # XXX Clean up / Simplify this update? + self.set_v(dp) + for body in self.bodies: + body.move(dt) + for joint in self.joints: + joint[0].move(dt) + # print('s2') + self.set_v(tmp_v) + + # self.find_contacts() # XXX Necessary to recheck contacts? + self.correct_penetrations() + self.times.append(self.t) + + self.t += dt + + def get_v(self): + return self.v + + def set_v(self, new_v): + if np.isnan(new_v).sum().item() > 0: + new_v[:] = 0.0 + self.v = new_v + for i, b in enumerate(self.bodies): + b.v = self.v[i * len(b.v):(i + 1) * len(b.v)] + + def set_p(self, new_p): + for i, b in enumerate(self.bodies): + b.set_p(new_p[i * self.vec_len:(i + 1) * self.vec_len]) + + def apply_forces(self, t): + return torch.cat([b.apply_forces(t) for b in self.bodies]) + + + def correct_penetrations(self): + # corrects 1 + for c in self.contacts: + # corrects for penetration in the fingers + if self.bodies[c[1]].name[0] == 'f': + if c[0][3].item() > 11: + pass + # norm = (c[0][1])/torch.norm(c[0][1]) + # pen = c[0][3].item() + + # corrected_p = self.bodies[c[1]].p + torch.cat([torch.tensor([0.0]).double(), c[0][1]]) + # self.bodies[c[1]].set_p(corrected_p) + + # corrects for penetration with the environment + if self.bodies[c[1]].name == 'env': + if c[0][3].item() > self.tol: + for b in self.bodies: + if b.name[0] == 'o': + corrected_p = b.p + torch.cat([torch.tensor([0.0]).double(), torch.tensor([0.0]).double(), -torch.tensor([c[0][3].item()]).double()]) + b.set_p(corrected_p) + + # object frame + irot = torch.tensor([[math.cos(self.bodies[0].p[0].item()), math.sin(self.bodies[0].p[0].item())],[-math.sin(self.bodies[0].p[0].item()),math.cos(self.bodies[0].p[0].item())]]).view(2,2) + pos = self.bodies[0].pos.view(2,1).float() + + # finds the fingers + for b in self.bodies: + if b.name[0] == 'f': + inside = False + # finger in objec frame + ray = torch.matmul(irot,torch.tensor([b.p[1],b.p[2]]).view(2,1).float()-pos).view(2,1) + for f in self.facets: + # checks is the ray in y+ intercepts the facet (Jordan Polygon Theorem) + if (ray[0] <= max(f[0][0],f[1][0])) and (ray[0] >= min(f[0][0],f[1][0])): + if f[0][0] == f[1][0]: + if (ray[1] <= f[0][1]): + inside = True - inside + if (ray[1] <= f[1][1]): + inside = True - inside + else: + if (ray[1] <= ((ray[0]-f[1][0])*(f[0][1]-f[1][1])/(f[0][0]-f[1][0]) + f[1][1])): + inside = True - inside + if inside: + # finds closest facet + dist = 1e6 + for f in self.facets: + v1 = torch.tensor(f[1].copy()).view(2,1)-torch.tensor(f[0].copy()).view(2,1) # AB + v2 = ray - torch.tensor(f[0].copy()).view(2,1) + d1 = torch.matmul(v1.view(1,2),v2)/torch.norm(v1) + + if d1.item() <= 0: + p1 = torch.tensor(f[0].copy()).view(2,1) + elif d1.item() >= torch.norm(torch.tensor(f[1].copy()-f[0].copy()).view(2,1)).item(): + p1 = torch.tensor(f[1].copy()).view(2,1) + else: + p1 = torch.tensor(f[0].copy()).view(2,1) + d1.view(1,1)*v1/torch.norm(v1) + + if torch.norm(p1 - ray).item() < dist: + dist = torch.norm(p1 - ray).item() + correction = (p1 - ray).view(2,1) + + # projects to closes facet and assigns + b.set_p(b.p + torch.cat([torch.tensor([0.0]), torch.matmul(irot.transpose(0, 1), correction.float()).view(2)])) + + def find_contacts(self): + # import time + # start_c1 = time.time() + self.contacts = [] + # ODE contact detection + # self.space.collide([self], self.contact_callback) + # pdb.set_trace() + for i, b1 in enumerate(self.bodies): + g1 = Namespace() + g1.no_contact = b1.no_contact + g1.body_ref = b1 + g1.body = i + for j, b2 in enumerate(self.bodies[:i]): + g2 = Namespace() + g2.no_contact = b2.no_contact + g2.body_ref = b2 + g2.body = j + self.contact_callback([self], g1, g2, self.gamma) + + # end_c1 = time.time() + # print("time per finding contacts: ") + # print(end_c1 - start_c1) + + + def restitutions(self): + restitutions = self._M.new_empty(len(self.contacts)) + for i, c in enumerate(self.contacts): + r1 = self.bodies[c[1]].restitution + r2 = self.bodies[c[2]].restitution + restitutions[i] = (r1 + r2) / 2 + # restitutions[i] = math.sqrt(r1 * r2) + return restitutions + + def M(self): + return self._M + + def Je(self): + Je = self._M.new_zeros(self.num_constraints, + self.vec_len * len(self.bodies)) + row = 0 + for joint in self.joints: + J1, J2 = joint[0].J() + i1 = joint[1] + i2 = joint[2] + Je[row:row + J1.size(0), + i1 * self.vec_len:(i1 + 1) * self.vec_len] = J1 + if J2 is not None: + Je[row:row + J2.size(0), + i2 * self.vec_len:(i2 + 1) * self.vec_len] = J2 + row += J1.size(0) + return Je + + def Jc(self): + Jc = self._M.new_zeros(len(self.contacts), self.vec_len * len(self.bodies)) + for i, contact in enumerate(self.contacts): + c = contact[0] # c = (normal, contact_pt_1, contact_pt_2, penetration_dist) + i1 = contact[1] + i2 = contact[2] + J1 = torch.cat([cross_2d(c[1], c[0]).reshape(1, 1), + c[0].unsqueeze(0)], dim=1) + J2 = -torch.cat([cross_2d(c[2], c[0]).reshape(1, 1), + c[0].unsqueeze(0)], dim=1) + Jc[i, i1 * self.vec_len:(i1 + 1) * self.vec_len] = J1 + Jc[i, i2 * self.vec_len:(i2 + 1) * self.vec_len] = J2 + return Jc + + def Jf(self): + Jf = self._M.new_zeros(len(self.contacts) * self.fric_dirs, + self.vec_len * len(self.bodies)) + for i, contact in enumerate(self.contacts): + c = contact[0] # c = (normal, contact_pt_1, contact_pt_2) + dir1 = left_orthogonal(c[0]) + dir2 = -dir1 + i1 = contact[1] # body 1 index + i2 = contact[2] # body 2 index + J1 = torch.cat([ + torch.cat([cross_2d(c[1], dir1).reshape(1, 1), + dir1.unsqueeze(0)], dim=1), + torch.cat([cross_2d(c[1], dir2).reshape(1, 1), + dir2.unsqueeze(0)], dim=1), + ], dim=0) + J2 = torch.cat([ + torch.cat([cross_2d(c[2], dir1).reshape(1, 1), + dir1.unsqueeze(0)], dim=1), + torch.cat([cross_2d(c[2], dir2).reshape(1, 1), + dir2.unsqueeze(0)], dim=1), + ], dim=0) + Jf[i * self.fric_dirs:(i + 1) * self.fric_dirs, + i1 * self.vec_len:(i1 + 1) * self.vec_len] = J1 + Jf[i * self.fric_dirs:(i + 1) * self.fric_dirs, + i2 * self.vec_len:(i2 + 1) * self.vec_len] = -J2 + return Jf + + def mu(self): + return self._memoized_mu(*[(c[1], c[2]) for c in self.contacts]) + + def _memoized_mu(self, *contacts): + # contacts is argument so that cacheing can be implemented at some point + mu = self._M.new_zeros(len(self.contacts)) + for i, contacts in enumerate(self.contacts): + i1 = contacts[1] + i2 = contacts[2] + # mu[i] = torch.sqrt(self.bodies[i1].fric_coeff * self.bodies[i2].fric_coeff) + mu[i] = 0.5 * (self.bodies[i1].fric_coeff + self.bodies[i2].fric_coeff) + return torch.diag(mu) + + def E(self): + return self._memoized_E(len(self.contacts)) + + def _memoized_E(self, num_contacts): + n = self.fric_dirs * num_contacts + E = self._M.new_zeros(n, num_contacts) + for i in range(num_contacts): + E[i * self.fric_dirs: (i + 1) * self.fric_dirs, i] += 1 + return E + + def save_state(self): + raise NotImplementedError + + def load_state(self, state_dict): + raise NotImplementedError + + def reset_engine(self): + raise NotImplementedError + def run_world(world, animation_dt=None, run_time=10, print_time=True, - screen=None, recorder=None, pixels_per_meter=1): - """Helper function to run a simulation forward once a world is created. - """ - # If in batched mode don't display simulation - if hasattr(world, 'worlds'): - screen = None - - if screen is not None: - import pygame - background = pygame.Surface(screen.get_size()) - background = background.convert() - background.fill((255, 255, 255)) - - if animation_dt is None: - animation_dt = float(world.dt) - elapsed_time = 0. - prev_frame_time = -animation_dt - start_time = time.time() - - while world.t < run_time: - world.step() - - if screen is not None: - for event in pygame.event.get(): - if event.type == pygame.QUIT: - return - - if elapsed_time - prev_frame_time >= animation_dt or recorder: - prev_frame_time = elapsed_time - - screen.blit(background, (0, 0)) - update_list = [] - for body in world.bodies: - update_list += body.draw(screen, pixels_per_meter=pixels_per_meter) - for joint in world.joints: - update_list += joint[0].draw(screen, pixels_per_meter=pixels_per_meter) - - # Visualize contact points and normal for debug - # (Uncomment contacts_debug line in contacts handler): - # if world.contacts_debug: - # for c in world.contacts_debug: - # (normal, p1, p2, penetration), b1, b2 = c - # b1_pos = world.bodies[b1].pos - # b2_pos = world.bodies[b2].pos - # p1 = p1 + b1_pos - # p2 = p2 + b2_pos - # pygame.draw.circle(screen, (0, 255, 0), p1.data.numpy().astype(int), 5) - # pygame.draw.circle(screen, (0, 0, 255), p2.data.numpy().astype(int), 5) - # pygame.draw.line(screen, (0, 255, 0), p1.data.numpy().astype(int), - # (p1.data.numpy() + normal.data.numpy() * 100).astype(int), 3) - - if not recorder: - # Don't refresh screen if recording - pygame.display.update(update_list) - # pygame.display.flip() # XXX - else: - recorder.record(world.t) - - elapsed_time = time.time() - start_time - if not recorder: - # Adjust frame rate dynamically to keep real time - wait_time = world.t - elapsed_time - if wait_time >= 0 and not recorder: - wait_time += animation_dt # XXX - time.sleep(max(wait_time - animation_dt, 0)) - # animation_dt -= 0.005 * wait_time - # elif wait_time < 0: - # animation_dt += 0.005 * -wait_time - # elapsed_time = time.time() - start_time - - elapsed_time = time.time() - start_time - if print_time: - print('\r ', '{} / {} {} '.format(int(world.t), int(elapsed_time), - 1 / animation_dt), end='') + screen=None, recorder=None, pixels_per_meter=1, traj = [Trajectory()], pos_f = [Reference()]): + """Helper function to run a simulation forward once a world is created. + """ + import math + # If in batched mode don't display simulation + if hasattr(world, 'worlds'): + screen = None + + if screen is not None: + import pygame + background = pygame.Surface(screen.get_size()) + background = background.convert() + background.fill((255, 255, 255)) + + if animation_dt is None: + animation_dt = float(world.dt) + elapsed_time = 0. + prev_frame_time = -animation_dt + start_time = time.time() + + world.idx = 0 + world.traj = traj + world.ref = pos_f + world.t_prev = -10.0 + # world.engine = get_instance(engines_module,'LemkeEngine') + while world.t < run_time: + if world.t - world.t_prev >= 0.099: + # print(world.t) + for body in world.bodies: + if body.name == "obj": + world.states.append(body.p) + if body.name == "f0": + world.fingers1.append(body.p) + if body.name == "f1": + world.fingers2.append(body.p) + world.idx += 1 + world.t_prev = round(world.t,2) + world.applied = True + + world.step() + # pdb.set_trace() + + if screen is not None: + for event in pygame.event.get(): + if event.type == pygame.QUIT: + return + + if elapsed_time - prev_frame_time >= animation_dt or recorder: + prev_frame_time = elapsed_time + + screen.blit(background, (0, 0)) + update_list = [] + for body in world.bodies: + update_list += body.draw(screen, pixels_per_meter=pixels_per_meter) + for joint in world.joints: + update_list += joint[0].draw(screen, pixels_per_meter=pixels_per_meter) + + # Visualize contact points and normal for debug + # (Uncomment contacts_debug line in contacts handler): + if world.contacts_debug: + for c in world.contacts_debug: + (normal, p1, p2, penetration), b1, b2 = c + b1_pos = world.bodies[b1].pos + b2_pos = world.bodies[b2].pos + p1 = p1 + b1_pos + p2 = p2 + b2_pos + pygame.draw.circle(screen, (0, 255, 0), p1.data.numpy().astype(int), 5) + pygame.draw.circle(screen, (0, 0, 255), p2.data.numpy().astype(int), 5) + pygame.draw.line(screen, (0, 255, 0), p1.data.numpy().astype(int), + (p1.data.numpy() + normal.data.numpy() * 100).astype(int), 3) + + if not recorder: + pass + # Don't refresh screen if recording + pygame.display.update(update_list) + pygame.display.flip() # XXX + else: + recorder.record(world.t) + + elapsed_time = time.time() - start_time + if not recorder: + # Adjust frame rate dynamically to keep real time + wait_time = world.t - elapsed_time + if wait_time >= 0 and not recorder: + wait_time += animation_dt # XXX + time.sleep(max(wait_time - animation_dt, 0)) + # animation_dt -= 0.005 * wait_time + # elif wait_time < 0: + # animation_dt += 0.005 * -wait_time + # elapsed_time = time.time() - start_time + + elapsed_time = time.time() - start_time + if print_time: + print('\r ', '{} / {} {} '.format(float(world.t), float(elapsed_time), + 1 / animation_dt), end='')