using UnityEngine; using MLAgents; [RequireComponent(typeof(JointDriveController))] // Required to set joint forces public class CrawlerAgent : Agent { [Header("Target To Walk Towards")][Space(10)] public Transform target; public Transform ground; public bool detectTargets; public bool targetIsStatic = false; public bool respawnTargetWhenTouched; public float targetSpawnRadius; [Header("Body Parts")][Space(10)] public Transform body; public Transform leg0Upper; public Transform leg0Lower; public Transform leg1Upper; public Transform leg1Lower; public Transform leg2Upper; public Transform leg2Lower; public Transform leg3Upper; public Transform leg3Lower; [Header("Joint Settings")][Space(10)] JointDriveController m_JdController; Vector3 m_DirToTarget; float m_MovingTowardsDot; float m_FacingDot; [Header("Reward Functions To Use")][Space(10)] public bool rewardMovingTowardsTarget; // Agent should move towards target public bool rewardFacingTarget; // Agent should face the target public bool rewardUseTimePenalty; // Hurry up [Header("Foot Grounded Visualization")][Space(10)] public bool useFootGroundedVisualization; public MeshRenderer foot0; public MeshRenderer foot1; public MeshRenderer foot2; public MeshRenderer foot3; public Material groundedMaterial; public Material unGroundedMaterial; bool m_IsNewDecisionStep; int m_CurrentDecisionStep; Quaternion m_LookRotation; Matrix4x4 m_TargetDirMatrix; public override void InitializeAgent() { m_JdController = GetComponent(); m_CurrentDecisionStep = 1; m_DirToTarget = target.position - body.position; //Setup each body part m_JdController.SetupBodyPart(body); m_JdController.SetupBodyPart(leg0Upper); m_JdController.SetupBodyPart(leg0Lower); m_JdController.SetupBodyPart(leg1Upper); m_JdController.SetupBodyPart(leg1Lower); m_JdController.SetupBodyPart(leg2Upper); m_JdController.SetupBodyPart(leg2Lower); m_JdController.SetupBodyPart(leg3Upper); m_JdController.SetupBodyPart(leg3Lower); } /// /// We only need to change the joint settings based on decision freq. /// public void IncrementDecisionTimer() { if (m_CurrentDecisionStep == agentParameters.numberOfActionsBetweenDecisions || agentParameters.numberOfActionsBetweenDecisions == 1) { m_CurrentDecisionStep = 1; m_IsNewDecisionStep = true; } else { m_CurrentDecisionStep++; m_IsNewDecisionStep = false; } } /// /// Add relevant information on each body part to observations. /// public void CollectObservationBodyPart(BodyPart bp) { var rb = bp.rb; AddVectorObs(bp.groundContact.touchingGround ? 1 : 0); // Whether the bp touching the ground var velocityRelativeToLookRotationToTarget = m_TargetDirMatrix.inverse.MultiplyVector(rb.velocity); AddVectorObs(velocityRelativeToLookRotationToTarget); var angularVelocityRelativeToLookRotationToTarget = m_TargetDirMatrix.inverse.MultiplyVector(rb.angularVelocity); AddVectorObs(angularVelocityRelativeToLookRotationToTarget); if (bp.rb.transform != body) { var localPosRelToBody = body.InverseTransformPoint(rb.position); AddVectorObs(localPosRelToBody); AddVectorObs(bp.currentXNormalizedRot); // Current x rot AddVectorObs(bp.currentYNormalizedRot); // Current y rot AddVectorObs(bp.currentZNormalizedRot); // Current z rot AddVectorObs(bp.currentStrength / m_JdController.maxJointForceLimit); } } public override void CollectObservations() { m_JdController.GetCurrentJointForces(); // Update pos to target m_DirToTarget = target.position - body.position; m_LookRotation = Quaternion.LookRotation(m_DirToTarget); m_TargetDirMatrix = Matrix4x4.TRS(Vector3.zero, m_LookRotation, Vector3.one); RaycastHit hit; if (Physics.Raycast(body.position, Vector3.down, out hit, 10.0f)) { AddVectorObs(hit.distance); } else AddVectorObs(10.0f); // Forward & up to help with orientation var bodyForwardRelativeToLookRotationToTarget = m_TargetDirMatrix.inverse.MultiplyVector(body.forward); AddVectorObs(bodyForwardRelativeToLookRotationToTarget); var bodyUpRelativeToLookRotationToTarget = m_TargetDirMatrix.inverse.MultiplyVector(body.up); AddVectorObs(bodyUpRelativeToLookRotationToTarget); foreach (var bodyPart in m_JdController.bodyPartsDict.Values) { CollectObservationBodyPart(bodyPart); } } /// /// Agent touched the target /// public void TouchedTarget() { AddReward(1f); if (respawnTargetWhenTouched) { GetRandomTargetPos(); } } /// /// Moves target to a random position within specified radius. /// public void GetRandomTargetPos() { var newTargetPos = Random.insideUnitSphere * targetSpawnRadius; newTargetPos.y = 5; target.position = newTargetPos + ground.position; } public override void AgentAction(float[] vectorAction, string textAction) { if (detectTargets) { foreach (var bodyPart in m_JdController.bodyPartsDict.Values) { if (bodyPart.targetContact && !IsDone() && bodyPart.targetContact.touchingTarget) { TouchedTarget(); } } } // If enabled the feet will light up green when the foot is grounded. // This is just a visualization and isn't necessary for function if (useFootGroundedVisualization) { foot0.material = m_JdController.bodyPartsDict[leg0Lower].groundContact.touchingGround ? groundedMaterial : unGroundedMaterial; foot1.material = m_JdController.bodyPartsDict[leg1Lower].groundContact.touchingGround ? groundedMaterial : unGroundedMaterial; foot2.material = m_JdController.bodyPartsDict[leg2Lower].groundContact.touchingGround ? groundedMaterial : unGroundedMaterial; foot3.material = m_JdController.bodyPartsDict[leg3Lower].groundContact.touchingGround ? groundedMaterial : unGroundedMaterial; } // Joint update logic only needs to happen when a new decision is made if (m_IsNewDecisionStep) { // The dictionary with all the body parts in it are in the jdController var bpDict = m_JdController.bodyPartsDict; var i = -1; // Pick a new target joint rotation bpDict[leg0Upper].SetJointTargetRotation(vectorAction[++i], vectorAction[++i], 0); bpDict[leg1Upper].SetJointTargetRotation(vectorAction[++i], vectorAction[++i], 0); bpDict[leg2Upper].SetJointTargetRotation(vectorAction[++i], vectorAction[++i], 0); bpDict[leg3Upper].SetJointTargetRotation(vectorAction[++i], vectorAction[++i], 0); bpDict[leg0Lower].SetJointTargetRotation(vectorAction[++i], 0, 0); bpDict[leg1Lower].SetJointTargetRotation(vectorAction[++i], 0, 0); bpDict[leg2Lower].SetJointTargetRotation(vectorAction[++i], 0, 0); bpDict[leg3Lower].SetJointTargetRotation(vectorAction[++i], 0, 0); // Update joint strength bpDict[leg0Upper].SetJointStrength(vectorAction[++i]); bpDict[leg1Upper].SetJointStrength(vectorAction[++i]); bpDict[leg2Upper].SetJointStrength(vectorAction[++i]); bpDict[leg3Upper].SetJointStrength(vectorAction[++i]); bpDict[leg0Lower].SetJointStrength(vectorAction[++i]); bpDict[leg1Lower].SetJointStrength(vectorAction[++i]); bpDict[leg2Lower].SetJointStrength(vectorAction[++i]); bpDict[leg3Lower].SetJointStrength(vectorAction[++i]); } // Set reward for this step according to mixture of the following elements. if (rewardMovingTowardsTarget) { RewardFunctionMovingTowards(); } if (rewardFacingTarget) { RewardFunctionFacingTarget(); } if (rewardUseTimePenalty) { RewardFunctionTimePenalty(); } IncrementDecisionTimer(); } /// /// Reward moving towards target & Penalize moving away from target. /// void RewardFunctionMovingTowards() { m_MovingTowardsDot = Vector3.Dot(m_JdController.bodyPartsDict[body].rb.velocity, m_DirToTarget.normalized); AddReward(0.03f * m_MovingTowardsDot); } /// /// Reward facing target & Penalize facing away from target /// void RewardFunctionFacingTarget() { m_FacingDot = Vector3.Dot(m_DirToTarget.normalized, body.forward); AddReward(0.01f * m_FacingDot); } /// /// Existential penalty for time-contrained tasks. /// void RewardFunctionTimePenalty() { AddReward(-0.001f); } /// /// Loop over body parts and reset them to initial conditions. /// public override void AgentReset() { if (m_DirToTarget != Vector3.zero) { transform.rotation = Quaternion.LookRotation(m_DirToTarget); } transform.Rotate(Vector3.up, Random.Range(0.0f, 360.0f)); foreach (var bodyPart in m_JdController.bodyPartsDict.Values) { bodyPart.Reset(bodyPart); } if (!targetIsStatic) { GetRandomTargetPos(); } m_IsNewDecisionStep = true; m_CurrentDecisionStep = 1; } }