How to use the tensorboard.plugins.custom_scalar.layout_pb2.Category function in tensorboard

def layout_dashboard(writer):
    from tensorboard import summary
    from tensorboard.plugins.custom_scalar import layout_pb2
    
    # This action does not have to be performed at every step, so the action is not
    # taken care of by an op in the graph. We only need to specify the layout once. 
    # We only need to specify the layout once (instead of per step).
    layout_summary = summary.custom_scalar_pb(layout_pb2.Layout(
        category=[
            layout_pb2.Category(
            title='losses',
            chart=[
                layout_pb2.Chart(
                    title='losses',
                    multiline=layout_pb2.MultilineChartContent(
                    tag=[r'loss.*'],
                )),
                layout_pb2.Chart(
                    title='baz',
                    margin=layout_pb2.MarginChartContent(
                    series=[
                        layout_pb2.MarginChartContent.Series(
                        value='loss/baz/scalar_summary',
                        lower='baz_lower/baz/scalar_summary',
                        upper='baz_upper/baz/scalar_summary'),
                    ],

title='losses',
                    multiline=layout_pb2.MultilineChartContent(
                    tag=[r'loss.*'],
                )),
                layout_pb2.Chart(
                    title='baz',
                    margin=layout_pb2.MarginChartContent(
                    series=[
                        layout_pb2.MarginChartContent.Series(
                        value='loss/baz/scalar_summary',
                        lower='baz_lower/baz/scalar_summary',
                        upper='baz_upper/baz/scalar_summary'),
                    ],
                )), 
            ]),
            layout_pb2.Category(
            title='trig functions',
            chart=[
                layout_pb2.Chart(
                    title='wave trig functions',
                    multiline=layout_pb2.MultilineChartContent(
                    tag=[r'trigFunctions/cosine', r'trigFunctions/sine'],
                )),
                # The range of tangent is different. Let's give it its own chart.
                layout_pb2.Chart(
                    title='tan',
                    multiline=layout_pb2.MultilineChartContent(
                    tag=[r'trigFunctions/tangent'],
                )),
            ],
        # This category we care less about. Let's make it initially closed.
        closed=True),

),
                            layout_pb2.Chart(
                                title="baz",
                                margin=layout_pb2.MarginChartContent(
                                    series=[
                                        layout_pb2.MarginChartContent.Series(
                                            value="loss/baz/scalar_summary",
                                            lower="loss/baz_lower/scalar_summary",
                                            upper="loss/baz_upper/scalar_summary",
                                        ),
                                    ],
                                ),
                            ),
                        ],
                    ),
                    layout_pb2.Category(
                        title="trig functions",
                        chart=[
                            layout_pb2.Chart(
                                title="wave trig functions",
                                multiline=layout_pb2.MultilineChartContent(
                                    tag=[
                                        r"trigFunctions/cosine",
                                        r"trigFunctions/sine",
                                    ],
                                ),
                            ),
                            # The range of tangent is different. Give it its own chart.
                            layout_pb2.Chart(
                                title="tan",
                                multiline=layout_pb2.MultilineChartContent(
                                    tag=[r"trigFunctions/tangent"],

[
                ("Policy Loss", [r"policy_loss"]),
                ("Value Loss", [r"value_loss"]),
                ("Policy Entropy", [r"policy_entropy"]),
                ("Explained Variance", [r"explained_variance"]),
                ("Approx KL", [r"approxkl"]),
                ("Clip Fraction", [r"clipfrac"]),
            ]
        ),
    )

    # Intentionally unused:
    # + serial_timesteps (just total_timesteps / num_envs)
    # + time_elapsed (TensorBoard already logs wall-clock time)
    # + nupdates (this is already logged as step)
    time = layout_pb2.Category(
        title="Time",
        chart=gen_multiline_charts([("Total Timesteps", [r"total_timesteps"]), ("FPS", [r"fps"])]),
    )

    categories = [episode_rewards, game_outcome, training, time]
    return summary_lib.custom_scalar_pb(layout_pb2.Layout(category=categories))

def layout_dashboard(writer):
    from tensorboard import summary
    from tensorboard.plugins.custom_scalar import layout_pb2
    
    # This action does not have to be performed at every step, so the action is not
    # taken care of by an op in the graph. We only need to specify the layout once. 
    # We only need to specify the layout once (instead of per step).
    layout_summary = summary.custom_scalar_pb(layout_pb2.Layout(
        category=[
            layout_pb2.Category(
            title='losses',
            chart=[
                layout_pb2.Chart(
                    title='losses',
                    multiline=layout_pb2.MultilineChartContent(
                    tag=[r'loss.*'],
                )),
                layout_pb2.Chart(
                    title='baz',
                    margin=layout_pb2.MarginChartContent(
                    series=[
                        layout_pb2.MarginChartContent.Series(
                        value='loss/baz/scalar_summary',
                        lower='baz_lower/baz/scalar_summary',
                        upper='baz_upper/baz/scalar_summary'),
                    ],

categories = []
    for k, v in layout.items():
        charts = []
        for chart_name, chart_meatadata in v.items():
            tags = chart_meatadata[1]
            if chart_meatadata[0] == 'Margin':
                assert len(tags) == 3
                mgcc = layout_pb2.MarginChartContent(series=[layout_pb2.MarginChartContent.Series(value=tags[0],
                                                                                                  lower=tags[1],
                                                                                                  upper=tags[2])])
                chart = layout_pb2.Chart(title=chart_name, margin=mgcc)
            else:
                mlcc = layout_pb2.MultilineChartContent(tag=tags)
                chart = layout_pb2.Chart(title=chart_name, multiline=mlcc)
            charts.append(chart)
        categories.append(layout_pb2.Category(title=k, chart=charts))

    layout = layout_pb2.Layout(category=categories)
    plugin_data = SummaryMetadata.PluginData(plugin_name='custom_scalars')
    smd = SummaryMetadata(plugin_data=plugin_data)
    tensor = TensorProto(dtype='DT_STRING',
                         string_val=[layout.SerializeToString()],
                         tensor_shape=TensorShapeProto())
    return Summary(value=[Summary.Value(tag='custom_scalars__config__', tensor=tensor, metadata=smd)])

with tf.name_scope("trigFunctions"):
        summary_lib.scalar("cosine", tf.cos(step))
        summary_lib.scalar("sine", tf.sin(step))
        summary_lib.scalar("tangent", tf.tan(step))

    merged_summary = tf.compat.v1.summary.merge_all()

    with tf.compat.v1.Session() as sess, tf.summary.FileWriter(
        LOGDIR
    ) as writer:
        # We only need to specify the layout once (instead of per step).
        layout_summary = summary_lib.custom_scalar_pb(
            layout_pb2.Layout(
                category=[
                    layout_pb2.Category(
                        title="losses",
                        chart=[
                            layout_pb2.Chart(
                                title="losses",
                                multiline=layout_pb2.MultilineChartContent(
                                    tag=[r"loss(?!.*margin.*)"],
                                ),
                            ),
                            layout_pb2.Chart(
                                title="baz",
                                margin=layout_pb2.MarginChartContent(
                                    series=[
                                        layout_pb2.MarginChartContent.Series(
                                            value="loss/baz/scalar_summary",
                                            lower="loss/baz_lower/scalar_summary",
                                            upper="loss/baz_upper/scalar_summary",