In this example you create a vectorized reinforcement learning environment by specifying your custom step, reset, and setup MATLAB® functions.
For this example, create a vectorized environment that represents a discrete-time double-integrator system with a continuous action space. A scalar version of this environment is available as a predefined environment. For more information, see Use Predefined Control System Environments.
The observations from the environment are the position and velocity of a point along a line. Create an observation specification for these signals.
The environment has a continuous action space where the agent can apply an acceleration value in the range [-1, 1] to the point. Create the action specification for this action.
Next, specify the setup, step, and reset functions. Unlike rlFunctionEnv, which requires only step and reset functions, rlFunctionVectorEnv also requires a setup function that initializes the environment data (typically the environment states and parameters). This data is shared across all environment instances. For this example, use the supplied functions vec_dblint_setupfcn.m, vec_dblint_resetfcn.m, and vec_dblint_stepfcn.m.
Setup Function
The setup function receives a structure containing a NumEnv field and returns an env_data structure that holds the state and parameters for NumEnv environment instances. For this example, vec_dblint_setup.m initializes a state matrix of size with two rows (position and velocity) and NumEnv columns (each column corresponds to a different environment instance). The function also stores the sampling time as a parameter.
Display the setup function.
function env_data = vec_dblint_setupfcn(info)
% Vectorized setup function for a double-integrator environment.
% This function initializes env_data.
% Create matrix of initial states.
env_data.x = zeros(2, info.NumEnv);
% Store the sample time as parameter.
env_data.Ts = 0.1;
Reset Function
The reset function receives both the env_data structure and a logical vector resetidx. The function then resets only the environments indicated by resetidx, and returns as outputs both the observation for the environments that were reset, and the updated env_data. In this function, you use matrix operations so that environment instances can be reset simultaneously.
Display the reset function.
function [obs, env_data] = vec_dblint_resetfcn(env_data, resetidx)
% Vectorized reset function for a double-integrator environment.
% This function resets the environment instances corresponding to the
% true elements of resetidx.
% Number of environment instances to reset.
num_reset_env = nnz(resetidx);
% Create vector of random numbers from -0.5 to 0.5 excluding zero.
r = rand(1,num_reset_env) - 0.5;
r(r == 0) = 1;
% Reset state for indicated instances.
% Set initial positions randomly to +4 or -4 and initial velocities to 0.
env_data.x(1, resetidx) = 4.*sign(r);
env_data.x(2, resetidx) = 0;
% Assign initial observation for the indicated instances.
obs{1} = env_data.x(:, resetidx);
Step Function
The step function receives the env_data structure and a cell array of actions. The function returns observations, rewards, an is-done flag vector, and the updated env_data. In this function, you write the dynamics in a fully vectorized fashion so that all environment instances can be stepped simultaneously.
Display the step function.
function [obs, rwd, isd, env_data] = vec_dblint_stepfcn(env_data, act)
% Vectorized step function for a double-integrator environment.
% Write the dynamics in matrix form to take advantage of vectorization.
% State and input matrices representing a double integrator system.
a = [0,1;0 0];
b = [0;1];
% The number of environments is the batch dimension of the action.
% The batch dimension is the dimension following the last specification
% dimension.
num_env = size(act{1}, 3);
% Extract the action and limit it between -1 and 1.
u = max(min(act{1}(:,:),1),-1);
% Extract the state matrix (2 by NumEnv) and the sample time.
x = env_data.x;
Ts = env_data.Ts;
% Update the state for all environment instances using the Euler method.
x = x + Ts.*(a*x + b*u);
% Set the next observation value.
obs{1} = x;
% Calculate reward for all environment instances (1 by NumEnv).
rwd = 1./(abs(x(1,:)) + 1);
% Set the isdone vector (1 by NumEnv) to zero.
isd = uint8(zeros(1, num_env));
% Store the updated state (2 by NumEnv).
env_data.x = x;
Create the custom vectorized environment using the defined observation and action specifications, the function handles, and the desired number of parallel environment instances to create. Here, 128 environment instances are created.
venv =
rlVectorEnv with properties:
NumEnv: 128
Use validateEnvironment to validate the environment.
This function does not result in any error, so the environment is valid.
You can now create an agent for venv and train or simulate it as you would for any other environment. Because the vectorized environment steps all instances at the same time, training with algorithms that benefit from large batches of experience can be significantly faster.
