Kerasを用いたDQN実装における、自作誤差関数のcompile時に発生するエラー

前提・実現したいこと

kerasを用いてDQNを実装しようとしています。その際にkerasでは扱っていない誤差関数であるhuber_lossを自作関数として定義して、modelにcompileしようとしています。

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

ValueError: Shapes () and (?, ?) are incompatible

該当のソースコード

python
1    def init_model(self):
2
3        # tensorflowにあるhuber_lossをラップ
4        def huber_loss_tensorflow_wrap(y_true, y_pred):
5            return tf.losses.huber_loss(labels=y_true, predictions=y_pred)
6
7        state_input = Input(shape=(4, 28, 28), name='state')
8        action_input = Input(shape=[None], name='action', dtype='int32')
9
10        x = Convolution2D(16, (4, 4), activation='relu', strides=(2, 2), padding='same')(state_input)
11        x = Convolution2D(32, (2, 2), activation='relu', strides=(1, 1), padding='same')(x)
12        x = Convolution2D(32, (2, 2), activation='relu', strides=(1, 1), padding='same')(x)
13        x = Flatten()(x)
14        x = Dense(128, activation='relu')(x)
15
16        y_pred = Dense(self.n_actions, activation='linear', name='main_input')(x)
17        y_true = Input(shape=(self.n_actions, ), name='y_true')
18
19        loss_out = Lambda(lambda y_s: huber_loss_tensorflow_wrap(y_s[0], y_s[1]), output_shape=(1, ), name='loss')([y_true, y_pred])
20
21        self.model = Model(inputs=[state_input, action_input, y_true], outputs=[loss_out, y_pred])
22
23        opt = RMSprop(lr=self.learning_rate)
24
25        # 以下compile時にエラー
26        self.model.compile(loss=huber_loss_tensorflow_wrap, optimizer=opt)
27
28        self.target_model = copy.copy(self.model)

試したこと

エラーをたどると
lib\site-packages\tensorflow\python\ops\losses\losses_impl.py　にて
predictions.get_shape().assert_is_compatible_with(labels.get_shape())箇所でエラーが出ます。
定義されているget_shape()によって返ってくる値がpredicitionsとlabelsで異なるのが原因かと思いますが、いかがでしょうか。