ml-agents/unity-environment/Assets/ML-Agents/Scripts/CoreBrainHeuristic.cs

using System.Collections;
using System.Collections.Generic;
using UnityEngine;

#if UNITY_EDITOR
using UnityEditor;
#endif

/// CoreBrain which decides actions using developer-provided Decision script.
public class CoreBrainHeuristic : ScriptableObject, CoreBrain
{
    [SerializeField]
    private bool broadcast = true;

    /**< Reference to the brain that uses this CoreBrainHeuristic */
    public Brain brain;

    MLAgents.Batcher brainBatcher;

    /**< Reference to the Decision component used to decide the actions */
    public Decision decision;

    /// Create the reference to the brain
    public void SetBrain(Brain b)
    {
        brain = b;
    }

    /// Create the reference to decision
    public void InitializeCoreBrain(MLAgents.Batcher brainBatcher)
    {
        decision = brain.gameObject.GetComponent<Decision>();

        if ((brainBatcher == null)
            || (!broadcast))
        {
            this.brainBatcher = null;
        }
        else
        {
            this.brainBatcher = brainBatcher; ;
            this.brainBatcher.SubscribeBrain(brain.gameObject.name);
        }
    }

    /// Uses the Decision Component to decide that action to take
    public void DecideAction(Dictionary<Agent, AgentInfo> agentInfo)
    {
        if (brainBatcher != null)
        {
            brainBatcher.SendBrainInfo(brain.gameObject.name, agentInfo);
        }

        if (decision == null)
        {
            throw new UnityAgentsException("The Brain is set to Heuristic, but no decision script attached to it");
        }

        foreach (Agent agent in agentInfo.Keys)
        {
            agent.UpdateVectorAction(decision.Decide(
                agentInfo[agent].stackedVectorObservation,
                agentInfo[agent].visualObservations,
                agentInfo[agent].reward,
                agentInfo[agent].done,
                agentInfo[agent].memories));
            
        }

        foreach (Agent agent in agentInfo.Keys)
        {
            agent.UpdateMemoriesAction(decision.MakeMemory(
                agentInfo[agent].stackedVectorObservation,
                agentInfo[agent].visualObservations,
                agentInfo[agent].reward,
                agentInfo[agent].done,
                agentInfo[agent].memories));
        }
    }

    /// Displays an error if no decision component is attached to the brain
    public void OnInspector()
    {
        #if UNITY_EDITOR
            EditorGUILayout.LabelField("", GUI.skin.horizontalSlider);
            broadcast = EditorGUILayout.Toggle(new GUIContent("Broadcast",
            "If checked, the brain will broadcast states and actions to Python."), broadcast);
            if (brain.gameObject.GetComponent<Decision>() == null)
            {
            EditorGUILayout.HelpBox("You need to add a 'Decision' component to this gameObject", MessageType.Error);
            }
        #endif
    }

}