diff --git a/agimus_controller/agimus_controller/factory/robot_model.py b/agimus_controller/agimus_controller/factory/robot_model.py index ec133a51..ae133145 100644 --- a/agimus_controller/agimus_controller/factory/robot_model.py +++ b/agimus_controller/agimus_controller/factory/robot_model.py @@ -10,9 +10,12 @@ @dataclass class RobotModelParameters: + full_q0: npt.NDArray[np.float64] = np.array( + [], dtype=np.float64 + ) # Initial full configuration of the robot q0: npt.NDArray[np.float64] = np.array( [], dtype=np.float64 - ) # Initial configuration of the robot + ) # Initial reduced configuration of the robot free_flyer: bool = False # True if the robot has a free flyer locked_joint_names: list[str] = field(default_factory=list) urdf_path: Path = Path() # Path to the URDF file @@ -33,10 +36,15 @@ class RobotModelParameters: ) # Red color for the collision model def __post_init__(self): - # Check q0 is not empty - if len(self.q0) == 0: - raise ValueError("q0 cannot be empty.") - + # Check full_q0 is not empty + if len(self.full_q0) == 0: + raise ValueError("full_q0 cannot be empty.") + + # Check q0 is not empty or if there is no reduced model, q0 is the full configuration + if len(self.q0) == 0 and not len(self.locked_joint_names) == 0: + raise ValueError("q0 cannot be empty while reducing the model.") + elif len(self.q0) == 0: + self.q0 = self.full_q0 # Handle armature: if self.armature.size == 0: # Use a default armature filled with 0s, based on the size of q0 @@ -73,6 +81,7 @@ def __init__(self, param: RobotModelParameters): self._collision_model = None self._visual_model = None self._q0 = deepcopy(self._params.q0) + self._full_q0 = deepcopy(self._params.full_q0) self.load_models() # Populate models @property @@ -108,6 +117,11 @@ def q0(self) -> np.array: """Initial configuration of the robot.""" return self._q0 + @property + def full_q0(self) -> np.array: + """Initial full configuration of the robot.""" + return self._full_q0 + def load_models(self) -> None: """Load and prepare robot models based on parameters.""" self._load_full_pinocchio_models() @@ -145,7 +159,7 @@ def _apply_locked_joints(self) -> None: raise ValueError(f"Joint {jn} not found in the robot model.") self._robot_model = pin.buildReducedModel( - self._full_robot_model, joints_to_lock, self._q0 + self._full_robot_model, joints_to_lock, self._full_q0 ) def _update_collision_model_to_capsules(self) -> None: diff --git a/agimus_controller/agimus_controller/mpc_data.py b/agimus_controller/agimus_controller/mpc_data.py index a01e0a94..a8ae345f 100644 --- a/agimus_controller/agimus_controller/mpc_data.py +++ b/agimus_controller/agimus_controller/mpc_data.py @@ -16,14 +16,13 @@ class OCPResults: @dataclass class OCPDebugData: - # Solver infos - problem_solved: bool = False - # Debug data result: list[TrajectoryPoint] references: list[TrajectoryPoint] kkt_norms: list[np.float64] collision_distance_residuals: list[dict[np.float64]] + # Solver infos + problem_solved: bool = False @dataclass diff --git a/agimus_controller/agimus_controller/ocp_base.py b/agimus_controller/agimus_controller/ocp_base.py index a0475b46..19f9a77c 100644 --- a/agimus_controller/agimus_controller/ocp_base.py +++ b/agimus_controller/agimus_controller/ocp_base.py @@ -1,11 +1,17 @@ from abc import ABC, abstractmethod + import numpy as np +import numpy.typing as npt from agimus_controller.mpc_data import OCPResults, OCPDebugData from agimus_controller.trajectory import WeightedTrajectoryPoint class OCPBase(ABC): + """Base class for the Optimal Control Problem (OCP) solver. + + This class defines the interface for the OCP solver.""" + def __init__(self) -> None: pass @@ -13,30 +19,72 @@ def __init__(self) -> None: def set_reference_horizon( self, reference_trajectory: list[WeightedTrajectoryPoint] ) -> None: - ... + """Set the reference trajectory and the weights of the costs for the OCP solver. This method should be implemented by the derived class.""" + pass - @abstractmethod @property + @abstractmethod def horizon_size() -> int: - ... + """Returns the horizon size of the OCP. + + Returns: + int: size of the horizon. + """ + pass - @abstractmethod @property - def dt() -> int: - ... + @abstractmethod + def dt() -> float: + """Returns the time step of the OCP in seconds. + + Returns: + int: time step of the OCP. + """ + pass @abstractmethod def solve( - self, x0: np.ndarray, x_init: list[np.ndarray], u_init: list[np.ndarray] + self, + x0: npt.NDArray[np.float64], + x_warmstart: list[npt.NDArray[np.float64]], + u_warmstart: list[npt.NDArray[np.float64]], ) -> None: - ... + """Solver for the OCP. This method should be implemented by the derived class. + The method should solve the OCP for the given initial state and warmstart values. + + Args: + x0 (npt.NDArray[np.float64]): current state of the robot. + x_warmstart (list[npt.NDArray[np.float64]]): Warmstart values for the state. This doesn't include the current state. + u_warmstart (list[npt.NDArray[np.float64]]): Warmstart values for the control inputs. + """ + pass - @abstractmethod @property + @abstractmethod def ocp_results(self) -> OCPResults: - ... + """Returns the results of the OCP solver. + The solve method should be called before calling this method. + + Returns: + OCPResults: Class containing the results of the OCP solver. + """ + pass + + @ocp_results.setter + def ocp_results(self, value: OCPResults) -> None: + """Set the output data structure of the OCP. + + Args: + value (OCPResults): New output data structure of the OCP. + """ + pass - @abstractmethod @property + @abstractmethod def debug_data(self) -> OCPDebugData: - ... + """Returns the debug data of the OCP solver. + + Returns: + OCPDebugData: Class containing the debug data of the OCP solver. + """ + pass diff --git a/agimus_controller/agimus_controller/ocp_base_croco.py b/agimus_controller/agimus_controller/ocp_base_croco.py new file mode 100644 index 00000000..fa45b033 --- /dev/null +++ b/agimus_controller/agimus_controller/ocp_base_croco.py @@ -0,0 +1,146 @@ +from abc import abstractmethod + +import crocoddyl +import mim_solvers +import numpy as np +import numpy.typing as npt + +from agimus_controller.factory.robot_model import RobotModels +from agimus_controller.mpc_data import OCPResults, OCPDebugData +from agimus_controller.ocp_base import OCPBase +from agimus_controller.ocp_param_base import OCPParamsBaseCroco +from agimus_controller.trajectory import TrajectoryPointWeights + + +class OCPBaseCroco(OCPBase): + def __init__( + self, + robot_models: RobotModels, + ocp_params: OCPParamsBaseCroco, + ) -> None: + """Defines common behavior for all OCP using croccodyl. This is an abstract class with some helpers to design OCPs in a more friendly way. + + Args: + robot_models (RobotModels): All models of the robot. + ocp_params (OCPParamsBaseCroco): Input data structure of the OCP. + """ + # Setting the robot model + self._robot_models = robot_models + self._collision_model = self._robot_models.collision_model + self._armature = self._robot_models._params.armature + + # Stat and actuation model + self._state = crocoddyl.StateMultibody(self._robot_models.robot_model) + self._actuation = crocoddyl.ActuationModelFull(self._state) + + # Setting the OCP parameters + self._ocp_params = ocp_params + self._ocp = None + self._ocp_results = None + self._running_model_list = [] + self._terminal_model = [] + self._problem = None + + @property + def horizon_size(self) -> int: + """Number of time steps in the horizon.""" + return self._ocp_params.horizon_size + + @property + def dt(self) -> float: + """Integration step of the OCP.""" + return self._ocp_params.dt + + @abstractmethod + def create_running_model_list(self) -> list[crocoddyl.ActionModelAbstract]: + """Create the list of running models.""" + pass + + @abstractmethod + def create_terminal_model(self) -> crocoddyl.ActionModelAbstract: + """Create the terminal model.""" + pass + + @abstractmethod + def update_crocoddyl_problem( + self, + x0: npt.NDArray[np.float64], + weighted_trajectory_points: list[TrajectoryPointWeights], + ) -> None: + """Update the Crocoddyl problem's references, weights and x0.""" + pass + + def solve( + self, + x0: npt.NDArray[np.float64], + x_warmstart: list[npt.NDArray[np.float64]], + u_warmstart: list[npt.NDArray[np.float64]], + ) -> None: + """Solves the OCP. + The results can be accessed through the ocp_results property. + + Args: + x0 (npt.NDArray[np.float64]): Current state of the robot. + x_warmstart (list[npt.NDArray[np.float64]]): Predicted states for the OCP. + u_warmstart (list[npt.NDArray[np.float64]]): Predicted control inputs for the OCP. + """ + ### Creation of the state and actuation models + + if self._ocp is None: + # Create the running models + self._running_model_list = self.create_running_model_list() + # Create the terminal model + self._terminal_model = self.create_terminal_model() + # Create the shooting problem + self._problem = crocoddyl.ShootingProblem( + x0, self._running_model_list, self._terminal_model + ) + # Create solver + callbacks + self._ocp = mim_solvers.SolverCSQP(self._problem) + + # Merit function + self._ocp.use_filter_line_search = self._ocp_params.use_filter_line_search + + # Parameters of the solver + self._ocp.termination_tolerance = self._ocp_params.termination_tolerance + self._ocp.max_qp_iters = self._ocp_params.qp_iters + self._ocp.eps_abs = self._ocp_params.eps_abs + self._ocp.eps_rel = self._ocp_params.eps_rel + self._ocp.with_callbacks = self._ocp_params.callbacks + + # Creating the warmstart lists for the solver + # Solve the OCP + self._ocp.solve([x0] + x_warmstart, u_warmstart, self._ocp_params.solver_iters) + + # Store the results + self._ocp_results = OCPResults( + states=self._ocp.xs, + ricatti_gains=self._ocp.K, + feed_forward_terms=self._ocp.us, + ) + + @property + def ocp_results(self) -> OCPResults: + """Output data structure of the OCP. + + Returns: + OCPResults: Output data structure of the OCP. It contains the states, Ricatti gains, and feed-forward terms. + """ + return self._ocp_results + + @ocp_results.setter + def ocp_results(self, value: OCPResults) -> None: + """Set the output data structure of the OCP. + + Args: + value (OCPResults): New output data structure of the OCP. + """ + self._ocp_results = value + + @property + def debug_data(self) -> OCPDebugData: + pass + + @debug_data.setter + def debug_data(self, value: OCPDebugData) -> None: + pass diff --git a/agimus_controller/agimus_controller/ocp_param_base.py b/agimus_controller/agimus_controller/ocp_param_base.py new file mode 100644 index 00000000..040d2691 --- /dev/null +++ b/agimus_controller/agimus_controller/ocp_param_base.py @@ -0,0 +1,21 @@ +from dataclasses import dataclass + +import numpy as np + + +@dataclass +class OCPParamsBaseCroco: + """Input data structure of the OCP.""" + + # Data relevant to solve the OCP + dt: float # Integration step of the OCP + horizon_size: int # Number of time steps in the horizon + solver_iters: int # Number of solver iterations + qp_iters: int = 200 # Number of QP iterations (must be a multiple of 25). + termination_tolerance: float = ( + 1e-3 # Termination tolerance (norm of the KKT conditions) + ) + eps_abs: float = 1e-6 # Absolute tolerance of the solver + eps_rel: float = 0.0 # Relative tolerance of the solver + callbacks: bool = False # Flag to enable/disable callbacks + use_filter_line_search = False # Flag to enable/disable the filter line searchs diff --git a/agimus_controller/tests/ressources/simple_ocp_croco_results.pkl b/agimus_controller/tests/ressources/simple_ocp_croco_results.pkl new file mode 100644 index 00000000..1ea7aae7 Binary files /dev/null and b/agimus_controller/tests/ressources/simple_ocp_croco_results.pkl differ diff --git a/agimus_controller/tests/test_ocp_croco_base.py b/agimus_controller/tests/test_ocp_croco_base.py new file mode 100644 index 00000000..6a523dd1 --- /dev/null +++ b/agimus_controller/tests/test_ocp_croco_base.py @@ -0,0 +1,261 @@ +from os.path import dirname +import unittest +from pathlib import Path +import crocoddyl +import example_robot_data as robex +import numpy as np +import pinocchio as pin +import pickle +from agimus_controller.ocp_base_croco import OCPBaseCroco +from agimus_controller.ocp_param_base import OCPParamsBaseCroco +from agimus_controller.factory.robot_model import RobotModels, RobotModelParameters + + +class TestOCPBaseCroco(unittest.TestCase): + @classmethod + def setUpClass(self): + # Mock the RobotModelFactory and OCPParamsCrocoBase + robot = robex.load("panda") + urdf_path = Path(robot.urdf) + srdf_path = Path(robot.urdf.replace("urdf", "srdf")) + urdf_meshes_dir = urdf_path.parent.parent.parent.parent.parent + free_flyer = False + locked_joint_names = ["panda_joint1", "panda_joint2"] + reduced_nq = robot.model.nq - len(locked_joint_names) + full_q0 = np.zeros(robot.model.nq) + armature = np.full(reduced_nq, 0.1) + q0 = np.zeros(robot.model.nq - len(locked_joint_names)) + # Store shared initial parameters + self.params = RobotModelParameters( + full_q0=full_q0, + q0=q0, + free_flyer=free_flyer, + locked_joint_names=locked_joint_names, + urdf_path=urdf_path, + srdf_path=srdf_path, + urdf_meshes_dir=urdf_meshes_dir, + collision_as_capsule=True, + self_collision=True, + armature=armature, + ) + + self.robot_models = RobotModels(self.params) + + self.ocp_params = OCPParamsBaseCroco( + dt=0.1, horizon_size=10, solver_iters=100, callbacks=True + ) + + # Create a concrete implementation of OCPBaseCroco + class TestOCPCroco(OCPBaseCroco): + def create_running_model_list(self): + return None + + def create_terminal_model(self): + return None + + def update_crocoddyl_problem(self, x0, trajectory_points_list): + return None + + def set_reference_horizon(self, horizon_size): + return None + + self.ocp = TestOCPCroco(self.robot_models, self.ocp_params) + + def test_horizon_size(self): + """Test the horizon_size property.""" + self.assertEqual(self.ocp.horizon_size, self.ocp_params.horizon_size) + + def test_dt(self): + """Test the dt property.""" + self.assertAlmostEqual(self.ocp.dt, self.ocp_params.dt) + + +class TestSimpleOCPCroco(unittest.TestCase): + class TestOCPCroco(OCPBaseCroco): + def create_running_model_list(self): + # Running cost model + running_cost_model = crocoddyl.CostModelSum(self._state) + + ### Creation of cost terms + # State Regularization cost + x_residual = crocoddyl.ResidualModelState( + self._state, + np.concatenate( + ( + pin.neutral(self._robot_models.robot_model), + np.zeros(self._robot_models.robot_model.nv), + ) + ), + ) + x_reg_cost = crocoddyl.CostModelResidual(self._state, x_residual) + # Control Regularization cost + u_residual = crocoddyl.ResidualModelControl(self._state) + u_reg_cost = crocoddyl.CostModelResidual(self._state, u_residual) + + # End effector frame cost + frame_placement_residual = crocoddyl.ResidualModelFramePlacement( + self._state, + self._robot_models.robot_model.getFrameId("panda_hand_tcp"), + pin.SE3(np.eye(3), np.array([1.0, 1.0, 1.0])), + ) + + goal_tracking_cost = crocoddyl.CostModelResidual( + self._state, frame_placement_residual + ) + running_cost_model.addCost("stateReg", x_reg_cost, 0.1) + running_cost_model.addCost("ctrlRegGrav", u_reg_cost, 0.0001) + running_cost_model.addCost("gripperPoseRM", goal_tracking_cost, 1.0) + # Create Differential Action Model (DAM), i.e. continuous dynamics and cost functions + running_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics( + self._state, + self._actuation, + running_cost_model, + ) + running_model = crocoddyl.IntegratedActionModelEuler( + running_DAM, + ) + running_model.differential.armature = self._robot_models.armature + running_model_list = [running_model] * (self._ocp_params.horizon_size - 1) + return running_model_list + + def create_terminal_model(self): + # Terminal cost models + terminal_cost_model = crocoddyl.CostModelSum(self._state) + + ### Creation of cost terms + # State Regularization cost + x_residual = crocoddyl.ResidualModelState( + self._state, + np.concatenate( + ( + pin.neutral(self._robot_models.robot_model), + np.zeros(self._robot_models.robot_model.nv), + ) + ), + ) + x_reg_cost = crocoddyl.CostModelResidual(self._state, x_residual) + + # End effector frame cost + frame_placement_residual = crocoddyl.ResidualModelFramePlacement( + self._state, + self._robot_models.robot_model.getFrameId("panda_hand_tcp"), + pin.SE3(np.eye(3), np.array([1.0, 1.0, 1.0])), + ) + + goal_tracking_cost = crocoddyl.CostModelResidual( + self._state, frame_placement_residual + ) + terminal_cost_model.addCost("stateReg", x_reg_cost, 0.1) + terminal_cost_model.addCost("gripperPose", goal_tracking_cost, 50) + + # Create Differential Action Model (DAM), i.e. continuous dynamics and cost functions + terminal_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics( + self._state, + self._actuation, + terminal_cost_model, + ) + + terminal_model = crocoddyl.IntegratedActionModelEuler(terminal_DAM, 0.0) + terminal_model.differential.armature = self._robot_models.armature + return terminal_model + + def set_reference_horizon(self, horizon_size): + ### Not implemented in this OCP example. + return None + + def update_crocoddyl_problem(self, x0, trajectory_points_list): + ### Not implemented in this OCP example. + return None + + @classmethod + def setUpClass(self): + # Mock the RobotModelFactory and OCPParamsCrocoBase + robot = robex.load("panda") + urdf_path = Path(robot.urdf) + srdf_path = Path(robot.urdf.replace("urdf", "srdf")) + urdf_meshes_dir = urdf_path.parent.parent.parent.parent.parent + free_flyer = False + locked_joint_names = ["panda_finger_joint1", "panda_finger_joint2"] + reduced_nq = robot.model.nq - len(locked_joint_names) + full_q0 = np.zeros(robot.model.nq) + q0 = np.zeros(reduced_nq) + armature = np.full(reduced_nq, 0.1) + # Store shared initial parameters + self.params = RobotModelParameters( + q0=q0, + full_q0=full_q0, + free_flyer=free_flyer, + locked_joint_names=locked_joint_names, + urdf_path=urdf_path, + srdf_path=srdf_path, + urdf_meshes_dir=urdf_meshes_dir, + collision_as_capsule=True, + self_collision=True, + armature=armature, + ) + + self.robot_models = RobotModels(self.params) + self.robot_model = self.robot_models.robot_model + self.collision_model = self.robot_models.collision_model + + # Set mock parameters + self.ocp_params = OCPParamsBaseCroco( + dt=0.1, horizon_size=10, solver_iters=100, callbacks=True + ) + self.state_reg = np.concatenate( + (pin.neutral(self.robot_model), np.zeros(self.robot_model.nv)) + ) + self.state_warmstart = [np.zeros(self.robot_model.nq + self.robot_model.nv)] * ( + self.ocp_params.horizon_size - 1 + ) # The first state is the current state + self.control_warmstart = [np.zeros(self.robot_model.nq)] * ( + self.ocp_params.horizon_size - 1 + ) + # Create a concrete implementation of OCPBaseCroco + self.ocp = self.TestOCPCroco(self.robot_models, self.ocp_params) + self.ocp.solve(self.state_reg, self.state_warmstart, self.control_warmstart) + self.save_results() + + @classmethod + def save_results(self): + results = { + "states": self.ocp.ocp_results.states.tolist(), + "ricatti_gains": self.ocp.ocp_results.ricatti_gains.tolist(), + "feed_forward_terms": self.ocp.ocp_results.feed_forward_terms.tolist(), + } + + with open("ressources/simple_ocp_croco_results.pkl", "wb") as handle: + pickle.dump(results, handle, protocol=pickle.HIGHEST_PROTOCOL) + + def test_check_results(self): + file_path = "ressources/simple_ocp_croco_results.pkl" + # Load the results + with open(file_path, "rb") as handle: + results = pickle.load(handle) + # Checking the states + for iter, state in enumerate(results["states"]): + np.testing.assert_array_almost_equal( + state, + self.ocp.ocp_results.states.tolist()[iter], + err_msg="States are not equal", + ) + + # Checking the ricatti gains + for iter, gain in enumerate(results["ricatti_gains"]): + np.testing.assert_array_almost_equal( + gain, + self.ocp.ocp_results.ricatti_gains.tolist()[iter], + err_msg="Ricatti gains are not equal", + ) + + # Checking the feed forward terms + for iter, term in enumerate(results["feed_forward_terms"]): + np.testing.assert_array_almost_equal( + term, + self.ocp.ocp_results.feed_forward_terms.tolist()[iter], + err_msg="Feed forward term are not equal", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/agimus_controller/tests/test_ocp_param_base.py b/agimus_controller/tests/test_ocp_param_base.py new file mode 100644 index 00000000..0855f7b4 --- /dev/null +++ b/agimus_controller/tests/test_ocp_param_base.py @@ -0,0 +1,39 @@ +import unittest + +from agimus_controller.ocp_param_base import OCPParamsBaseCroco + + +class TestOCPParamsCrocoBase(unittest.TestCase): + """ + TestOCPParamsCrocoBase unittests parameters settters and getters of OCPParamsBaseCroco class. + """ + + def __init__(self, methodName="runTest"): + super().__init__(methodName) + + def test_initialization(self): + """ + Test the initialization of the OCPParamsBaseCroco class. + """ + params = { + "dt": 0.01, + "horizon_size": 100, + "solver_iters": 50, + "qp_iters": 200, + "termination_tolerance": 1e-3, + "eps_abs": 1e-6, + "eps_rel": 0, + "callbacks": False, + } + ocp_param_base_croco = OCPParamsBaseCroco(**params) + for key, val in params.items(): + res = getattr(ocp_param_base_croco, key) + self.assertEqual( + res, + val, + f"Value missmatch for parameter '{key}'. Expected: '{val}', got: '{res}'", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/agimus_controller/tests/test_robot_models.py b/agimus_controller/tests/test_robot_models.py index c8ab1813..fd792c2e 100644 --- a/agimus_controller/tests/test_robot_models.py +++ b/agimus_controller/tests/test_robot_models.py @@ -18,9 +18,12 @@ def test_valid_initialization(self): srdf_path = Path(robot.urdf.replace("urdf", "srdf")) urdf_meshes_dir = urdf_path.parent.parent.parent.parent.parent free_flyer = False - q0 = np.zeros(robot.model.nq) + full_q0 = np.zeros(robot.model.nq) locked_joint_names = ["panda_joint1", "panda_joint2"] + reduced_nq = robot.model.nq - len(locked_joint_names) + q0 = np.zeros(reduced_nq) params = RobotModelParameters( + full_q0=full_q0, q0=q0, free_flyer=free_flyer, locked_joint_names=locked_joint_names, @@ -29,9 +32,9 @@ def test_valid_initialization(self): urdf_meshes_dir=urdf_meshes_dir, collision_as_capsule=True, self_collision=True, - armature=np.linspace(0.1, 0.9, 9), + armature=np.linspace(0.1, 0.9, reduced_nq), ) - + self.assertTrue(np.array_equal(params.full_q0, full_q0)) self.assertTrue(np.array_equal(params.q0, q0)) self.assertEqual(params.free_flyer, free_flyer) self.assertEqual(params.locked_joint_names, ["panda_joint1", "panda_joint2"]) @@ -40,7 +43,9 @@ def test_valid_initialization(self): self.assertEqual(params.urdf_meshes_dir, urdf_meshes_dir) self.assertTrue(params.collision_as_capsule) self.assertTrue(params.self_collision) - self.assertTrue(np.array_equal(params.armature, np.linspace(0.1, 0.9, 9))) + self.assertTrue( + np.array_equal(params.armature, np.linspace(0.1, 0.9, reduced_nq)) + ) def test_empty_q0_raises_error(self): """Test that an empty q0 raises a ValueError.""" @@ -55,8 +60,11 @@ def test_armature_default_value(self): urdf_meshes_dir = urdf_path.parent.parent.parent.parent.parent locked_joint_names = ["panda_joint1", "panda_joint2"] free_flyer = False - q0 = np.zeros(robot.model.nq) + full_q0 = np.zeros(robot.model.nq) + locked_joint_names = ["panda_joint1", "panda_joint2"] + q0 = np.zeros(robot.model.nq - len(locked_joint_names)) params = RobotModelParameters( + full_q0=full_q0, q0=q0, free_flyer=free_flyer, locked_joint_names=locked_joint_names, @@ -70,27 +78,27 @@ def test_armature_default_value(self): def test_armature_mismatched_shape_raises_error(self): """Test that a mismatched armature raises a ValueError.""" - q0 = np.array([0.0, 1.0, 2.0]) + full_q0 = np.array([0.0, 1.0, 2.0]) armature = np.array([0.1, 0.2]) # Wrong shape with self.assertRaises(ValueError): - RobotModelParameters(q0=q0, armature=armature) + RobotModelParameters(full_q0=full_q0, armature=armature) def test_invalid_urdf_path_raises_error(self): """Test that an invalid URDF path raises a ValueError.""" - q0 = np.array([0.0, 1.0, 2.0]) + full_q0 = np.array([0.0, 1.0, 2.0]) with self.assertRaises(ValueError): - RobotModelParameters(q0=q0, urdf_path=Path("invalid_path")) + RobotModelParameters(full_q0=full_q0, urdf_path=Path("invalid_path")) def test_invalid_srdf_path_type_raises_error(self): """Test that a non-string SRDF path raises a ValueError.""" - q0 = np.array([0.0, 1.0, 2.0]) + full_q0 = np.array([0.0, 1.0, 2.0]) with self.assertRaises(ValueError): - RobotModelParameters(q0=q0, srdf_path=Path("invalid_path")) + RobotModelParameters(full_q0=full_q0, srdf_path=Path("invalid_path")) class TestRobotModels(unittest.TestCase): @classmethod - def setUpClass(self): + def setUpClass(cls): """ This method sets up the shared environment for all test cases in the class. """ @@ -100,10 +108,12 @@ def setUpClass(self): srdf_path = Path(robot.urdf.replace("urdf", "srdf")) urdf_meshes_dir = urdf_path.parent.parent.parent.parent.parent free_flyer = False - q0 = np.zeros(robot.model.nq) + full_q0 = np.zeros(robot.model.nq) locked_joint_names = ["panda_joint1", "panda_joint2"] - # Store shared initial parameters - self.shared_params = RobotModelParameters( + reduced_nq = robot.model.nq - len(locked_joint_names) + q0 = np.zeros(reduced_nq) + cls.params = RobotModelParameters( + full_q0=full_q0, q0=q0, free_flyer=free_flyer, locked_joint_names=locked_joint_names, @@ -112,15 +122,15 @@ def setUpClass(self): urdf_meshes_dir=urdf_meshes_dir, collision_as_capsule=True, self_collision=True, - armature=np.linspace(0.1, 0.9, robot.model.nq), + armature=np.linspace(0.1, 0.9, reduced_nq), ) def setUp(self): """ This method ensures that a fresh RobotModelParameters and RobotModels instance - are created for each test case while still benefiting from shared setup computations. + are created for each test case. """ - self.params = deepcopy(self.shared_params) + self.params = deepcopy(self.params) self.robot_models = RobotModels(self.params) def test_initial_configuration(self):