ML-Agents, I can't use heuristic only. So I can't use ML-Agents

Hello, I was trying for MLAgents it worked on 2D but it doesn’t work on 3D could someone please help me ?

9842790--1416315--3.PNG

using System.Collections.Generic;
using UnityEngine;
using Unity.MLAgents;
using Unity.MLAgents.Actuators;
using Unity.MLAgents.Sensors;

public class MoveToTargetAgent : Agent
{
    //[SerializeField] private Transform environment; // Environment object
    [SerializeField] private Transform target; // Target object
    [SerializeField] private Material rewardMaterial; // Reward material
    [SerializeField] private Material penaltyMaterial; // Penalty material

    public override void OnEpisodeBegin()
    {
        // Randomly set the positions of the target and the agent
        target.localPosition = new Vector3(Random.Range(-4.4f, 4.665f), 0.22f, Random.Range(4.516f, -4.648f));
        transform.localPosition = new Vector3(Random.Range(0f, 4.665f), 0.22f, Random.Range(4.516f, -4.648f));

        // Rotate the environment (optional)
        // environment.rotation = Quaternion.Euler(0f, Random.Range(0f, 360f), 0f);
    }

    public override void CollectObservations(VectorSensor sensor)
    {
        // Add the positions of the agent and the target as observations
        sensor.AddObservation(transform.localPosition);
        sensor.AddObservation(target.localPosition);
    }

    public override void OnActionReceived(ActionBuffers actions)
    {
        float moveX = actions.ContinuousActions[0]; // Movement along the X-axis (right-left)
        float moveZ = actions.ContinuousActions[1]; // Movement along the Z-axis (forward-backward)

        float movementSpeed = 10f;

        // Create a movement vector
        Vector3 movement = new Vector3(moveX, 0f, moveZ) * Time.deltaTime * movementSpeed;

        // Update the agent's position
        transform.localPosition += movement;
    }

    public override void Heuristic(in ActionBuffers actionsOut)
    {
        ActionSegment<float> continuousActions = actionsOut.ContinuousActions;

        // Log detected inputs
        Debug.Log("Horizontal Input: " + Input.GetAxis("Horizontal"));
        Debug.Log("Vertical Input: " + Input.GetAxis("Vertical"));

        // Map keyboard inputs to continuous actions
        continuousActions[0] = Input.GetAxis("Horizontal"); // X-axis (right-left)
        continuousActions[1] = 0; // No movement along the Y-axis (up-down)
        continuousActions[2] = Input.GetAxis("Vertical"); // Z-axis (forward-backward)
    }

    private void OnTriggerEnter(Collider other)
    {
        if (other.CompareTag("Target")) // Compare with the target object's tag
        {
            AddReward(10f);
            ChangeMaterial(rewardMaterial);
            EndEpisode();
        }
        else if (other.CompareTag("Wall")) // Compare with the wall object's tag
        {
            AddReward(-2f);
            ChangeMaterial(penaltyMaterial);
            EndEpisode();
        }
    }

    private void ChangeMaterial(Material material)
    {
        // Change the material of the target (or another visual element)
        target.GetComponent<MeshRenderer>().material = material;
    }
}

you only have 2 branches set in the inspector and your code is trying to access 3 branches (0, 1 & 2)