//Put this script on your blue cube.
using System;
using System.Collections;
using UnityEngine;
using Unity.MLAgents;
using Unity.MLAgents.Actuators;
using Random = UnityEngine.Random;
public class PushAgentBasic : Agent
{
///
/// The ground. The bounds are used to spawn the elements.
///
public GameObject ground;
public GameObject area;
///
/// The area bounds.
///
[HideInInspector]
public Bounds areaBounds;
PushBlockSettings m_PushBlockSettings;
///
/// The block to be pushed to the goal.
///
public GameObject block;
///
/// Detects when the block touches the goal.
///
[HideInInspector]
public GoalDetect goalDetect;
Rigidbody m_BlockRb; //cached on initialization
Rigidbody m_AgentRb; //cached on initialization
Material m_GroundMaterial; //cached on Awake()
///
/// We will be changing the ground material based on success/failue
///
Renderer m_GroundRenderer;
EnvironmentParameters m_ResetParams;
void Awake()
{
m_PushBlockSettings = FindObjectOfType();
goalDetect = block.GetComponent();
goalDetect.agent = this;
// Cache the agent rigidbody
m_AgentRb = GetComponent();
// Cache the block rigidbody
m_BlockRb = block.GetComponent();
// Get the ground's bounds
areaBounds = ground.GetComponent().bounds;
// Get the ground renderer so we can change the material when a goal is scored
m_GroundRenderer = ground.GetComponent();
// Starting material
m_GroundMaterial = m_GroundRenderer.material;
}
public override void Initialize()
{
m_ResetParams = Academy.Instance.EnvironmentParameters;
SetResetParameters();
}
///
/// Use the ground's bounds to pick a random spawn position.
///
public Vector3 GetRandomSpawnPos()
{
var foundNewSpawnLocation = false;
var randomSpawnPos = Vector3.zero;
var tries = 0;
while (foundNewSpawnLocation == false && tries < 50)
{
var randomPosX = Random.Range(-areaBounds.extents.x * m_PushBlockSettings.spawnAreaMarginMultiplier,
areaBounds.extents.x * m_PushBlockSettings.spawnAreaMarginMultiplier);
var randomPosZ = Random.Range(-areaBounds.extents.z * m_PushBlockSettings.spawnAreaMarginMultiplier,
areaBounds.extents.z * m_PushBlockSettings.spawnAreaMarginMultiplier);
randomSpawnPos = ground.transform.position + new Vector3(randomPosX, 1f, randomPosZ);
if (Physics.CheckBox(randomSpawnPos, new Vector3(2.5f, 0.01f, 2.5f)) == false)
{
foundNewSpawnLocation = true;
}
tries++;
}
return randomSpawnPos;
}
///
/// Called when the agent moves the block into the goal.
///
public void ScoredAGoal()
{
// We use a reward of 5.
AddReward(5f);
// By marking an agent as done AgentReset() will be called automatically.
EndEpisode();
// Swap ground material for a bit to indicate we scored.
StartCoroutine(GoalScoredSwapGroundMaterial(m_PushBlockSettings.goalScoredMaterial, 0.5f));
}
///
/// Swap ground material, wait time seconds, then swap back to the regular material.
///
IEnumerator GoalScoredSwapGroundMaterial(Material mat, float time)
{
m_GroundRenderer.material = mat;
yield return new WaitForSeconds(time); // Wait for 2 sec
m_GroundRenderer.material = m_GroundMaterial;
}
///
/// Called every step of the engine. Here the agent takes an action.
///
public override void OnActionReceived(ActionBuffers actionBuffers)
{
// Penalty given each step to encourage agent to finish task quickly.
AddReward(-1f / MaxStep);
}
///
/// Resets the block position and velocities.
///
void ResetBlock()
{
// Get a random position for the block.
block.transform.position = GetRandomSpawnPos();
// Reset block velocity back to zero.
m_BlockRb.velocity = Vector3.zero;
// Reset block angularVelocity back to zero.
m_BlockRb.angularVelocity = Vector3.zero;
}
///
/// In the editor, if "Reset On Done" is checked then AgentReset() will be
/// called automatically anytime we mark done = true in an agent script.
///
public override void OnEpisodeBegin()
{
var rotation = Random.Range(0, 4);
var rotationAngle = rotation * 90f;
area.transform.Rotate(new Vector3(0f, rotationAngle, 0f));
ResetBlock();
transform.position = GetRandomSpawnPos();
m_AgentRb.velocity = Vector3.zero;
m_AgentRb.angularVelocity = Vector3.zero;
SetResetParameters();
}
public void SetGroundMaterialFriction()
{
var groundCollider = ground.GetComponent();
groundCollider.material.dynamicFriction = m_ResetParams.GetWithDefault("dynamic_friction", 0);
groundCollider.material.staticFriction = m_ResetParams.GetWithDefault("static_friction", 0);
}
public void SetBlockProperties()
{
var scale = m_ResetParams.GetWithDefault("block_scale", 2);
//Set the scale of the block
m_BlockRb.transform.localScale = new Vector3(scale, 0.75f, scale);
// Set the drag of the block
m_BlockRb.drag = m_ResetParams.GetWithDefault("block_drag", 0.5f);
}
void SetResetParameters()
{
SetGroundMaterialFriction();
SetBlockProperties();
}
}