mlagent: reset agent pose when falling out of the plane

Question

Just got into ml-agent,I tried to make my own training scene. In order to reset the training scene,I modified the code from the rollerball demo for starter,down below is the reset part of the code:

public class forklift : Agent
{
    Rigidbody rBody;
    void Start () {
        rBody = GetComponent<Rigidbody>();
    }

    public Transform Target;
    public override void OnEpisodeBegin()
    {
       // If the Agent fell, zero its momentum
        if (this.transform.localPosition.y < -0.2f)
        {
            //this.rBody.angularVelocity = Vector3.zero;
            this.rBody.velocity = Vector3.zero;
            this.transform.localPosition = new Vector3( 0, 0.2f, 0);
        }
        
        Debug.Log("agent y: " + this.transform.localPosition.y);
        
        // Move the target to a new spot (between(8,8)&(-8,-8),leaving width 4.5 free space)
        Target.localPosition = new Vector3(Random.value * 1.6f - 0.8f,
                                           0f,
                                           Random.value * 1.6f - 0.8f);
        //Target.transform.localRotation = Quaternion.Euler(0.0f, Random.value *360, 0.0f);
    }

However, the code doesn't work properly if the agent fall off the plane. It seems to just move the falling agent to the reset coordinate, which causes the agent to remain the falling pose.Is there a code to reset the pose, too?