Using roll-forward window to create a training set for ML based on multivariate time series

Based on the simplifed sample dataframe

import pandas as pd
import numpy as np
timestamps = pd.date_range(start='2017-01-01', end='2017-01-5', inclusive='left')
values = np.arange(0,len(timestamps))
df = pd.DataFrame({'A': values ,'B' : values*2},
                       index = timestamps )
print(df)

            A  B
2017-01-01  0  0
2017-01-02  1  2
2017-01-03  2  4
2017-01-04  3  6

I want to use a roll-forward window of size 2 with a stride of 1 to create a resulting dataframe like

     timestep_1  timestep_2  target_A  
0  A 0           1           2         
   B 0           2           2         
1  A 1           2           3
   B 2           4           3

To create a dataframe for the training of an ML model that predicts target values based on the values of n previous timesteps, where n=2, in the example above.

I.e., for each window step, a data item is created with the two values of A and B in this window and the A value immediately to the right of the window as target_A, where the index is the number of the data item.

My first idea was to use pandas

https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.rolling.html

But that seems to only work in combination with aggregate functions such as sum, which is a completely different use case.

Any ideas on how to implement this rolling-window-based sampling approach?