# Create Custom Environment Using Step and Reset Functions

This example shows how to create a custom environment by writing your own MATLAB® `step` and `reset` functions.

Using the `rlFunctionEnv` function, you can create a MATLAB reinforcement learning environment from an observation specification, an action specification, and `step` and `reset` functions that you supply. You can then train a reinforcement learning agent in this environment. For this example, the necessary `step` and `reset` functions are already defined.

Creating an environment using custom functions is especially useful for simpler environments that do not need many helper functions and have no special visualization requirements. For more complex environments, you can create an environment object using a template class. For more information, see Create Custom Environment from Class Template.

For more information on creating reinforcement learning environments, see Reinforcement Learning Environments and Create Custom Simulink Environments.

### DIscrete Action Space Cart-Pole MATLAB Environment

The cart-pole environment is a pole attached to an unactuated joint on a cart, which moves along a frictionless track. The training goal is to make the pendulum stand upright without falling over. For this environment:

• The upward balanced pendulum position is 0 radians and the downward hanging position is $\pi$ radians.

• The pendulum starts upright with an initial angle that is between –0.05 radians and 0.05 radians.

• The force action signal from the agent to the environment is either –10 N or 10 N.

• The observations from the environment are the cart position, cart velocity, pendulum angle, and pendulum angular velocity.

• The episode terminates if the pole is more than 12 degrees from vertical or the cart moves more than 2.4 m from the original position.

• A reward of +1 is provided for every time step that the pole remains upright. A penalty of –10 is applied when the pendulum falls.

### Observation and Action Specifications

The observations from the environment are the cart position, cart velocity, pendulum angle, and pendulum angle derivative.

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

The environment has a discrete action space where the agent can apply one of two possible force values to the cart: `-10` or `10` N.

```ActInfo = rlFiniteSetSpec([-10 10]); ActInfo.Name = "CartPole Action";```

For more information on specifying environment actions and observations, see `rlNumericSpec` and `rlFiniteSetSpec`.

### Define Environment Reset and Step Functions

To define a custom environment, first specify the custom `step` and `reset` functions. These functions must be in your current working folder or on the MATLAB path.

The `reset` function sets the initial state of the environment. This function must have the following signature.

```[InitialObservation,Info] = myResetFunction() ```

The first output argument is the initial observation. The second output argument can be any useful environment information that you want to pass from one step to the next, such as for example the environment state, or a structure containing state and parameters.

At the beginning of the training (or simulation) episode, `train` (or `sim`) calls your reset function and uses its second output argument to initialize the `Info` property of your custom environment. During a training (or simulation) step, `train` (or `sim`) supplies the current value of `Info` as the second input argument of `StepFcn`, and then uses the fourth output argument returned by `StepFcn` to update the value of `Info`.

For this example, use the second argument to store the initial states of the cart-pole environment: the position and velocity of the cart, the pendulum angle, and the pendulum angle derivative. The `reset` function sets the cart angle to a random value each time the environment is reset.

For this example, use the custom reset function defined in `myResetFunction.m`.

`type myResetFunction.m`
```function [InitialObservation, InitialState] = myResetFunction() % Reset function to place custom cart-pole environment into a random % initial state. % Theta (randomize) T0 = 2 * 0.05 * rand() - 0.05; % Thetadot Td0 = 0; % X X0 = 0; % Xdot Xd0 = 0; % Return initial environment state variables as logged signals. InitialState = [X0;Xd0;T0;Td0]; InitialObservation = InitialState; end ```

The `step` function specifies how the environment advances to the next state based on a given action. This function must have the following signature.

```[NextObservation,Reward,IsDone,UpdatedInfo] = myStepFunction(Action,Info) ```

To calculate the new state, the step function applies the dynamic equation to the current state stored in `Info`. The function then returns the updated state in `UpdatedInfo`. At the next training (or simulation) step, `train` or `sim` takes the fourth output argument obtained during the previous step, `UpdatedInfo`, and supplies it to the step function as the second input argument, `Info`.

For this example, use the custom step function defined in `myStepFunction.m`. For implementation simplicity, this function redefines physical constants, such as the cart mass, every time `step` is executed. An alternative is to define the physical constants in the reset function, define `Info` as a structure containing both state and parameters, and therefore use `Info` to store the physical constants as well as the environment states. This alternative implementation lets you easily change some of the parameters during the simulation or training if you need to.

`type myStepFunction.m`
```function [NextObs,Reward,IsDone,NextState] = myStepFunction(Action,State) % Custom step function to construct cart-pole environment for the function % name case. % % This function applies the given action to the environment and evaluates % the system dynamics for one simulation step. % Define the environment constants. % Acceleration due to gravity in m/s^2 Gravity = 9.8; % Mass of the cart CartMass = 1.0; % Mass of the pole PoleMass = 0.1; % Half the length of the pole HalfPoleLength = 0.5; % Max force the input can apply MaxForce = 10; % Sample time Ts = 0.02; % Pole angle at which to fail the episode AngleThreshold = 12 * pi/180; % Cart distance at which to fail the episode DisplacementThreshold = 2.4; % Reward each time step the cart-pole is balanced RewardForNotFalling = 1; % Penalty when the cart-pole fails to balance PenaltyForFalling = -10; % Check if the given action is valid. if ~ismember(Action,[-MaxForce MaxForce]) error('Action must be %g for going left and %g for going right.',... -MaxForce,MaxForce); end Force = Action; % Unpack the state vector from the logged signals. XDot = State(2); Theta = State(3); ThetaDot = State(4); % Cache to avoid recomputation. CosTheta = cos(Theta); SinTheta = sin(Theta); SystemMass = CartMass + PoleMass; temp = (Force + PoleMass*HalfPoleLength*ThetaDot*ThetaDot*SinTheta)/SystemMass; % Apply motion equations. ThetaDotDot = (Gravity*SinTheta - CosTheta*temp) / ... (HalfPoleLength*(4.0/3.0 - PoleMass*CosTheta*CosTheta/SystemMass)); XDotDot = temp - PoleMass*HalfPoleLength*ThetaDotDot*CosTheta/SystemMass; % Perform Euler integration to calculate next state. NextState = State + Ts.*[XDot;XDotDot;ThetaDot;ThetaDotDot]; % Copy next state to next observation. NextObs = NextState; % Check terminal condition. X = NextObs(1); Theta = NextObs(3); IsDone = abs(X) > DisplacementThreshold || abs(Theta) > AngleThreshold; % Calculate reward. if ~IsDone Reward = RewardForNotFalling; else Reward = PenaltyForFalling; end end ```

Construct the custom environment using the observation and action specification, and the names of your `step` and `reset` functions.

`env = rlFunctionEnv(ObsInfo,ActInfo,"myStepFunction","myResetFunction");`

To verify the operation of your environment, `rlFunctionEnv` automatically calls `validateEnvironment` after creating the environment.

### Pass Additional Arguments Using Anonymous Functions

While the custom reset and step functions that you must pass to `rlFuntionEnv` must have exactly zero and two arguments, respectively, you can avoid this limitation by using anonymous functions. Specifically, you define the `reset` and `step` functions to be passed to `rlFuntionEnv` as anonymous functions (with zero and two arguments, respectively). In turn, these anonymous functions, call your custom functions that have additional arguments.

For example, to pass the additional arguments `arg1` and `arg2` to both the `step` and `reset` function, you can write the following functions.

```[InitialObservation,Info] = myResetFunction(arg1,arg2) [Observation,Reward,IsDone,Info] = myStepFunction(Action,Info,arg1,arg2) ```

Then, with `arg1` and `arg2` in the MATLAB workspace, define the following handles to anonymous `reset` and `step` functions with zero and two arguments, respectively.

```ResetHandle = @() myResetFunction(arg1,arg2); StepHandle = @(Action,Info) myStepFunction(Action,Info,arg1,arg2); ```

If `arg1` and `arg2` are available at the time that `ResetHandle` and `StepHandle` are created, the workspaces of both anonymous functions include those values. The values persist within the function workspaces even if you clear the variables from the MATLAB workspace. When `ResetHandle` and `StepHandle` are evaluated, they invoke `myResetFunction` and `myStepFunction` and pass them a copy of `arg1` and `arg2`. For more information, see Anonymous Functions.

Using additional input arguments can create a more efficient environment implementation. For example, `myStepFunction2.m` is a custom step function that takes the environment constants as its third input argument (`envPars`). By doing so, this function avoids redefining the environment constants at each step.

`type myStepFunction2.m`
```function [NextObs,Reward,IsDone,NextState] = myStepFunction2(Action,State,EnvPars) % Custom step function to construct cart-pole environment for the function % handle case. % % This function applies the given action to the environment and evaluates % the system dynamics for one simulation step. % Check if the given action is valid. if ~ismember(Action,[-EnvPars.MaxForce EnvPars.MaxForce]) error('Action must be %g for going left and %g for going right.',... -EnvPars.MaxForce,EnvPars.MaxForce); end Force = Action; % Unpack the state vector from the logged signals. XDot = State(2); Theta = State(3); ThetaDot = State(4); % Cache to avoid recomputation. CosTheta = cos(Theta); SinTheta = sin(Theta); MassPole = EnvPars.MassPole; HalfLen = EnvPars.HalfLength; SystemMass = EnvPars.MassCart + MassPole; temp = (Force + MassPole*HalfLen*ThetaDot*ThetaDot*SinTheta)/SystemMass; % Apply motion equations. ThetaDotDot = (EnvPars.Gravity*SinTheta - CosTheta*temp)... / (HalfLen*(4.0/3.0 - MassPole*CosTheta*CosTheta/SystemMass)); XDotDot = temp - MassPole*HalfLen*ThetaDotDot*CosTheta/SystemMass; % Perform Euler integration. NextState = State + EnvPars.Ts.*[XDot;XDotDot;ThetaDot;ThetaDotDot]; % Copy next state to next observation. NextObs = NextState; % Check terminal condition. X = NextObs(1); Theta = NextObs(3); IsDone = abs(X) > EnvPars.XThreshold || ... abs(Theta) > EnvPars.ThetaThresholdRadians; % Calculate reward. if ~IsDone Reward = EnvPars.RewardForNotFalling; else Reward = EnvPars.PenaltyForFalling; end end ```

Create the structure that contains the environment parameters.

```% Acceleration due to gravity in m/s^2 envPars.Gravity = 9.8; % Mass of the cart envPars.MassCart = 1.0; % Mass of the pole envPars.MassPole = 0.1; % Half the length of the pole envPars.HalfLength = 0.5; % Max force the input can apply envPars.MaxForce = 10; % Sample time envPars.Ts = 0.02; % Angle at which to fail the episode envPars.ThetaThresholdRadians = 12 * pi/180; % Distance at which to fail the episode envPars.XThreshold = 2.4; % Reward each time step the cart-pole is balanced envPars.RewardForNotFalling = 1; % Penalty when the cart-pole fails to balance envPars.PenaltyForFalling = -5;```

Create an anonymous function that calls your custom step function, passing `envPars` as an additional input argument.

`StepHandle = @(Action,Info) myStepFunction2(Action,Info,envPars);`

Because `envPars` is available at the time that `StepHandle` is created, the anonymous function workspace includes a copy of `envPars`. When `StepHandle` is evaluated, it calls `myStepFunction2` passing its copy of `envPars`.

Use the same `reset` function, specifying it as a function handle rather than by using its name.

`ResetHandle = @() myResetFunction;`

Create the environment using the handles to the anonymous functions.

`env2 = rlFunctionEnv(ObsInfo,ActInfo,StepHandle,ResetHandle);`

### Visually Inspect Outputs of Custom Functions

While `rlFunctionEnv` automatically calls `validateEnvironment` after creating the environment, it might be useful to visually inspect the output of your functions to further confirm that their behavior conforms to your expectations. To do so, initialize your environment using the `reset` function and run one simulation step using the `step` function. For reproducibility, set the random generator seed before validation.

Validate the environment created using function names.

```rng(0); InitialObs = reset(env)```
```InitialObs = 4×1 0 0 0.0315 0 ```
```[NextObs,Reward,IsDone,Info] = step(env,10); NextObs```
```NextObs = 4×1 0 0.1947 0.0315 -0.2826 ```

Validate the environment created using function handles.

```rng(0); InitialObs2 = reset(env2)```
```InitialObs2 = 4×1 0 0 0.0315 0 ```
```[NextObs2,Reward2,IsDone2,Info2] = step(env2,10); NextObs2```
```NextObs2 = 4×1 0 0.1947 0.0315 -0.2826 ```

Both environments initialize and simulate successfully, producing the same state values in `NextObs`.