using System; using System.Collections.Generic; using UnityEngine; using System.IO; namespace Unity.MLAgents.SideChannels { /// /// Collection of static utilities for managing the registering/unregistering of /// and the sending/receiving of messages for all the channels. /// public static class SideChannelsManager { static Dictionary s_RegisteredChannels = new Dictionary(); struct CachedSideChannelMessage { public Guid ChannelId; public byte[] Message; } static readonly Queue s_CachedMessages = new Queue(); /// /// Register a side channel to begin sending and receiving messages. This method is /// available for environments that have custom side channels. All built-in side /// channels within the ML-Agents Toolkit are managed internally and do not need to /// be explicitly registered/unregistered. A side channel may only be registered once. /// /// The side channel to register. public static void RegisterSideChannel(SideChannel sideChannel) { var channelId = sideChannel.ChannelId; if (s_RegisteredChannels.ContainsKey(channelId)) { throw new UnityAgentsException( $"A side channel with id {channelId} is already registered. " + "You cannot register multiple side channels of the same id."); } // Process any messages that we've already received for this channel ID. var numMessages = s_CachedMessages.Count; for (var i = 0; i < numMessages; i++) { var cachedMessage = s_CachedMessages.Dequeue(); if (channelId == cachedMessage.ChannelId) { sideChannel.ProcessMessage(cachedMessage.Message); } else { s_CachedMessages.Enqueue(cachedMessage); } } s_RegisteredChannels.Add(channelId, sideChannel); } /// /// Unregister a side channel to stop sending and receiving messages. This method is /// available for environments that have custom side channels. All built-in side /// channels within the ML-Agents Toolkit are managed internally and do not need to /// be explicitly registered/unregistered. Unregistering a side channel that has already /// been unregistered (or never registered in the first place) has no negative side effects. /// Note that unregistering a side channel may not stop the Python side /// from sending messages, but it does mean that sent messages with not result in a call /// to . Furthermore, /// those messages will not be buffered and will, in essence, be lost. /// /// The side channel to unregister. public static void UnregisterSideChannel(SideChannel sideChannel) { if (s_RegisteredChannels.ContainsKey(sideChannel.ChannelId)) { s_RegisteredChannels.Remove(sideChannel.ChannelId); } } /// /// Unregisters all the side channels from the communicator. /// internal static void UnregisterAllSideChannels() { s_RegisteredChannels = new Dictionary(); } /// /// Returns the SideChannel of Type T if there is one registered, or null if it doesn't. /// If there are multiple SideChannels of the same type registered, the returned instance is arbitrary. /// /// /// internal static T GetSideChannel() where T : SideChannel { foreach (var sc in s_RegisteredChannels.Values) { if (sc.GetType() == typeof(T)) { return (T)sc; } } return null; } /// /// Grabs the messages that the registered side channels will send to Python at the current step /// into a singe byte array. /// /// internal static byte[] GetSideChannelMessage() { return GetSideChannelMessage(s_RegisteredChannels); } /// /// Grabs the messages that the registered side channels will send to Python at the current step /// into a singe byte array. /// /// A dictionary of channel type to channel. /// internal static byte[] GetSideChannelMessage(Dictionary sideChannels) { using (var memStream = new MemoryStream()) { using (var binaryWriter = new BinaryWriter(memStream)) { foreach (var sideChannel in sideChannels.Values) { var messageList = sideChannel.MessageQueue; foreach (var message in messageList) { binaryWriter.Write(sideChannel.ChannelId.ToByteArray()); binaryWriter.Write(message.Length); binaryWriter.Write(message); } sideChannel.MessageQueue.Clear(); } return memStream.ToArray(); } } } /// /// Separates the data received from Python into individual messages for each registered side channel. /// /// The byte array of data received from Python. internal static void ProcessSideChannelData(byte[] dataReceived) { ProcessSideChannelData(s_RegisteredChannels, dataReceived); } /// /// Separates the data received from Python into individual messages for each registered side channel. /// /// A dictionary of channel type to channel. /// The byte array of data received from Python. internal static void ProcessSideChannelData(Dictionary sideChannels, byte[] dataReceived) { while (s_CachedMessages.Count != 0) { var cachedMessage = s_CachedMessages.Dequeue(); if (sideChannels.ContainsKey(cachedMessage.ChannelId)) { sideChannels[cachedMessage.ChannelId].ProcessMessage(cachedMessage.Message); } else { Debug.Log(string.Format( "Unknown side channel data received. Channel Id is " + ": {0}", cachedMessage.ChannelId)); } } if (dataReceived.Length == 0) { return; } using (var memStream = new MemoryStream(dataReceived)) { using (var binaryReader = new BinaryReader(memStream)) { while (memStream.Position < memStream.Length) { Guid channelId = Guid.Empty; byte[] message = null; try { channelId = new Guid(binaryReader.ReadBytes(16)); var messageLength = binaryReader.ReadInt32(); message = binaryReader.ReadBytes(messageLength); } catch (Exception ex) { throw new UnityAgentsException( "There was a problem reading a message in a SideChannel. Please make sure the " + "version of MLAgents in Unity is compatible with the Python version. Original error : " + ex.Message); } if (sideChannels.ContainsKey(channelId)) { sideChannels[channelId].ProcessMessage(message); } else { // Don't recognize this ID, but cache it in case the SideChannel that can handle // it is registered before the next call to ProcessSideChannelData. s_CachedMessages.Enqueue(new CachedSideChannelMessage { ChannelId = channelId, Message = message }); } } } } } } }