From 7441515c0ed7e1005c85184ab614f3b77ad3a471 Mon Sep 17 00:00:00 2001
From: Erwin Coumans <erwin.coumans@gmail.com>
Date: Mon, 19 Jun 2017 13:15:05 -0700
Subject: [PATCH] Preliminary version of pybullet.createMultiBody including
 links connected to parent by a joint. See createMultiBodyLinks.py example.

---
 examples/SharedMemory/PhysicsClientC_API.cpp  | 28 +++---
 examples/SharedMemory/PhysicsClientC_API.h    |  2 +
 .../PhysicsServerCommandProcessor.cpp         | 85 +++++++++++++++++--
 examples/SharedMemory/SharedMemoryCommands.h  |  1 -
 .../pybullet/examples/createMultiBodyLinks.py | 57 +++++++++++++
 examples/pybullet/pybullet.c                  | 33 +++++--
 6 files changed, 181 insertions(+), 25 deletions(-)
 create mode 100644 examples/pybullet/examples/createMultiBodyLinks.py

diff --git a/examples/SharedMemory/PhysicsClientC_API.cpp b/examples/SharedMemory/PhysicsClientC_API.cpp
index dde3d2ed8..84f856467 100644
--- a/examples/SharedMemory/PhysicsClientC_API.cpp
+++ b/examples/SharedMemory/PhysicsClientC_API.cpp
@@ -993,6 +993,12 @@ int b3CreateMultiBodyBase(b3SharedMemoryCommandHandle commandHandle, double mass
 			command->m_createMultiBodyArgs.m_linkVisualShapeUniqueIds[baseLinkIndex] = visualShapeUniqueId;
 
 			command->m_createMultiBodyArgs.m_linkMasses[baseLinkIndex] = mass;
+
+			command->m_createMultiBodyArgs.m_linkParentIndices[baseLinkIndex] = -2;//no parent
+			command->m_createMultiBodyArgs.m_linkJointAxis[baseLinkIndex+0]=0;
+			command->m_createMultiBodyArgs.m_linkJointAxis[baseLinkIndex+1]=0;
+			command->m_createMultiBodyArgs.m_linkJointAxis[baseLinkIndex+2]=0;
+			command->m_createMultiBodyArgs.m_linkJointTypes[baseLinkIndex]=-1;
 			command->m_createMultiBodyArgs.m_numLinks++;
 		}
 		return numLinks;
@@ -1004,6 +1010,8 @@ int b3CreateMultiBodyLink(b3SharedMemoryCommandHandle commandHandle, double link
 									double linkVisualShapeIndex, 
 									double linkPosition[3], 
 									double linkOrientation[4],
+									double linkInertialFramePosition[3], 
+									double linkInertialFrameOrientation[4],
 									int linkParentIndex,
 									int linkJointType,
 									double linkJointAxis[3])
@@ -1033,23 +1041,23 @@ int b3CreateMultiBodyLink(b3SharedMemoryCommandHandle commandHandle, double link
 			command->m_createMultiBodyArgs.m_linkInertias[linkIndex*3+2] = linkMass;
 
 
-			command->m_createMultiBodyArgs.m_linkInertialFramePositions[linkIndex*3+0] = 0;
-			command->m_createMultiBodyArgs.m_linkInertialFramePositions[linkIndex*3+1] = 0;
-			command->m_createMultiBodyArgs.m_linkInertialFramePositions[linkIndex*3+2] = 0;
+			command->m_createMultiBodyArgs.m_linkInertialFramePositions[linkIndex*3+0] = linkInertialFramePosition[0];
+			command->m_createMultiBodyArgs.m_linkInertialFramePositions[linkIndex*3+1] = linkInertialFramePosition[1];
+			command->m_createMultiBodyArgs.m_linkInertialFramePositions[linkIndex*3+2] = linkInertialFramePosition[2];
 
-			command->m_createMultiBodyArgs.m_linkInertialFrameOrientations[linkIndex*4+0] = 0;
-			command->m_createMultiBodyArgs.m_linkInertialFrameOrientations[linkIndex*4+1] = 0;
-			command->m_createMultiBodyArgs.m_linkInertialFrameOrientations[linkIndex*4+2] = 0;
-			command->m_createMultiBodyArgs.m_linkInertialFrameOrientations[linkIndex*4+3] = 1;
+			command->m_createMultiBodyArgs.m_linkInertialFrameOrientations[linkIndex*4+0] = linkInertialFrameOrientation[0];
+			command->m_createMultiBodyArgs.m_linkInertialFrameOrientations[linkIndex*4+1] = linkInertialFrameOrientation[1];
+			command->m_createMultiBodyArgs.m_linkInertialFrameOrientations[linkIndex*4+2] = linkInertialFrameOrientation[2];
+			command->m_createMultiBodyArgs.m_linkInertialFrameOrientations[linkIndex*4+3] = linkInertialFrameOrientation[3];
 
 			command->m_createMultiBodyArgs.m_linkCollisionShapeUniqueIds[linkIndex]= linkCollisionShapeIndex;
 			command->m_createMultiBodyArgs.m_linkVisualShapeUniqueIds[linkIndex] = linkVisualShapeIndex;
 			
 			command->m_createMultiBodyArgs.m_linkParentIndices[linkIndex] = linkParentIndex;
 			command->m_createMultiBodyArgs.m_linkJointTypes[linkIndex] = linkJointType;
-			command->m_createMultiBodyArgs.m_linkJointAxis[0] = linkJointAxis[0];
-			command->m_createMultiBodyArgs.m_linkJointAxis[1] = linkJointAxis[1];
-			command->m_createMultiBodyArgs.m_linkJointAxis[2] = linkJointAxis[2];
+			command->m_createMultiBodyArgs.m_linkJointAxis[3*linkIndex+0] = linkJointAxis[0];
+			command->m_createMultiBodyArgs.m_linkJointAxis[3*linkIndex+1] = linkJointAxis[1];
+			command->m_createMultiBodyArgs.m_linkJointAxis[3*linkIndex+2] = linkJointAxis[2];
 
 			command->m_createMultiBodyArgs.m_linkMasses[linkIndex] = linkMass;
 			command->m_createMultiBodyArgs.m_numLinks++;
diff --git a/examples/SharedMemory/PhysicsClientC_API.h b/examples/SharedMemory/PhysicsClientC_API.h
index 24a3ddd77..9b32a0daf 100644
--- a/examples/SharedMemory/PhysicsClientC_API.h
+++ b/examples/SharedMemory/PhysicsClientC_API.h
@@ -348,6 +348,8 @@ int b3CreateMultiBodyLink(b3SharedMemoryCommandHandle commandHandle, double link
 									double linkVisualShapeIndex, 
 									double linkPosition[3], 
 									double linkOrientation[4],
+									double linkInertialFramePosition[3],
+									double linkInertialFrameOrientation[4],
 									int linkParentIndex,
 									int linkJointType,
 									double linkJointAxis[3]);
diff --git a/examples/SharedMemory/PhysicsServerCommandProcessor.cpp b/examples/SharedMemory/PhysicsServerCommandProcessor.cpp
index ee0412a9b..d3a62ec52 100644
--- a/examples/SharedMemory/PhysicsServerCommandProcessor.cpp
+++ b/examples/SharedMemory/PhysicsServerCommandProcessor.cpp
@@ -1547,7 +1547,11 @@ struct ProgrammaticUrdfInterface : public URDFImporterInterface
     ///pure virtual interfaces, precondition is a valid linkIndex (you can assert/terminate if the linkIndex is out of range)
     virtual std::string getLinkName(int linkIndex) const
 	{
-		return "link";
+		std::string linkName = "link";
+		char numstr[21]; // enough to hold all numbers up to 64-bits
+		sprintf(numstr, "%d", linkIndex);
+		linkName = linkName + numstr;
+		return linkName;
 	}
 
 	//various derived class in internal source code break with new pure virtual methods, so provide some default implementation
@@ -1587,8 +1591,11 @@ struct ProgrammaticUrdfInterface : public URDFImporterInterface
 
     virtual std::string getJointName(int linkIndex) const
 	{
-		b3Assert(0);
-		return "joint";
+		std::string jointName = "joint";
+		char numstr[21]; // enough to hold all numbers up to 64-bits
+		sprintf(numstr, "%d", linkIndex);
+		jointName = jointName + numstr;
+		return jointName;
 	}
 	
     //fill mass and inertial data. If inertial data is missing, please initialize mass, inertia to sensitive values, and inertialFrame to identity.
@@ -1623,7 +1630,7 @@ struct ProgrammaticUrdfInterface : public URDFImporterInterface
 	{
 		for (int i=0;i<m_createBodyArgs.m_numLinks;i++)
 		{
-			if (m_createBodyArgs.m_linkParents[i] == urdfLinkIndex)
+			if (m_createBodyArgs.m_linkParentIndices[i] == urdfLinkIndex)
 			{
 				childLinkIndices.push_back(i);
 			}
@@ -1638,10 +1645,72 @@ struct ProgrammaticUrdfInterface : public URDFImporterInterface
 
 	virtual bool getJointInfo2(int urdfLinkIndex, btTransform& parent2joint, btTransform& linkTransformInWorld, btVector3& jointAxisInJointSpace, int& jointType, btScalar& jointLowerLimit, btScalar& jointUpperLimit, btScalar& jointDamping, btScalar& jointFriction, btScalar& jointMaxForce, btScalar& jointMaxVelocity) const 
 	{
-		//backwards compatibility for custom file importers
-		jointMaxForce = 0;
-		jointMaxVelocity = 0;
-		return getJointInfo(urdfLinkIndex, parent2joint, linkTransformInWorld, jointAxisInJointSpace, jointType, jointLowerLimit, jointUpperLimit, jointDamping, jointFriction);
+		bool isValid = false;
+
+		int jointTypeOrg = m_createBodyArgs.m_linkJointTypes[urdfLinkIndex];
+
+		switch (jointTypeOrg)
+		{
+		case eRevoluteType:
+		{
+			isValid = true;
+			jointType = URDFRevoluteJoint;
+			break;
+		}
+		case	ePrismaticType:
+		{
+			isValid = true;
+			jointType = URDFPrismaticJoint;
+			break;
+		}
+		case	eFixedType:
+		{
+			isValid = true;
+			jointType = URDFFixedJoint;
+			break;
+		}
+		//case	eSphericalType:
+		//case	ePlanarType:
+		//case	eFixedType:
+		//case ePoint2PointType:
+		//case eGearType:
+		default:
+		{
+		}
+		};
+
+		if (isValid)
+		{
+			//backwards compatibility for custom file importers
+			jointMaxForce = 0;
+			jointMaxVelocity = 0;
+			jointFriction = 0;
+			jointDamping = 0;
+			jointLowerLimit = 1;
+			jointUpperLimit = -1;
+
+			parent2joint.setOrigin(btVector3(
+				m_createBodyArgs.m_linkPositions[urdfLinkIndex*3+0],
+				m_createBodyArgs.m_linkPositions[urdfLinkIndex*3+1],
+				m_createBodyArgs.m_linkPositions[urdfLinkIndex*3+2]));
+			parent2joint.setRotation(btQuaternion(
+				m_createBodyArgs.m_linkOrientations[urdfLinkIndex*4+0],
+				m_createBodyArgs.m_linkOrientations[urdfLinkIndex*4+1],
+				m_createBodyArgs.m_linkOrientations[urdfLinkIndex*4+2],
+				m_createBodyArgs.m_linkOrientations[urdfLinkIndex*4+3]
+			));
+
+		
+			linkTransformInWorld.setIdentity();
+
+			jointAxisInJointSpace.setValue(
+				m_createBodyArgs.m_linkJointAxis[3*urdfLinkIndex+0],
+				m_createBodyArgs.m_linkJointAxis[3*urdfLinkIndex+1],
+				m_createBodyArgs.m_linkJointAxis[3*urdfLinkIndex+2]);
+
+		
+		}
+		return isValid;
 	};
     
     virtual bool getRootTransformInWorld(btTransform& rootTransformInWorld) const
diff --git a/examples/SharedMemory/SharedMemoryCommands.h b/examples/SharedMemory/SharedMemoryCommands.h
index 260c6cd7d..9658d6e6c 100644
--- a/examples/SharedMemory/SharedMemoryCommands.h
+++ b/examples/SharedMemory/SharedMemoryCommands.h
@@ -854,7 +854,6 @@ struct b3CreateMultiBodyArgs
 	double m_linkOrientations[4*MAX_CREATE_MULTI_BODY_LINKS];
 
 	int m_numLinks;
-	int m_linkParents[MAX_CREATE_MULTI_BODY_LINKS];
 	double m_linkMasses[MAX_CREATE_MULTI_BODY_LINKS];
 	double m_linkInertias[MAX_CREATE_MULTI_BODY_LINKS*3];
 
diff --git a/examples/pybullet/examples/createMultiBodyLinks.py b/examples/pybullet/examples/createMultiBodyLinks.py
new file mode 100644
index 000000000..7cfc03272
--- /dev/null
+++ b/examples/pybullet/examples/createMultiBodyLinks.py
@@ -0,0 +1,57 @@
+import pybullet as p
+import time
+
+
+p.connect(p.GUI)
+p.createCollisionShape(p.GEOM_PLANE)
+p.createMultiBody(0,0)
+
+sphereRadius = 0.05
+colSphereId = p.createCollisionShape(p.GEOM_SPHERE,radius=sphereRadius)
+colBoxId = p.createCollisionShape(p.GEOM_BOX,halfExtents=[sphereRadius,sphereRadius,sphereRadius])
+
+mass = 1
+visualShapeId = -1
+
+#"linkMasses","linkCollisionShapeIndices", "linkVisualShapeIndices",
+#"linkPositions", "linkOrientations","linkParentIndices", "linkJointTypes","linkJointAxis", 
+
+linkMasses=[1]
+linkCollisionShapeIndices=[colBoxId]
+linkVisualShapeIndices=[-1]
+linkPositions=[[0,0,0.11]]
+linkOrientations=[[0,0,0,1]]
+linkInertialFramePositions=[[0,0,0]]
+linkInertialFrameOrientations=[[0,0,0,1]]
+indices=[0]
+jointTypes=[p.JOINT_REVOLUTE]
+axis=[[0,0,1]]
+
+for i in range (3):
+	for j in range (3):
+		for k in range (3):
+			basePosition = [1+i*5*sphereRadius,1+j*5*sphereRadius,1+k*5*sphereRadius+1]
+			baseOrientation = [0,0,0,1]
+			if (k&2):
+				sphereUid = p.createMultiBody(mass,colSphereId,visualShapeId,basePosition,baseOrientation)
+			else:
+				sphereUid = p.createMultiBody(mass,colBoxId,visualShapeId,basePosition,baseOrientation,linkMasses,linkCollisionShapeIndices,linkVisualShapeIndices,linkPositions,linkOrientations,linkInertialFramePositions, linkInertialFrameOrientations,linkParentIndices=indices,linkJointTypes=jointTypes,linkJointAxis=axis)			
+			
+			p.changeDynamics(sphereUid,-1,spinningFriction=0.001, rollingFriction=0.001,linearDamping=0.0)
+			for joint in range (p.getNumJoints(sphereUid)):
+				p.setJointMotorControl2(sphereUid,joint,p.VELOCITY_CONTROL,targetVelocity=1,force=10)
+
+
+p.setGravity(0,0,-10)
+#p.setRealTimeSimulation(1)
+
+p.getNumJoints(sphereUid)
+for i in range (p.getNumJoints(sphereUid)):
+	p.getJointInfo(sphereUid,i)
+	
+while (1):
+	events = p.getKeyboardEvents()
+	if (len(events)):
+		p.stepSimulation()
+	time.sleep(0.01)
+	
\ No newline at end of file
diff --git a/examples/pybullet/pybullet.c b/examples/pybullet/pybullet.c
index 478c92bd7..fdbf401a0 100644
--- a/examples/pybullet/pybullet.c
+++ b/examples/pybullet/pybullet.c
@@ -5093,17 +5093,19 @@ static PyObject* pybullet_createMultiBody(PyObject* self, PyObject* args, PyObje
 	PyObject* linkParentIndicesObj=0;
 	PyObject* linkJointTypesObj=0;
 	PyObject* linkJointAxisObj=0;
-	
+	PyObject* linkInertialFramePositionObj=0;
+	PyObject* linkInertialFrameOrientationObj=0;
 
 	static char* kwlist[] = {"baseMass","baseCollisionShapeIndex","baseVisualShapeIndex","basePosition", "baseOrientation",
 		"linkMasses","linkCollisionShapeIndices", "linkVisualShapeIndices",
-		"linkPositions", "linkOrientations","linkParentIndices", "linkJointTypes","linkJointAxis", 
+		"linkPositions", "linkOrientations","linkInertialFramePosition","linkInertialFrameOrientation", "linkParentIndices", "linkJointTypes","linkJointAxis", 
 		"useMaximalCoordinates","physicsClientId",
 
 		NULL};
-	if (!PyArg_ParseTupleAndKeywords(args, keywds, "|diiOOOOOOOOOOii", kwlist, &baseMass,&baseCollisionShapeIndex,&baseVisualShapeIndex,
+	if (!PyArg_ParseTupleAndKeywords(args, keywds, "|diiOOOOOOOOOOOOii", kwlist, &baseMass,&baseCollisionShapeIndex,&baseVisualShapeIndex,
 				&basePosObj, &baseOrnObj,
 		&linkMassesObj, &linkCollisionShapeIndicesObj, &linkVisualShapeIndicesObj, &linkPositionsObj, &linkOrientationsObj,
+		&linkInertialFramePositionObj, &linkInertialFrameOrientationObj,
 		&linkParentIndicesObj, &linkJointTypesObj,&linkJointAxisObj,
 		&useMaximalCoordinates, &physicsClientId))
 	{
@@ -5127,7 +5129,8 @@ static PyObject* pybullet_createMultiBody(PyObject* self, PyObject* args, PyObje
 		int numLinkParentIndices = linkParentIndicesObj?PySequence_Size(linkParentIndicesObj):0;
 		int numLinkJointTypes = linkJointTypesObj?PySequence_Size(linkJointTypesObj):0;
 		int numLinkJoinAxis = linkJointAxisObj? PySequence_Size(linkJointAxisObj):0;
-
+		int numLinkInertialFramePositions = linkInertialFramePositionObj? PySequence_Size(linkInertialFramePositionObj) : 0;
+		int numLinkInertialFrameOrientations = linkInertialFrameOrientationObj? PySequence_Size(linkInertialFrameOrientationObj) : 0;
 
 		PyObject* seqLinkMasses = linkMassesObj?PySequence_Fast(linkMassesObj, "expected a sequence"):0;
 		PyObject* seqLinkCollisionShapes = linkCollisionShapeIndicesObj?PySequence_Fast(linkCollisionShapeIndicesObj, "expected a sequence"):0;
@@ -5137,6 +5140,8 @@ static PyObject* pybullet_createMultiBody(PyObject* self, PyObject* args, PyObje
 		PyObject* seqLinkParentIndices = linkParentIndicesObj?PySequence_Fast(linkParentIndicesObj, "expected a sequence"):0;
 		PyObject* seqLinkJointTypes = linkJointTypesObj?PySequence_Fast(linkJointTypesObj, "expected a sequence"):0;
 		PyObject* seqLinkJoinAxis = linkJointAxisObj?PySequence_Fast(linkJointAxisObj, "expected a sequence"):0;
+		PyObject* seqLinkInertialFramePositions = linkInertialFramePositionObj?PySequence_Fast(linkInertialFramePositionObj, "expected a sequence"):0;
+		PyObject* seqLinkInertialFrameOrientations = linkInertialFrameOrientationObj?PySequence_Fast(linkInertialFrameOrientationObj, "expected a sequence"):0;
 
 		if ((numLinkMasses==numLinkCollisionShapes) &&
 			(numLinkMasses==numLinkVisualShapes) &&
@@ -5144,7 +5149,9 @@ static PyObject* pybullet_createMultiBody(PyObject* self, PyObject* args, PyObje
 			(numLinkMasses==numLinkOrientations) &&
 			(numLinkMasses==numLinkParentIndices) &&
 			(numLinkMasses==numLinkJointTypes) &&
-			(numLinkMasses==numLinkJoinAxis))
+			(numLinkMasses==numLinkJoinAxis) &&
+			(numLinkMasses==numLinkInertialFramePositions) &&
+			(numLinkMasses==numLinkInertialFrameOrientations))
 		{
 			b3SharedMemoryStatusHandle statusHandle;
 			int statusType;
@@ -5165,7 +5172,11 @@ static PyObject* pybullet_createMultiBody(PyObject* self, PyObject* args, PyObje
 				double linkPosition[3];
 				double linkOrientation[4];
 				double linkJointAxis[3];
-
+				double linkInertialFramePosition[3];
+				double linkInertialFrameOrientation[4];
+				
+				pybullet_internalGetVector3FromSequence(seqLinkInertialFramePositions,i,linkInertialFramePosition);
+				pybullet_internalGetVector4FromSequence(linkInertialFrameOrientationObj,i,linkInertialFrameOrientation);
 				pybullet_internalGetVector3FromSequence(seqLinkPositions,i,linkPosition);
 				pybullet_internalGetVector4FromSequence(seqLinkOrientations,i,linkOrientation);
 				pybullet_internalGetVector3FromSequence(seqLinkJoinAxis,i,linkJointAxis);
@@ -5178,6 +5189,8 @@ static PyObject* pybullet_createMultiBody(PyObject* self, PyObject* args, PyObje
 									linkVisualShapeIndex, 
 									linkPosition, 
 									linkOrientation,
+									linkInertialFramePosition,
+									linkInertialFrameOrientation,
 									linkParentIndex,
 									linkJointType,
 									linkJointAxis
@@ -5201,6 +5214,10 @@ static PyObject* pybullet_createMultiBody(PyObject* self, PyObject* args, PyObje
 				Py_DECREF(seqLinkJointTypes);
 			if (seqLinkJoinAxis)
 				Py_DECREF(seqLinkJoinAxis);
+			if (seqLinkInertialFramePositions)
+				Py_DECREF(seqLinkInertialFramePositions);
+			if (seqLinkInertialFrameOrientations)
+				Py_DECREF(seqLinkInertialFrameOrientations);
 
 			if (useMaximalCoordinates>0)
 			{
@@ -5233,6 +5250,10 @@ static PyObject* pybullet_createMultiBody(PyObject* self, PyObject* args, PyObje
 				Py_DECREF(seqLinkJointTypes);
 			if (seqLinkJoinAxis)
 				Py_DECREF(seqLinkJoinAxis);
+			if (seqLinkInertialFramePositions)
+				Py_DECREF(seqLinkInertialFramePositions);
+			if (seqLinkInertialFrameOrientations)
+				Py_DECREF(seqLinkInertialFrameOrientations);
 
 			PyErr_SetString(SpamError, "All link arrays need to be same size.");
 			return NULL;