Patent Description:
This specification relates to reinforcement learning.

In a reinforcement learning system, an agent interacts with an environment by performing actions that are selected by the reinforcement learning system in response to receiving observations that characterize the current state of the environment.

Some reinforcement learning systems select the action to be performed by the agent in response to receiving a given observation in accordance with an output of a neural network.

Some neural networks are deep neural networks that include one or more hidden layers in addition to an output layer. Examples of relevant prior art are: <NPL>, and<NPL>.

This specification describes a reinforcement learning system that controls an agent interacting with an environment by, at each of multiple time steps, processing data characterizing the current state of the environment at the time step (i.e., an "observation") to select an action to be performed by the agent from a set of actions.

At each time step, the state of the environment at the time step depends on the state of the environment at the previous time step and the action performed by the agent at the previous time step.

A model of the environment may be constructed such that the environment model is causally correct. The environment model may be used in a reinforcement learning system as part of a planning subsystem for determining actions for the agent to perform, for example, to achieve a particular task or goal. Generally, such a system may receive the current observation and performs a plurality of planning iterations. The system then selects the action to be performed in response to the current observation based on the results of the planning iterations. At each planning iteration, the system may generate a sequence of actions that progress the environment to new states starting from the state represented by the current observation.

Unlike conventional systems, the system does not require performing the planning iterations using a full simulator of the environment, i.e., does not use a simulator of the environment to determine which state the environment will transition into as a result of a given action being performed in a given state. Instead, the system uses an environment model that is configured to receive an input action selected from the set of actions and to update its hidden state to simulate a predicted next environment state that the environment would transition into if the agent performed the input action when the environment is in the input environment state. That is, the environment model does not need to model the environment fully, only the necessary parts that have an influence on planning whilst still fulfilling the condition of causal correctness. Each hidden state may be a lower-dimensional representation of an observation tThe environment model is configured to update its hidden state based on a latent representation computed from the hidden state. The updating may not require processing an observation of an environment state. Thus, the system can perform planning using only these hidden states without ever being required to reconstruct the full state of the environment or even a full observation characterizing a state.

According to an aspect, there is provided a computer-implemented method of using an environment model to simulate state transitions of an environment being interacted with by an agent that is controlled using a policy neural network. The policy neural network is configured to receive an observation characterizing a state of the environment, update a belief representation of the state of the environment, generate a latent representation from the belief representation, and generate an output specifying an action to be performed by the agent from the latent representation. The method comprises initializing an internal representation of a state of the environment at a current time point. The method further comprises repeatedly performing the following operations: receiving an action to be performed by the agent; generating, based on the internal representation, a predicted latent representation that is a prediction of a latent representation that would have been generated by the policy neural network by processing an observation characterizing the state of the environment corresponding to the internal representation; and updating the internal representation to simulate a state transition caused by the agent performing the received action by processing the predicted latent representation and the received action using the environment model.

The use of an intermediate latent representation between the belief representation and the action output in the policy neural network and the use of a predicted latent representation by the environment model ensures casual correctness as the latent representation provides a variable to block the influence of confounding variables that may arise if the environment model does not capture all aspects of the environment. Given the latent representation, the action performed by an agent is conditionally independent of the state.

The method may further comprise generating, from the internal representation of the state of the environment, a target to be provided for use in controlling the agent. The target may be a reward and/or an expected return.

The method comprises based on a result of repeatedly performing the operations, an action to be performed by the agent in the environment at the current time point. The target may be used in selecting an action for the agent to perform. The method may further comprise controlling the agent to perform the selected action.

Initializing an internal representation of a state of the environment at a current time point may comprise receiving, by the policy neural network, an observation characterizing the state of the environment at the current time point. The initializing may further comprise updating, by the policy neural network and based on processing the received observation, a belief representation of the state of the environment. The initializing may further comprise initializing the internal representation based on the belief representation of the state of the environment.

Updating the internal representation may not include processing the observation to be provided to the policy neural network that characterizes the state of the environment. Updating the internal representation may be based upon the initial representation of a state of the environment, one or more past actions and the received action. For example, the updating may not require processing of past states of the internal representation caused by the past actions. As such, updating the internal representation may be based upon the initial representation of a state of the environment, one or more past actions and the received action only.

The method may further comprise processing, by the policy neural network, the belief representation of the state of the environment and the action that is performed by the agent to update the belief representation of the state of the environment at a future time point that is after the current time point. Updating the belief representation of the state of the environment at the future time point may further comprise processing an observation that characterizes the state of the environment at the future time point.

The latent representation may correspond to one or more layers of the policy neural network after updating the belief representation of the state of the environment. The one or more layers may comprise an input layer of the policy neural network after updating the belief representation of the state of the environment. The belief representation may be an internal state of the policy neural network.

The latent representation may correspond to respective probabilities generated by the policy neural network for controlling the agent to perform different actions.

The latent representation may correspond to an intended action to be performed by the agent before selecting actions under exploration. That is, an exploration technique, such as epsilon-greedy exploration, may be used to change an intended action to a different action in order to cause the agent to explore other actions to gain more information regarding the environment. Thus, the executed action may be different from the intended action in this regard. The policy neural network may be recurrent neural network. It will be appreciated that the policy neural network has a plurality of network parameters.

The environment model may comprise a neural network having a plurality of network parameters. The environment model may comprise a recurrent neural network. The environment model may further comprise a generative model conditioned on the states of the recurrent neural network to generate the target to be provided for use in controlling the agent. The generative model may comprise one or more neural network layers.

The environment specified in the latent representation of the state of the environment may correspond to a partial view of the environment being interacted with by the agent. For example, a partial view may comprise a subset of the full data characterizing the environment. The partial view may be based upon past observations and/or past actions.

Generating the latent representation from the belief representation may comprise sampling from a distribution of a plurality of variables that describe the latent representation, the distribution being generated by the policy neural network and being conditioned on the belief representation. Generating the predicted latent representation that is a prediction of the latent representation may comprise sampling from a distribution of a plurality of variables that describe the latent representation, the distribution being generated by the environment model and being conditioned on the internal representation.

The method may further comprise iteratively training the environment model on training data to determine trained values of the model parameters, wherein the training data includes observation received by the agent during interaction with the environment.

Training the environment model may comprise, at each training iteration generating, by the environment model, a training predicted latent representation; evaluating an objective function measuring a difference between the training predicted latent representation and the actual latent representation that is generated by the policy neural network; and updating, based on a computed gradient of the objective function, corresponding values of the environment model parameters.

According to another aspect, there is provided a system comprising one or more computers and one or more storage devices storing instructions that are operable, when executed by the one or more computers, to cause the one or more computers to perform the operations of the above method aspect.

According to a further aspect, there is provided a computer storage medium encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform the operations of the method aspect.

It will be appreciated that features described in the context of one aspect may be combined with features described in the context of another aspect.

Oftentimes, when predicting state transitions of an environment that is being interacted with by the agent, an environment model fails to make correct predictions due to issues related to incorrect causal reasoning. That is, the environment model incorrectly attributes at least part of a reward received by the agent to an action that is performed by the agent.

The techniques described in this specification, however, provide the environment model with additional information derived from latent representations that corresponds to frontdoor or backdoor variables in causal reasoning (see <NPL>.

This allows the environment model to generate simulations that are causally correct and are robust to changes in the policy that is being followed in selecting actions to be performed by the agent. This improves the effectiveness of the reinforcement learning algorithm and can avoid dangerous or otherwise risky situations that would otherwise result from the environment model not accurately accounting for stochasticity or other external factors within the environment. As one example, the techniques allow for the system to control an agent navigating through an environment to avoid taking a path with promising immediate rewards but substantially negative punishments in the long run.

In addition, the techniques prevent the environment model from being unrealistically optimistic about the environment, especially when simulating target rewards to be received by the agent, which in turn results in improved performance of the policy neural network when used in selecting actions to be performed by the agent.

In some implementations, the techniques use an environment model to simulate environments that correspond to a partial view of the entire environment. Such partial models are neither conditioned on, nor configured to generate the full set of observable data (which is typically very high dimensional) to be provided to the agent. For example, the observable data may include images of the environment. Conditioning the environment model based upon the image data may require modeling and generating images. Given the high-dimensionality of image data, such modeling is computationally expensive and may be intractable. In addition, the modeling may capture aspects of visual details that are unnecessary for the purposes of action planning. However, simply choosing not to model certain aspects of the environment may cause the environment model to be causally incorrect as the unmodelled aspects may become confounding variables. This may result in suboptimal actions being taken by the agent. Therefore, controlling an agent using a casually correct partial view environment model as described herein is able to achieve a high level of accuracy in simulating (at least a partial view of) the environment while requiring less computational resources (e.g., memory, computing power, or both) than systems that are required to predict high-dimensional observations in order to perform planning.

This specification describes a reinforcement learning system that controls an agent interacting with an environment by, at each of multiple time steps, processing data characterizing the current state of the environment at the time step (i.e., an "observation") to select an action to be performed by the agent.

The environment is a real-world environment and the agent is a mechanical agent interacting with the real-world environment, e.g., a robot or an autonomous or semi-autonomous land, air, or sea vehicle navigating through the environment.

In these implementations, the observations may include, e.g., one or more of: images, object position data, and sensor data to capture observations as the agent as it interacts with the environment, for example sensor data from an image, distance, or position sensor or from an actuator.

For example in the case of a robot the observations may include data characterizing the current state of the robot, e.g., one or more of: joint position, joint velocity, joint force, torque or acceleration, e.g., gravity-compensated torque feedback, and global or relative pose of an item held by the robot.

In the case of a robot or other mechanical agent or vehicle the observations may similarly include one or more of the position, linear or angular velocity, force, torque or acceleration, and global or relative pose of one or more parts of the agent. The observations may be defined in <NUM>, <NUM> or <NUM> dimensions, and may be absolute and/or relative observations.

The observations may also include, for example, sensed electronic signals such as motor current or a temperature signal; and/or image or video data for example from a camera or a LIDAR sensor, e.g., data from sensors of the agent or data from sensors that are located separately from the agent in the environment.

In the case of an electronic agent the observations may include data from one or more sensors monitoring part of a plant or service facility such as current, voltage, power, temperature and other sensors and/or electronic signals representing the functioning of electronic and/or mechanical items of equipment.

In these implementations, the actions may be control inputs to control the robot, e.g., torques for the joints of the robot or higher-level control commands, or the autonomous or semi-autonomous land, air, sea vehicle, e.g., torques to the control surface or other control elements of the vehicle or higher-level control commands.

In other words, the actions can include for example, position, velocity, or force/torque/acceleration data for one or more joints of a robot or parts of another mechanical agent. Action data may additionally or alternatively include electronic control data such as motor control data, or more generally data for controlling one or more electronic devices within the environment the control of which has an effect on the observed state of the environment. For example in the case of an autonomous or semi-autonomous land or air or sea vehicle the actions may include actions to control navigation e.g. steering, and movement e.g., braking and/or acceleration of the vehicle.

In some implementations the environment is a simulated environment and the agent is implemented as one or more computers interacting with the simulated environment. Training an agent in a simulated environment may enable the agent to learn from large amounts of simulated training data while avoiding risks associated with training the agent in a real world environment, e.g., damage to the agent due to performing poorly chosen actions. An agent trained in a simulated environment may thereafter be deployed in a real-world environment.

The simulated environment may be a motion simulation environment, e.g., a driving simulation or a flight simulation, and the agent may be a simulated vehicle navigating through the motion simulation. In these implementations, the actions may be control inputs to control the simulated user or simulated vehicle.

In another example, the simulated environment may be a video game and the agent may be a simulated user playing the video game. Generally, in the case of a simulated environment, the observations may include simulated versions of one or more of the previously described observations or types of observations and the actions may include simulated versions of one or more of the previously described actions or types of actions.

In a further example the environment may be a chemical synthesis or a protein folding environment such that each state is a respective state of a protein chain or of one or more intermediates or precursor chemicals and the agent is a computer system for determining how to fold the protein chain or synthesize the chemical. In this example, the actions are possible folding actions for folding the protein chain or actions for assembling precursor chemicals/intermediates and the result to be achieved may include, e.g., folding the protein so that the protein is stable and so that it achieves a particular biological function or providing a valid synthetic route for the chemical. As another example, the agent may be a mechanical agent that performs or controls the protein folding actions selected by the system automatically without human interaction. The observations may include direct or indirect observations of a state of the protein and/or may be derived from simulation.

In a similar way the environment may be a drug design environment such that each state is a respective state of a potential pharma chemical drug and the agent is a computer system for determining elements of the pharma chemical drug and/or a synthetic pathway for the pharma chemical drug. The drug/synthesis may be designed based on a reward derived from a target for the drug, for example in simulation. As another example, the agent may be a mechanical agent that performs or controls synthesis of the drug.

In some applications the agent may be a static or mobile software agent i.e. a computer programs configured to operate autonomously and/or with other software agents or people to perform a task. For example the environment may be an integrated circuit routing environment and the system may be configured to learn to perform a routing task for routing interconnection lines of an integrated circuit such as an ASIC. The rewards (or costs) may then be dependent on one or more routing metrics such as an interconnect resistance, capacitance, impedance, loss, speed or propagation delay, physical line parameters such as width, thickness or geometry, and design rules. The observations may be observations of component positions and interconnections; the actions may comprise component placing actions e.g. to define a component position or orientation and/or interconnect routing actions e.g. interconnect selection and/or placement actions. The routing task may thus comprise placing components i.e. determining positions and/or orientations of components of the integrated circuit, and/or determining a routing of interconnections between the components. Once the routing task has been completed an integrated circuit, e.g. ASIC, may be fabricated according to the determined placement and/or routing. Or the environment may be a data packet communications network environment, and the agent be a router to route packets of data over the communications network based on observations of the network.

Generally, in the case of a simulated environment, the observations may include simulated versions of one or more of the previously described observations or types of observations and the actions may include simulated versions of one or more of the previously described actions or types of actions.

In some other applications the agent may control actions in a real-world environment including items of equipment, for example in a data center, in a power/water distribution system, or in a manufacturing plant or service facility. The observations may then relate to operation of the plant or facility. For example the observations may include observations of power or water usage by equipment, or observations of power generation or distribution control, or observations of usage of a resource or of waste production. The actions may include actions controlling or imposing operating conditions on items of equipment of the plant/facility, and/or actions that result in changes to settings in the operation of the plant/facility e.g. to adjust or turn on/off components of the plant/facility.

In some further applications, the environment is a real-world environment and the agent manages distribution of tasks across computing resources e.g. on a mobile device and/or in a data center. In these implementations, the actions may include assigning tasks to particular computing resources.

<FIG> shows an example reinforcement learning system <NUM>. The reinforcement learning system <NUM> is an example of a system implemented as computer programs on one or more computers in one or more locations in which the systems, components, and techniques described below are implemented.

The system <NUM> controls an agent <NUM> interacting with an environment <NUM> by using a planning engine <NUM> which in turn includes a policy neural network <NUM> and an environment model <NUM> to select actions <NUM> to be performed by the agent <NUM> that cause the state of the environment to transition into new states.

The system <NUM> includes the planning engine <NUM> and a set of model parameters <NUM> for the components of the planning engine <NUM>, including parameters of the policy neural network <NUM> and parameters of the environment model <NUM>. The system <NUM> may also include a training engine <NUM>.

Briefly, the environment model <NUM> is a model which, given information at a given time step, is able to make a prediction about at least one future time. As will be described in more detail below, the environment model <NUM> is used in the system <NUM> to make a prediction about multiple time steps after the given time step. This is referred to as a rollout. It represents an imagined trajectory of the environment at times after the given time step, assuming that the agent performs certain actions.

At each of the multiple time steps, the policy neural network <NUM> is configured to process an input that includes the current observation <NUM> characterizing the current state of the environment <NUM> in accordance with the model parameters <NUM> to generate an action selection output <NUM> ("action selection policy") that could be used to control the agent.

For example, the policy neural network <NUM> can be a recurrent neural network, e.g., a LSTM network, that can receive an input including an observation of the environment and to process the input in accordance with a set of policy neural network parameters to generate an action selection output <NUM> that can be used to determine an action <NUM> to be performed by the agent <NUM> at each of multiple time steps.

A few examples of the action selection output <NUM> are described next.

In one example, the action selection output <NUM> may include a respective numerical probability value for each action in a set of possible actions that can be performed by the agent. If being used to control the agent, the system <NUM> could select the action to be performed by the agent, e.g., by sampling an action in accordance with the probability values for the actions, or by selecting the action with the highest probability value.

In another example, the action selection output <NUM> may directly define the action to be performed by the agent, e.g., by defining the values of torques that should be applied to the joints of a robotic agent.

In another example, the action selection output <NUM> may include a respective Q-value for each action in the set of possible actions that can be performed by the agent. If being used to directly control the agent, the system <NUM> could process the Q-values (e.g., using a soft-max function) to generate a respective probability value for each possible action, which can be used to select the action to be performed by the agent (as described earlier). The system <NUM> could also select the action with the highest Q-value as the action to be performed by the agent.

The Q value for an action is an estimate of a "return" that would result from the agent performing the action in response to the current observation <NUM> and thereafter selecting future actions performed by the agent <NUM> in accordance with current values of the policy neural network parameters.

A return refers to a cumulative measure of "rewards" <NUM> received by the agent, for example, a time-discounted sum of rewards. The agent can receive a respective reward <NUM> at each time step, where the reward <NUM> is specified by a scalar numerical value and characterizes, e.g., a progress of the agent towards completing an assigned task.

Rather than directly using the policy neural network <NUM> to control the agent <NUM>, however, the system <NUM> instead uses the environment model <NUM> to perform a plurality of planning iterations. The system then selects the action <NUM> to be performed in response to the current observation based on the results of the planning iterations. At each planning iteration, the system <NUM> can generate a sequence of actions <NUM> that progress the environment <NUM> to new states starting from the state represented by the current observation <NUM>. Typically the environment model <NUM> is used to produce multiple trajectories starting from the current observation <NUM>. This can aid in determining more effective action selection policies to maximize expected cumulative reward for the agent <NUM>.

Specifically, the system <NUM> can use the environment model <NUM> together with an MDP-based planning algorithm, e.g., a tree-based search or other look-ahead planning methods, to achieve high quality agent performance in a range of challenging and visually complex domains, without any knowledge of their underlying dynamics. For example, the MDP-based planning algorithm may be a Monte Carlo tree search (MCTS) algorithm. At each time step, the system <NUM> makes use of an action selection policy, the reward estimate, and, when relevant, the value estimate generated by the environment model in accordance with current model parameters. Each value estimate, when considered, specifies a value of the environment being in the predicted next environment state to performing the task. The system runs the MCTS algorithm using these data to determine an action selection output and, in some cases, an estimated value, based on which a next action to be performed by the agent can be selected.

Examples of trajectory planning over a series of internal planning iterations and how to use planning iterations to control an agent are described in <NPL>, and in<NPL>.

More specifically, the environment model <NUM> can first be initialized using an initial action, e.g., a candidate action selected by the system for a given state of the environment according to a currently adopted action selection policy, and an initial hidden state of the policy neural network <NUM>. To perform a rollout of multiple state transitions forward from the given state of the environment, the environment model <NUM> is then configured to receive, over multiple time steps, actions to be performed by the agent and generates as output respective targets (e.g., in terms of rewards to be received by the agent at each of multiple time steps) for each of the multiple trajectories starting from the given environment state that the system <NUM> can provide to the policy neural network <NUM>. The system <NUM> then uses these targets to select action to be performed by the agent <NUM>.

Simulating state transitions of the environment being interacted with by the agent using the environment model <NUM> is described in more detail below with reference to <FIG>.

The exact architectures of the environment model <NUM> may vary, but typically, the environment model <NUM> can be a recurrent neural network that is configured to receive as input an action performed by the agent and to process the input in accordance with a set of environment model parameters to generate as output a target (e.g., in terms of rewards to be received by the agent as a result of performing the action) and to update its hidden state to simulate a state transition of the environment caused by the action. For example, the architecture of the environment model <NUM> may include a sequence of one or more layers (e.g., convolutional layers or fully-connected layers), followed by one or more recurrent layers (e.g., long short-term memory (LSTM) layers) and an output layer that generates the environment model output including the observation.

Because the environment model <NUM> is specifically configured to predict target rewards for use in action selection, i.e., instead of a full set of observable data that characterizes the entire environment, it may be referred to as a partial environment model. The output data of the environment model generally has lower dimension, simpler modality, or both than the actual observation data that could have been received by the system during the interaction of the agent with the environment.

<FIG> is an example illustration of dynamics of an environment being interacted with by a reinforcement learning agent. The environment is modeled using multiple states et that can each transition into a subsequent state et+<NUM> in accordance with a transition probability of form p(et+<NUM>|et, at).

In the example of <FIG>, at each step t, the system can use the environment model <NUM> to output a current target yt (e.g., in terms of rewards to be received by the agent), sample a latent representation zt from a given distribution p(zt|ht) that is conditioned on the current hidden state ht of the environment model (which corresponds to an "internal representation" of a state of the environment et. ), select an action at according to a given action selection policy and based on the current hidden state ht and on the sampled latent representation zt, and determine an updated environment model hidden state ht+<NUM> using (i) the current environment model hidden state ht, (ii) the current latent representation zt, and (iii) the selected action at. Specifically, the selected action corresponds to the current action to be performed by the agent at the current step which can cause the environment to transition into a subsequent state.

In particular, the given distribution p(zt|ht) can be parameterized by the output of a given layer, e.g., an input layer or an intermediate layer, of the environment model <NUM> or a combination of the outputs of multiple layers of the environment model <NUM>.

The latent representation zt sampled from p(zt|ht) will be referred to as a predicted latent representation of a state of the environment computed by the environment model, i.e., a prediction of a target latent representation that would have been derived from a hidden state of the policy neural network <NUM> based on using the policy neural network to process a policy network input including an observation characterizing the state of the environment. Deriving the target latent representation from the policy neural network hidden states is described in more detail below.

<FIG> shows an example of selecting actions to be performed by the agent using the policy neural network <NUM>. In the example of <FIG>, at each time step t, the policy neural network <NUM> receives an observation yt and processes the observation to update its hidden state st (which corresponds to a "belief representation" of a state of the environment et) from which a latent representation zt can be derived, and to generate an action selection output π(at|zt) that is conditioned on the latent representation zt. The system can then cause the agent to interact with the environment by taking an action at according to the action selection output at the current time step, e.g., by passing a control signal to a control system for the agent,.

Computing the latent representation zt usually involves sampling from a given distribution m(zt|st) that is conditioned on a hidden state of the policy neural network <NUM>. Similarly, the given distribution m(zt|st) can be parameterized by the output of a given layer, e.g., an output layer or an intermediate layer preceded by the one or more recurrent layers, of the policy neural network <NUM> or a combination of the outputs of multiple layers of the policy neural network <NUM>. In some implementations, each such output is generated by the policy neural network based on its current hidden state. For example, the policy neural network <NUM> generates, at a given layer and for each of one or more pre-defined latent factors, an output that parameterizes a distribution, e.g., a Gaussian distribution, over a set of possible values for the latent factor. The system can then sample a value for each latent factor from the distribution.

For example, the hidden states of the policy neural network can be used as parameterizing the distribution m(zt|st) from which the latent representation zt is determined. Each hidden state is generally defined by an ordered collection of numeric values that has a fixed number of values.

As another example, the action selection probabilities can be used as parameterizing the distribution m(zt|st) from which the latent representation zt is determined. For example, in cases where the action selection output includes a respective numerical probability value for each action in a discrete set of possible actions, the latent representation zt can be a concatenated representation, e.g., a vector, of the probabilities specified by the action selection outputs. Alternatively, the latent representation zt can be the actual action identified by the action selection output (i.e., the action assigned with the highest probability or an action sampled from the probability distribution) without applying any exploration policy, e.g., ε-greedy exploration policy.

The training engine <NUM> trains the policy neural network <NUM> and the environment model <NUM> by using a reinforcement learning technique to iteratively adjust the values of the set of parameters of the policy neural network <NUM> and the environment model <NUM> based on the interactions of the agent with the environment. An example of a suitable reinforcement learning technique is described in <NPL>.

In particular, during training, trajectories generated as a result of the agent interacting with the environment are stored in a trajectory buffer <NUM>. Each trajectory can include (i) a sequence of observations y received by the agent that characterize respective states of the environment and that specify respective rewards issued to the agent by the environment and (ii) a sequence of actions a performed by the agent in response to the observations.

The training engine <NUM> updates the parameters of the policy neural network <NUM> to encourage it to generate policy outputs that maximize the expected cumulative reward received by the system <NUM>. The training engine <NUM> updates the model parameters of the environment model <NUM> to encourage it to more accurately simulate state transitions of the environment caused by the actions that would be performed by the agent including predicting accurate target rewards for these actions. More specifically, the training engine <NUM> can encourage the environment model <NUM> to generate predicted latent representations that emulate the actual latent representations generated by the policy neural network. For example, at each of training iteration: the training engine <NUM> can use the environment model <NUM> to generate a training predicted latent representation for a state of the environment, evaluate an objective function measuring a difference between the training predicted latent representation and the actual latent representation that is generated by the policy neural network <NUM> for the same state of the environment, and determine an update to current values of the environment model parameters based on a computed gradient of the objective function.

Training the environment model <NUM> together with the policy neural network <NUM> on the trajectory data tunes the values of the set of parameters of the policy neural network <NUM> to cause the policy neural network <NUM> to derive more useful information from the observations <NUM> which aid in causal reasoning of state transitions within the environment and long-term trajectory planning and cause the policy neural network <NUM> to determine more effective policy outputs to maximize expected cumulative reward, e.g., a long-term time-discounted sum of rewards received by the system <NUM>, even when the observations <NUM> do not accurately account for stochasticity or other external factors within the environment <NUM>.

Once trained, the system <NUM> can be used, for example, to select actions <NUM> to be performed by the agent <NUM>. For example, if the reward <NUM> includes a value rating the success of the interaction of the agent <NUM> with the environment <NUM>, e.g. a value representing the amount of time it takes for the agent to accomplish a task starting from a current state of the environment, then the action <NUM> of the agent <NUM> may be selected as an action that that is predicted by the system <NUM>, i.e., as a result of the multiple planning iterations, to optimize the component of the reward <NUM> corresponding to the value.

<FIG> is a flow diagram of an example process <NUM> for simulating state transitions of an environment. For convenience, the process <NUM> will be described as being performed by a system of one or more computers located in one or more locations. For example, a reinforcement learning system, e.g., the reinforcement learning system <NUM> of <FIG>, appropriately programmed, can perform the process <NUM>.

The system can repeatedly perform the process <NUM> to generate multiple trajectories including simulated state transitions caused by actions that would be performed by the agent that, for example, can be used as training data for the system.

The system initializes a hidden state of the environment model (<NUM>). In general, each hidden state of the environment model can be generated from processing (i) the preceding environment model hidden state ht-<NUM>, (ii) the preceding latent representation zt-<NUM>, and (iii) the preceding action at-<NUM> that has been previously performed by the agent in the trajectory. Thus hidden state captures information determined (or understood) by the environment model for the environment at a corresponding state of the environment.

If the action is the first action in the trajectory, the corresponding hidden state can be an initial hidden state computed by performing the following steps <NUM>-<NUM>.

The system receives, an observation characterizing the state of the environment at the current time point (<NUM>). The observation can be a first observation in the sequence of observations included in the trajectory.

The system updates a hidden state of the policy neural network (<NUM>). That is, the system uses the policy neural network to process a policy network input including the received first observation to determine the initial policy network hidden state, and to output an action selection output ψ(a<NUM>|s<NUM>) that the system can use to select a first action a<NUM> to be performed by the agent at the beginning of the trajectory.

The system initializes the environment model hidden state based on the initial hidden state of the policy neural network and on the selected first action (<NUM>). For example, referring back to the example of <FIG>, the system can evaluate a predetermined initialization function h<NUM> = g(s<NUM>, a<NUM>) using the initial policy network hidden state and the action to determine an output that the system can use as the initial hidden state of the environment model.

The system then repeatedly performs the steps <NUM>-<NUM> of the process <NUM> to simulate a plurality of state transitions within the environment. In other words, the system performs steps <NUM>-<NUM> at each of multiple future trajectory steps forward from the first step. For convenience, each of the steps <NUM>-<NUM> will be described as being performed at a "current" trajectory step.

The system receives an action to be performed by the agent (<NUM>) at the current trajectory step. In general, the action can be an arbitrary action that has been selected, by using the policy neural network or another action selection component of the system implementing any of a variety of action selection policies.

For example, the system can select, according to a given action selection policy and based on a predicted latent representation (as generated at step <NUM>) and, optionally, the environment model hidden state, the action to be performed by the agent.

The system generates a predicted latent representation (<NUM>) based on the current hidden state of the environment model. The predicted latent representation zt is a prediction of a latent representation that would have been generated by using the policy neural network to process a policy network input including an observation characterizing the state of the environment that corresponds to the current environment model hidden state. In particular, the system can generate the predicted latent representation zt, which can take any of a variety of forms as described above, by sampling from a given distribution p(zt|ht) that is conditioned on the current environment model hidden state ht.

Optionally, the system also uses the environment model to output a current target yt based on the current hidden state ht that the system can use in action selection during planning. That is, the target is conditioned on the current environment model hidden state ht. The exact format of the targets depend on actual choice of the planning algorithm and thus may vary, but typically, the target specifies a reward that the system can use in action selection starting from the first trajectory step.

The system updates the environment model hidden state (<NUM>) to simulate a state transition caused by the agent performing the received action. The system can determine the update to the current environment model hidden state, i.e., modify the current hidden state of the environment model to transition into a subsequent hidden state, by processing (i) the current environment model hidden state ht, (ii) the predicted latent representation zt, and (iii) the received action at using the environment model and in accordance with current values of the environment model parameters.

Unlike conventional RL planning techniques, e.g., action-conditional prediction using an autoregressive generative model or option-conditional prediction using a value prediction network, which can predict outcomes of a sequence of actions on the environment in an overly redundant or causally incorrect manner, simulating state transitions caused by the agent performing various action as described above is only conditioned on, and therefore dependent upon, an minimally required amount of information to ensure causal correctness. In particular, the system can update the environment model hidden state by conditioning on the action and the minimally required amount of information to reproduce the action distribution in the model training data. Given such information, the action is independent of the other inputs and thus the model becomes less likely to be confounded by any information within the observations that aren't considered, or accounted for, by the environment model when simulating the actual, e.g., real-world, environment. By using such environment model, the system can simulate causally correct state transitions while being less affected by changes to the action selection policy and, when used in training the policy neural network, can assist in the learning of an optimal action selection policy that the system can use in controlling the agent to maximize an expected cumulative reward received by the agent by learning robust and causally correct action selection policies.

An example algorithm for using the environment model to generate a simulated trajectory under a given action selection policy ψ is shown below.

In the example algorithm shown above, at each trajectory step, the system uses the environment model to determine a current target yt ("wanted targets") for use in planning. In this example, the action to be performed by the agent is selected according to the action selection policy ψ, where the policy relies upon both the current hidden state of the environment model and the predicted latent representation when used in selecting the actions. In particular, the system determines the predicted latent representation zt ("partial view") by sampling from the distribution p(zt|ht) that is conditioned on the current hidden state of the environment model.

Claim 1:
A computer-implemented method of using an environment model to simulate state transitions of a real-world environment being interacted with by a mechanical or electronic agent that is controlled using a policy neural network, wherein the policy neural network is configured to receive an observation characterizing a state of the environment, update a belief representation of the state of the environment, generate a latent representation from the belief representation, and generate an output specifying an action to be performed by the agent from the latent representation, and wherein the method comprises:
initializing an internal representation of a state of the environment at a current time point;
repeatedly performing the following operations:
receiving an action to be performed by the agent;
generating, based on the internal representation, a predicted latent representation that is a prediction of a latent representation that would have been generated by the policy neural network by processing an observation characterizing the state of the environment corresponding to the internal representation; and
updating the internal representation to simulate a state transition caused by the agent performing the received action by processing the predicted latent representation and the received action using the environment model;
selecting, based on a result of repeatedly performing the operations, the action to be performed by the agent in the environment at the current time point and causing the agent to perform the action.