using System.IO;
using Google.Protobuf;
using System.Collections.Generic;
using Unity.MLAgents.Sensors;
using Unity.MLAgents.Policies;
namespace Unity.MLAgents.Demonstrations
{
///
/// Responsible for writing demonstration data to stream (typically a file stream).
///
///
public class DemonstrationWriter
{
///
/// Number of bytes reserved for the at the start of the demo file.
///
internal const int MetaDataBytes = 32;
DemonstrationMetaData m_MetaData;
Stream m_Writer;
float m_CumulativeReward;
ObservationWriter m_ObservationWriter = new ObservationWriter();
///
/// Create a DemonstrationWriter that will write to the specified stream.
/// The stream must support writes and seeking.
///
///
public DemonstrationWriter(Stream stream)
{
m_Writer = stream;
}
///
/// Writes the initial data to the stream.
///
/// Base name of the demonstration file(s).
/// The name of the Brain the agent is attached to.
/// The parameters of the Brain the agent is attached to.
internal void Initialize(
string demonstrationName, BrainParameters brainParameters, string brainName)
{
if (m_Writer == null)
{
// Already closed
return;
}
m_MetaData = new DemonstrationMetaData { demonstrationName = demonstrationName };
var metaProto = m_MetaData.ToProto();
metaProto.WriteDelimitedTo(m_Writer);
WriteBrainParameters(brainName, brainParameters);
}
///
/// Writes meta-data. Note that this is called at the *end* of recording, but writes to the
/// beginning of the file.
///
void WriteMetadata()
{
if (m_Writer == null)
{
// Already closed
return;
}
var metaProto = m_MetaData.ToProto();
var metaProtoBytes = metaProto.ToByteArray();
m_Writer.Write(metaProtoBytes, 0, metaProtoBytes.Length);
m_Writer.Seek(0, 0);
metaProto.WriteDelimitedTo(m_Writer);
}
///
/// Writes brain parameters to file.
///
/// The name of the Brain the agent is attached to.
/// The parameters of the Brain the agent is attached to.
void WriteBrainParameters(string brainName, BrainParameters brainParameters)
{
if (m_Writer == null)
{
// Already closed
return;
}
// Writes BrainParameters to file.
m_Writer.Seek(MetaDataBytes + 1, 0);
var brainProto = brainParameters.ToProto(brainName, false);
brainProto.WriteDelimitedTo(m_Writer);
}
///
/// Write AgentInfo experience to file.
///
/// for the agent being recorded.
/// List of sensors to record for the agent.
internal void Record(AgentInfo info, List sensors)
{
if (m_Writer == null)
{
// Already closed
return;
}
// Increment meta-data counters.
m_MetaData.numberSteps++;
m_CumulativeReward += info.reward;
if (info.done)
{
EndEpisode();
}
// Generate observations and add AgentInfo to file.
var agentProto = info.ToInfoActionPairProto();
foreach (var sensor in sensors)
{
agentProto.AgentInfo.Observations.Add(sensor.GetObservationProto(m_ObservationWriter));
}
agentProto.WriteDelimitedTo(m_Writer);
}
///
/// Performs all clean-up necessary.
///
public void Close()
{
if (m_Writer == null)
{
// Already closed
return;
}
EndEpisode();
m_MetaData.meanReward = m_CumulativeReward / m_MetaData.numberEpisodes;
WriteMetadata();
m_Writer.Close();
m_Writer = null;
}
///
/// Performs necessary episode-completion steps.
///
void EndEpisode()
{
m_MetaData.numberEpisodes += 1;
}
}
}