float xInput = Input.GetAxis("Horizontal")
float zInput = Input.GetAxis("Vertical");

Vector3 dir = new Vector3(xInput, 0, zInput) * movementSpeed;
dir.y = rb.velocity.y;

rb.velocity = dir;