File size: 2,834 Bytes
05c9ac2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
using System;
using System.Collections.Generic;
using System.IO;
using Unity.MLAgents.CommunicatorObjects;
using UnityEditor;
using UnityEngine;
#if UNITY_2020_2_OR_NEWER
using UnityEditor.AssetImporters;
#else
using UnityEditor.Experimental.AssetImporters;
#endif
using Unity.MLAgents.Demonstrations;
namespace Unity.MLAgents.Editor
{
/// <summary>
/// Asset Importer used to parse demonstration files.
/// </summary>
[ScriptedImporter(1, new[] { "demo" })]
internal class DemonstrationImporter : ScriptedImporter
{
const string k_IconPath = "Packages/com.unity.ml-agents/Editor/Icons/DemoIcon.png";
public override void OnImportAsset(AssetImportContext ctx)
{
var inputType = Path.GetExtension(ctx.assetPath);
if (inputType == null)
{
throw new Exception("Demonstration import error.");
}
try
{
// Read first three proto objects containing metadata, brain parameters, and observations.
Stream reader = File.OpenRead(ctx.assetPath);
var metaDataProto = DemonstrationMetaProto.Parser.ParseDelimitedFrom(reader);
var metaData = metaDataProto.ToDemonstrationMetaData();
reader.Seek(DemonstrationWriter.MetaDataBytes + 1, 0);
var brainParamsProto = BrainParametersProto.Parser.ParseDelimitedFrom(reader);
var brainParameters = brainParamsProto.ToBrainParameters();
// Read the first AgentInfoActionPair so that we can get the observation sizes.
List<ObservationSummary> observationSummaries;
try
{
var agentInfoActionPairProto = AgentInfoActionPairProto.Parser.ParseDelimitedFrom(reader);
observationSummaries = agentInfoActionPairProto.GetObservationSummaries();
}
catch
{
// Just in case there weren't any AgentInfoActionPair or they couldn't be read.
observationSummaries = new List<ObservationSummary>();
}
reader.Close();
var demonstrationSummary = ScriptableObject.CreateInstance<DemonstrationSummary>();
demonstrationSummary.Initialize(brainParameters, metaData, observationSummaries);
userData = demonstrationSummary.ToString();
var texture = (Texture2D)
AssetDatabase.LoadAssetAtPath(k_IconPath, typeof(Texture2D));
ctx.AddObjectToAsset(ctx.assetPath, demonstrationSummary, texture);
ctx.SetMainObject(demonstrationSummary);
}
catch
{
// ignored
}
}
}
}
|