unityのml-agentでエピソード実行中はアクションを起こさないようにしたい

Question

前提

unityでml-agentを用いた実装をしています。
ゴールは下の図にあるように、上の石が下の石にうまく乗るように学習させることです。

実現したいこと

エピソード中は動かずにエピソードが始まるタイミングで位置と角度が決まるようにしたいです。

該当のソースコード

このコードで行動を決めるとエピソード中に石が動いて落ちないように学習してしまいます。
そうではなく、エピソード中に石は動かず、初期位置を学習してうまく乗る部分を探し出してほしいです。

public class RockAgent : Agent
{
    public GameObject rock;
    public GameObject rockB;

    Rigidbody m_Rb_rock;
    Rigidbody m_Rb_rockB;

    public override void Initialize()
    {
        m_Rb_rock = rock.GetComponent<Rigidbody>();
        m_Rb_rockB = rockB.GetComponent<Rigidbody>();
    }

    public override void OnEpisodeBegin()
    {
        rock.transform.localPosition = new Vector3(0f, 5f, 0f);
        rock.transform.rotation = Quaternion.Euler(0f, 0f, 0f);
        m_Rb_rock.velocity = Vector3.zero;
        m_Rb_rock.angularVelocity = Vector3.zero;

        rockB.transform.localPosition = new Vector3(0f, 0.1f, 0f);
        rockB.transform.rotation = Quaternion.Euler(0f, 0f, 0f);
        m_Rb_rockB.velocity = Vector3.zero;
        m_Rb_rockB.angularVelocity = Vector3.zero;
    }

    public override void CollectObservations(VectorSensor sensor)
    {
        sensor.AddObservation(rock.transform.localPosition);
        sensor.AddObservation(rock.transform.localEulerAngles);
    }

    public override void OnActionReceived(ActionBuffers actionBuffers)
    {
        float turnX = 0f;
        if (actionBuffers.DiscreteActions[0] == 1f)
        {
            turnX = -1f;
        }
        else if (actionBuffers.DiscreteActions[0] == 2f)
        {
            turnX = 1f;
        }
        rock.transform.Rotate(200f * turnX * Time.deltaTime, 0, 0, Space.World);

        float turnY = 0f;
        if (actionBuffers.DiscreteActions[1] == 1f)
        {
            turnY = -1f;
        }
        else if (actionBuffers.DiscreteActions[1] == 2f)
        {
            turnY = 1f;
        }
        rock.transform.Rotate(0, 200f * turnY * Time.deltaTime, 0, Space.World);

        float turnZ = 0f;
        if (actionBuffers.DiscreteActions[2] == 1f)
        {
            turnZ = -1f;
        }
        else if (actionBuffers.DiscreteActions[2] == 2f)
        {
            turnZ = 1f;
        }
        rock.transform.Rotate(0, 0, 200f * turnZ * Time.deltaTime, Space.World);

        if (MaxStep > 0) AddReward(1f / MaxStep);

        if (rock.transform.position.y < 0f)
        {
            SetReward(-1.0f);
            EndEpisode();
        }
    }

Comments

[Bongo]

ご提示いただいたソースコードが、中身が空のOnActionReceivedだけのようなのですが、投稿時のミスでしょうか?

[545498_ho]

ありがとうございます。全体のコードを載せました。

Answer 1

エージェントにDecisionRequesterはアタッチされているでしょうか?あれがあると設定した周期で自動的に行動決定要求を出してくるので、今回の目的には邪魔になるかと思います。もしDecisionRequesterがあれば、ひとまずそれは削除してしまうのがいいでしょう。
その上で、下記のようにOnEpisodeBeginで一発だけRequestDecisionを実行してはいかがでしょうか。

using Unity.MLAgents;
using Unity.MLAgents.Actuators;
using Unity.MLAgents.Sensors;
using UnityEngine;

public class RockAgent : Agent
{
	// 後述のrockに付けるタグ
	// 別途Tags & Layersでこの名前のタグを登録しておく必要があります
	static readonly string RockTag = "Rock";

	// 後述の1エピソードの最大時間
	static readonly float EpisodeTime = 10.0f;

	public GameObject rock;
	public GameObject rockB;

	// 後述のrockBが置かれている地面オブジェクト
	public GameObject ground;

	Rigidbody m_Rb_rock;
	Rigidbody m_Rb_rockB;

	// 後述の岩モデルの初期位置
	Vector3 initialPosition;
	Vector3 initialPositionB;

	// 後述のエピソード開始フラグ
	bool episodeBegun;

	// 後述のエピソード中断二乗距離
	float distanceThreshold;

	// 後述の残り時間
	float timeLeft;

	// rockがrockBを押しのけるようにして地面の上に乗り、地面から落ちることなく
	// いつまでもエピソードが終わらないケースがあるようだったので、地面に
	// 下記CollisionDetectorをアタッチすることでrockが地面に触れたのを検出し
	// エピソードを終了させることにしました
	[DisallowMultipleComponent]
	[RequireComponent(typeof(Collider))]
	private class CollisionDetector : MonoBehaviour
	{
		public RockAgent Agent { get; set; }

		void OnCollisionEnter(Collision collision)
		{
			if (collision.gameObject.CompareTag(RockTag))
			{
				Agent.StopEpisode();
			}
		}
	}

	// CollisionDetectorからもエピソードの停止が行えるよう
	// 停止処理を単独のメソッドとして分離しました
	void StopEpisode()
	{
		episodeBegun = false;
		SetReward(-1.0f);
		EndEpisode();
	}

	public override void Initialize()
	{
		m_Rb_rock = rock.GetComponent<Rigidbody>();
		m_Rb_rockB = rockB.GetComponent<Rigidbody>();

		// 前述のCollisionDetectorを地面にアタッチしておきます
		var detector = ground.AddComponent<CollisionDetector>();
		detector.Agent = this;

		// rockにはタグを付けておきます
		rock.tag = RockTag;

		// 今回の実験では、ご質問者さんの岩モデルとは異なる岩モデルを代用品として
		// 用意したため、rockの初期位置が(0, 5, 0)、rockBが(0, 0.1, 0)で決め打ちだと
		// 私の岩モデルでは少々不都合でした
		// そこで、シーン上に配置した時の位置を初期位置とするようにしました
		initialPosition = rock.transform.localPosition;
		initialPositionB = rockB.transform.localPosition;

		// rockとrockBの初期距離の4倍を限界としました
		distanceThreshold = ((initialPosition - initialPositionB) * 4.0f).sqrMagnitude;
	}

	public override void OnEpisodeBegin()
	{
		rock.transform.localPosition = initialPosition;
		rock.transform.rotation = Quaternion.Euler(0f, 0f, 0f);
		m_Rb_rock.velocity = Vector3.zero;
		m_Rb_rock.angularVelocity = Vector3.zero;

		rockB.transform.localPosition = initialPositionB;
		rockB.transform.rotation = Quaternion.Euler(0f, 0f, 0f);
		m_Rb_rockB.velocity = Vector3.zero;
		m_Rb_rockB.angularVelocity = Vector3.zero;

		// 残り時間を初期値に設定
		timeLeft = EpisodeTime;

		// エピソード開始時に行動決定およびアクションを要求
		RequestDecision();
	}

	public override void CollectObservations(VectorSensor sensor)
	{
		// 念のため申し上げますと、今回の状況設定ではエピソード開始時の
		// 岩の姿勢は常に一定ですので、それをエピソード開始時の1回だけ
		// 観測したところで変化はなく、行動決定の役には立たなそうな気がします
		sensor.AddObservation(rock.transform.localPosition);
		sensor.AddObservation(rock.transform.localEulerAngles);
	}

	public override void OnActionReceived(ActionBuffers actionBuffers)
	{
		// エピソード開始時の回転一発で目標姿勢に向けるのであれば、
		// 離散アクションより連続アクションの方が適しているんじゃないかと思い
		// ContinuousActionsを使うよう変更しました
		var rotation = Quaternion.identity;
		for (var i = 0; i < 4; i++)
		{
			rotation[i] = actionBuffers.ContinuousActions[i];
		}
		rotation.Normalize();
		rock.transform.localRotation = rotation;

		// 行動決定はエピソード開始時の1回しか行わないので
		// OnActionReceived内で結果を評価するわけにはいかないため
		// ここではエピソード開始フラグを立てるだけにしました
		episodeBegun = true;
	}

	// 結果を評価し報酬を与えるのはFixedUpdate内で行いました
	void FixedUpdate()
	{
		if (!episodeBegun)
		{
			return;
		}

		// rockがおおむね静止している間、残り時間をカウントダウンしていき...
		if (m_Rb_rock.velocity.sqrMagnitude < 0.01f)
		{
			AddReward(1f / EpisodeTime);
			timeLeft -= Time.deltaTime;
		}

		// EpisodeTime秒持ちこたえたら、ペナルティなしで
		// エピソードを終えることにしました
		if (timeLeft <= 0.0f)
		{
			EndEpisode();
		}

		// ご質問者さんの条件に加えて、何らかの理由で岩が吹っ飛んで行方不明に
		// なった場合に備え、rockとrockBが限界距離を超えて遠ざかった場合にも
		// エピソードを中止するようにしました
		// また、前述のようにCollisionDetectorが衝突を検出した場合も
		// CollisionDetectorによってエピソードが中断されることになります
		if (rock.transform.position.y < 0f || (rock.transform.localPosition - rockB.transform.localPosition).sqrMagnitude > distanceThreshold)
		{
			StopEpisode();
		}
	}
}

なおコード中のコメントでも申し上げましたが、行動を離散型から連続型に変更（XYZ軸それぞれについて200°/秒で回転するかしないか...の代わりに、目標姿勢を表すクォータニオンの4成分を出力させる）しましたので、エージェントのインスペクター上でContinuous Actionsを4、Discrete Branchesを0に設定しています。

50万回の訓練を行ったところ、下図のように次第に成績が向上し...

確実に成功とまではいきませんでしたが、高確率で安定な姿勢で落とせるようになりました。

念のため申し上げますと、今回の実験ではrockの初期姿勢しか判断材料に加えていませんので（その初期姿勢にしても、コード中のコメントで申し上げましたように役に立っているとは思えず、実質判断材料なしの手探り）、rockやrockBの形を変えたり、あるいはrockBの位置をずらしただけでも成功率が落ちそうな気がします。
いろいろな岩に対応させるには、何らかの手段で岩の形をエージェントに教えてやる必要があるでしょう。たとえば岩の周りからさまざまな角度でカメラセンサーで撮影する...とかでしょうかね?

高い塔を目指した結果

using Unity.MLAgents;
using Unity.MLAgents.Actuators;
using Unity.MLAgents.Sensors;
using UnityEngine;

public class RockAgent : Agent
{
	static readonly string RockTag = "Rock";
	static readonly float EpisodeTime = 10.0f;

	public GameObject rock;
	public GameObject rockB;
	public GameObject ground;

	float distanceThreshold;
	bool episodeBegun;
	Vector3 initialPosition;
	Vector3 initialPositionB;
	Rigidbody m_Rb_rock;
	Rigidbody m_Rb_rockB;
	float timeLeft;

	void StopEpisode()
	{
		// 失敗時のペナルティを増やしました
		episodeBegun = false;
		SetReward(-4.0f);
		EndEpisode();
	}

	[DisallowMultipleComponent]
	[RequireComponent(typeof(Collider))]
	private class CollisionDetector : MonoBehaviour
	{
		public RockAgent Agent { get; set; }

		private void OnCollisionEnter(Collision collision)
		{
			if (collision.gameObject.CompareTag(RockTag))
			{
				Agent.StopEpisode();
			}
		}
	}

	public override void Initialize()
	{
		m_Rb_rock = rock.GetComponent<Rigidbody>();
		m_Rb_rockB = rockB.GetComponent<Rigidbody>();
		var detector = ground.AddComponent<CollisionDetector>();
		detector.Agent = this;
		rock.tag = RockTag;
		initialPosition = rock.transform.localPosition;
		initialPositionB = rockB.transform.localPosition;
		distanceThreshold = ((initialPosition - initialPositionB) * 4.0f).sqrMagnitude;
	}

	public override void OnEpisodeBegin()
	{
		rock.transform.localPosition = initialPosition;
		rock.transform.rotation = Quaternion.Euler(0f, 0f, 0f);
		m_Rb_rock.velocity = Vector3.zero;
		m_Rb_rock.angularVelocity = Vector3.zero;
		rockB.transform.localPosition = initialPositionB;
		rockB.transform.rotation = Quaternion.Euler(0f, 0f, 0f);
		m_Rb_rockB.velocity = Vector3.zero;
		m_Rb_rockB.angularVelocity = Vector3.zero;
		timeLeft = EpisodeTime;
		RequestDecision();
	}

	public override void CollectObservations(VectorSensor sensor)
	{
		sensor.AddObservation(rock.transform.localPosition);
		sensor.AddObservation(rock.transform.localEulerAngles);
	}

	public override void OnActionReceived(ActionBuffers actionBuffers)
	{
		// 慣性モーメントが最も小さい方角が長軸だと仮定し、その軸を縦に置いた姿勢から
		// 一定角度（さしあたり15°）以内の範囲で回転量を決めさせることにしました
		var actions = actionBuffers.ContinuousActions;
		var inertiaTensor = m_Rb_rock.inertiaTensor;
		var axis = Vector3.zero;
		if (inertiaTensor.x < inertiaTensor.y)
		{
			if (inertiaTensor.x < inertiaTensor.z)
			{
				axis.x = 1.0f;
			}
			else
			{
				axis.z = 1.0f;
			}
		}
		else
		{
			if (inertiaTensor.y < inertiaTensor.z)
			{
				axis.y = 1.0f;
			}
			else
			{
				axis.z = 1.0f;
			}
		}
		var rotation = Quaternion.FromToRotation(
			m_Rb_rock.inertiaTensorRotation * axis,
			actions[3] > 0.0f ? Vector3.up : Vector3.down);
		rotation = Quaternion.AngleAxis(actions[3], Vector3.up) * rotation;
		var targetDirection = new Vector3(actions[0], actions[1], actions[2]).normalized;
		if (targetDirection.sqrMagnitude > 0.0f)
		{
			rotation = Quaternion.RotateTowards(
				Quaternion.identity,
				Quaternion.FromToRotation(Vector3.up, targetDirection),
				15.0f) * rotation;
		}
		rock.transform.localRotation = rotation;

		// 回転だけできれいに乗る位置を決めさせるのは酷なように思い、
		// 位置を少しずらすことも許すようにしました
		// ただし、ずれに応じていくらか減点しています
		var shift = Vector3.ClampMagnitude(
			new Vector3(actions[4], 0.0f, actions[5]),
			0.0625f);
		rock.transform.localPosition += shift;
		AddReward(-shift.sqrMagnitude * 16.0f);

		episodeBegun = true;
	}

	void FixedUpdate()
	{
		if (!episodeBegun)
		{
			return;
		}

		if (m_Rb_rock.velocity.sqrMagnitude < 0.01f)
		{
			// 静止時の獲得報酬が多すぎたように感じ、引き下げました
			AddReward((4.0f * Time.deltaTime) / EpisodeTime);
			timeLeft -= Time.deltaTime;
		}

		if (timeLeft <= 0.0f)
		{
			// 適当な高さ（さしあたりinitialPosition.yの2倍）から適当なサイズのBoxCastを行い...
			var origin = ground.transform.position;
			if (Physics.BoxCast(
					new Vector3(origin.x, initialPosition.y * 2.0f, origin.z),
					new Vector3(1.0f, 1.0f, 0.01f),
					Vector3.down,
					out var hitInfo))
			{
				// 塔の高さを測定して、高さに応じたボーナス点を与えることにしました
				// このボーナス点の式は勘で決めたもので、根拠があるわけではありません
				AddReward(Mathf.Pow(Mathf.Clamp01((hitInfo.point.y - origin.y) / initialPosition.y), 8.0f) * 32.0f);
			}

			EndEpisode();
		}

		if ((rock.transform.position.y < 0f) ||
			((rock.transform.localPosition - rockB.transform.localPosition).sqrMagnitude >
			 distanceThreshold))
		{
			StopEpisode();
		}
	}
}

Comments

[545498_ho]

再現までしてくださってありがとうございます。
一度自分でもやってみます。

[545498_ho]

石がなるべく高くなるような傾きで乗ってほしい場合、どのように報酬を与えればよいでしょうか？

[Bongo]

返信が遅くなりすみませんでした。試してみたところですと、当初は出来上がった塔の高さを測定し報酬に加味すればいいかと思ったのですが、もしかするとそれだけでは大して積み方が変わらないかもしれません。おそらく、高いけれど不安定な積み方は要求される位置・回転がシビアで、学習の過程で正しい向きを発見できる確率が低すぎるんじゃないかと感じました。

自由な回転を許した上で高さが最大の塔を作れるような報酬の与え方ができればいいのでしょうが、あいにくいまいちいい案が思いつかず、ちょっと不本意ですが慣性モーメントが小さい軸を垂直に持ってきた上で、そこから小さな角度以内に回転を制限してやることで高い塔の発見確率を上げてやることにしました。また、回転だけでは条件が厳しすぎるかと思い、落とす位置を少しずらすことも許すようにしてみました。なおこの変更に伴い、Behavior ParametersのContinuous Actionsを6に増やしました。

回答に変更したエージェントのコードと結果の図を追記しましたが、一応は縦に積めそうな位置・回転を見つけられたようです。ですが当初の回答で提示しました横置きの結果と比べるとだいぶ不安定になってしまいました。
あれだけお膳立てをすればもっと成功率が上がるんじゃないかと期待してはいたのですが、100万ステップぐらいで成績は上げ止まり、現状ではあれ以上成功率は改善しないような気がします。もっとゲームのルールを甘くしてやるべきかもしれませんね。未確認ですが、上からポトッと落とすのではなく、静かに乗っけてやることを許せば成功率が上がるかもしれません。

[545498_ho]

お忙しい中、丁寧にご回答いただき誠にありがとうございます。
また疑問に思う部分が出てきたら宜しくお願い致します。