using System;
using System.Collections.Generic;
using UnityEngine;
using Barracuda;
using System.IO;
using MLAgents;
namespace MLAgentsExamples
{
///
/// Utility class to allow the NNModel file for an agent to be overriden during inference.
/// This is used internally to validate the file after training is done.
/// The behavior name to override and file path are specified on the commandline, e.g.
/// player.exe --mlagents-override-model behavior1 /path/to/model1.nn --mlagents-override-model behavior2 /path/to/model2.nn
///
/// Additionally, a number of episodes to run can be specified; after this, the application will quit.
/// Note this will only work with example scenes that have 1:1 Agent:Behaviors. More complicated scenes like WallJump
/// probably won't override correctly.
///
public class ModelOverrider : MonoBehaviour
{
const string k_CommandLineModelOverrideFlag = "--mlagents-override-model";
const string k_CommandLineQuitAfterEpisodesFlag = "--mlagents-quit-after-episodes";
// The attached Agent
Agent m_Agent;
// Assets paths to use, with the behavior name as the key.
Dictionary m_BehaviorNameOverrides = new Dictionary();
// Cached loaded NNModels, with the behavior name as the key.
Dictionary m_CachedModels = new Dictionary();
// Max episodes to run. Only used if > 0
// Will default to 1 if override models are specified, otherwise 0.
int m_MaxEpisodes;
int m_NumSteps;
///
/// Get the asset path to use from the commandline arguments.
///
///
void GetAssetPathFromCommandLine()
{
m_BehaviorNameOverrides.Clear();
var maxEpisodes = 0;
var args = Environment.GetCommandLineArgs();
for (var i = 0; i < args.Length - 1; i++)
{
if (args[i] == k_CommandLineModelOverrideFlag && i < args.Length-2)
{
var key = args[i + 1].Trim();
var value = args[i + 2].Trim();
m_BehaviorNameOverrides[key] = value;
}
else if (args[i] == k_CommandLineQuitAfterEpisodesFlag)
{
Int32.TryParse(args[i + 1], out maxEpisodes);
}
}
if (m_BehaviorNameOverrides.Count > 0)
{
// If overriding models, set maxEpisodes to 1 or the command line value
m_MaxEpisodes = maxEpisodes > 0 ? maxEpisodes : 1;
Debug.Log($"setting m_MaxEpisodes to {maxEpisodes}");
}
}
void OnEnable()
{
m_Agent = GetComponent();
GetAssetPathFromCommandLine();
if (m_BehaviorNameOverrides.Count > 0)
{
OverrideModel();
}
}
void FixedUpdate()
{
if (m_MaxEpisodes > 0)
{
if (m_NumSteps > m_MaxEpisodes * m_Agent.maxStep)
{
// Stop recording so that we don't write partial rewards to the timer info.
TimerStack.Instance.Recording = false;
Application.Quit(0);
}
}
m_NumSteps++;
}
NNModel GetModelForBehaviorName(string behaviorName)
{
if (m_CachedModels.ContainsKey(behaviorName))
{
return m_CachedModels[behaviorName];
}
if (!m_BehaviorNameOverrides.ContainsKey(behaviorName))
{
Debug.Log($"No override for behaviorName {behaviorName}");
return null;
}
var assetPath = m_BehaviorNameOverrides[behaviorName];
byte[] model = null;
try
{
model = File.ReadAllBytes(assetPath);
}
catch(IOException)
{
Debug.Log($"Couldn't load file {assetPath}", this);
// Cache the null so we don't repeatedly try to load a missing file
m_CachedModels[behaviorName] = null;
return null;
}
var asset = ScriptableObject.CreateInstance();
asset.modelData = ScriptableObject.CreateInstance();
asset.modelData.Value = model;
asset.name = "Override - " + Path.GetFileName(assetPath);
m_CachedModels[behaviorName] = asset;
return asset;
}
///
/// Load the NNModel file from the specified path, and give it to the attached agent.
///
void OverrideModel()
{
m_Agent.LazyInitialize();
var bp = m_Agent.GetComponent();
var nnModel = GetModelForBehaviorName(bp.behaviorName);
Debug.Log($"Overriding behavior {bp.behaviorName} for agent with model {nnModel?.name}");
// This might give a null model; that's better because we'll fall back to the Heuristic
m_Agent.GiveModel($"Override_{bp.behaviorName}", nnModel);
}
}
}