AnnaMats's picture
Second Push
05c9ac2
using Unity.MLAgents.Actuators;
namespace Unity.MLAgents.Tests.Actuators
{
internal class TestActuator : IActuator
{
public ActionBuffers LastActionBuffer;
public int[][] Masks;
public bool m_HeuristicCalled;
public int m_DiscreteBufferSize;
public TestActuator(ActionSpec actuatorSpace, string name)
{
ActionSpec = actuatorSpace;
Name = name;
}
public void OnActionReceived(ActionBuffers actionBuffers)
{
LastActionBuffer = actionBuffers;
}
public void WriteDiscreteActionMask(IDiscreteActionMask actionMask)
{
for (var i = 0; i < Masks.Length; i++)
{
foreach (var actionIndex in Masks[i])
{
actionMask.SetActionEnabled(i, actionIndex, false);
}
}
}
public ActionSpec ActionSpec { get; }
public string Name { get; }
public void ResetData()
{
}
public void Heuristic(in ActionBuffers actionBuffersOut)
{
m_HeuristicCalled = true;
m_DiscreteBufferSize = actionBuffersOut.DiscreteActions.Length;
}
}
}