78c163f37be165673b040cff49e1ee98eb89ca0f,tf_agents/bandits/agents/greedy_reward_prediction_agent_test.py,AgentTest,testTrainAgentWithConstraint,#AgentTest#,253

Before Change

        constraints=[neural_constraint])
    observations = np.array([[1, 2], [3, 4]], dtype=np.float32)
    actions = np.array([0, 1], dtype=np.int32)
    rewards = np.array([[0.5, 6.0], [3.0, 4.0]], dtype=np.float32)
    initial_step, final_step = _get_initial_and_final_steps(
        observations, rewards)
    action_step = _get_action_step(actions)
    experience = _get_experience(initial_step, action_step, final_step)

After Change

        self._action_spec,
        constraint_network=constraint_net)

    reward_spec = {
        "reward": tensor_spec.TensorSpec(
            shape=(), dtype=tf.float32, name="reward"),
        "constraint": tensor_spec.TensorSpec(
            shape=(), dtype=tf.float32, name="constraint")
    }
    self._time_step_spec = ts.time_step_spec(self._obs_spec, reward_spec)

    agent = greedy_agent.GreedyRewardPredictionAgent(
        self._time_step_spec,
        self._action_spec,
        reward_network=reward_net,
        optimizer=optimizer,
        constraints=[neural_constraint])
    observations = np.array([[1, 2], [3, 4]], dtype=np.float32)
    actions = np.array([0, 1], dtype=np.int32)
    rewards = {
        "reward": np.array([0.5, 3.0], dtype=np.float32),
        "constraint": np.array([6.0, 4.0], dtype=np.float32)
    }
    initial_step, final_step = _get_initial_and_final_steps_nested_rewards(
        observations, rewards)
    action_step = _get_action_step(actions)
    experience = _get_experience(initial_step, action_step, final_step)

Instances