arc-agents/lib/modules/imitation-learning-agent.js

A library for creating and deploying gaming agents at scale

const ProbabilisticAgentWrapper = require('../agent-wrappers/probabilistic-agent-wrapper')
const BaseAgent = require('./base-agent')
const { modelsMapping } = require('./agent-core')

const { default: axios } = require('axios')

class ImitationLearningAgent extends BaseAgent {
  async sendDataToPlatform() {
    try {
      await axios.post(`${BaseAgent.backend}platform/upload-training-data`, {
          gameId: BaseAgent.gameId,
          userId: this.userId,
          architectureId: this.architectureId,
          trainingData: this.dataCollector.trainingData
        },
        { headers: BaseAgent.getHeaders() }
      )
      return true
    }
    catch (err) {
      console.log(err.response?.data)
      return false
    }
  }

  async train(trainingData, config) {
    const configMapping = {
      "simple": { updatableCells: [], multiplier: 1 },
      "neural-network": { epochs: 2500, batchSize: 32 }
    }
    try {
      const trainingResponse = await axios.post(`${BaseAgent.backend}training/train-model`, {
          gameId: BaseAgent.gameId,
          userId: this.userId,
          architectureId: this.architectureId,
          modelData: this.model.modelData,
          metadata: this.metadata,
          trainingData: trainingData,
          trainingConfig: config ?? configMapping[this.metadata.modelType]
        },
        { headers: BaseAgent.getHeaders() }
      )
      const { trainedModelData, memory } = trainingResponse.data
      this.trainedModel = new modelsMapping[this.metadata.modelType](trainedModelData)
      this.trainedModel.addMemory(memory)
      this.trainedModelData = this.trainedModel.modelData
      this.trainedAgent = new ProbabilisticAgentWrapper(this.trainedModel, this.agentConfig)
      return true
    }
    catch (err) {
      console.log(err.response?.data)
      return false
    }
  }

  async save(newModelBool=false) {
    this.model.formatModelToSave()
    if (!newModelBool && this.trainedModelData) {
      this.trainedModelData.timesUpdated += 1
    }
    
    await super.save(newModelBool)

    if (!newModelBool) {
      this.model = new modelsMapping[this.metadata.modelType](this.trainedModelData)
      this.modelData = this.trainedModelData
      this.discardTraining(true)
    }
  }

  discardTraining(discardData=false) {
    this.trainedModelData = null
    this.trainedModel = null
    this.trainedAgent = null
    if (discardData) {
      this.dataCollector.reset()
    }
  }
}

module.exports = ImitationLearningAgent