def test_input_gradient(self):
        """
        """
        us = np.ones((1, 2))
        cost_grad =\
            TwoWheeledConfigModule.gradient_cost_fn_with_input(None, us)

        # numeric grad
        eps = 1e-4
        expected_grad = np.zeros((1, 2))
        for i in range(2): 
            tmp_u = us.copy()
            tmp_u[0, i] = us[0, i] + eps
            forward = \
                TwoWheeledConfigModule.input_cost_fn(tmp_u)
            forward = np.sum(forward)
            tmp_u = us.copy()
            tmp_u[0, i] = us[0, i] - eps
            backward = \
                TwoWheeledConfigModule.input_cost_fn(tmp_u)
            backward = np.sum(backward)

            expected_grad[0, i] = (forward - backward) / (2. * eps)
        
        assert cost_grad == pytest.approx(expected_grad)
    def test_calc_costs(self):
        # make config
        config = TwoWheeledConfigModule()
        # set
        pred_len = 5
        state_size = 3
        input_size = 2
        pop_size = 2
        pred_xs = np.ones((pop_size, pred_len, state_size))
        g_xs = np.ones((pop_size, pred_len, state_size)) * 0.5
        input_samples = np.ones((pop_size, pred_len, input_size)) * 0.5

        costs = config.input_cost_fn(input_samples)
        expected_costs = np.ones((pop_size, pred_len, input_size))*0.5

        assert costs == pytest.approx(expected_costs**2 * np.diag(config.R))

        costs = config.state_cost_fn(pred_xs, g_xs)
        expected_costs = np.ones((pop_size, pred_len, state_size))*0.5
        
        assert costs == pytest.approx(expected_costs**2 * np.diag(config.Q))

        costs = config.terminal_state_cost_fn(pred_xs[:, -1, :],\
                                              g_xs[:, -1, :])
        expected_costs = np.ones((pop_size, state_size))*0.5
        
        assert costs == pytest.approx(expected_costs**2 * np.diag(config.Sf))