Examples
Important
These examples are only to demonstrate the use of the GBRL library and its functions. Specific algorithm implementations can be found for stable_baselines3 in the GBRL_SB3 repository.
The full tutorial is also available as a jupyter notebook
Basic Usage: Training, Saving, Loading, Copying
In the following example, we will train, save, and load a GBRL model incrementally. We will use the base GradientBoostingTrees class and get familiarized with the basic usage of the GBRL library. We will train a GBRL model as in supervised learning on the Diabetes dataset from sklearn.
Basic imports and preprocessing
import numpy as np
import torch as th
import gymnasium as gym
from sklearn import datasets
from torch.nn.functional import mse_loss
from torch.distributions import Categorical
from gbrl import cuda_available
from gbrl.models import GBTModel, ParametricActor
Pre-process data
# CUDA is not deterministic
device = 'cuda' if cuda_available() else 'cpu'
# incremental learning dataset
X_numpy, y_numpy = datasets.load_diabetes(return_X_y=True, as_frame=False, scaled=False)
# Reshape target as GBRL works with 2D arrays
out_dim = 1 if len(y_numpy.shape) == 1 else y_numpy.shape[1]
input_dim = X_numpy.shape[1]
X, y = th.tensor(X_numpy, dtype=th.float32, device=device), th.tensor(y_numpy, dtype=th.float32, device=device)
Setting up a GBRL model
# initializing model parameters
tree_struct = {'max_depth': 4,
'n_bins': 256,
'min_data_in_leaf': 0,
'par_th': 2,
'grow_policy': 'oblivious'}
optimizer = {'algo': 'SGD',
'lr': 1.0,
'start_idx': 0,
'stop_idx': out_dim}
gbrl_params = {
"split_score_func": "Cosine",
"generator_type": "Quantile"
}
# setting up model
gbt_model = GBTModel(
input_dim=input_dim,
output_dim=out_dim,
tree_struct=tree_struct,
optimizers=optimizer,
params=gbrl_params,
verbose=0,
device=device)
gbt_model.set_bias_from_targets(y)
Incremental learning
# training for 10 epochs
n_epochs = 10
for _ in range(n_epochs):
# forward pass - setting requires_grad=True is mandatory for training
# y_pred is a torch tensor by default
y_pred = gbt_model(X, requires_grad=True)
# calculate loss - we must scale pytorch's mse loss function by 0.5 to get the correct MSE gradient
loss = 0.5 * mse_loss(y_pred, y)
loss.backward()
# perform a boosting step
gbt_model.step(X)
print(f"Boosting iteration: {gbt_model.get_iteration()} RMSE loss: {loss.sqrt()}")
GBT work with per-sample gradients but pytorch typically calculates the expected loss. GBRL internally multiplies the gradients with the number of samples when calling the step function. Therefore, when working with pytorch losses and multi-output targets one should take this into consideration. For example: when using a summation reduction
gbt_model = GBTModel(
input_dim=input_dim,
output_dim=out_dim,
tree_struct=tree_struct,
optimizers=optimizer,
params=gbrl_params,
verbose=0,
device=device)
gbt_model.set_bias_from_targets(y)
# continuing training 10 epochs using a sum reduction
n_epochs = 10
for _ in range(n_epochs):
y_pred = gbt_model(X, requires_grad=True)
# we divide the loss by the number of samples to compensate for GBRL's built-in multiplication by the same value
loss = 0.5 * mse_loss(y_pred, y, reduction='sum') / len(y_pred)
loss.backward()
# perform a boosting step
gbt_model.step(X)
print(f"Boosting iteration: {gbt_model.get_iteration()} RMSE loss: {loss.sqrt()}")
or when working with multi-dimensional outputs
y_multi = th.concat([y, y], dim=1)
out_dim = y_multi.shape[1]
gbt_model = GBTModel(
input_dim=input_dim,
output_dim=out_dim,
tree_struct=tree_struct,
optimizers=optimizer,
params=gbrl_params,
verbose=0,
device=device)
gbt_model.set_bias_from_targets(y_multi)
# continuing training 10 epochs using a sum reduction
n_epochs = 10
for _ in range(n_epochs):
y_pred = gbt_model(X, requires_grad=True)
# we multiply the loss by the output dimension to compensate for pytorch's mean reduction for MSE loss that averages across all dimensions.
# this step is necessary to get the correct loss gradient - however the loss value itself is correct
loss = 0.5 * mse_loss(y_pred, y_multi) * out_dim
loss.backward()
# perform a boosting step
gbt_model.step(X)
print(f"Boosting iteration: {gbt_model.get_iteration()} RMSE loss: {(loss / out_dim).sqrt()}")
Saving, loading, and copying a GBRL Model
# Call the save_learner method of a GBRL class
# GBRL will automatically save the file with the .gbrl_model ending
# The file will be saved in the current working directory
# Provide the absolute path to save the file in a different directory.
gbt_model.save_learner('gbt_model_tutorial')
# Loading a saved model is similar and is done by calling the specific class instance.
loaded_gbt_model = GBTModel.load_learner('gbt_model_tutorial', device=device)
# Copying a model is straighforward
copied_model = gbt_model.copy()
Manually Calculated Gradients
Alternatively, GBRL can use manually calculated gradients. Calling the predict method instead of the __call__ method, returns a numpy array instead of a PyTorch tensor. Autograd libraries or manual calculations can be used to calculate gradients. Fitting manually calculated gradients is done using the _model.step method that receives numpy arrays.
# initializing model parameters
tree_struct = {'max_depth': 4,
'n_bins': 256,
'min_data_in_leaf': 0,
'par_th': 2,
'grow_policy': 'oblivious'}
optimizer = { 'algo': 'SGD',
'lr': 1.0,
'start_idx': 0,
'stop_idx': 1}
gbrl_params = {
"split_score_func": "Cosine",
"generator_type": "Quantile"}
# setting up model
gbt_model = GBTModel(
input_dim=input_dim,
output_dim=1,
tree_struct=tree_struct,
optimizers=optimizer,
params=gbrl_params,
verbose=0,
device=device)
# works with numpy arrays as well as PyTorch tensors
gbt_model.set_bias_from_targets(y_numpy)
# training for 10 epochs
n_epochs = 10
for _ in range(n_epochs):
# y_pred is a numpy array
# set tensor = False to output a numpy array instead of a tensor
y_pred = gbt_model(X_numpy, tensor=False)
loss = np.sqrt(0.5 * ((y_pred - y_numpy)**2).mean())
grads = y_pred - y_numpy
# perform a boosting step with manual gradients
gbt_model.step(X_numpy, grads)
print(f"Boosting iteration: {gbt_model.get_iteration()} RMSE loss: {loss}")
Multiple boosting iterations
GBRL supports training multiple boosting iterations with targets similar to other GBT libraries. This is done using the fit method.
Important
Only the RMSE loss function is supported for the fit method
gbt_model = GBTModel(
input_dim=input_dim,
output_dim=1,
tree_struct=tree_struct,
optimizers=optimizer,
params=gbrl_params,
verbose=1,
device=device)
final_loss = gbt_model.fit(X_numpy, y_numpy, iterations=10)
RL using GBRL
Now that we have seen how GBRL is trained using incremental learning and PyTorch, we can use it within an RL training loop.
Important
When collecting a rollout, often the observations are flattened. As GBRL works with 2D arrays, GBRL automatically assumes that the flattened inputs are a single sample and reshapes accordingly. In case of a flattened array that represents multiple samples and a single input dimension, the user must reshape the array manually.
Let’s start by training a simple Reinforce algorithm.
def calculate_returns(rewards, gamma):
returns = []
running_g = 0.0
for reward in rewards[::-1]:
running_g = reward + gamma * running_g
returns.insert(0, running_g)
return returns
env = gym.make("CartPole-v1")
wrapped_env = gym.wrappers.RecordEpisodeStatistics(env, 50) # Records episode-reward
num_episodes = 1000
gamma = 0.99
optimizer = { 'algo': 'SGD',
'lr': 0.05,
'start_idx': 0,
'stop_idx': env.action_space.n}
bias = np.zeros(env.action_space.n, dtype=np.single)
agent = ParametricActor(
input_dim=env.observation_space.shape[0],
output_dim=env.action_space.n,
tree_struct=tree_struct,
policy_optimizer=optimizer,
params=gbrl_params,
verbose=0,
bias=bias,
device='cpu')
update_every = 20
rollout_buffer = {'actions': [], 'obs': [], 'returns': []}
for episode in range(num_episodes):
# gymnasium v26 requires users to set seed while resetting the environment
obs, info = wrapped_env.reset(seed=0)
rollout_buffer['rewards'] = []
done = False
while not done:
# obs is a flattened array representing a single sample and multiple input dimensions
# hence GBRL reshapes obs automatically to a 2D-array.
action_logits = agent(obs)
action = Categorical(logits=action_logits).sample()
action_numpy = action.cpu().numpy()
rollout_buffer['obs'].append(obs)
obs, reward, terminated, truncated, info = wrapped_env.step(action_numpy.squeeze())
rollout_buffer['rewards'].append(reward)
rollout_buffer['actions'].append(action)
done = terminated or truncated
rollout_buffer['returns'].extend(calculate_returns(rollout_buffer['rewards'], gamma))
if episode % update_every == 0 and episode > 0:
returns = th.tensor(rollout_buffer['returns'], device=device)
actions = th.cat(rollout_buffer['actions']).to(device)
# input to model can be either a torch tensor or a numpy ndarray
observations = np.stack(rollout_buffer['obs'])
# model update
action_logits = agent(observations, requires_grad=True)
dist = Categorical(logits=action_logits)
log_probs = dist.log_prob(actions)
# calculate reinforce loss with subtracted baseline
loss = -(log_probs * (returns - returns.mean())).mean()
loss.backward()
grads = agent.step(observations)
rollout_buffer = {'actions': [], 'obs': [], 'returns': []}
if episode % 100 == 0:
print(f"Episode {episode} - boosting iteration: {agent.get_iteration()} episodic return: {np.mean(wrapped_env.return_queue)}")
Explainability
GBRL implements SHAP value calculation. SHAP values can be calculated over the entire ensemble as well as for individual trees. GBRL returns SHAP values with shap: [n_samples, n_features, n_actions].
# per tree shap values
tree_shap = agent.tree_shap(0, obs)
# for the entire ensemble
shap_values = agent.shap(obs)
SHAP values are calculated internally and can be plotted using the SHAP library.
import shap
import matplotlib.pyplot as plt
plt.close('all')
explainable_values_action_1 = shap.Explanation(tree_shap.squeeze()[: , 0])
explainable_values_action_2 = shap.Explanation(tree_shap.squeeze()[: , 1])
fig, ax = plt.subplots()
shap.plots.bar(explainable_values_action_1, ax=ax)
ax.set_title("SHAP values Action 1")
fig, ax = plt.subplots()
shap.plots.bar(explainable_values_action_2, ax=ax)
ax.set_title("SHAP values Action 2")
plt.show()
GBRL’s SHAP values are optimizer-aware. Pass return_base=True to receive
both the feature attributions and a base value alongside them.
shap() explains the entire ensemble:
phi, base = agent.shap(obs, return_base=True)
# phi — shape (n_samples, n_features, output_dim)
# Each entry is the contribution of one feature to the prediction
# relative to the expected prediction.
# base — shape (n_samples, output_dim)
# The expected model output (bias + expected tree contributions).
# For SGD this is one value shared by every sample. For Adam it is
# sample-specific: it is the expected output of that sample's
# frozen-state local surrogate, not a single global E[f(X)].
#
# Identity: base + phi.sum(axis=1) == predict(obs) for every sample.
tree_shap() explains one tree at a time. Its base and phi_t are
different objects that describe only that tree’s contribution, not the whole model:
phi_t, base_t = agent.tree_shap(0, obs, return_base=True)
# phi_t — shape (n_samples, n_features, output_dim)
# Feature contributions for tree 0 only.
# base_t — shape (n_samples, output_dim)
# Expected contribution of tree 0 alone.
#
# Identity: base_t + phi_t.sum(axis=1) == contribution of tree 0 per sample.
# This does NOT reconstruct predict(obs); for that, use shap().
Both identities above hold for SGD and Adam.
With SGD a tree’s contribution is just its leaf value scaled by the learning rate,
and base is the same for every sample by construction.
With Adam a tree’s contribution also depends on the optimizer state accumulated by
all earlier trees. GBRL replays that state tree by tree and holds it fixed while
explaining the alternative leaves of that tree, so base is per-sample. The
returned values reconstruct that sample’s own prediction, but they are not the
exact Shapley values of the full Adam prediction function, because alternative
earlier-tree histories are not recomputed.
Learning Rate Schedulers
GBRL supports learning rate scheduling to control the learning rate throughout training. Two schedulers are available:
Constant (default): Fixed learning rate throughout training
Linear: Linearly interpolates between an initial and final learning rate
Note
Linear scheduler on GPU is only supported for oblivious trees (grow_policy='oblivious').
Constant Scheduler (Default)
# Constant learning rate (default behavior)
optimizer = {
'algo': 'SGD',
'lr': 0.1, # Fixed learning rate
'start_idx': 0,
'stop_idx': out_dim
}
Linear Scheduler
The linear scheduler interpolates the learning rate from lr (initial) to stop_lr (final) over T trees:
where \(t\) is the current tree index (0-indexed, so \(t \in [0, T-1]\)). The schedule covers trees 0 through T-1, and at tree T and beyond, the learning rate is held constant at stop_lr. This means:
At tree 0: \(lr_0 = lr\) (initial learning rate)
At tree T-1: \(lr_{T-1} = lr + \frac{T-1}{T} \times (stop\_lr - lr)\) (approaching final learning rate)
At tree T and beyond: \(lr_t = stop\_lr\) (held constant)
Edge Case (T=1): When T=1, the schedule contains only tree 0 which uses lr (since \(lr_0 = lr + 0/1 \times (stop\_lr - lr) = lr\)). The interpolation phase is skipped, so tree 1 and all subsequent trees immediately use stop_lr.
Parameter Constraints:
Tmust be a positive integer (minimum 1). It should equal the number of trees you expect to build.lrandstop_lrmust be positive floats.stop_lrcan be greater thanlrfor warming schedules.At tree T and for all subsequent trees, the scheduler holds at
stop_lr.
# Linear learning rate decay from 0.1 to 0.01 over 100 trees
optimizer = {
'algo': 'SGD',
'lr': 0.1, # Initial learning rate
'stop_lr': 0.01, # Final learning rate
'T': 100, # Number of trees for the schedule
'scheduler': 'Linear',
'start_idx': 0,
'stop_idx': out_dim
}
tree_struct = {
'max_depth': 4,
'n_bins': 256,
'min_data_in_leaf': 0,
'par_th': 2,
'grow_policy': 'oblivious' # Required for GPU linear scheduler
}
gbt_model = GBTModel(
input_dim=input_dim,
output_dim=out_dim,
tree_struct=tree_struct,
optimizers=optimizer,
params=gbrl_params,
verbose=1,
device=device
)
Monotonic Constraints
Monotonic constraints enforce that the model output is monotonically increasing or decreasing with respect to specific input features. This is useful for incorporating domain knowledge or ensuring interpretable behavior.
Note
Monotonic constraints are only supported for oblivious trees (grow_policy='oblivious')
optimized with SGD. Constraints apply to the output dimensions defined by start_idx to
stop_idx-1 in the optimizer configuration. For GBTModel, this typically covers all
outputs. For actor-critic models, constraints affect only the policy outputs (not value
function outputs).
Warning
Adam is not supported for constrained outputs, and GBRL raises a ValueError if you
try. Constraints are enforced by ordering the leaf gradients, but Adam’s update
is non-linear in the leaf gradient \(g\), and \(m\) and \(v\) depend on the path each sample took through the earlier trees. Two inputs differing in one feature therefore accumulate different optimizer state, while leaf values are shared — so no ordering of leaf values makes every sample’s contribution monotone. With SGD the update is a fixed signed scale, which preserves the ordering.
This applies to the whole model: if any optimizer uses Adam, monotonic constraints are rejected. Use SGD for a model that needs constrained outputs.
Note
MultiGBTLearner does not support monotonic constraints at all and raises a ValueError
if any are supplied.
How Constraints Are Enforced:
Monotonic constraints are enforced through two mechanisms:
During tree growing: Incompatible splits are rejected or pruned to preserve monotonicity. The constraint-aware scoring function pools the left and right child means when a split would violate the monotonic ordering, effectively reducing the score of such splits.
After each tree is built: the leaf values for the affected output indices (
start_idxtostop_idx-1) are projected onto the monotone cone — the closest set of leaf values satisfying the constraints.The projection uses Dykstra’s cyclic projection method. For each constrained depth, leaf pairs differing only in that depth’s split are averaged when they violate the constraint; because those pairs are disjoint, this is already the exact projection for that depth. With two or more constrained features the per-depth constraint sets overlap, so simply cycling between them would land somewhere monotone but not at the closest monotone point. Dykstra carries a correction term per depth, which recovers the true isotonic projection. A short cleanup phase then clears any residual violation left by Dykstra’s asymptotic convergence.
Dykstra runs up to a fixed pass limit. The returned leaf values are always monotone; they are the closest monotone values provided that phase converged, which it does in the large majority of cases. If the limit is reached, GBRL emits a warning and the result is monotone but only approximately nearest.
Practical Trade-offs:
Split search may be slower due to constraint checking and mean pooling during scoring
Convergence may be affected for
GBTModeland actor-critic models (policy outputs only) since some gradient directions are restrictedThe projection moves leaf values as little as possible while satisfying the constraints, but a constrained model will still generally fit worse than an unconstrained one
Setting Monotonic Constraints
Constraints are specified as a dictionary mapping feature indices to constraint specifications:
constraints = {
feature_index: (direction, output_indices),
...
}
Where:
feature_index: The input feature to constrain (int or numpy integer type)direction: The constraint direction. All three forms are accepted as valid inputs for each direction:'increasing','+', or1for increasing constraints, and'decreasing','-', or-1for decreasing constraints.output_indices: Single int or list of output dimensions to apply the constraint to
Only specify the features you want to constrain - unlisted features have no constraints.
from gbrl.models import GBTModel
input_dim = 4
out_dim = 2
# Feature 0: increasing for output 0
# Feature 1: decreasing for outputs 0 and 1
# Features 2-3: no constraints (not listed)
monotonic_constraints = {
0: ("increasing", 0),
1: ("decreasing", [0, 1]),
}
tree_struct = {
'max_depth': 4,
'n_bins': 256,
'min_data_in_leaf': 0,
'par_th': 2,
'grow_policy': 'oblivious' # Required for monotonic constraints
}
# Specify which output dimensions to optimize (constraints apply to indices start_idx to stop_idx-1)
optimizer = {
'algo': 'SGD',
'lr': 0.1,
'start_idx': 0,
'stop_idx': out_dim # constraints apply to output indices 0 to stop_idx-1
}
gbrl_params = {
'split_score_func': 'Cosine',
'generator_type': 'Quantile'
}
gbt_model = GBTModel(
input_dim=input_dim,
output_dim=out_dim,
tree_struct=tree_struct,
optimizers=optimizer,
params=gbrl_params,
monotonic_constraints=monotonic_constraints,
verbose=1,
device='cuda' # GPU supported for oblivious trees
)
Combining Schedulers and Constraints
Monotonic constraints and linear schedulers can be used together:
monotonic_constraints = {
0: ("increasing", 0),
1: ("decreasing", [0, 1]),
}
optimizer = {
'algo': 'SGD',
'lr': 0.1,
'stop_lr': 0.01,
'T': 100,
'scheduler': 'Linear',
'start_idx': 0,
'stop_idx': out_dim
}
tree_struct = {
'max_depth': 4,
'n_bins': 256,
'min_data_in_leaf': 0,
'par_th': 2,
'grow_policy': 'oblivious'
}
gbt_model = GBTModel(
input_dim=input_dim,
output_dim=out_dim,
tree_struct=tree_struct,
optimizers=optimizer,
params=gbrl_params,
monotonic_constraints=monotonic_constraints,
verbose=1,
device='cuda'
)