TRAINING A NEURAL NETWORK TO CONTROL AN AGENT USING TASK-RELEVANT ADVERSARIAL IMITATION LEARNING

A method is proposed of training a neural network to generate action data for controlling an agent to perform a task in an environment. The method includes obtaining, for each of a plurality of performances of the task, one or more first tuple datasets, each first tuple dataset comprising state data characterizing a state of the environment at a corresponding time during the performance of the task; and a concurrent process of training the neural network and a discriminator network. The training process comprises a plurality of neural network update steps and a plurality of discriminator network update steps. Each neural network update step comprises: receiving state data characterizing a current state of the environment; using the neural network and the state data to generate action data indicative of an action to be performed by the agent; forming a second tuple dataset comprising the state data; using the second tuple dataset to generate a reward value, wherein the reward value comprises an imitation value generated by the discriminator network based on the second tuple dataset; and updating one or more parameters of the neural network based on the reward value. Each discriminator network update step comprises updating the discriminator network based on a plurality of the first tuple datasets and a plurality of the second tuple datasets, the update being to increase respective imitation values which the discriminator network generates upon receiving any of the plurality of the first tuple datasets compared to respective imitation values which the discriminator network generates upon receiving any of the plurality of the second tuple datasets. The updating process is performed subject to a constraint that the updated discriminator network, upon receiving any of at least a certain proportion of a first subset of the first tuple datasets and/or any of at least a certain proportion of a second subset of the second tuple datasets, does not generate imitation values which correctly indicate that those tuple datasets are first or second tuple datasets.