Create Custom Vectorized Environment Using Step, Reset, and Setup Functions

R2026b

In this example you create a vectorized reinforcement learning environment by supplying your custom step, reset and setup MATLAB® functions.

For this example, create a vectorized environment that represents a system for balancing a pole on a cart. The observations from the environment are the cart position, cart velocity, pendulum angle, and pendulum angular velocity. For additional details about the scalar version of this environment, see Create Custom Environment Using Step and Reset Functions. Create an observation specification for these signals.

obsInfo = rlNumericSpec([4 1]);
obsInfo.Name = "CartPole States";
obsInfo.Description = 'x, dx, theta, dtheta';

The environment has a continuous action space where the agent can apply a force value in the range [-10, 10] N to the cart. Create the action specification for this action.

actInfo = rlNumericSpec([1 1]);
actInfo.Name = "CartPole Action";
actInfo.UpperLimit =  10;
actInfo.LowerLimit = -10;

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) shared across all environment instances. For this example, use the supplied functions vec_cartpole_setup.m, vec_cartpole_reset.m, and vec_cartpole_step.m.

Setup Function

Display the 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 all environment instances. For this example, vec_cartpole_setup.m initializes a state matrix of size with four rows (each row corresponds to one of the four states) and NumEnv columns (each column corresponds to a different environment instance).

Display the setup function.

type("vec_cartpole_setup.m")

function env_data = vec_cartpole_setup(setup_info)

% Copyright 2026 The MathWorks, Inc.

% Setup 4 states for each environment instance.
env_data.State = zeros(4,setup_info.NumEnv);

Reset Function

The reset function receives the env_data structure and a logical or integer index resetidx indicating which environments to reset. It returns observations for the reset environments and the updated env_data. In vec_cartpole_reset.m, only the environments indicated by resetidx are reset. Here, you use matrix operations so that environment instances are reset simultaneously.

Display the reset function.

type("vec_cartpole_reset.m")

function [obs,env_data] = vec_cartpole_reset(env_data,resetidx)

% Copyright 2026 The MathWorks, Inc.

% Vectorized reset function for the continuous cart-pole environment. 
% Notice how the states of the system are reset in a vectorized fashion.

% Randomize theta +- 0.5 rads, set all other states to 0.
env_data.State(:,resetidx) = 0.0;
env_data.State(3,resetidx) = 0.05.*(2.*rand(1,nnz(resetidx)) - 1);

% Form the observation.
obs = {env_data.State(:,resetidx)};

Step Function

The step function receives the env_data structure and a cell array of actions, and returns observations, rewards, an is-done flag vector, and the updated env_data. In vec_cartpole_step.m, the cart-pole dynamics are written in a fully vectorized fashion so that all environment instances are stepped simultaneously.

Display the step function.

type("vec_cartpole_step.m")

function [obs,rwd,isd,env_data] = vec_cartpole_step(env_data,act)

% Copyright 2026 The MathWorks, Inc.

% Vectorized reset function for the continuous cart-pole environment. 
% Notice how the dynamics are written to take advantage of vectorization.

% Get force applied on cart-pole.
f = max(min(act{1}(:,:),10),-10);            

% Unpack state vector.
dx     = env_data.State(2,:);
theta  = env_data.State(3,:);
dtheta = env_data.State(4,:);

s = sin(theta);
c = cos(theta);

g = 9.81;
mc = 1.0;
mp = 0.1;
l = 0.5;
ts = 0.02;

ddthetaNum = g.*s + c.*(-f - mp.*l.*dtheta.^2.*s)./(mc+mp);
ddthetaDen = l.*(4/3 - mp.*c.^2./(mc+mp));
ddtheta =  ddthetaNum./ddthetaDen;

ddxNum = f + mp.*l.*(dtheta.^2.*s - ddtheta.*c);
ddxDen = mc + mp;
ddx = ddxNum./ddxDen;

% Perform Euler integration.
env_data.State = env_data.State + ts.*[dx;ddx;dtheta;ddtheta];

% Form the observation.
obs = {env_data.State};

% Compute the terminal condition.
x           = env_data.State(1,:);
theta       = env_data.State(3,:);
xlim        = 2.4;
thetalim    = 12*pi/180;
isd = uint8(abs(x) > xlim | abs(theta) > thetalim);

% Compute the reward.
distReward = 1 - abs(x)/2.4;
rwd = 0.5 + 0.5 * distReward;

% Apply penalty for falling.
rwd(isd > 0) = -50.0;

Create the custom vectorized environment using the defined observation and action specifications, the vectorized function handles, and the desired number of parallel environments. Here, 128 environment instances are created.

venv = rlFunctionVectorEnv(obsInfo,actInfo,...
    @vec_cartpole_step,...
    @vec_cartpole_reset,...
    @vec_cartpole_setup,...
    NumEnv=128)

venv = 
  rlVectorEnv with properties:

    NumEnv: 128

Use validateEnvironment to validate the environment.

validateEnvironment(venv)

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.