code: def calculate_step_return(
      rewards,
      discounts,
      importance_weights,
      importance_weights_one_earlier,
      estimated_state_values,
      estimated_q_values,
      j_step,
  ):
      trajectory_length = len(rewards[0])
      num_trajectories = len(rewards)
      j_step = int(min(j_step, trajectory_length - 1))
  
      weighted_discounts = np.multiply(discounts, importance_weights)
      weighted_discounts_one_earlier = np.multiply(
          discounts, importance_weights_one_earlier
      )
  
      importance_sampled_cumulative_reward = np.sum(
          np.multiply(
              weighted_discounts[:, : j_step + 1],
              rewards[:, : j_step + 1],
          ),
          axis=1,
      )
  
      if j_step < trajectory_length - 1:
          direct_method_value = (
              weighted_discounts_one_earlier[:, j_step + 1]
              * estimated_state_values[:, j_step + 1]
          )
      else:
          direct_method_value = np.zeros([num_trajectories])
  
      control_variate = np.sum(
          np.multiply(
              weighted_discounts[:, : j_step + 1],
              estimated_q_values[:, : j_step + 1],
          )
          - np.multiply(
              weighted_discounts_one_earlier[:, : j_step + 1],
              estimated_state_values[:, : j_step + 1],
          ),
          axis=1,
      )
  
      j_step_return = (
          importance_sampled_cumulative_reward
          + direct_method_value
          - control_variate
      )
  
      return j_step_return
  
file_path: ReAgent/reagent/evaluation/weighted_sequential_doubly_robust_estimator.py
imports: import itertools
  import logging
  import numpy
  import scipy
  from reagent.evaluation.cpe import CpeEstimate
  from reagent.evaluation.evaluation_data_page import EvaluationDataPage
label_desc: Use image node_id 6 for calling the WeightedSequentialDoublyRobustEstimator obj's underlying member method code with example usage: obj.calculate_step_return(rewards, discounts, importance_weights, importance_weights_one_earlier, estimated_state_values, estimated_q_values, j_step) and returns: j_step_return