Multivariate LSTM cross feature dependencies

Question

I was working myself through handson-ml2, and chapter 15 in particular.

I want to generalize the multiple steps ahead approach to multiple features and one target. In order to test my understanding I create some series, which are either following a sin wave or a cos wave with some frequency. `

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

start = -10
stop = 10
n_steps = 2*(stop-start) #40
forecast_horizon = 10
num_features = 4

x_axis = np.linspace(start, stop, n_steps+forecast_horizon)
series_sin = np.stack([np.sin(np.random.rand() * x_axis) for i in range(10000)])
series_cos = np.stack([np.cos(np.random.rand() * x_axis) for i in range(10000)])

rand = np.random.rand(10000, n_steps + forecast_horizon).round()

target = np.where(rand, series_sin, series_cos)
series = np.stack([target, rand, series_sin, series_cos], axis = 2)

X_train = series[:7000,:n_steps]
X_valid = series[7000:9000,:n_steps]
X_test = series[9000:,:n_steps]
Y = np.empty([10000, n_steps, forecast_horizon])
for step_ahead in range(1, n_steps + 1):
  Y[:, step_ahead - 1, :] = \
  target[:, step_ahead:step_ahead + forecast_horizon] \
  .reshape(10000, forecast_horizon)
y_train = Y[:7000]
y_valid = Y[7000:9000]
y_test = Y[9000:]

` When plotting the target, sin and cos waves for some sample one sees that according to the rand array the target is either the sin wave or cos wave. time series plot

Now I want to train a neural network forecasting the time series.

model = tf.keras.models.Sequential([
  tf.keras.layers.Dense(512, activation = "linear"),
  tf.keras.layers.Bidirectional(
      tf.keras.layers.LSTM(
          128, 
          input_shape=(n_steps,num_features), 
          return_sequences=True, 
          activation="sigmoid")),
  tf.keras.layers.Bidirectional(
      tf.keras.layers.LSTM(
          64, 
          input_shape=(n_steps,num_features), 
          return_sequences=True, 
          activation="sigmoid")),
  tf.keras.layers.Dense(256, activation = "linear"),
  tf.keras.layers.Dense(
      forecast_horizon, 
      activation = "linear")   
])

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.01), 
    loss=tf.keras.losses.MeanSquaredError()
    )
model.fit(X_train, y_train, 
    validation_data=(X_valid, y_valid),
    epochs=2)

My assumption would be that it is very easy for the model to learn the task since feature 2,3,4 are perfect predictors for the target. But as one can see on the plot below the model does not learn any cross feature dependencies.

time series plot with forecast

Any ideas?

cheers, Felix