tensorflowでValueError: Dimensions must be equal, but are 4 and 2 for 'loss/output_1_loss/loss

前提・実現したいこと

前回の質問同様，オープンソースのニューラルネットワークのコード（Variational Integrator Networks）を用いて，二重振り子の学習をしたいと考えています．
論文で行われているideal single pendulumの学習の再現は実行済みです．

発生している問題・エラーメッセージ

二重振り子のシステムをdoublependulumのシステム名で追記し，以下のコード

python run_exp.py experiments VIN_VV doublependulum noisy 1 60 1

を実行したところ，以下のエラーが出てしまいました．

ValueError: Dimensions must be equal, but are 4 and 2 for 'loss/output_1_loss/loss_output_1_loss_Normal/log_prob/truediv' (op: 'RealDiv') with input shapes: [?,21,4], [?,?,2].
  In call to configurable 'main' (<function main at 0x182f01488>)

該当のソースコード

Python
1run_exp.pyのmainの内容は以下の通りです．文字数の関係で一部省略しています．
2
3
4import matplotlib
5matplotlib.use('TKAgg')
6
7import sys
8import os
9
10import tensorflow as tf
11import tensorflow.keras as tfk
12import numpy as np
13
14from tensorflow_probability import distributions as tfd
15
16import gin
17import build_utils
18import plot_utils
19
20import matplotlib.pyplot as plt
21import seaborn as sns
22
23@gin.configurable
24def main(exp_dir, model_name, system_name, observations, num_train_traj, num_train_steps, seed,
25        train_flag=True, eval_flag=True, save_every=100):
26
27    exp_dir = os.path.join(
28        root_dir, system_name,
29        observations + '-{}'.format(num_train_traj*num_train_steps),
30        model_name, str(seed))
31
32    checkpoint_dir = os.path.join(exp_dir, 'training/')
33    eval_dir = os.path.join(exp_dir, 'eval/')
34
35    if not os.path.isdir(eval_dir):
36        os.makedirs(eval_dir)
37
38    model = build_utils.create_model(observations, model_name)
39    system = build_utils.create_system(system_name, observations, seed)
40
41    if train_flag:
42
43        checkpoint_path = checkpoint_dir + 'cp-{epoch:04d}.ckpt'
44        cp_callback = tfk.callbacks.ModelCheckpoint(
45            filepath=checkpoint_path, 
46            verbose=1, 
47            save_weights_only=True,
48            save_freq=save_every)
49
50        model.save_weights(checkpoint_path.format(epoch=0))
51
52        train_dataset, valid_dataset, callbacks, num_epochs, obs, states =\
53            build_utils.create_dataset(system, model.horizon, num_train_traj, num_train_steps)
54
55
56        np.savez(eval_dir+'/obs_states.npz', obs=obs, states=states)
57
58        if callbacks is not None:
59            callbacks += [cp_callback]
60        else:
61            callbacks = [cp_callback]
62
63        model.fit(train_dataset, validation_data=valid_dataset,
64            epochs=num_epochs, callbacks=callbacks + [cp_callback])
65
66    if eval_flag:
67
68        latest = tf.train.latest_checkpoint(checkpoint_dir)
69        model.load_weights(latest)
70        eval_model(eval_dir, model, system, model_name)
71
72

Python
1また，doublependulumは以下のように記述しました．
2
3class DoublePendulum(System):
4
5    def _init_system(self, mass1=1, mass2=1, length1=1, length2=1, g=9.82, friction=0., **kwargs):
6
7        def doublependulum_ODE(t, y):
8            t1, t2, w1, w2 = y
9            a1 = (length2 / length1) * (mass2 / (mass1 + mass2)) * np.cos(t1 - t2)
10            a2 = (length1 / length2) * np.cos(t1 - t2)
11            f1 = -(length2 / length1) * (mass2 / (mass1 + mass2)) * (w2**2) * np.sin(t1 - t2) - \
12                (g / length1) * np.sin(t1)
13            f2 = (length1 / length2) * (w1**2) * np.sin(t1 - t2) - (g / length2) * np.sin(t2)
14            g1 = (f1 - a1 * f2) / (1 - a1 * a2)
15            g2 = (f2 - a2 * f1) / (1 - a1 * a2)
16
17            return np.stack([w1, w2, g1, g2])
18
19        self._ODE = doublependulum_ODE
20
21    def _random_init_state(self):
22
23        sign1 = [-1., 1.]
24        sign2 = [-1., 1.]
25        q10 = np.pi - self.RNG.uniform(low=0.1, high=1./2.*np.pi)
26        q10 = self.RNG.choice(sign1)*q10
27        q1dot0 = self.RNG.uniform(low=-1., high=1.)
28        q20 = np.pi - self.RNG.uniform(low=0.1, high=1./2.*np.pi)
29        q20 = self.RNG.choice(sign2)*q20
30        q2dot0 = self.RNG.uniform(low=-1., high=1.)
31        y0 = np.hstack([q10, q1dot0,q20,q2dot0])
32        return y0
33
34    def compute_energy(self, qqdot):
35
36        mass1 = self.system_param["mass"]
37        mass2 = self.system_param["mass"]
38        length1 = self.system_param["length"]
39        length2 = self.system_param["length"]
40        g = self.system_param["g"]
41
42        q, qdot = np.split(qqdot, 2, axis=-1)
43
44        t1, t2 = q     # theta 1 and theta 2
45        w1, w2 = qdot # omega 1 and omega 2
46
47        # kinetic energy (T)
48        K1 = 0.5 * mass1 * (length1 * w1)**2
49        K2 = 0.5 * mass2 * ((length1 * w1)**2 + (length2 * w2)**2 + 2 * length1 * length2 * w1 * w2 * np.cos(t1 - t2))
50        K = K1 + K2
51        
52        # potential energy (V)
53        y1 = -length1 * np.cos(t1)
54        y2 = y1 - length2 * np.cos(t2)
55        U = mass1 * g * y1 + mass2 * g * y2
56
57        E = K + U
58
59        return E

試したこと

singlependulumと比べて，doublependulumでは振り子の数が２つに増えるため，変数の数が２倍になることでエラーが出ているのだと思うのですが，どこを直せばいいのかが見つけられず途方に暮れています．些細なことでもいいのでアドバイスいただけると助かります．

Tracebackは以下の通りです．こちらも文字数の関係で一部省略しています．

Traceback (most recent call last):
  File "run_exp.py", line 244, in <module>
    main(root_dir, model_name, system_name, observations, num_train_traj, num_train_steps, seed)
   File "/Users/ootakemomoko/.pyenv/versions/3.6.0/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py", line 3322, in _create_op_internal
    op_def=op_def)
  File "/Users/ootakemomoko/.pyenv/versions/3.6.0/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py", line 1786, in __init__
    control_input_ops)
  File "/Users/ootakemomoko/.pyenv/versions/3.6.0/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py", line 1622, in _create_c_op
    raise ValueError(str(e))

補足情報（FW/ツールのバージョンなど）

0S： MacOS Catalina ver10.15.7
仮想環境名：VIN-Python3.6

pip==21.0.1
Pytho==3.6.0
jax==0.1.60
numpy==1.19.5
matplotlib==3.2.0
tensorflow==2.1.0
tensorflow-probability==0.9.0
tensorflow-estimator==2.1.0

行動規範の内容に同意します

回答1件

自己解決

質問内容で明記していたdoublependulum_ODEを使って初期値を作成する関数integrate_ODEにおいて，解が(t1,　t2,　w1,　w2)であるにもかかわらず，sol = np.empty((len(t),2)としていたことが原因でした．下記の通りsol = ((lent(t),4)とすることで解決しました．

Python
1def integrate_ODE(self, num_steps, step_size, y0=None, rtol=1e-12):
2
3        T = num_steps * step_size
4        t = np.linspace(0.0, T, num_steps)
5
6        if y0 is None:
7            y0 = self._random_init_state()
8
9        solver = ode(self._ODE).set_integrator('dop853', rtol=rtol)
10        sol = np.empty((len(t), 4))
11        sol[0] = y0
12        solver.set_initial_value(y0)
13        k = 1
14        while solver.successful() and solver.t < T:
15            solver.integrate(t[k])
16            sol[k] = solver.y
17            k += 1
18        
19        return sol