using System.Collections.Generic;
using Unity.MLAgents.Sensors;
using Unity.Barracuda;
namespace Unity.MLAgents.Inference
{
///
/// Mapping between Tensor names and generators.
/// A TensorGenerator implements a Dictionary of strings (node names) to an Action.
/// The Action take as argument the tensor, the current batch size and a Dictionary of
/// Agent to AgentInfo corresponding to the current batch.
/// Each Generator reshapes and fills the data of the tensor based of the data of the batch.
/// When the TensorProxy is an Input to the model, the shape of the Tensor will be modified
/// depending on the current batch size and the data of the Tensor will be filled using the
/// Dictionary of Agent to AgentInfo.
/// When the TensorProxy is an Output of the model, only the shape of the Tensor will be
/// modified using the current batch size. The data will be pre-filled with zeros.
///
internal class TensorGenerator
{
public interface IGenerator
{
///
/// Modifies the data inside a Tensor according to the information contained in the
/// AgentInfos contained in the current batch.
///
/// The tensor the data and shape will be modified.
/// The number of agents present in the current batch.
///
/// List of AgentInfos containing the information that will be used to populate
/// the tensor's data.
///
void Generate(
TensorProxy tensorProxy, int batchSize, IList infos);
}
readonly Dictionary m_Dict = new Dictionary();
///
/// Returns a new TensorGenerators object.
///
/// The seed the Generators will be initialized with.
/// Tensor allocator.
/// Dictionary of AgentInfo.id to memory for use in the inference model.
///
public TensorGenerator(
int seed,
ITensorAllocator allocator,
Dictionary> memories,
object barracudaModel = null)
{
// Generator for Inputs
m_Dict[TensorNames.BatchSizePlaceholder] =
new BatchSizeGenerator(allocator);
m_Dict[TensorNames.SequenceLengthPlaceholder] =
new SequenceLengthGenerator(allocator);
m_Dict[TensorNames.RecurrentInPlaceholder] =
new RecurrentInputGenerator(allocator, memories);
if (barracudaModel != null)
{
var model = (Model)barracudaModel;
for (var i = 0; i < model.memories.Count; i++)
{
m_Dict[model.memories[i].input] =
new BarracudaRecurrentInputGenerator(i, allocator, memories);
}
}
m_Dict[TensorNames.PreviousActionPlaceholder] =
new PreviousActionInputGenerator(allocator);
m_Dict[TensorNames.ActionMaskPlaceholder] =
new ActionMaskInputGenerator(allocator);
m_Dict[TensorNames.RandomNormalEpsilonPlaceholder] =
new RandomNormalInputGenerator(seed, allocator);
// Generators for Outputs
m_Dict[TensorNames.ActionOutput] = new BiDimensionalOutputGenerator(allocator);
m_Dict[TensorNames.RecurrentOutput] = new BiDimensionalOutputGenerator(allocator);
m_Dict[TensorNames.ValueEstimateOutput] = new BiDimensionalOutputGenerator(allocator);
}
public void InitializeObservations(List sensors, ITensorAllocator allocator)
{
// Loop through the sensors on a representative agent.
// For vector observations, add the index to the (single) VectorObservationGenerator
// For visual observations, make a VisualObservationInputGenerator
var visIndex = 0;
VectorObservationGenerator vecObsGen = null;
for (var sensorIndex = 0; sensorIndex < sensors.Count; sensorIndex++)
{
var sensor = sensors[sensorIndex];
var shape = sensor.GetObservationShape();
// TODO generalize - we currently only have vector or visual, but can't handle "2D" observations
var isVectorSensor = (shape.Length == 1);
if (isVectorSensor)
{
if (vecObsGen == null)
{
vecObsGen = new VectorObservationGenerator(allocator);
}
vecObsGen.AddSensorIndex(sensorIndex);
}
else
{
m_Dict[TensorNames.VisualObservationPlaceholderPrefix + visIndex] =
new VisualObservationInputGenerator(sensorIndex, allocator);
visIndex++;
}
}
if (vecObsGen != null)
{
m_Dict[TensorNames.VectorObservationPlaceholder] = vecObsGen;
}
}
///
/// Populates the data of the tensor inputs given the data contained in the current batch
/// of agents.
///
/// Enumerable of tensors that will be modified.
/// The number of agents present in the current batch
///
/// List of AgentsInfos and Sensors that contains the
/// data that will be used to modify the tensors
/// One of the tensor does not have an
/// associated generator.
public void GenerateTensors(
IReadOnlyList tensors, int currentBatchSize, IList infos)
{
for (var tensorIndex = 0; tensorIndex < tensors.Count; tensorIndex++)
{
var tensor = tensors[tensorIndex];
if (!m_Dict.ContainsKey(tensor.name))
{
throw new UnityAgentsException(
$"Unknown tensorProxy expected as input : {tensor.name}");
}
m_Dict[tensor.name].Generate(tensor, currentBatchSize, infos);
}
}
}
}