/// <summary>
/// Utility method for writing a PoseExtractor to an ObservationWriter.
/// </summary>
/// <param name="writer"></param>
/// <param name="settings"></param>
/// <param name="poseExtractor"></param>
/// <param name="baseOffset">The offset into the ObservationWriter to start writing at.</param>
/// <returns>The number of observations written.</returns>
public static int WritePoses(this ObservationWriter writer, PhysicsSensorSettings settings, PoseExtractor poseExtractor, int baseOffset = 0)
{
var offset = baseOffset;
if (settings.UseModelSpace)
{
foreach (var pose in poseExtractor.GetEnabledModelSpacePoses())
{
if (settings.UseModelSpaceTranslations)
{
writer.Add(pose.position, offset);
offset += 3;
}
if (settings.UseModelSpaceRotations)
{
writer.Add(pose.rotation, offset);
offset += 4;
}
}
foreach (var vel in poseExtractor.GetEnabledModelSpaceVelocities())
{
if (settings.UseModelSpaceLinearVelocity)
{
writer.Add(vel, offset);
offset += 3;
}
}
}
if (settings.UseLocalSpace)
{
foreach (var pose in poseExtractor.GetEnabledLocalSpacePoses())
{
if (settings.UseLocalSpaceTranslations)
{
writer.Add(pose.position, offset);
offset += 3;
}
if (settings.UseLocalSpaceRotations)
{
writer.Add(pose.rotation, offset);
offset += 4;
}
}
foreach (var vel in poseExtractor.GetEnabledLocalSpaceVelocities())
{
if (settings.UseLocalSpaceLinearVelocity)
{
writer.Add(vel, offset);
offset += 3;
}
}
}
return(offset - baseOffset);
}
public PhysicsBodySensor(ArticulationBody rootBody, PhysicsSensorSettings settings, string sensorName = null)
{
var poseExtractor = new ArticulationBodyPoseExtractor(rootBody);
m_PoseExtractor = poseExtractor;
m_SensorName = string.IsNullOrEmpty(sensorName) ? $"ArticulationBodySensor:{rootBody?.name}" : sensorName;
m_Settings = settings;
var numJointExtractorObservations = 0;
var articBodies = poseExtractor.Bodies;
if (articBodies != null)
{
m_JointExtractors = new IJointExtractor[articBodies.Length - 1]; // skip the root
for (var i = 1; i < articBodies.Length; i++)
{
var jointExtractor = new ArticulationBodyJointExtractor(articBodies[i]);
numJointExtractorObservations += jointExtractor.NumObservations(settings);
m_JointExtractors[i - 1] = jointExtractor;
}
}
else
{
m_JointExtractors = new IJointExtractor[0];
}
var numTransformObservations = m_PoseExtractor.GetNumPoseObservations(settings);
m_Shape = new[] { numTransformObservations + numJointExtractorObservations };
}
/// <summary>
/// Utility method for writing a PoseExtractor to an ObservationWriter.
/// </summary>
/// <param name="writer"></param>
/// <param name="settings"></param>
/// <param name="poseExtractor"></param>
/// <param name="baseOffset">The offset into the ObservationWriter to start writing at.</param>
/// <returns>The number of observations written.</returns>
public static int WritePoses(this ObservationWriter writer, PhysicsSensorSettings settings, PoseExtractor poseExtractor, int baseOffset = 0)
{
var offset = baseOffset;
if (settings.UseModelSpace)
{
var poses = poseExtractor.ModelSpacePoses;
var vels = poseExtractor.ModelSpaceVelocities;
for (var i = 0; i < poseExtractor.NumPoses; i++)
{
var pose = poses[i];
if (settings.UseModelSpaceTranslations)
{
writer.Add(pose.position, offset);
offset += 3;
}
if (settings.UseModelSpaceRotations)
{
writer.Add(pose.rotation, offset);
offset += 4;
}
if (settings.UseModelSpaceLinearVelocity)
{
writer.Add(vels[i], offset);
offset += 3;
}
}
}
if (settings.UseLocalSpace)
{
var poses = poseExtractor.LocalSpacePoses;
var vels = poseExtractor.LocalSpaceVelocities;
for (var i = 0; i < poseExtractor.NumPoses; i++)
{
var pose = poses[i];
if (settings.UseLocalSpaceTranslations)
{
writer.Add(pose.position, offset);
offset += 3;
}
if (settings.UseLocalSpaceRotations)
{
writer.Add(pose.rotation, offset);
offset += 4;
}
if (settings.UseLocalSpaceLinearVelocity)
{
writer.Add(vels[i], offset);
offset += 3;
}
}
}
return(offset - baseOffset);
}
public PhysicsBodySensor(ArticulationBody rootBody, PhysicsSensorSettings settings, string sensorName = null)
{
m_PoseExtractor = new ArticulationBodyPoseExtractor(rootBody);
m_SensorName = string.IsNullOrEmpty(sensorName) ? $"ArticulationBodySensor:{rootBody?.name}" : sensorName;
m_Settings = settings;
var numTransformObservations = settings.TransformSize(m_PoseExtractor.NumPoses);
m_Shape = new[] { numTransformObservations };
}
/// <summary>
/// Construct a new PhysicsBodySensor
/// </summary>
/// <param name="rootBody"></param>
/// <param name="settings"></param>
/// <param name="sensorName"></param>
public PhysicsBodySensor(Rigidbody rootBody, GameObject rootGameObject, PhysicsSensorSettings settings, string sensorName = null)
{
m_PoseExtractor = new RigidBodyPoseExtractor(rootBody, rootGameObject);
m_SensorName = string.IsNullOrEmpty(sensorName) ? $"PhysicsBodySensor:{rootBody?.name}" : sensorName;
m_Settings = settings;
var numTransformObservations = settings.TransformSize(m_PoseExtractor.NumPoses);
m_Shape = new[] { numTransformObservations };
}
/// <summary>
/// Compute the number of floats needed to represent the poses for the given PhysicsSensorSettings.
/// </summary>
/// <param name="settings"></param>
/// <returns></returns>
public int GetNumPoseObservations(PhysicsSensorSettings settings)
{
int obsPerPose = 0;
obsPerPose += settings.UseModelSpaceTranslations ? 3 : 0;
obsPerPose += settings.UseModelSpaceRotations ? 4 : 0;
obsPerPose += settings.UseLocalSpaceTranslations ? 3 : 0;
obsPerPose += settings.UseLocalSpaceRotations ? 4 : 0;
obsPerPose += settings.UseModelSpaceLinearVelocity ? 3 : 0;
obsPerPose += settings.UseLocalSpaceLinearVelocity ? 3 : 0;
return(NumEnabledPoses * obsPerPose);
}
public static int NumObservations(Rigidbody body, Joint joint, PhysicsSensorSettings settings)
{
if (body == null || joint == null)
{
return(0);
}
var numObservations = 0;
if (settings.UseJointForces)
{
// 3 force and 3 torque values
numObservations += 6;
}
return(numObservations);
}
public int Write(PhysicsSensorSettings settings, ObservationWriter writer, int offset)
{
if (m_Body == null || m_Body.isRoot)
{
return(0);
}
var currentOffset = offset;
// Write joint positions
if (settings.UseJointPositionsAndAngles)
{
switch (m_Body.jointType)
{
case ArticulationJointType.RevoluteJoint:
case ArticulationJointType.SphericalJoint:
// All joint positions are angular
for (var dofIndex = 0; dofIndex < m_Body.dofCount; dofIndex++)
{
var jointRotationRads = m_Body.jointPosition[dofIndex];
writer[currentOffset++] = Mathf.Sin(jointRotationRads);
writer[currentOffset++] = Mathf.Cos(jointRotationRads);
}
break;
case ArticulationJointType.FixedJoint:
// No observations
break;
case ArticulationJointType.PrismaticJoint:
writer[currentOffset++] = GetPrismaticValue();
break;
}
}
if (settings.UseJointForces)
{
for (var dofIndex = 0; dofIndex < m_Body.dofCount; dofIndex++)
{
// take tanh to keep in [-1, 1]
writer[currentOffset++] = (float)System.Math.Tanh(m_Body.jointForce[dofIndex]);
}
}
return(currentOffset - offset);
}
public PhysicsBodySensor(ArticulationBody rootBody, PhysicsSensorSettings settings, string sensorName = null)
{
var poseExtractor = new ArticulationBodyPoseExtractor(rootBody);
m_PoseExtractor = poseExtractor;
m_SensorName = string.IsNullOrEmpty(sensorName) ? $"ArticulationBodySensor:{rootBody?.name}" : sensorName;
m_Settings = settings;
var numJointExtractorObservations = 0;
m_JointExtractors = new List <IJointExtractor>(poseExtractor.NumEnabledPoses);
foreach (var articBody in poseExtractor.GetEnabledArticulationBodies())
{
var jointExtractor = new ArticulationBodyJointExtractor(articBody);
numJointExtractorObservations += jointExtractor.NumObservations(settings);
m_JointExtractors.Add(jointExtractor);
}
var numTransformObservations = m_PoseExtractor.GetNumPoseObservations(settings);
m_Shape = new[] { numTransformObservations + numJointExtractorObservations };
}
public int Write(PhysicsSensorSettings settings, ObservationWriter writer, int offset)
{
if (m_Body == null || m_Joint == null)
{
return(0);
}
var currentOffset = offset;
if (settings.UseJointForces)
{
// Take tanh of the forces and torques to ensure they're in [-1, 1]
writer[currentOffset++] = (float)System.Math.Tanh(m_Joint.currentForce.x);
writer[currentOffset++] = (float)System.Math.Tanh(m_Joint.currentForce.y);
writer[currentOffset++] = (float)System.Math.Tanh(m_Joint.currentForce.z);
writer[currentOffset++] = (float)System.Math.Tanh(m_Joint.currentTorque.x);
writer[currentOffset++] = (float)System.Math.Tanh(m_Joint.currentTorque.y);
writer[currentOffset++] = (float)System.Math.Tanh(m_Joint.currentTorque.z);
}
return(currentOffset - offset);
}
public static int NumObservations(ArticulationBody body, PhysicsSensorSettings settings)
{
if (body == null || body.isRoot)
{
return(0);
}
var totalCount = 0;
if (settings.UseJointPositionsAndAngles)
{
switch (body.jointType)
{
case ArticulationJointType.RevoluteJoint:
case ArticulationJointType.SphericalJoint:
// Both RevoluteJoint and SphericalJoint have all angular components.
// We use sine and cosine of the angles for the observations.
totalCount += 2 * body.dofCount;
break;
case ArticulationJointType.FixedJoint:
// Since FixedJoint can't moved, there aren't any interesting observations for it.
break;
case ArticulationJointType.PrismaticJoint:
// One linear component
totalCount += body.dofCount;
break;
}
}
if (settings.UseJointForces)
{
totalCount += body.dofCount;
}
return(totalCount);
}
/// <summary>
/// Construct a new PhysicsBodySensor
/// </summary>
/// <param name="poseExtractor"></param>
/// <param name="settings"></param>
/// <param name="sensorName"></param>
public PhysicsBodySensor(
RigidBodyPoseExtractor poseExtractor,
PhysicsSensorSettings settings,
string sensorName
)
{
m_PoseExtractor = poseExtractor;
m_SensorName = sensorName;
m_Settings = settings;
var numJointExtractorObservations = 0;
m_JointExtractors = new List <IJointExtractor>(poseExtractor.NumEnabledPoses);
foreach (var rb in poseExtractor.GetEnabledRigidbodies())
{
var jointExtractor = new RigidBodyJointExtractor(rb);
numJointExtractorObservations += jointExtractor.NumObservations(settings);
m_JointExtractors.Add(jointExtractor);
}
var numTransformObservations = m_PoseExtractor.GetNumPoseObservations(settings);
m_Shape = new[] { numTransformObservations + numJointExtractorObservations };
}
/// <summary>
/// Construct a new PhysicsBodySensor
/// </summary>
/// <param name="rootBody">The root Rigidbody. This has no Joints on it (but other Joints may connect to it).</param>
/// <param name="rootGameObject">Optional GameObject used to find Rigidbodies in the hierarchy.</param>
/// <param name="virtualRoot">Optional GameObject used to determine the root of the poses,
/// <param name="settings"></param>
/// <param name="sensorName"></param>
public PhysicsBodySensor(
Rigidbody rootBody,
GameObject rootGameObject,
GameObject virtualRoot,
PhysicsSensorSettings settings,
string sensorName = null
)
{
var poseExtractor = new RigidBodyPoseExtractor(rootBody, rootGameObject, virtualRoot);
m_PoseExtractor = poseExtractor;
m_SensorName = string.IsNullOrEmpty(sensorName) ? $"PhysicsBodySensor:{rootBody?.name}" : sensorName;
m_Settings = settings;
var numJointExtractorObservations = 0;
var rigidBodies = poseExtractor.Bodies;
if (rigidBodies != null)
{
m_JointExtractors = new IJointExtractor[rigidBodies.Length - 1]; // skip the root
for (var i = 1; i < rigidBodies.Length; i++)
{
var jointExtractor = new RigidBodyJointExtractor(rigidBodies[i]);
numJointExtractorObservations += jointExtractor.NumObservations(settings);
m_JointExtractors[i - 1] = jointExtractor;
}
}
else
{
m_JointExtractors = new IJointExtractor[0];
}
var numTransformObservations = m_PoseExtractor.GetNumPoseObservations(settings);
m_Shape = new[] { numTransformObservations + numJointExtractorObservations };
}
public int NumObservations(PhysicsSensorSettings settings)
{
return(NumObservations(m_Body, m_Joint, settings));
}
