working on memory allocation running, but not working properly

python/quadruped_reactive_walking/WB_MPC/CrocoddylOCP.py

@@ -11,43 +11,45 @@ class OCP:
     def __init__(self, pd:ProblemData, target:Target):
         self.pd = pd
         self.target = target
-        self.state = crocoddyl.StateMultibody(self.pd.model)
+        
         self.results = OcpResult()
-        if pd.useFixedBase == 0:
-            self.actuation = crocoddyl.ActuationModelFloatingBase(self.state)
-        else:
-            self.actuation = crocoddyl.ActuationModelFull(self.state)
+        self.initialize_models()
+    
+    def initialize_models(self):
+        self.models = []
+        for _ in range(self.pd.T+1):
+            self.models.append(Model(self.pd))
+        
 
     def make_crocoddyl_ocp(self, x0):
         """ Create a shooting problem for a simple walking gait.
 
         :param x0: initial state
         """
-        self.control = crocoddyl.ControlParametrizationModelPolyZero(
-            self.actuation.nu)
+        
 
         # Compute the current foot positions
-        q0 = x0[:self.state.nq]
+        q0 = x0[:self.pd.nq]
         pin.forwardKinematics(self.pd.model, self.pd.rdata, q0)
         pin.updateFramePlacements(self.pd.model, self.pd.rdata)
 
-        model = []
         for t in range(self.pd.T):
             target = self.target.evaluate_in_t(t)
             freeIds = [idf for idf in self.pd.allContactIds if idf not in self.target.contactSequence[t]]
             contactIds = self.target.contactSequence[t]
-            model += self.createFootstepModels(target, contactIds, freeIds)
+            self.appendTargetToModel(self.models[t], target, freeIds, False)
 
         freeIds = [idf for idf in self.pd.allContactIds if idf not in self.target.contactSequence[self.pd.T]]
-        contactIds = self.target.contactSequence[self.pd.T]
-        model += self.createFootstepModels(self.target.evaluate_in_t(self.pd.T), contactIds, freeIds, True)
+        #contactIds = self.target.contactSequence[self.pd.T]
+        self.appendTargetToModel(self.models[self.pd.T], self.target.evaluate_in_t(self.pd.T), freeIds, True)
 
-        problem = crocoddyl.ShootingProblem(x0, model[:-1], model[-1])
+        problem = crocoddyl.ShootingProblem(x0, 
+                                            [m.model for m in self.models[:-1]], 
+                                            self.models[-1].model)
 
         return problem
 
-    def createFootstepModels(self,  target, supportFootIds,
-                             swingFootIds, isTerminal=False):
+    def appendTargetToModel(self, model, target, swingFootIds, isTerminal=False):
         """ Action models for a footstep phase.
         :param numKnots: number of knots for the footstep phase
         :param supportFootIds: Ids of the supporting feet
@@ -55,7 +57,6 @@ class OCP:
         :return footstep action models
         """
         # Action models for the foot swing
-        footSwingModel = []
         swingFootTask = []
         for i in swingFootIds:
             try:
@@ -63,72 +64,9 @@ class OCP:
                 swingFootTask += [[i, pin.SE3(np.eye(3), tref)]]
             except:
                 pass
+        
+        model.tracking_cost(swingFootTask)
 
-        footSwingModel += [self.createSwingFootModel(supportFootIds, swingFootTask=swingFootTask, isTerminal=isTerminal)]
-        return footSwingModel
-
-    def createSwingFootModel(self, supportFootIds, swingFootTask=None, isTerminal=False):
-        """ Action model for a swing foot phase.
-
-        :param timeStep: step duration of the action model
-        :param supportFootIds: Ids of the constrained feet
-        :param comTask: CoM task
-        :param swingFootTask: swinging foot task
-        :return action model for a swing foot phase
-        """
-        # Creating a 3D multi-contact model, and then including the supporting
-        # foot
-        nu = self.actuation.nu
-        contactModel = crocoddyl.ContactModelMultiple(self.state, nu)
-        for i in supportFootIds:
-            supportContactModel = crocoddyl.ContactModel3D(self.state, i, np.array([0., 0., 0.]), nu,
-                                                           np.array([0., 0.]))
-            contactModel.addContact(self.pd.model.frames[i].name + "_contact", supportContactModel)
-
-        # Creating the cost model for a contact phase
-        costModel = crocoddyl.CostModelSum(self.state, nu)
-
-        if not isTerminal:
-
-            for i in supportFootIds:
-                cone = crocoddyl.FrictionCone(self.pd.Rsurf, self.pd.mu, 4, False)
-                coneResidual = crocoddyl.ResidualModelContactFrictionCone(self.state, i, cone, nu)
-                coneActivation = crocoddyl.ActivationModelQuadraticBarrier(crocoddyl.ActivationBounds(cone.lb, cone.ub))
-                frictionCone = crocoddyl.CostModelResidual(self.state, coneActivation, coneResidual)
-                costModel.addCost(self.pd.model.frames[i].name + "_frictionCone", frictionCone, self.pd.friction_cone_w)
-            if swingFootTask is not None:
-                for i in swingFootTask:
-                    frameTranslationResidual = crocoddyl.ResidualModelFrameTranslation(self.state, i[0], i[1].translation,nu)
-                    footTrack = crocoddyl.CostModelResidual(self.state, frameTranslationResidual)
-                    costModel.addCost(self.pd.model.frames[i[0]].name + "_footTrack", footTrack, self.pd.foot_tracking_w)
-
-            ctrlResidual = crocoddyl.ResidualModelControl(self.state, self.pd.uref)
-            ctrlReg = crocoddyl.CostModelResidual(self.state, ctrlResidual)
-            costModel.addCost("ctrlReg", ctrlReg, self.pd.control_reg_w)
-
-            ctrl_bound_residual = crocoddyl.ResidualModelControl(self.state, nu)
-            ctrl_bound_activation = crocoddyl.ActivationModelQuadraticBarrier(crocoddyl.ActivationBounds(-self.pd.effort_limit, self.pd.effort_limit))
-            ctrl_bound = crocoddyl.CostModelResidual(self.state, ctrl_bound_activation, ctrl_bound_residual)
-            costModel.addCost("ctrlBound", ctrl_bound, self.pd.control_bound_w)
-
-        stateResidual = crocoddyl.ResidualModelState(self.state, self.pd.xref, nu)
-        stateActivation = crocoddyl.ActivationModelWeightedQuad(self.pd.state_reg_w**2)
-        stateReg = crocoddyl.CostModelResidual(self.state, stateActivation, stateResidual)
-        costModel.addCost("stateReg", stateReg, 1)
-
-        if isTerminal:
-            stateResidual = crocoddyl.ResidualModelState(self.state, self.pd.xref, nu)
-            stateActivation = crocoddyl.ActivationModelWeightedQuad(self.pd.terminal_velocity_w**2)
-            stateReg = crocoddyl.CostModelResidual(self.state, stateActivation, stateResidual)
-            costModel.addCost("terminalVelocity", stateReg, 1)
-
-        dmodel = crocoddyl.DifferentialActionModelContactFwdDynamics(self.state, self.actuation, contactModel,
-                                                                     costModel, 0., True)
-
-        model = crocoddyl.IntegratedActionModelEuler(
-            dmodel, self.control, self.pd.dt)
-
-        return model
 
 # Solve
     def solve(self, x0, guess=None):
@@ -148,7 +86,6 @@ class OCP:
         self.ddp.solve(xs, us, 1, False)
         print("Solver time: ", time()- start_time, "\n")
 
-
     def get_results(self):
         self.results.x = self.ddp.xs.tolist()
         self.results.a = self.get_croco_acc()
@@ -189,4 +126,89 @@ class OCP:
          for m in self.ddp.problem.runningDatas]
         return acc
 
-    
+    
\ No newline at end of file
+
+class Model:
+    def __init__(self, pd, supportFootIds=[], isTerminal=False):
+        self.pd = pd
+        self.supportFootIds=supportFootIds
+        self.isTerminal=isTerminal
+
+        self.state = crocoddyl.StateMultibody(self.pd.model)
+        if pd.useFixedBase == 0:
+            self.actuation = crocoddyl.ActuationModelFloatingBase(self.state)
+        else:
+            self.actuation = crocoddyl.ActuationModelFull(self.state)
+        self.control = crocoddyl.ControlParametrizationModelPolyZero(self.actuation.nu)
+        self.nu = self.actuation.nu
+
+        self.createStandardModel()
+
+    def createStandardModel(self):
+        """ Action model for a swing foot phase.
+
+        :param timeStep: step duration of the action model
+        :param supportFootIds: Ids of the constrained feet
+        :param comTask: CoM task
+        :param swingFootTask: swinging foot task
+        :return action model for a swing foot phase
+        """
+        # Creating a 3D multi-contact model, and then including the supporting
+        # foot
+
+        
+        self.contactModel = crocoddyl.ContactModelMultiple(self.state, self.nu)
+        for i in self.supportFootIds:
+            supportContactModel = crocoddyl.ContactModel3D(self.state, i, np.array([0., 0., 0.]), self.nu,
+                                                           np.array([0., 0.]))
+            self.contactModel.addContact(self.pd.model.frames[i].name + "_contact", supportContactModel)
+
+        # Creating the cost model for a contact phase
+        costModel = crocoddyl.CostModelSum(self.state, self.nu)
+
+        if not self.isTerminal:
+
+            for i in self.supportFootIds:
+                cone = crocoddyl.FrictionCone(self.pd.Rsurf, self.pd.mu, 4, False)
+                coneResidual = crocoddyl.ResidualModelContactFrictionCone(self.state, i, cone, self.nu)
+                coneActivation = crocoddyl.ActivationModelQuadraticBarrier(crocoddyl.ActivationBounds(cone.lb, cone.ub))
+                frictionCone = crocoddyl.CostModelResidual(self.state, coneActivation, coneResidual)
+                costModel.addCost(self.pd.model.frames[i].name + "_frictionCone", frictionCone, self.pd.friction_cone_w)
+
+            ctrlResidual = crocoddyl.ResidualModelControl(self.state, self.pd.uref)
+            ctrlReg = crocoddyl.CostModelResidual(self.state, ctrlResidual)
+            costModel.addCost("ctrlReg", ctrlReg, self.pd.control_reg_w)
+
+            ctrl_bound_residual = crocoddyl.ResidualModelControl(self.state, self.nu)
+            ctrl_bound_activation = crocoddyl.ActivationModelQuadraticBarrier(crocoddyl.ActivationBounds(-self.pd.effort_limit, self.pd.effort_limit))
+            ctrl_bound = crocoddyl.CostModelResidual(self.state, ctrl_bound_activation, ctrl_bound_residual)
+            costModel.addCost("ctrlBound", ctrl_bound, self.pd.control_bound_w)
+
+        stateResidual = crocoddyl.ResidualModelState(self.state, self.pd.xref, self.nu)
+        stateActivation = crocoddyl.ActivationModelWeightedQuad(self.pd.state_reg_w**2)
+        stateReg = crocoddyl.CostModelResidual(self.state, stateActivation, stateResidual)
+        costModel.addCost("stateReg", stateReg, 1)
+
+        if self.isTerminal:
+            stateResidual = crocoddyl.ResidualModelState(self.state, self.pd.xref, self.nu)
+            stateActivation = crocoddyl.ActivationModelWeightedQuad(self.pd.terminal_velocity_w**2)
+            stateReg = crocoddyl.CostModelResidual(self.state, stateActivation, stateResidual)
+            costModel.addCost("terminalVelocity", stateReg, 1)
+
+        self.costModel = costModel
+
+        self.dmodel = crocoddyl.DifferentialActionModelContactFwdDynamics(self.state, self.actuation, self.contactModel,
+                                                                    self.costModel, 0., True)
+        self.model = crocoddyl.IntegratedActionModelEuler(self.dmodel, self.control, self.pd.dt)
+
+    def tracking_cost(self, swingFootTask):
+        if swingFootTask is not None:
+            for i in swingFootTask:
+                frameTranslationResidual = crocoddyl.ResidualModelFrameTranslation(self.state, i[0], i[1].translation,self.nu)
+                footTrack = crocoddyl.CostModelResidual(self.state, frameTranslationResidual)
+                if self.pd.model.frames[i[0]].name + "_footTrack" in self.dmodel.costs.active.tolist():
+                    self.dmodel.costs.removeCost(self.pd.model.frames[i[0]].name + "_footTrack")
+                self.costModel.addCost(self.pd.model.frames[i[0]].name + "_footTrack", footTrack, self.pd.foot_tracking_w)
+
+
+