Efficient Real-World RL for Legged Locomotion via Adaptive Policy Regularization
from functools import partial

import jax
import jax.numpy as jnp
import numpy as np
from flax import struct
from flax.training.train_state import TrainState

from jaxrl5.types import PRNGKey


@partial(jax.jit, static_argnames="apply_fn")
def _sample_actions(rng, apply_fn, params, observations: np.ndarray, **kwargs) -> np.ndarray:
    key, rng = jax.random.split(rng)
    dist = apply_fn({"params": params}, observations, **kwargs)
    return dist.sample(seed=key), rng


@partial(jax.jit, static_argnames="apply_fn")
def _eval_actions(apply_fn, params, observations: np.ndarray) -> np.ndarray:
    dist = apply_fn({"params": params}, observations)
    return dist.mode()


class Agent(struct.PyTreeNode):
    actor: TrainState
    rng: PRNGKey

    def eval_actions(self, observations: np.ndarray, **kwargs) -> np.ndarray:
        actions = _eval_actions(self.actor.apply_fn, self.actor.params, observations, **kwargs)
        return np.asarray(actions)

    def sample_actions(self, observations: np.ndarray, **kwargs) -> np.ndarray:
        actions, new_rng = _sample_actions(
            self.rng, self.actor.apply_fn, self.actor.params, observations, **kwargs
        )
        return np.asarray(actions), self.replace(rng=new_rng)
jaxrl5/agents/agent.py
realquantumcookie/APRL
"""Implementations of algorithms for continuous control."""

from functools import partial
from typing import Dict, Optional, Sequence, Tuple

import flax
import gym
import jax
import jax.numpy as jnp
import numpy as np
import optax
from flax import struct
from flax.training.train_state import TrainState

from jaxrl5.agents.agent import Agent
from jaxrl5.agents.sac.temperature import Temperature
from jaxrl5.data.dataset import DatasetDict
from jaxrl5.distributions import TanhNormal
from jaxrl5.networks import (
    MLP,
    Ensemble,
    MLPResNetV2,
    StateActionValue,
    subsample_ensemble,
)
from flax.core.frozen_dict import unfreeze

def convert_to_numpy_array(param):
    if isinstance(param, list):
        return jax.tree_map(convert_to_numpy_array, param)
    else:
        return jnp.array(param)

def tree_multimap(func, tree1, tree2):
    """Apply a function element-wise to two trees."""
    return jax.tree_map(lambda x, y: func(x, y), tree1, tree2)

def compute_critic_param_change_norm(before_params, after_params):
    param_squares = tree_multimap(lambda p1, p2: jnp.sum((p2 - p1) ** 2), before_params, after_params)
    param_change_norm = jnp.sqrt(jnp.sum(jnp.array(jax.tree_util.tree_leaves(param_squares))))
    return param_change_norm

def compute_gradient_norm(grads):
    grad_squares = jax.tree_map(lambda g: jnp.sum(g ** 2), grads)
    grad_norm = jnp.sqrt(jnp.sum(jnp.array(jax.tree_util.tree_leaves(grad_squares))))
    return grad_norm

# From https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/text_classification_flax.ipynb#scrollTo=ap-zaOyKJDXM
def decay_mask_fn(params):
    flat_params = flax.traverse_util.flatten_dict(params)
    flat_mask = {path: path[-1] != "bias" for path in flat_params}
    return flax.core.FrozenDict(flax.traverse_util.unflatten_dict(flat_mask))


class SACLearner(Agent):
    critic: TrainState
    target_critic: TrainState
    temp: TrainState
    tau: float
    discount: float
    target_entropy: float
    num_qs: int = struct.field(pytree_node=False)
    num_min_qs: Optional[int] = struct.field(
        pytree_node=False
    )  # See M in RedQ https://arxiv.org/abs/2101.05982
    backup_entropy: bool = struct.field(pytree_node=False)
    interior_linear_c: float = struct.field(pytree_node=False)
    interior_quadratic_c: float = struct.field(pytree_node=False)
    exterior_linear_c: float = struct.field(pytree_node=False)
    exterior_quadratic_c: float = struct.field(pytree_node=False)
    max_gradient_norm: float = struct.field(pytree_node=False)

    @classmethod
    def create(
        cls,
        seed: int,
        observation_space: gym.Space,
        action_space: gym.Space,
        actor_lr: float = 3e-4,
        critic_lr: float = 3e-4,
        temp_lr: float = 3e-4,
        hidden_dims: Sequence[int] = (256, 256),
        discount: float = 0.99,
        tau: float = 0.005,
        num_qs: int = 2,
        num_min_qs: Optional[int] = None,
        critic_dropout_rate: Optional[float] = None,
        critic_weight_decay: Optional[float] = None,
        max_gradient_norm: Optional[float] = None,
        critic_layer_norm: bool = False,
        target_entropy: Optional[float] = None,
        init_temperature: float = 1.0,
        backup_entropy: bool = True,
        use_pnorm: bool = False,
        use_critic_resnet: bool = False,
        interior_linear_c: float = 0.0,
        interior_quadratic_c: float = 0.0,
        exterior_linear_c: float = 0.0,
        exterior_quadratic_c: float = 0.0,
    ):
        """
        An implementation of the version of Soft-Actor-Critic described in https://arxiv.org/abs/1812.05905
        """
        action_dim = action_space.shape[-1]
        observations = observation_space.sample()
        actions = action_space.sample()

        if target_entropy is None:
            target_entropy = -action_dim / 2

        rng = jax.random.PRNGKey(seed)
        rng, actor_key, critic_key, temp_key = jax.random.split(rng, 4)

        actor_base_cls = partial(
            MLP, hidden_dims=hidden_dims, activate_final=True, use_pnorm=use_pnorm
        )
        actor_def = TanhNormal(actor_base_cls, action_dim)
        actor_params = actor_def.init(actor_key, observations)["params"]
        actor = TrainState.create(
            apply_fn=actor_def.apply,
            params=actor_params,
            tx=optax.adam(learning_rate=actor_lr),
        )

        if use_critic_resnet:
            critic_base_cls = partial(
                MLPResNetV2,
                num_blocks=1,
            )
        else:
            critic_base_cls = partial(
                MLP,
                hidden_dims=hidden_dims,
                activate_final=True,
                dropout_rate=critic_dropout_rate,
                use_layer_norm=critic_layer_norm,
                use_pnorm=use_pnorm,
            )
        critic_cls = partial(StateActionValue, base_cls=critic_base_cls)
        critic_def = Ensemble(critic_cls, num=num_qs)
        critic_params = critic_def.init(critic_key, observations, actions)["params"]
        if critic_weight_decay is not None:
            if max_gradient_norm is not None:
                tx = optax.chain(
                    optax.clip_by_global_norm(max_gradient_norm),
                    optax.adamw(
                        learning_rate=critic_lr,
                        weight_decay=critic_weight_decay,
                        mask=decay_mask_fn,
                    )
                )
            else:
                tx = optax.adamw(
                    learning_rate=critic_lr,
                    weight_decay=critic_weight_decay,
                    mask=decay_mask_fn,
                )
        else:
            if max_gradient_norm is not None:
                tx = optax.chain(
                        optax.clip_by_global_norm(max_gradient_norm),
                        optax.adam(learning_rate=critic_lr)
                )
            else:
                tx = optax.adam(learning_rate=critic_lr)

        critic = TrainState.create(
            apply_fn=critic_def.apply,
            params=critic_params,
            tx=tx,
        )
        target_critic_def = Ensemble(critic_cls, num=num_min_qs or num_qs)
        target_critic = TrainState.create(
            apply_fn=target_critic_def.apply,
            params=critic_params,
            tx=optax.GradientTransformation(lambda _: None, lambda _: None),
        )

        temp_def = Temperature(init_temperature)
        temp_params = temp_def.init(temp_key)["params"]
        temp = TrainState.create(
            apply_fn=temp_def.apply,
            params=temp_params,
            tx=optax.adam(learning_rate=temp_lr),
        )

        return cls(
            rng=rng,
            actor=actor,
            critic=critic,
            target_critic=target_critic,
            temp=temp,
            target_entropy=target_entropy,
            tau=tau,
            discount=discount,
            num_qs=num_qs,
            num_min_qs=num_min_qs,
            backup_entropy=backup_entropy,
            interior_linear_c=interior_linear_c,
            interior_quadratic_c=interior_quadratic_c,
            exterior_linear_c=exterior_linear_c,
            exterior_quadratic_c=exterior_quadratic_c,
            max_gradient_norm=max_gradient_norm,
        )

    def reset_actor(self,
                    seed: int,
                    observation_space: gym.Space,
                    action_space: gym.Space,
                    actor_lr: float = 3e-4,
                    critic_lr: float = 3e-4,
                    temp_lr: float = 3e-4,
                    hidden_dims: Sequence[int] = (256, 256),
                    discount: float = 0.99,
                    tau: float = 0.005,
                    num_qs: int = 2,
                    num_min_qs: Optional[int] = None,
                    critic_dropout_rate: Optional[float] = None,
                    critic_weight_decay: Optional[float] = None,
                    critic_layer_norm: bool = False,
                    target_entropy: Optional[float] = None,
                    init_temperature: float = 1.0,
                    backup_entropy: bool = True,
                    use_pnorm: bool = False,
                    use_critic_resnet: bool = False,
                    interior_linear_c: float = 0.0,
                    interior_quadratic_c: float = 0.0,
                    exterior_linear_c: float = 0.0,
                    exterior_quadratic_c: float = 0.0,):
        action_dim = action_space.shape[-1]
        key, rng = jax.random.split(self.rng)

        # Get a fresh set of random actor parameters
        actor_base_cls = partial(
            MLP, hidden_dims=hidden_dims, activate_final=True, use_pnorm=use_pnorm
        )
        actor_def = TanhNormal(actor_base_cls, action_dim)
        actor_params = actor_def.init(key, self.batch["observations"])["params"]
        actor = TrainState.create(
            apply_fn=actor_def.apply,
            params=actor_params,
            tx=optax.adam(learning_rate=actor_lr),
        )

        # Replace the actor with the new random parameters
        return self.replace(actor=actor, rng=rng)
    
    def reset_critic(self,
                    seed: int,
                    observation_space: gym.Space,
                    action_space: gym.Space,
                    actor_lr: float = 3e-4,
                    critic_lr: float = 3e-4,
                    temp_lr: float = 3e-4,
                    hidden_dims: Sequence[int] = (256, 256),
                    discount: float = 0.99,
                    tau: float = 0.005,
                    num_qs: int = 2,
                    num_min_qs: Optional[int] = None,
                    critic_dropout_rate: Optional[float] = None,
                    critic_weight_decay: Optional[float] = None,
                    critic_layer_norm: bool = False,
                    target_entropy: Optional[float] = None,
                    init_temperature: float = 1.0,
                    backup_entropy: bool = True,
                    use_pnorm: bool = False,
                    use_critic_resnet: bool = False,
                    interior_linear_c: float = 0.0,
                    interior_quadratic_c: float = 0.0,
                    exterior_linear_c: float = 0.0,
                    exterior_quadratic_c: float = 0.0,
                    max_gradient_norm: Optional[float] = None):
        key, rng = jax.random.split(self.rng)

        if use_critic_resnet:
            critic_base_cls = partial(
                MLPResNetV2,
                num_blocks=1,
            )
        else:
            critic_base_cls = partial(
                MLP,
                hidden_dims=hidden_dims,
                activate_final=True,
                dropout_rate=critic_dropout_rate,
                use_layer_norm=critic_layer_norm,
                use_pnorm=use_pnorm,
            )

        critic_cls = partial(StateActionValue, base_cls=critic_base_cls)
        critic_def = Ensemble(critic_cls, num=self.num_qs)
        critic_params = critic_def.init(key, self.batch["observations"], self.batch["actions"])["params"]
        
        if critic_weight_decay is not None:
            if max_gradient_norm is not None:
                tx = optax.chain(
                    optax.clip_by_global_norm(max_gradient_norm),
                    optax.adamw(
                        learning_rate=critic_lr,
                        weight_decay=critic_weight_decay,
                        mask=decay_mask_fn,
                    )
                )
            else:
                tx = optax.adamw(
                    learning_rate=critic_lr,
                    weight_decay=critic_weight_decay,
                    mask=decay_mask_fn,
                )
        else:
            if max_gradient_norm is not None:
                tx = optax.chain(
                        optax.clip_by_global_norm(max_gradient_norm),
                        optax.adam(learning_rate=critic_lr)
                )
            else:
                tx = optax.adam(learning_rate=critic_lr)

        critic = TrainState.create(
            apply_fn=critic_def.apply,
            params=critic_params,
            tx=tx,
        )

        target_critic_def = Ensemble(critic_cls, num=self.num_min_qs or self.num_qs)
        target_critic = TrainState.create(
            apply_fn=target_critic_def.apply,
            params=critic_params,
            tx=optax.GradientTransformation(lambda _: None, lambda _: None),
        )

        # Replace the critic and target critic with the new random parameters
        return self.replace(critic=critic, target_critic=target_critic, rng=rng)

    def update_actor(self, batch: DatasetDict, output_range: Optional[Tuple[jnp.ndarray, jnp.ndarray]]) -> Tuple[Agent, Dict[str, float]]:
        key, rng = jax.random.split(self.rng)
        key2, rng = jax.random.split(rng)

        def actor_loss_fn(actor_params) -> Tuple[jnp.ndarray, Dict[str, float]]:
            dist = self.actor.apply_fn({"params": actor_params}, batch["observations"])
            actions = dist.sample(seed=key)
            log_probs = dist.log_prob(actions)
            qs = self.critic.apply_fn(
                {"params": self.critic.params},
                batch["observations"],
                actions,
                True,
                rngs={"dropout": key2},
            )  # training=True
            q = qs.mean(axis=0)
            
            # create a mask if the action is out of range (0 if out of range, 1 if in range)
            mask = 1 - jnp.logical_or(jnp.any(actions < output_range[0], axis=-1), jnp.any(actions > output_range[1], axis=-1))
            
            exterior_actions = jnp.where(mask[:, None], jnp.zeros_like(actions), actions)
            interior_actions = jnp.where(mask[:, None], actions, jnp.zeros_like(actions))
            
            exterior_l1_penalty = jnp.sum(jnp.abs(exterior_actions), axis=-1)
            exterior_l2_penalty = jnp.sum(exterior_actions ** 2, axis=-1)
            exterior_l2_penalty = jnp.sqrt(exterior_l2_penalty)
            
            interior_linear_penalty = jnp.sum(jnp.abs(interior_actions), axis=-1)
            interior_quadratic_penalty = jnp.sum(interior_actions ** 4, axis=-1)
            
            penalty_function = self.interior_quadratic_c * (self.exterior_linear_c / (3 * output_range[1] ** 3)) * interior_quadratic_penalty + self.exterior_linear_c * exterior_l1_penalty + self.exterior_quadratic_c * exterior_l2_penalty
            # penalty_function = self.ctrl_weight * oob_penalty + (self.ctrl_weight / (2 * output_range[1])) * ib_penalty
            
            actor_loss = (
                log_probs * self.temp.apply_fn({"params": self.temp.params}) - q + penalty_function
            ).mean()
            
            # compute how many actions are out of range
            oor_actions = jnp.mean(1 - mask)
            info_dict = {"actor_loss": actor_loss, 
                         "entropy": -log_probs.mean(), 
                         "oor_actions": oor_actions,
                         "penalty_function": penalty_function.mean()}
            
            return actor_loss, info_dict

        grads, actor_info = jax.grad(actor_loss_fn, has_aux=True)(self.actor.params)
        actor = self.actor.apply_gradients(grads=grads)

        return self.replace(actor=actor, rng=rng), actor_info

    def update_temperature(self, entropy: float) -> Tuple[Agent, Dict[str, float]]:
        def temperature_loss_fn(temp_params):
            temperature = self.temp.apply_fn({"params": temp_params})
            temp_loss = temperature * (entropy - self.target_entropy).mean()
            return temp_loss, {
                "temperature": temperature,
                "temperature_loss": temp_loss,
            }

        grads, temp_info = jax.grad(temperature_loss_fn, has_aux=True)(self.temp.params)
        temp = self.temp.apply_gradients(grads=grads)

        return self.replace(temp=temp), temp_info

    def update_critic(self, batch: DatasetDict, max_grad_norm: float = np.inf) -> Tuple[TrainState, Dict[str, float]]:

        dist = self.actor.apply_fn(
            {"params": self.actor.params}, batch["next_observations"]
        )

        rng = self.rng

        key, rng = jax.random.split(rng)
        next_actions = dist.sample(seed=key)

        # Used only for REDQ.
        key, rng = jax.random.split(rng)
        target_params = subsample_ensemble(
            key, self.target_critic.params, self.num_min_qs, self.num_qs
        )

        key, rng = jax.random.split(rng)
        next_qs = self.target_critic.apply_fn(
            {"params": target_params},
            batch["next_observations"],
            next_actions,
            True,
            rngs={"dropout": key},
        )  # training=True
        next_q = next_qs.min(axis=0)

        target_q = batch["rewards"] + self.discount * batch["masks"] * next_q

        if self.backup_entropy:
            next_log_probs = dist.log_prob(next_actions)
            target_q -= (
                self.discount
                * batch["masks"]
                * self.temp.apply_fn({"params": self.temp.params})
                * next_log_probs
            )

        key, rng = jax.random.split(rng)

        def critic_loss_fn(critic_params) -> Tuple[jnp.ndarray, Dict[str, float]]:
            qs = self.critic.apply_fn(
                {"params": critic_params},
                batch["observations"],
                batch["actions"],
                True,
                rngs={"dropout": key},
            )  # training=True
            critic_loss = ((qs - target_q) ** 2).mean()
            return critic_loss, {"critic_loss": critic_loss, "q": qs.mean()}

        grads, info = jax.grad(critic_loss_fn, has_aux=True)(self.critic.params)

        # Compute the gradient magnitudes
        critic_grad_magnitudes = compute_gradient_norm(grads) 

        critic_params_before = self.critic.params
        critic = self.critic.apply_gradients(grads=grads)
        critic_params_after = critic.params

        critic_params_change = compute_critic_param_change_norm(critic_params_before, critic_params_after)
        
        target_critic_params = optax.incremental_update(
            critic.params, self.target_critic.params, self.tau
        )
        target_critic = self.target_critic.replace(params=target_critic_params)
        
        info["critic_grad_magnitudes"] = critic_grad_magnitudes
        info["critic_weight_change"] = critic_params_change

        return self.replace(critic=critic, target_critic=target_critic, rng=rng), info
    

    def compute_td_error(self, batch: DatasetDict):

        dist = self.actor.apply_fn(
            {"params": self.actor.params}, batch["next_observations"]
        )

        rng = self.rng

        key, rng = jax.random.split(rng)
        next_actions = dist.sample(seed=key)

        # Used only for REDQ.
        key, rng = jax.random.split(rng)
        target_params = subsample_ensemble(
            key, self.target_critic.params, self.num_min_qs, self.num_qs
        )

        key, rng = jax.random.split(rng)
        next_qs = self.target_critic.apply_fn(
            {"params": target_params},
            batch["next_observations"],
            next_actions,
            True,
            rngs={"dropout": key},
        )  # training=True
        next_q = next_qs.min(axis=0)

        target_q = batch["rewards"] + self.discount * batch["masks"] * next_q

        if self.backup_entropy:
            next_log_probs = dist.log_prob(next_actions)
            target_q -= (
                self.discount
                * batch["masks"]
                * self.temp.apply_fn({"params": self.temp.params})
                * next_log_probs
            )

        key, rng = jax.random.split(rng)
        
        qs = self.critic.apply_fn(
            {"params": self.critic.params},
            batch["observations"],
            batch["actions"],
            True,
            rngs={"dropout": key},
        )  # training=True
        
        td_error = (qs - target_q) ** 2
        return td_error


    @partial(jax.jit, static_argnames=("utd_ratio", "actor_delay"))
    def update_with_delay(self, batch: DatasetDict, utd_ratio: int, output_range: Optional[Tuple[jnp.ndarray, jnp.ndarray]] = None, actor_delay: int = 1):
        new_agent = self
        for actor_i in range(utd_ratio//actor_delay):
            for critic_i in range(actor_delay):
                def slice(x):
                    assert x.shape[0] % utd_ratio == 0
                    batch_size = x.shape[0] // utd_ratio
                    return x[batch_size * (critic_i + actor_i * actor_delay) : batch_size * (critic_i + actor_i * actor_delay + 1)]

                mini_batch = jax.tree_util.tree_map(slice, batch)
                new_agent, critic_info = new_agent.update_critic(mini_batch)
            
            new_agent, actor_info = new_agent.update_actor(mini_batch, output_range=output_range)
            new_agent, temp_info = new_agent.update_temperature(actor_info["entropy"])
    
        return new_agent, {**actor_info, **critic_info, **temp_info}
    
    @partial(jax.jit, static_argnames="utd_ratio")
    def update(self, batch: DatasetDict, utd_ratio: int, output_range: Optional[Tuple[jnp.ndarray, jnp.ndarray]] = None):
        new_agent = self
        for i in range(utd_ratio):

            def slice(x):
                assert x.shape[0] % utd_ratio == 0
                batch_size = x.shape[0] // utd_ratio
                return x[batch_size * i : batch_size * (i + 1)]

            mini_batch = jax.tree_util.tree_map(slice, batch)
            new_agent, critic_info = new_agent.update_critic(mini_batch)
        
        new_agent, actor_info = new_agent.update_actor(mini_batch, output_range=output_range)
        new_agent, temp_info = new_agent.update_temperature(actor_info["entropy"])
            
        return new_agent, {**actor_info, **critic_info, **temp_info}
    
    @jax.jit    
    def get_value(self, observation: np.ndarray) -> float:
        # Convert the observation to a JAX numpy array
        jax_observation = convert_to_numpy_array(observation)
        
        rng = self.rng
        key, rng = jax.random.split(rng)
        key2, rng = jax.random.split(rng)
        
        # Use the actor network to generate an action
        dist = self.actor.apply_fn({"params": self.actor.params}, jax_observation[None])
        action = dist.sample(seed=key)[0]  # Generate a single action
        
        # Use the critic network to estimate the state-action value
        value = self.critic.apply_fn(
            {"params": self.critic.params},
            jax_observation[None],  # Add a batch dimension to the observation
            action[None],           # Add a batch dimension to the action
            True,
            rngs={"dropout": key2} 
        )
        # average over the ensemble
        value = value.mean(axis=0)[0] # Remove the batch dimension
        return jnp.array(value, float)

    @jax.jit    
    def self_perception(self, batch: DatasetDict) -> float:
        rng = self.rng
        key, rng = jax.random.split(rng)
        key2, rng = jax.random.split(rng)
        
        # Use the actor network to generate an action
        dist = self.actor.apply_fn({"params": self.actor.params}, batch["observations"])
        actions = dist.sample(seed=key)
        
        # Use the critic network to estimate the state-action value
        value = self.critic.apply_fn(
            {"params": self.critic.params},
            batch["observations"],
            actions,
            True,
            rngs={"dropout": key2} 
        )
        # average over the ensemble
        value = value.mean(axis=0)
        return jnp.array(value, float)[0]
jaxrl5/agents/sac/sac_learner.py
realquantumcookie/APRL
"""Implementations of algorithms for continuous control."""

from functools import partial
from typing import Dict, Optional, Sequence, Tuple

import flax
import gym
import jax
import jax.numpy as jnp
import numpy as np
import optax
from flax import struct
from flax.training.train_state import TrainState

from jaxrl5.agents.agent import Agent
from jaxrl5.agents.sac.temperature import Temperature
from jaxrl5.data.dataset import DatasetDict
from jaxrl5.distributions import TanhNormal
from jaxrl5.networks import (
    MLP,
    Ensemble,
    MLPResNetV2,
    StateActionValue,
    StateActionNextState,
    subsample_ensemble,
)
from flax.core.frozen_dict import unfreeze

def convert_to_numpy_array(param):
    if isinstance(param, list):
        return jax.tree_map(convert_to_numpy_array, param)
    else:
        return jnp.array(param)

def tree_multimap(func, tree1, tree2):
    """Apply a function element-wise to two trees."""
    return jax.tree_map(lambda x, y: func(x, y), tree1, tree2)

def compute_critic_param_change_norm(before_params, after_params):
    param_squares = tree_multimap(lambda p1, p2: jnp.sum((p2 - p1) ** 2), before_params, after_params)
    param_change_norm = jnp.sqrt(jnp.sum(jnp.array(jax.tree_util.tree_leaves(param_squares))))
    return param_change_norm

def compute_gradient_norm(grads):
    grad_squares = jax.tree_map(lambda g: jnp.sum(g ** 2), grads)
    grad_norm = jnp.sqrt(jnp.sum(jnp.array(jax.tree_util.tree_leaves(grad_squares))))
    return grad_norm

# From https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/text_classification_flax.ipynb#scrollTo=ap-zaOyKJDXM
def decay_mask_fn(params):
    flat_params = flax.traverse_util.flatten_dict(params)
    flat_mask = {path: path[-1] != "bias" for path in flat_params}
    return flax.core.FrozenDict(flax.traverse_util.unflatten_dict(flat_mask))


class SACLearnerWithDynamics(Agent):
    critic: TrainState
    target_critic: TrainState
    temp: TrainState
    dynamics_model: TrainState
    tau: float
    discount: float
    target_entropy: float
    num_qs: int = struct.field(pytree_node=False)
    num_min_qs: Optional[int] = struct.field(
        pytree_node=False
    )  # See M in RedQ https://arxiv.org/abs/2101.05982
    backup_entropy: bool = struct.field(pytree_node=False)
    interior_linear_c: float = struct.field(pytree_node=False)
    interior_quadratic_c: float = struct.field(pytree_node=False)
    exterior_linear_c: float = struct.field(pytree_node=False)
    exterior_quadratic_c: float = struct.field(pytree_node=False)
    max_gradient_norm: float = struct.field(pytree_node=False)

    @classmethod
    def create(
        cls,
        seed: int,
        observation_space: gym.Space,
        action_space: gym.Space,
        actor_lr: float = 3e-4,
        critic_lr: float = 3e-4,
        temp_lr: float = 3e-4,
        hidden_dims: Sequence[int] = (256, 256),
        discount: float = 0.99,
        tau: float = 0.005,
        num_qs: int = 2,
        num_min_qs: Optional[int] = None,
        critic_dropout_rate: Optional[float] = None,
        critic_weight_decay: Optional[float] = None,
        max_gradient_norm: Optional[float] = None,
        critic_layer_norm: bool = False,
        target_entropy: Optional[float] = None,
        init_temperature: float = 1.0,
        backup_entropy: bool = True,
        use_pnorm: bool = False,
        use_critic_resnet: bool = False,
        interior_linear_c: float = 0.0,
        interior_quadratic_c: float = 0.0,
        exterior_linear_c: float = 0.0,
        exterior_quadratic_c: float = 0.0,
    ):
        """
        An implementation of the version of Soft-Actor-Critic described in https://arxiv.org/abs/1812.05905
        """
        action_dim = action_space.shape[-1]
        obs_dim = observation_space.shape[-1]
        observations = observation_space.sample()
        actions = action_space.sample()

        if target_entropy is None:
            target_entropy = -action_dim / 2

        rng = jax.random.PRNGKey(seed)
        rng, actor_key, critic_key, temp_key, dynamics_key = jax.random.split(rng, 5)

        actor_base_cls = partial(
            MLP, hidden_dims=hidden_dims, activate_final=True, use_pnorm=use_pnorm
        )
        actor_def = TanhNormal(actor_base_cls, action_dim)
        actor_params = actor_def.init(actor_key, observations)["params"]
        actor = TrainState.create(
            apply_fn=actor_def.apply,
            params=actor_params,
            tx=optax.adam(learning_rate=actor_lr),
        )

        if use_critic_resnet:
            critic_base_cls = partial(
                MLPResNetV2,
                num_blocks=1,
            )
        else:
            critic_base_cls = partial(
                MLP,
                hidden_dims=hidden_dims,
                activate_final=True,
                dropout_rate=critic_dropout_rate,
                use_layer_norm=critic_layer_norm,
                use_pnorm=use_pnorm,
            )
        critic_cls = partial(StateActionValue, base_cls=critic_base_cls)
        critic_def = Ensemble(critic_cls, num=num_qs)
        critic_params = critic_def.init(critic_key, observations, actions)["params"]
        if critic_weight_decay is not None:
            if max_gradient_norm is not None:
                tx = optax.chain(
                    optax.clip_by_global_norm(max_gradient_norm),
                    optax.adamw(
                        learning_rate=critic_lr,
                        weight_decay=critic_weight_decay,
                        mask=decay_mask_fn,
                    )
                )
            else:
                tx = optax.adamw(
                    learning_rate=critic_lr,
                    weight_decay=critic_weight_decay,
                    mask=decay_mask_fn,
                )
        else:
            if max_gradient_norm is not None:
                tx = optax.chain(
                        optax.clip_by_global_norm(max_gradient_norm),
                        optax.adam(learning_rate=critic_lr)
                )
            else:
                tx = optax.adam(learning_rate=critic_lr)

        critic = TrainState.create(
            apply_fn=critic_def.apply,
            params=critic_params,
            tx=tx,
        )
        target_critic_def = Ensemble(critic_cls, num=num_min_qs or num_qs)
        target_critic = TrainState.create(
            apply_fn=target_critic_def.apply,
            params=critic_params,
            tx=optax.GradientTransformation(lambda _: None, lambda _: None),
        )

        temp_def = Temperature(init_temperature)
        temp_params = temp_def.init(temp_key)["params"]
        temp = TrainState.create(
            apply_fn=temp_def.apply,
            params=temp_params,
            tx=optax.adam(learning_rate=temp_lr),
        )
        
        # Create the dynamics model
        dynamics_base_cls = partial(
            MLP,
            hidden_dims=hidden_dims,
            activate_final=True,
        )
        dynamics_def = StateActionNextState(dynamics_base_cls, obs_dim)
        dynamics_params = dynamics_def.init(
            dynamics_key, observations, actions
        )["params"]
        dynamics_model = TrainState.create(
            apply_fn=dynamics_def.apply, 
            params=dynamics_params, 
            tx=optax.adam(learning_rate=critic_lr)
        )

        return cls(
            rng=rng,
            actor=actor,
            critic=critic,
            target_critic=target_critic,
            temp=temp,
            target_entropy=target_entropy,
            tau=tau,
            discount=discount,
            num_qs=num_qs,
            num_min_qs=num_min_qs,
            backup_entropy=backup_entropy,
            interior_linear_c=interior_linear_c,
            interior_quadratic_c=interior_quadratic_c,
            exterior_linear_c=exterior_linear_c,
            exterior_quadratic_c=exterior_quadratic_c,
            max_gradient_norm=max_gradient_norm,
            dynamics_model=dynamics_model,
        )

    def reset_actor(self,
                    seed: int,
                    observation_space: gym.Space,
                    action_space: gym.Space,
                    actor_lr: float = 3e-4,
                    critic_lr: float = 3e-4,
                    temp_lr: float = 3e-4,
                    hidden_dims: Sequence[int] = (256, 256),
                    discount: float = 0.99,
                    tau: float = 0.005,
                    num_qs: int = 2,
                    num_min_qs: Optional[int] = None,
                    critic_dropout_rate: Optional[float] = None,
                    critic_weight_decay: Optional[float] = None,
                    critic_layer_norm: bool = False,
                    target_entropy: Optional[float] = None,
                    init_temperature: float = 1.0,
                    backup_entropy: bool = True,
                    use_pnorm: bool = False,
                    use_critic_resnet: bool = False,
                    interior_linear_c: float = 0.0,
                    interior_quadratic_c: float = 0.0,
                    exterior_linear_c: float = 0.0,
                    exterior_quadratic_c: float = 0.0,):
        action_dim = action_space.shape[-1]
        key, rng = jax.random.split(self.rng)

        # Get a fresh set of random actor parameters
        actor_base_cls = partial(
            MLP, hidden_dims=hidden_dims, activate_final=True, use_pnorm=use_pnorm
        )
        actor_def = TanhNormal(actor_base_cls, action_dim)
        actor_params = actor_def.init(key, self.batch["observations"])["params"]
        actor = TrainState.create(
            apply_fn=actor_def.apply,
            params=actor_params,
            tx=optax.adam(learning_rate=actor_lr),
        )

        # Replace the actor with the new random parameters
        return self.replace(actor=actor, rng=rng)
    
    def reset_critic(self,
                    seed: int,
                    observation_space: gym.Space,
                    action_space: gym.Space,
                    actor_lr: float = 3e-4,
                    critic_lr: float = 3e-4,
                    temp_lr: float = 3e-4,
                    hidden_dims: Sequence[int] = (256, 256),
                    discount: float = 0.99,
                    tau: float = 0.005,
                    num_qs: int = 2,
                    num_min_qs: Optional[int] = None,
                    critic_dropout_rate: Optional[float] = None,
                    critic_weight_decay: Optional[float] = None,
                    critic_layer_norm: bool = False,
                    target_entropy: Optional[float] = None,
                    init_temperature: float = 1.0,
                    backup_entropy: bool = True,
                    use_pnorm: bool = False,
                    use_critic_resnet: bool = False,
                    interior_linear_c: float = 0.0,
                    interior_quadratic_c: float = 0.0,
                    exterior_linear_c: float = 0.0,
                    exterior_quadratic_c: float = 0.0,
                    max_gradient_norm: Optional[float] = None):
        key, rng = jax.random.split(self.rng)

        if use_critic_resnet:
            critic_base_cls = partial(
                MLPResNetV2,
                num_blocks=1,
            )
        else:
            critic_base_cls = partial(
                MLP,
                hidden_dims=hidden_dims,
                activate_final=True,
                dropout_rate=critic_dropout_rate,
                use_layer_norm=critic_layer_norm,
                use_pnorm=use_pnorm,
            )

        critic_cls = partial(StateActionValue, base_cls=critic_base_cls)
        critic_def = Ensemble(critic_cls, num=self.num_qs)
        critic_params = critic_def.init(key, self.batch["observations"], self.batch["actions"])["params"]
        
        if critic_weight_decay is not None:
            if max_gradient_norm is not None:
                tx = optax.chain(
                    optax.clip_by_global_norm(max_gradient_norm),
                    optax.adamw(
                        learning_rate=critic_lr,
                        weight_decay=critic_weight_decay,
                        mask=decay_mask_fn,
                    )
                )
            else:
                tx = optax.adamw(
                    learning_rate=critic_lr,
                    weight_decay=critic_weight_decay,
                    mask=decay_mask_fn,
                )
        else:
            if max_gradient_norm is not None:
                tx = optax.chain(
                        optax.clip_by_global_norm(max_gradient_norm),
                        optax.adam(learning_rate=critic_lr)
                )
            else:
                tx = optax.adam(learning_rate=critic_lr)

        critic = TrainState.create(
            apply_fn=critic_def.apply,
            params=critic_params,
            tx=tx,
        )

        target_critic_def = Ensemble(critic_cls, num=self.num_min_qs or self.num_qs)
        target_critic = TrainState.create(
            apply_fn=target_critic_def.apply,
            params=critic_params,
            tx=optax.GradientTransformation(lambda _: None, lambda _: None),
        )

        # Replace the critic and target critic with the new random parameters
        return self.replace(critic=critic, target_critic=target_critic, rng=rng)

    def update_actor(self, batch: DatasetDict, output_range: Optional[Tuple[jnp.ndarray, jnp.ndarray]]) -> Tuple[Agent, Dict[str, float]]:
        key, rng = jax.random.split(self.rng)
        key2, rng = jax.random.split(rng)

        def actor_loss_fn(actor_params) -> Tuple[jnp.ndarray, Dict[str, float]]:
            dist = self.actor.apply_fn({"params": actor_params}, batch["observations"])
            actions = dist.sample(seed=key)
            log_probs = dist.log_prob(actions)
            qs = self.critic.apply_fn(
                {"params": self.critic.params},
                batch["observations"],
                actions,
                True,
                rngs={"dropout": key2},
            )  # training=True
            q = qs.mean(axis=0)
            
            # create a mask if the action is out of range (0 if out of range, 1 if in range)
            mask = 1 - jnp.logical_or(jnp.any(actions < output_range[0], axis=-1), jnp.any(actions > output_range[1], axis=-1))
            
            exterior_actions = jnp.where(mask[:, None], jnp.zeros_like(actions), actions)
            interior_actions = jnp.where(mask[:, None], actions, jnp.zeros_like(actions))
            
            exterior_l1_penalty = jnp.sum(jnp.abs(exterior_actions), axis=-1)
            exterior_l2_penalty = jnp.sum(exterior_actions ** 2, axis=-1)
            exterior_l2_penalty = jnp.sqrt(exterior_l2_penalty)
            
            interior_linear_penalty = jnp.sum(jnp.abs(interior_actions), axis=-1)
            interior_quadratic_penalty = jnp.sum(interior_actions ** 4, axis=-1)
            
            penalty_function = self.interior_quadratic_c * (self.exterior_linear_c / (3 * output_range[1] ** 3)) * interior_quadratic_penalty + self.exterior_linear_c * exterior_l1_penalty + self.exterior_quadratic_c * exterior_l2_penalty
            # penalty_function = self.ctrl_weight * oob_penalty + (self.ctrl_weight / (2 * output_range[1])) * ib_penalty
            
            actor_loss = (
                log_probs * self.temp.apply_fn({"params": self.temp.params}) - q + penalty_function
            ).mean()
            
            # compute how many actions are out of range
            oor_actions = jnp.mean(1 - mask)
            info_dict = {"actor_loss": actor_loss, 
                         "entropy": -log_probs.mean(), 
                         "oor_actions": oor_actions,
                         "penalty_function": penalty_function.mean()}
            
            return actor_loss, info_dict

        grads, actor_info = jax.grad(actor_loss_fn, has_aux=True)(self.actor.params)
        actor = self.actor.apply_gradients(grads=grads)

        return self.replace(actor=actor, rng=rng), actor_info

    def update_temperature(self, entropy: float) -> Tuple[Agent, Dict[str, float]]:
        def temperature_loss_fn(temp_params):
            temperature = self.temp.apply_fn({"params": temp_params})
            temp_loss = temperature * (entropy - self.target_entropy).mean()
            return temp_loss, {
                "temperature": temperature,
                "temperature_loss": temp_loss,
            }

        grads, temp_info = jax.grad(temperature_loss_fn, has_aux=True)(self.temp.params)
        temp = self.temp.apply_gradients(grads=grads)

        return self.replace(temp=temp), temp_info

    def update_critic(self, batch: DatasetDict, max_grad_norm: float = np.inf) -> Tuple[TrainState, Dict[str, float]]:

        dist = self.actor.apply_fn(
            {"params": self.actor.params}, batch["next_observations"]
        )

        rng = self.rng

        key, rng = jax.random.split(rng)
        next_actions = dist.sample(seed=key)

        # Used only for REDQ.
        key, rng = jax.random.split(rng)
        target_params = subsample_ensemble(
            key, self.target_critic.params, self.num_min_qs, self.num_qs
        )

        key, rng = jax.random.split(rng)
        next_qs = self.target_critic.apply_fn(
            {"params": target_params},
            batch["next_observations"],
            next_actions,
            True,
            rngs={"dropout": key},
        )  # training=True
        next_q = next_qs.min(axis=0)

        target_q = batch["rewards"] + self.discount * batch["masks"] * next_q

        if self.backup_entropy:
            next_log_probs = dist.log_prob(next_actions)
            target_q -= (
                self.discount
                * batch["masks"]
                * self.temp.apply_fn({"params": self.temp.params})
                * next_log_probs
            )

        key, rng = jax.random.split(rng)

        def critic_loss_fn(critic_params) -> Tuple[jnp.ndarray, Dict[str, float]]:
            qs = self.critic.apply_fn(
                {"params": critic_params},
                batch["observations"],
                batch["actions"],
                True,
                rngs={"dropout": key},
            )  # training=True
            critic_loss = ((qs - target_q) ** 2).mean()
            return critic_loss, {"critic_loss": critic_loss, "q": qs.mean()}

        grads, info = jax.grad(critic_loss_fn, has_aux=True)(self.critic.params)

        # Compute the gradient magnitudes
        critic_grad_magnitudes = compute_gradient_norm(grads) 

        critic_params_before = self.critic.params
        critic = self.critic.apply_gradients(grads=grads)
        critic_params_after = critic.params

        critic_params_change = compute_critic_param_change_norm(critic_params_before, critic_params_after)
        
        target_critic_params = optax.incremental_update(
            critic.params, self.target_critic.params, self.tau
        )
        target_critic = self.target_critic.replace(params=target_critic_params)
        
        info["critic_grad_magnitudes"] = critic_grad_magnitudes
        info["critic_weight_change"] = critic_params_change

        return self.replace(critic=critic, target_critic=target_critic, rng=rng), info
    
    @jax.jit
    def compute_td_error(self, batch: DatasetDict) -> float:

        dist = self.actor.apply_fn(
            {"params": self.actor.params}, batch["next_observations"]
        )

        rng = self.rng

        key, rng = jax.random.split(rng)
        next_actions = dist.sample(seed=key)

        # Used only for REDQ.
        key, rng = jax.random.split(rng)
        target_params = subsample_ensemble(
            key, self.target_critic.params, self.num_min_qs, self.num_qs
        )

        key, rng = jax.random.split(rng)
        next_qs = self.target_critic.apply_fn(
            {"params": target_params},
            batch["next_observations"],
            next_actions,
            True,
            rngs={"dropout": key},
        )  # training=True
        next_q = next_qs.min(axis=0)

        target_q = batch["rewards"] + self.discount * batch["masks"] * next_q

        key, rng = jax.random.split(rng)

        qs = self.critic.apply_fn(
            {"params": self.critic.params},
            batch["observations"],
            batch["actions"],
            True,
            rngs={"dropout": key},
        )  # training=True
        td_error = ((qs - target_q) ** 2).mean()
        return jnp.array(td_error, float)


    @partial(jax.jit, static_argnames=("utd_ratio", "actor_delay"))
    def update_with_delay(self, batch: DatasetDict, utd_ratio: int, output_range: Optional[Tuple[jnp.ndarray, jnp.ndarray]] = None, actor_delay: int = 1):
        new_agent = self
        for actor_i in range(utd_ratio//actor_delay):
            for critic_i in range(actor_delay):
                def slice(x):
                    assert x.shape[0] % utd_ratio == 0
                    batch_size = x.shape[0] // utd_ratio
                    return x[batch_size * (critic_i + actor_i * actor_delay) : batch_size * (critic_i + actor_i * actor_delay + 1)]

                mini_batch = jax.tree_util.tree_map(slice, batch)
                new_agent, critic_info = new_agent.update_critic(mini_batch)
            
            new_agent, actor_info = new_agent.update_actor(mini_batch, output_range=output_range)
            new_agent, temp_info = new_agent.update_temperature(actor_info["entropy"])
            new_agent, dynamics_info = new_agent.update_dynamics(mini_batch)
    
        return new_agent, {**actor_info, **critic_info, **temp_info, **dynamics_info}
      
    @partial(jax.jit, static_argnames="utd_ratio")
    def update(self, batch: DatasetDict, utd_ratio: int, output_range: Optional[Tuple[jnp.ndarray, jnp.ndarray]] = None):
        new_agent = self
        for i in range(utd_ratio):

            def slice(x):
                assert x.shape[0] % utd_ratio == 0
                batch_size = x.shape[0] // utd_ratio
                return x[batch_size * i : batch_size * (i + 1)]

            mini_batch = jax.tree_util.tree_map(slice, batch)
            new_agent, critic_info = new_agent.update_critic(mini_batch)
        
        new_agent, actor_info = new_agent.update_actor(mini_batch, output_range=output_range)
        new_agent, temp_info = new_agent.update_temperature(actor_info["entropy"])
            
        return new_agent, {**actor_info, **critic_info, **temp_info}
    
    @jax.jit    
    def get_value(self, observation: np.ndarray) -> float:
        # Convert the observation to a JAX numpy array
        jax_observation = convert_to_numpy_array(observation)
        
        rng = self.rng
        key, rng = jax.random.split(rng)
        key2, rng = jax.random.split(rng)
        
        # Use the actor network to generate an action
        dist = self.actor.apply_fn({"params": self.actor.params}, jax_observation[None])
        action = dist.sample(seed=key)[0]  # Generate a single action
        
        # Use the critic network to estimate the state-action value
        value = self.critic.apply_fn(
            {"params": self.critic.params},
            jax_observation[None],  # Add a batch dimension to the observation
            action[None],           # Add a batch dimension to the action
            True,
            rngs={"dropout": key2} 
        )
        # average over the ensemble
        value = value.mean(axis=0)[0] # Remove the batch dimension
        return jnp.array(value, float)

    @jax.jit    
    def self_perception(self, batch: DatasetDict) -> float:
        rng = self.rng
        key, rng = jax.random.split(rng)
        key2, rng = jax.random.split(rng)
        
        # Use the actor network to generate an action
        dist = self.actor.apply_fn({"params": self.actor.params}, batch["observations"])
        actions = dist.sample(seed=key)
        
        # Use the critic network to estimate the state-action value
        value = self.critic.apply_fn(
            {"params": self.critic.params},
            batch["observations"],
            actions,
            True,
            rngs={"dropout": key2} 
        )
        # average over the ensemble
        value = value.mean(axis=0)
        return jnp.array(value, float)[0]
    
    
    def update_dynamics(self, batch: DatasetDict) -> Tuple[Agent, Dict[str, float]]:
        key, rng = jax.random.split(self.rng)

        def dynamics_loss_fn(dynamics_params) -> Tuple[jnp.ndarray, Dict[str, float]]:
            next_observations_pred = self.dynamics_model.apply_fn(
                {"params": dynamics_params},
                batch["observations"],
                batch["actions"],
            )

            dynamics_loss = ((next_observations_pred - batch["next_observations"]) ** 2).mean()

            return dynamics_loss, {
                "dynamics_loss": dynamics_loss,
            }

        grads, dynamics_info = jax.grad(dynamics_loss_fn, has_aux=True)(self.dynamics_model.params)
        dynamics_model = self.dynamics_model.apply_gradients(grads=grads)

        return self.replace(dynamics_model=dynamics_model, rng=rng), dynamics_info

    @jax.jit
    def compute_dynamics_surprise(self, batch):
        observations = batch["observations"]
        actions = batch["actions"]
        actual_next_observations = batch["next_observations"]

        next_observations_pred = self.dynamics_model.apply_fn(
                {"params": self.dynamics_model.params},
                batch["observations"],
                batch["actions"],
        )

        # Compute the mean squared error (MSE) for next observations and rewards
        next_observations_error = jnp.mean((next_observations_pred - actual_next_observations) ** 2)
        return next_observations_error

jaxrl5/agents/sac/sac_learner_wdynamics.py
realquantumcookie/APRL
import flax.linen as nn
import jax.numpy as jnp


class Temperature(nn.Module):
    initial_temperature: float = 1.0

    @nn.compact
    def __call__(self) -> jnp.ndarray:
        log_temp = self.param(
            "log_temp",
            init_fn=lambda key: jnp.full((), jnp.log(self.initial_temperature)),
        )
        return jnp.exp(log_temp)

jaxrl5/agents/sac/temperature.py
realquantumcookie/APRL
from jaxrl5.data.memory_efficient_replay_buffer import MemoryEfficientReplayBuffer
from jaxrl5.data.replay_buffer import ReplayBuffer

jaxrl5/data/__init__.py
realquantumcookie/APRL
from functools import partial
from random import sample
from typing import Dict, Iterable, Optional, Tuple, Union

import jax
import jax.numpy as jnp
import numpy as np
from flax.core import frozen_dict
from gym.utils import seeding

from jaxrl5.types import DataType

DatasetDict = Dict[str, DataType]


def _check_lengths(dataset_dict: DatasetDict, dataset_len: Optional[int] = None) -> int:
    for v in dataset_dict.values():
        if isinstance(v, dict):
            dataset_len = dataset_len or _check_lengths(v, dataset_len)
        elif isinstance(v, np.ndarray):
            item_len = len(v)
            dataset_len = dataset_len or item_len
            assert dataset_len == item_len, "Inconsistent item lengths in the dataset."
        else:
            raise TypeError("Unsupported type.")
    return dataset_len


def _subselect(dataset_dict: DatasetDict, index: np.ndarray) -> DatasetDict:
    new_dataset_dict = {}
    for k, v in dataset_dict.items():
        if isinstance(v, dict):
            new_v = _subselect(v, index)
        elif isinstance(v, np.ndarray):
            new_v = v[index]
        else:
            raise TypeError("Unsupported type.")
        new_dataset_dict[k] = new_v
    return new_dataset_dict


def _sample(
    dataset_dict: Union[np.ndarray, DatasetDict], indx: np.ndarray
) -> DatasetDict:
    if isinstance(dataset_dict, np.ndarray):
        return dataset_dict[indx]
    elif isinstance(dataset_dict, dict):
        batch = {}
        for k, v in dataset_dict.items():
            batch[k] = _sample(v, indx)
    else:
        raise TypeError("Unsupported type.")
    return batch


class Dataset(object):
    def __init__(self, dataset_dict: DatasetDict, seed: Optional[int] = None):
        self.dataset_dict = dataset_dict
        self.dataset_len = _check_lengths(dataset_dict)

        # Seeding similar to OpenAI Gym:
        # https://github.com/openai/gym/blob/master/gym/spaces/space.py#L46
        self._np_random = None
        self._seed = None
        if seed is not None:
            self.seed(seed)

    @property
    def np_random(self) -> np.random.RandomState:
        if self._np_random is None:
            self.seed()
        return self._np_random

    def seed(self, seed: Optional[int] = None) -> list:
        self._np_random, self._seed = seeding.np_random(seed)
        return [self._seed]

    def __len__(self) -> int:
        return self.dataset_len

    def sample(
        self,
        batch_size: int,
        keys: Optional[Iterable[str]] = None,
        indx: Optional[np.ndarray] = None,
    ) -> frozen_dict.FrozenDict:
        if indx is None:
            if hasattr(self.np_random, "integers"):
                indx = self.np_random.integers(len(self), size=batch_size)
            else:
                indx = self.np_random.randint(len(self), size=batch_size)

        batch = dict()

        if keys is None:
            keys = self.dataset_dict.keys()

        for k in keys:
            if isinstance(self.dataset_dict[k], dict):
                batch[k] = _sample(self.dataset_dict[k], indx)
            else:
                batch[k] = self.dataset_dict[k][indx]

        return frozen_dict.freeze(batch)

    def sample_jax(self, batch_size: int, keys: Optional[Iterable[str]] = None):
        if not hasattr(self, "rng"):
            self.rng = jax.random.PRNGKey(self._seed or 42)

            if keys is None:
                keys = self.dataset_dict.keys()

            jax_dataset_dict = {k: self.dataset_dict[k] for k in keys}
            jax_dataset_dict = jax.device_put(jax_dataset_dict)

            @jax.jit
            def _sample_jax(rng):
                key, rng = jax.random.split(rng)
                indx = jax.random.randint(
                    key, (batch_size,), minval=0, maxval=len(self)
                )
                return rng, jax.tree_map(
                    lambda d: jnp.take(d, indx, axis=0), jax_dataset_dict
                )

            self._sample_jax = _sample_jax

        self.rng, sample = self._sample_jax(self.rng)
        return sample

    def split(self, ratio: float) -> Tuple["Dataset", "Dataset"]:
        assert 0 < ratio and ratio < 1
        train_index = np.index_exp[: int(self.dataset_len * ratio)]
        test_index = np.index_exp[int(self.dataset_len * ratio) :]

        index = np.arange(len(self), dtype=np.int32)
        self.np_random.shuffle(index)
        train_index = index[: int(self.dataset_len * ratio)]
        test_index = index[int(self.dataset_len * ratio) :]

        train_dataset_dict = _subselect(self.dataset_dict, train_index)
        test_dataset_dict = _subselect(self.dataset_dict, test_index)
        return Dataset(train_dataset_dict), Dataset(test_dataset_dict)

    def _trajectory_boundaries_and_returns(self) -> Tuple[list, list, list]:
        episode_starts = [0]
        episode_ends = []

        episode_return = 0
        episode_returns = []

        for i in range(len(self)):
            episode_return += self.dataset_dict["rewards"][i]

            if self.dataset_dict["dones"][i]:
                episode_returns.append(episode_return)
                episode_ends.append(i + 1)
                if i + 1 < len(self):
                    episode_starts.append(i + 1)
                episode_return = 0.0

        return episode_starts, episode_ends, episode_returns

    def filter(
        self, take_top: Optional[float] = None, threshold: Optional[float] = None
    ):
        assert (take_top is None and threshold is not None) or (
            take_top is not None and threshold is None
        )

        (
            episode_starts,
            episode_ends,
            episode_returns,
        ) = self._trajectory_boundaries_and_returns()

        if take_top is not None:
            threshold = np.percentile(episode_returns, 100 - take_top)

        bool_indx = np.full((len(self),), False, dtype=bool)

        for i in range(len(episode_returns)):
            if episode_returns[i] >= threshold:
                bool_indx[episode_starts[i] : episode_ends[i]] = True

        self.dataset_dict = _subselect(self.dataset_dict, bool_indx)

        self.dataset_len = _check_lengths(self.dataset_dict)

    def normalize_returns(self, scaling: float = 1000):
        (_, _, episode_returns) = self._trajectory_boundaries_and_returns()
        self.dataset_dict["rewards"] /= np.max(episode_returns) - np.min(
            episode_returns
        )
        self.dataset_dict["rewards"] *= scaling

jaxrl5/data/dataset.py
realquantumcookie/APRL
import copy
from typing import Iterable, Optional, Tuple

import gym
import numpy as np
from flax.core import frozen_dict
from gym.spaces import Box

from jaxrl5.data.dataset import DatasetDict, _sample
from jaxrl5.data.replay_buffer import ReplayBuffer


class MemoryEfficientReplayBuffer(ReplayBuffer):
    def __init__(
        self,
        observation_space: gym.Space,
        action_space: gym.Space,
        capacity: int,
        pixel_keys: Tuple[str, ...] = ("pixels",),
    ):
        self.pixel_keys = pixel_keys

        observation_space = copy.deepcopy(observation_space)
        self._num_stack = None
        for pixel_key in self.pixel_keys:
            pixel_obs_space = observation_space.spaces[pixel_key]
            if self._num_stack is None:
                self._num_stack = pixel_obs_space.shape[-1]
            else:
                assert self._num_stack == pixel_obs_space.shape[-1]
            self._unstacked_dim_size = pixel_obs_space.shape[-2]
            low = pixel_obs_space.low[..., 0]
            high = pixel_obs_space.high[..., 0]
            unstacked_pixel_obs_space = Box(
                low=low, high=high, dtype=pixel_obs_space.dtype
            )
            observation_space.spaces[pixel_key] = unstacked_pixel_obs_space

        next_observation_space_dict = copy.deepcopy(observation_space.spaces)
        for pixel_key in self.pixel_keys:
            next_observation_space_dict.pop(pixel_key)
        next_observation_space = gym.spaces.Dict(next_observation_space_dict)

        self._first = True
        self._is_correct_index = np.full(capacity, False, dtype=bool)

        super().__init__(
            observation_space,
            action_space,
            capacity,
            next_observation_space=next_observation_space,
        )

    def insert(self, data_dict: DatasetDict):
        if self._insert_index == 0 and self._capacity == len(self) and not self._first:
            indxs = np.arange(len(self) - self._num_stack, len(self))
            for indx in indxs:
                element = super().sample(1, indx=indx)
                self._is_correct_index[self._insert_index] = False
                super().insert(element)

        data_dict = data_dict.copy()
        data_dict["observations"] = data_dict["observations"].copy()
        data_dict["next_observations"] = data_dict["next_observations"].copy()

        obs_pixels = {}
        next_obs_pixels = {}
        for pixel_key in self.pixel_keys:
            obs_pixels[pixel_key] = data_dict["observations"].pop(pixel_key)
            next_obs_pixels[pixel_key] = data_dict["next_observations"].pop(pixel_key)

        if self._first:
            for i in range(self._num_stack):
                for pixel_key in self.pixel_keys:
                    data_dict["observations"][pixel_key] = obs_pixels[pixel_key][..., i]

                self._is_correct_index[self._insert_index] = False
                super().insert(data_dict)

        for pixel_key in self.pixel_keys:
            data_dict["observations"][pixel_key] = next_obs_pixels[pixel_key][..., -1]

        self._first = data_dict["dones"]

        self._is_correct_index[self._insert_index] = True
        super().insert(data_dict)

        for i in range(self._num_stack):
            indx = (self._insert_index + i) % len(self)
            self._is_correct_index[indx] = False

    def sample(
        self,
        batch_size: int,
        keys: Optional[Iterable[str]] = None,
        indx: Optional[np.ndarray] = None,
        pack_obs_and_next_obs: bool = False,
    ) -> frozen_dict.FrozenDict:
        """Samples from the replay buffer.

        Args:
            batch_size: Minibatch size.
            keys: Keys to sample.
            indx: Take indices instead of sampling.
            pack_obs_and_next_obs: whether to pack img and next_img into one image.
                It's useful when they have overlapping frames.

        Returns:
            A frozen dictionary.
        """

        if indx is None:
            if hasattr(self.np_random, "integers"):
                indx = self.np_random.integers(len(self), size=batch_size)
            else:
                indx = self.np_random.randint(len(self), size=batch_size)

            for i in range(batch_size):
                while not self._is_correct_index[indx[i]]:
                    if hasattr(self.np_random, "integers"):
                        indx[i] = self.np_random.integers(len(self))
                    else:
                        indx[i] = self.np_random.randint(len(self))
        else:
            raise NotImplementedError()

        if keys is None:
            keys = self.dataset_dict.keys()
        else:
            assert "observations" in keys

        keys = list(keys)
        keys.remove("observations")

        batch = super().sample(batch_size, keys, indx)
        batch = batch.unfreeze()

        obs_keys = self.dataset_dict["observations"].keys()
        obs_keys = list(obs_keys)
        for pixel_key in self.pixel_keys:
            obs_keys.remove(pixel_key)

        batch["observations"] = {}
        for k in obs_keys:
            batch["observations"][k] = _sample(
                self.dataset_dict["observations"][k], indx
            )

        for pixel_key in self.pixel_keys:
            obs_pixels = self.dataset_dict["observations"][pixel_key]
            obs_pixels = np.lib.stride_tricks.sliding_window_view(
                obs_pixels, self._num_stack + 1, axis=0
            )
            obs_pixels = obs_pixels[indx - self._num_stack]

            if pack_obs_and_next_obs:
                batch["observations"][pixel_key] = obs_pixels
            else:
                batch["observations"][pixel_key] = obs_pixels[..., :-1]
                if "next_observations" in keys:
                    batch["next_observations"][pixel_key] = obs_pixels[..., 1:]

        return frozen_dict.freeze(batch)

jaxrl5/data/memory_efficient_replay_buffer.py
realquantumcookie/APRL
import collections
from typing import Optional, Union

import gym
import gym.spaces
import jax
import numpy as np

from jaxrl5.data.dataset import Dataset, DatasetDict


def _init_replay_dict(
    obs_space: gym.Space, capacity: int
) -> Union[np.ndarray, DatasetDict]:
    if isinstance(obs_space, gym.spaces.Box):
        return np.empty((capacity, *obs_space.shape), dtype=obs_space.dtype)
    elif isinstance(obs_space, gym.spaces.Dict):
        data_dict = {}
        for k, v in obs_space.spaces.items():
            data_dict[k] = _init_replay_dict(v, capacity)
        return data_dict
    else:
        raise TypeError()


def _insert_recursively(
    dataset_dict: DatasetDict, data_dict: DatasetDict, insert_index: int
):
    if isinstance(dataset_dict, np.ndarray):
        dataset_dict[insert_index] = data_dict
    elif isinstance(dataset_dict, dict):
        assert dataset_dict.keys() == data_dict.keys()
        for k in dataset_dict.keys():
            _insert_recursively(dataset_dict[k], data_dict[k], insert_index)
    else:
        raise TypeError()


class ReplayBuffer(Dataset):
    def __init__(
        self,
        observation_space: gym.Space,
        action_space: gym.Space,
        capacity: int,
        next_observation_space: Optional[gym.Space] = None,
    ):
        if next_observation_space is None:
            next_observation_space = observation_space

        observation_data = _init_replay_dict(observation_space, capacity)
        next_observation_data = _init_replay_dict(next_observation_space, capacity)
        dataset_dict = dict(
            observations=observation_data,
            next_observations=next_observation_data,
            actions=np.empty((capacity, *action_space.shape), dtype=action_space.dtype),
            rewards=np.empty((capacity,), dtype=np.float32),
            masks=np.empty((capacity,), dtype=np.float32),
            dones=np.empty((capacity,), dtype=bool),
        )

        super().__init__(dataset_dict)

        self._size = 0
        self._capacity = capacity
        self._insert_index = 0

    def __len__(self) -> int:
        return self._size

    def insert(self, data_dict: DatasetDict):
        _insert_recursively(self.dataset_dict, data_dict, self._insert_index)

        self._insert_index = (self._insert_index + 1) % self._capacity
        self._size = min(self._size + 1, self._capacity)

    def initialize_with_dataset(self, dataset: Dataset,
                                num_samples: Optional[int]):
        assert self._insert_index == 0, 'Can insert a batch online in an empty replay buffer.'

        dataset_size = len(dataset.dataset_dict['observations'])

        eps = 1e-6
        lim = 1 - eps
        dataset.dataset_dict["actions"] = np.clip(dataset.dataset_dict["actions"], -lim, lim)

        if num_samples is None:
            num_samples = dataset_size
        else:
            num_samples = min(dataset_size, num_samples)
        assert self._capacity >= num_samples, 'Dataset cannot be larger than the replay buffer capacity.'

        for k in self.dataset_dict.keys():
            self.dataset_dict[k][:num_samples] = dataset.dataset_dict[k][:num_samples]


        
        self._insert_index = num_samples
        self._size = num_samples

    def get_iterator(self, queue_size: int = 2, sample_args: dict = {}):
        # See https://flax.readthedocs.io/en/latest/_modules/flax/jax_utils.html#prefetch_to_device
        # queue_size = 2 should be ok for one GPU.

        queue = collections.deque()

        def enqueue(n):
            for _ in range(n):
                data = self.sample(**sample_args)
                queue.append(jax.device_put(data))

        enqueue(queue_size)
        while queue:
            yield queue.popleft()
            enqueue(1)

jaxrl5/data/replay_buffer.py
realquantumcookie/APRL
from jaxrl5.distributions.tanh_deterministic import TanhDeterministic
from jaxrl5.distributions.tanh_normal import Normal, TanhNormal

jaxrl5/distributions/__init__.py
realquantumcookie/APRL
from typing import Type

import flax.linen as nn
import jax.numpy as jnp

from jaxrl5.networks import default_init


class TanhDeterministic(nn.Module):
    base_cls: Type[nn.Module]
    action_dim: int

    @nn.compact
    def __call__(self, inputs, *args, **kwargs) -> jnp.ndarray:
        x = self.base_cls()(inputs, *args, **kwargs)

        means = nn.Dense(
            self.action_dim, kernel_init=default_init(), name="OutputDenseMean"
        )(x)

        means = nn.tanh(means)

        return means

jaxrl5/distributions/tanh_deterministic.py
realquantumcookie/APRL
import functools
from typing import Optional, Type

import tensorflow_probability

from jaxrl5.distributions.tanh_transformed import TanhTransformedDistribution

tfp = tensorflow_probability.substrates.jax
tfd = tfp.distributions

import flax.linen as nn
import jax.numpy as jnp

from jaxrl5.networks import default_init


class Normal(nn.Module):
    base_cls: Type[nn.Module]
    action_dim: int
    log_std_min: Optional[float] = -20
    log_std_max: Optional[float] = 2
    state_dependent_std: bool = True
    squash_tanh: bool = False

    @nn.compact
    def __call__(self, inputs, *args, **kwargs) -> tfd.Distribution:
        x = self.base_cls()(inputs, *args, **kwargs)

        means = nn.Dense(
            self.action_dim, kernel_init=default_init(), name="OutputDenseMean"
        )(x)
        if self.state_dependent_std:
            log_stds = nn.Dense(
                self.action_dim, kernel_init=default_init(), name="OutputDenseLogStd"
            )(x)
        else:
            log_stds = self.param(
                "OutpuLogStd", nn.initializers.zeros, (self.action_dim,), jnp.float32
            )

        log_stds = jnp.clip(log_stds, self.log_std_min, self.log_std_max)

        distribution = tfd.MultivariateNormalDiag(
            loc=means, scale_diag=jnp.exp(log_stds)
        )

        if self.squash_tanh:
            return TanhTransformedDistribution(distribution)
            # if output_range is not None:
            #     lows, highs = output_range
            # else:
            #     lows, highs = (-1.0, 1.0)

            # return TanhTransformedDistribution(distribution, low=lows, high=highs)
        else:
            return distribution
            


TanhNormal = functools.partial(Normal, squash_tanh=True)

jaxrl5/distributions/tanh_normal.py
realquantumcookie/APRL
from typing import Any, Optional

import tensorflow_probability

tfp = tensorflow_probability.substrates.jax
tfd = tfp.distributions
tfb = tfp.bijectors

import jax
import jax.numpy as jnp

# Inspired by
# https://github.com/deepmind/acme/blob/300c780ffeb88661a41540b99d3e25714e2efd20/acme/jax/networks/distributional.py#L163
# but modified to only compute a mode.


class TanhTransformedDistribution(tfd.TransformedDistribution):
    def __init__(self, distribution: tfd.Distribution, validate_args: bool = False):
        super().__init__(
            distribution=distribution, bijector=tfb.Tanh(), validate_args=validate_args
         )
    # def __init__(self, distribution: tfd.Distribution, validate_args: bool = False, low=-1.0, high=1.0):
    #     self.shift = (high + low) / 2.0
    #     self.scale = (high - low) / 2.0
    #     bijector = tfb.Chain([
    #         tfb.Shift(self.shift),
    #         tfb.Scale(self.scale),
    #         tfb.Tanh(),
    #         tfb.Scale(1.0 / self.scale),
    #         tfb.Shift(-self.shift),
    #     ], name="ranged_tanh")

    #     super().__init__(
    #         distribution=distribution, bijector=bijector, validate_args=validate_args
    #     )

    def mode(self) -> jnp.ndarray:
        return self.bijector.forward(self.distribution.mode())

    @classmethod
    def _parameter_properties(cls, dtype: Optional[Any], num_classes=None):
        td_properties = super()._parameter_properties(dtype, num_classes=num_classes)
        del td_properties["bijector"]
        return td_properties
    
jaxrl5/distributions/tanh_transformed.py
realquantumcookie/APRL
from jaxrl5.networks.ensemble import Ensemble, subsample_ensemble
from jaxrl5.networks.mlp import MLP, default_init
from jaxrl5.networks.mlp_resnet import MLPResNetV2
from jaxrl5.networks.state_action_value import StateActionValue
from jaxrl5.networks.state_action_nextstate import StateActionNextState

jaxrl5/networks/__init__.py
realquantumcookie/APRL
from jaxrl5.networks.encoders.d4pg_encoder import D4PGEncoder
from jaxrl5.networks.encoders.ln_resnet_encoder import ResNetV2Encoder

jaxrl5/networks/encoders/__init__.py
realquantumcookie/APRL
from typing import Sequence

import flax.linen as nn
import jax.numpy as jnp

from jaxrl5.networks import default_init


class D4PGEncoder(nn.Module):
    features: Sequence[int] = (32, 32, 32, 32)
    filters: Sequence[int] = (2, 1, 1, 1)
    strides: Sequence[int] = (2, 1, 1, 1)
    padding: str = "VALID"

    @nn.compact
    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        assert len(self.features) == len(self.strides)

        for features, filter_, stride in zip(self.features, self.filters, self.strides):
            x = nn.Conv(
                features,
                kernel_size=(filter_, filter_),
                strides=(stride, stride),
                kernel_init=default_init(),
                padding=self.padding,
            )(x)
            x = nn.relu(x)

        return x.reshape((*x.shape[:-3], -1))

jaxrl5/networks/encoders/d4pg_encoder.py
realquantumcookie/APRL
# Based on:
# https://github.com/google/flax/blob/main/examples/imagenet/models.py
# and
# https://github.com/google-research/big_transfer/blob/master/bit_jax/models.py
from functools import partial
from typing import Any, Callable, Tuple, Type

import flax.linen as nn
import jax.numpy as jnp
from flax import linen as nn

from jaxrl5.networks import default_init


class ResNetV2Block(nn.Module):
    """ResNet block."""

    filters: int
    conv_cls: Type[nn.Module]
    norm_cls: Type[nn.Module]
    act: Callable
    strides: Tuple[int, int] = (1, 1)

    @nn.compact
    def __call__(self, x):
        residual = x
        y = self.norm_cls()(x)
        y = self.act(y)
        y = self.conv_cls(self.filters, (3, 3), self.strides)(y)
        y = self.norm_cls()(y)
        y = self.act(y)
        y = self.conv_cls(self.filters, (3, 3))(y)

        if residual.shape != y.shape:
            residual = self.conv_cls(self.filters, (1, 1), self.strides)(residual)

        return residual + y


class MyGroupNorm(nn.GroupNorm):
    def __call__(self, x):
        if x.ndim == 3:
            x = x[jnp.newaxis]
            x = super().__call__(x)
            return x[0]
        else:
            return super().__call__(x)


class ResNetV2Encoder(nn.Module):
    """ResNetV2."""

    stage_sizes: Tuple[int]
    num_filters: int = 64
    dtype: Any = jnp.float32
    act: Callable = nn.relu

    @nn.compact
    def __call__(self, x):
        conv_cls = partial(
            nn.Conv, use_bias=False, dtype=self.dtype, kernel_init=default_init()
        )
        norm_cls = partial(MyGroupNorm, num_groups=4, epsilon=1e-5, dtype=self.dtype)

        if x.shape[-2] == 224:
            x = conv_cls(self.num_filters, (7, 7), (2, 2), padding=[(3, 3), (3, 3)])(x)
            x = nn.max_pool(x, (3, 3), strides=(2, 2), padding="SAME")
        else:
            x = conv_cls(self.num_filters, (3, 3))(x)

        for i, block_size in enumerate(self.stage_sizes):
            for j in range(block_size):
                strides = (2, 2) if i > 0 and j == 0 else (1, 1)
                x = ResNetV2Block(
                    self.num_filters * 2**i,
                    strides=strides,
                    conv_cls=conv_cls,
                    norm_cls=norm_cls,
                    act=self.act,
                )(x)

        x = norm_cls()(x)
        x = self.act(x)
        return x.reshape((*x.shape[:-3], -1))

jaxrl5/networks/encoders/ln_resnet_encoder.py
realquantumcookie/APRL
from typing import Type

import flax.linen as nn
import jax
import jax.numpy as jnp


class Ensemble(nn.Module):
    net_cls: Type[nn.Module]
    num: int = 2

    @nn.compact
    def __call__(self, *args):
        ensemble = nn.vmap(
            self.net_cls,
            variable_axes={"params": 0},
            split_rngs={"params": True, "dropout": True},
            in_axes=None,
            out_axes=0,
            axis_size=self.num,
        )
        return ensemble()(*args)


def subsample_ensemble(key: jax.random.PRNGKey, params, num_sample: int, num_qs: int):
    if num_sample is not None:
        all_indx = jnp.arange(0, num_qs)
        indx = jax.random.choice(key, a=all_indx, shape=(num_sample,), replace=False)

        if "Ensemble_0" in params:
            ens_params = jax.tree_util.tree_map(
                lambda param: param[indx], params["Ensemble_0"]
            )
            params = params.copy(add_or_replace={"Ensemble_0": ens_params})
        else:
            params = jax.tree_util.tree_map(lambda param: param[indx], params)
    return params

jaxrl5/networks/ensemble.py
realquantumcookie/APRL
from typing import Callable, Optional, Sequence

import flax.linen as nn
import jax.numpy as jnp

# default_init = nn.initializers.xavier_uniform
def default_init(scale: Optional[float] = jnp.sqrt(2)):
    return nn.initializers.orthogonal(scale)

class MLP(nn.Module):
    hidden_dims: Sequence[int]
    activations: Callable[[jnp.ndarray], jnp.ndarray] = nn.relu
    activate_final: bool = False
    use_layer_norm: bool = False
    scale_final: Optional[float] = None
    dropout_rate: Optional[float] = None
    use_pnorm: bool = False

    @nn.compact
    def __call__(self, x: jnp.ndarray, training: bool = False) -> jnp.ndarray:

        for i, size in enumerate(self.hidden_dims):
            if i + 1 == len(self.hidden_dims) and self.scale_final is not None:
                x = nn.Dense(size, kernel_init=default_init(self.scale_final))(x)
            else:
                x = nn.Dense(size, kernel_init=default_init())(x)

            if i + 1 < len(self.hidden_dims) or self.activate_final:
                if self.dropout_rate is not None and self.dropout_rate > 0:
                    x = nn.Dropout(rate=self.dropout_rate)(
                        x, deterministic=not training
                    )
                if self.use_layer_norm:
                    x = nn.LayerNorm()(x)
                x = self.activations(x)
        if self.use_pnorm:
            x /= jnp.linalg.norm(x, axis=-1, keepdims=True).clip(1e-10)
        return x

jaxrl5/networks/mlp.py
realquantumcookie/APRL
from typing import Any, Callable, Sequence

import flax.linen as nn
import jax.numpy as jnp
from flax import linen as nn

ModuleDef = Any


class MLPResNetV2Block(nn.Module):
    """MLPResNet block."""

    features: int
    act: Callable

    @nn.compact
    def __call__(self, x):
        residual = x
        y = nn.LayerNorm()(x)
        y = self.act(y)
        y = nn.Dense(self.features)(y)
        y = nn.LayerNorm()(y)
        y = self.act(y)
        y = nn.Dense(self.features)(y)

        if residual.shape != y.shape:
            residual = nn.Dense(self.features)(residual)

        return residual + y


class MLPResNetV2(nn.Module):
    """MLPResNetV2."""

    num_blocks: int
    features: int = 256
    dtype: Any = jnp.float32
    act: Callable = nn.relu

    @nn.compact
    def __call__(self, x, training=False):
        x = nn.Dense(self.features)(x)
        for _ in range(self.num_blocks):
            x = MLPResNetV2Block(self.features, act=self.act)(x)
        x = nn.LayerNorm()(x)
        x = self.act(x)
        return x

jaxrl5/networks/mlp_resnet.py
realquantumcookie/APRL
import flax.linen as nn
import jax.numpy as jnp
# from jax.nn import tanh

# from jaxrl5.networks import default_init
my_init = nn.initializers.xavier_uniform

class StateActionNextState(nn.Module):
    base_cls: nn.Module
    obs_dim: int
    
    @nn.compact
    def __call__(
        self, observations: jnp.ndarray, actions: jnp.ndarray, *args, **kwargs
    ) -> jnp.ndarray:
        inputs = jnp.concatenate([observations, actions], axis=-1)
        outputs = self.base_cls()(inputs, *args, **kwargs)
        
        residual = nn.Dense(self.obs_dim, kernel_init=my_init())(outputs)
        next_state = observations + residual

        return next_state

jaxrl5/networks/state_action_nextstate.py
realquantumcookie/APRL
import flax.linen as nn
import jax.numpy as jnp

from jaxrl5.networks import default_init


class StateActionValue(nn.Module):
    base_cls: nn.Module

    @nn.compact
    def __call__(
        self, observations: jnp.ndarray, actions: jnp.ndarray, *args, **kwargs
    ) -> jnp.ndarray:
        inputs = jnp.concatenate([observations, actions], axis=-1)
        outputs = self.base_cls()(inputs, *args, **kwargs)

        value = nn.Dense(1, kernel_init=default_init())(outputs)

        return jnp.squeeze(value, -1)

jaxrl5/networks/state_action_value.py
realquantumcookie/APRL
from typing import Any, Dict, Union

import flax
import numpy as np

DataType = Union[np.ndarray, Dict[str, "DataType"]]
PRNGKey = Any
Params = flax.core.FrozenDict[str, Any]

jaxrl5/types.py
realquantumcookie/APRL
import gym
from gym.wrappers.flatten_observation import FlattenObservation

from jaxrl5.wrappers.single_precision import SinglePrecision
from jaxrl5.wrappers.universal_seed import UniversalSeed
from jaxrl5.wrappers.wandb_video import WANDBVideo

def wrap_gym(env: gym.Env, rescale_actions: bool = True) -> gym.Env:
    env = SinglePrecision(env)
    env = UniversalSeed(env)
    if rescale_actions:
        env = gym.wrappers.RescaleAction(env, -1, 1)

    if isinstance(env.observation_space, gym.spaces.Dict):
        env = FlattenObservation(env)

    env = gym.wrappers.ClipAction(env)

    return env

jaxrl5/wrappers/__init__.py
realquantumcookie/APRL
import collections

import gym
import numpy as np
from gym.spaces import Box


class FrameStack(gym.Wrapper):
    def __init__(self, env, num_stack: int, stacking_key: str = "pixels"):
        super().__init__(env)
        self._num_stack = num_stack
        self._stacking_key = stacking_key

        assert stacking_key in self.observation_space.spaces
        pixel_obs_spaces = self.observation_space.spaces[stacking_key]

        self._env_dim = pixel_obs_spaces.shape[-1]

        low = np.repeat(pixel_obs_spaces.low[..., np.newaxis], num_stack, axis=-1)
        high = np.repeat(pixel_obs_spaces.high[..., np.newaxis], num_stack, axis=-1)
        new_pixel_obs_spaces = Box(low=low, high=high, dtype=pixel_obs_spaces.dtype)
        self.observation_space.spaces[stacking_key] = new_pixel_obs_spaces

        self._frames = collections.deque(maxlen=num_stack)

    def reset(self):
        obs = self.env.reset()
        for i in range(self._num_stack):
            self._frames.append(obs[self._stacking_key])
        obs[self._stacking_key] = self.frames
        return obs

    @property
    def frames(self):
        return np.stack(self._frames, axis=-1)

    def step(self, action):
        obs, reward, done, info = self.env.step(action)
        self._frames.append(obs[self._stacking_key])
        obs[self._stacking_key] = self.frames
        return obs, reward, done, info

jaxrl5/wrappers/frame_stack.py
realquantumcookie/APRL
import gym
import numpy as np


class RepeatAction(gym.Wrapper):
    def __init__(self, env, action_repeat=4):
        super().__init__(env)
        self._action_repeat = action_repeat

    def step(self, action: np.ndarray):
        total_reward = 0.0
        done = None
        combined_info = {}

        for _ in range(self._action_repeat):
            obs, reward, done, info = self.env.step(action)
            total_reward += reward
            combined_info.update(info)
            if done:
                break

        return obs, total_reward, done, combined_info

jaxrl5/wrappers/repeat_action.py
realquantumcookie/APRL
import copy

import gym
import numpy as np
from gym.spaces import Box, Dict


def _convert_space(obs_space):
    if isinstance(obs_space, Box):
        obs_space = Box(obs_space.low, obs_space.high, obs_space.shape)
    elif isinstance(obs_space, Dict):
        for k, v in obs_space.spaces.items():
            obs_space.spaces[k] = _convert_space(v)
        obs_space = Dict(obs_space.spaces)
    else:
        raise NotImplementedError
    return obs_space


def _convert_obs(obs):
    if isinstance(obs, np.ndarray):
        if obs.dtype == np.float64:
            return obs.astype(np.float32)
        else:
            return obs
    elif isinstance(obs, dict):
        obs = copy.copy(obs)
        for k, v in obs.items():
            obs[k] = _convert_obs(v)
        return obs


class SinglePrecision(gym.ObservationWrapper):
    def __init__(self, env):
        super().__init__(env)

        obs_space = copy.deepcopy(self.env.observation_space)
        self.observation_space = _convert_space(obs_space)

    def observation(self, observation):
        return _convert_obs(observation)

jaxrl5/wrappers/single_precision.py
realquantumcookie/APRL
import gym


class UniversalSeed(gym.Wrapper):
    def seed(self, seed: int):
        seeds = self.env.seed(seed)
        self.env.observation_space.seed(seed)
        self.env.action_space.seed(seed)
        return seeds

jaxrl5/wrappers/universal_seed.py
realquantumcookie/APRL
from typing import Optional

import gym
import numpy as np
import wandb


class WANDBVideo(gym.Wrapper):
    def __init__(
        self,
        env: gym.Env,
        name: str = "video",
        pixel_hw: int = 84,
        render_kwargs={},
        max_videos: Optional[int] = None,
    ):
        super().__init__(env)

        self._name = name
        self._pixel_hw = pixel_hw
        self._render_kwargs = render_kwargs
        self._max_videos = max_videos
        self._video = []

    def _add_frame(self, obs):
        if self._max_videos is not None and self._max_videos <= 0:
            return
        if isinstance(obs, dict) and "pixels" in obs:
            if obs["pixels"].ndim == 4:
                self._video.append(obs["pixels"][..., -1])
            else:
                self._video.append(obs["pixels"])
        else:
            self._video.append(
                self.render(
                    height=self._pixel_hw,
                    width=self._pixel_hw,
                    mode="rgb_array",
                    **self._render_kwargs
                )
            )

    def reset(self, **kwargs):
        self._video.clear()
        obs = super().reset(**kwargs)
        self._add_frame(obs)
        return obs

    def step(self, action: np.ndarray):

        obs, reward, done, info = super().step(action)
        self._add_frame(obs)

        if done and len(self._video) > 0:
            if self._max_videos is not None:
                self._max_videos -= 1
            video = np.moveaxis(np.stack(self._video), -1, 1)
            if video.shape[1] == 1:
                video = np.repeat(video, 3, 1)
            video = wandb.Video(video, fps=20, format="mp4")
            wandb.log({self._name: video}, commit=False)

        return obs, reward, done, info

jaxrl5/wrappers/wandb_video.py
realquantumcookie/APRL
from .robots import *
from .tasks import *
from .utils import *
from .floors import *
rail_mujoco_walker/__init__.py
realquantumcookie/APRL
from .heightfield_arena import *
rail_mujoco_walker/floors/__init__.py
realquantumcookie/APRL
from .heightfield_arena_goals import HEIGHTFIELD_ARENA_GOALS
from .heightfield_arena import HeightFieldArena
from .crouching_heightfield_arena import CrouchingHeightfieldArena
rail_mujoco_walker/floors/heightfield_arena/__init__.py
realquantumcookie/APRL
from .heightfield_arena import HeightFieldArena
import numpy as np

CROUCHING_LOCATION = np.array([0.05, 0.18, 1.01])
CROUCHING_SIZE = np.array([0.03, 0.05, 0.5])
CROUCHING_YAW = 10.0/180 * np.pi

def _helper_is_in_croching_range(pos_2d : np.ndarray, crouching_location_2d : np.ndarray, crouching_size_2d : np.ndarray, crouching_yaw : float):
    assert pos_2d.shape == (2,) and crouching_location_2d.shape == (2,) and crouching_size_2d.shape == (2,)
    delta_pos = pos_2d - crouching_location_2d
    x_axis_unit = [np.cos(crouching_yaw), np.sin(crouching_yaw)]
    y_axis_unit = [-np.sin(crouching_yaw), np.cos(crouching_yaw)]
    x_scalar = np.inner(delta_pos, x_axis_unit)
    y_scalar = np.inner(delta_pos, y_axis_unit)
    # https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-geom-size
    # it is safe to do so because mujoco has size specified by half size
    return np.abs(x_scalar) < crouching_size_2d[0] and np.abs(y_scalar) < crouching_size_2d[1]


class CrouchingHeightfieldArena(
    HeightFieldArena
):
    def __init__(self, scale, name="crouching_heightfield_arena"):
        super().__init__(scale, name=name)
    
    def _build(self, name):
        super()._build(name)

        size = [self.scale * s for s in CROUCHING_SIZE[:2]] + [CROUCHING_SIZE[2]]
        pos = [self.scale * p for p in CROUCHING_LOCATION[:2]] + [CROUCHING_LOCATION[2]]
        euler = (0.0, 0.0, CROUCHING_YAW)
        
        self._crouch_overhead = self._mjcf_root.worldbody.add(
            'geom',
            name='crouch_overhead',
            type='box',
            rgba=(1.0, 1.0, 1.0, 0.3),
            pos=pos,
            size=size,
            euler=euler
        )

    def should_crouch(self, robot_pos : np.ndarray):
        assert robot_pos.ndim == 1 and robot_pos.shape[0] >= 2
        return _helper_is_in_croching_range(
            robot_pos[:2],
            CROUCHING_LOCATION[:2] * self.scale,
            CROUCHING_SIZE[:2] * self.scale + 1.0, # 1.0 to account for the robot's size
            CROUCHING_YAW
        )

    @staticmethod
    def get_should_crouch_callback(scale : float):
        def is_in_crouching_range(
            robot_pos : np.ndarray
        ) -> bool:
            assert robot_pos.ndim == 1 and robot_pos.shape[0] >= 2
            return _helper_is_in_croching_range(
                robot_pos[:2],
                CROUCHING_LOCATION[:2] * scale,
                CROUCHING_SIZE[:2] * scale + 0.5, # 0.5 to account for the robot's size
                CROUCHING_YAW
            )
        return is_in_crouching_range
rail_mujoco_walker/floors/heightfield_arena/crouching_heightfield_arena.py
realquantumcookie/APRL
import typing
import numpy as np
import random

GOAL_THRESHOLD = 1.0
def normalize_angle(angle):
    return (angle + np.pi) % (2 * np.pi) - np.pi

class GoalGraph():
    """
    A graph of goals that the car can drive to. Once the car arrives at a goal,
    the goal will be changed to one of its successors.
    """
    def __init__(self, scale, goal_config='corners', goal_threshold=GOAL_THRESHOLD):
        if goal_config == 'full':
            self.goals = [
                (0.05, 0.05),
                (0.33, -0.15),
                (0.3, -0.45),
                (0.65, -0.1),
                (0.525, 0.2),
                (0.6, 0.75),
                (0.3, 0.6),
                (-0.2, 0.6),
                (-0.5, 0.57),
                (-0.6, 0.05),
                (-0.7, -0.6),
                (0.0, -0.8),
            ]

            self.graph = {
                0: [1, 6, 7, 9, 10, 11],
                1: [0, 2, 3, 4],
                2: [1, 3, 11],
                3: [1, 2, 4],
                4: [1, 3, 6],
                5: [4, 6],
                6: [0, 4, 5, 7],
                7: [0, 6, 8],
                8: [7, 9],
                9: [0, 8, 10],
                10: [0, 9, 11],
                11: [0, 2, 10],
            }

            self.start_headings = [(0, 2 * np.pi)] * len(self.goals)
        elif goal_config == 'edges':
            self.goals = [
                (0.9, 0.0),
                (0.0, 0.9),
                (-0.9, 0.0),
                (0.0, -0.9),
            ]

            self.start_headings = [
                (a - 0.1, a + 0.1) for a in [
                    np.deg2rad(90),
                    np.deg2rad(180),
                    np.deg2rad(-90),
                    np.deg2rad(0),
                ]
            ]

            self.graph = {
                0: [1],
                1: [2],
                2: [3],
                3: [0],
            }
        elif goal_config == 'corners':
            self.goals = [
                (0.9, 0.9),
                (-0.9, 0.9),
                (-0.9, -0.9),
                (0.9, -0.9),
            ]

            self.start_headings = [
                (a - 0.1, a + 0.1) for a in [
                    np.deg2rad(135),
                    -np.deg2rad(135),
                    -np.deg2rad(45),
                    np.deg2rad(45),
                ]
            ]

            self.graph = {
                0: [1],
                1: [2],
                2: [3],
                3: [0],
            }
        elif goal_config == 'ring_dense':
            self.goals = [
                (0.9, 0.9),
                (0.0, 0.9),
                (-0.9, 0.9),
                (-0.9, 0.0),
                (-0.9, -0.9),
                (0.0, -0.9),
                (0.9, -0.9),
                (0.9, 0.0),
            ]

            self.start_headings = [
                (a - 0.1, a + 0.1) for a in [
                    np.deg2rad(135),
                    np.deg2rad(-180),
                    -np.deg2rad(135),
                    np.deg2rad(-90),
                    -np.deg2rad(45),
                    np.deg2rad(0),
                    -np.deg2rad(45),
                    np.deg2rad(90),
                ]
            ]

            self.graph = {
                0: [1, 2],
                1: [2, 3],
                2: [3, 4],
                3: [4, 5],
                4: [5, 6],
                5: [6, 7],
                6: [7, 0],
                7: [0, 1],
            }
        elif goal_config == 'ring_small_inner':
            self.goals = [
                (-0.10888671875, 0.609375),
                (0.07421875, 0.072265625),
                (-0.072265625, -0.220703125),
                (-0.5849609375, 0.072265625),
                (-0.49951171875, 0.67041015625),
            ]

            self.start_headings = [
                (a - 0.1, a + 0.1) for a in [
                    np.deg2rad(-45),
                    np.deg2rad(-110),
                    np.deg2rad(-120),
                    np.deg2rad(110),
                    np.deg2rad(45),
                ]
            ]

            self.graph = {
                0: [1],
                1: [2],
                2: [3],
                3: [4],
                4: [0],
            }
        elif goal_config == 'small_inner_1fork':
            self.goals = [
                (-0.10888671875, 0.609375),
                (0.07421875, 0.072265625),
                (-0.072265625, -0.220703125),
                (-0.5849609375, 0.072265625),
                (-0.49951171875, 0.67041015625),
                (-0.2919921875, -0.5380859375),
                (-0.6826171875, -0.46484375)
            ]

            self.start_headings = [
                (a - 0.1, a + 0.1) for a in [
                    np.deg2rad(-45),
                    np.deg2rad(-110),
                    np.deg2rad(-120),
                    np.deg2rad(110),
                    np.deg2rad(45),
                    np.deg2rad(-160),
                    np.deg2rad(90),
                ]
            ]

            self.graph = {
                0: [1],
                1: [2],
                2: [3, 5],
                3: [4],
                4: [0],
                5: [6],
                6: [3],
            }
        elif goal_config == 'small_inner_graph':
            self.goals = [
                (-0.10888671875, 0.609375),
                (0.0498046875, 0.1943359375),
                (-0.072265625, -0.220703125),
                (-0.5849609375, 0.072265625),
                (-0.49951171875, 0.67041015625),
                (-0.2919921875, -0.5380859375),
                (-0.6826171875, -0.46484375),
                (0.2939453125, -0.0498046875),
                (0.2939453125, -0.4892578125),
                (0.025390625, -0.806640625),
                (0.28, 0.53),
            ]

            self.start_headings = [
                (a - 0.3, a + 0.3) for a in [
                    np.deg2rad(-45),
                    np.deg2rad(-90),
                    np.deg2rad(-120),
                    np.deg2rad(110),
                    np.deg2rad(45),
                    np.deg2rad(-160),
                    np.deg2rad(90),
                    np.deg2rad(-30),
                    np.deg2rad(-110),
                    np.deg2rad(120),
                    np.deg2rad(90),
                ]
            ]

            self.graph = {
                0: [1],
                1: [2, 7, 10],
                2: [3, 5],
                3: [4],
                4: [0],
                5: [6],
                6: [3],
                7: [8, 10],
                8: [9],
                9: [5],
                10: [0],
            }
        elif goal_config == 'small_inner_graph_reversed':
            self.goals = [
                (-0.10888671875, 0.609375),
                (0.0498046875, 0.1943359375),
                (-0.072265625, -0.220703125),
                (-0.5849609375, 0.072265625),
                (-0.49951171875, 0.67041015625),
                (-0.2919921875, -0.5380859375),
                (-0.6826171875, -0.46484375),
                (0.2939453125, -0.0498046875),
                (0.2939453125, -0.4892578125),
                (0.025390625, -0.806640625),
                (0.28, 0.53),
            ]

            self.start_headings = [
                (a - 0.3 + np.pi, a + 0.3 + np.pi) for a in [
                    np.deg2rad(-45),
                    np.deg2rad(-90),
                    np.deg2rad(-120),
                    np.deg2rad(110),
                    np.deg2rad(45),
                    np.deg2rad(-160),
                    np.deg2rad(90),
                    np.deg2rad(-30),
                    np.deg2rad(-110),
                    np.deg2rad(120),
                    np.deg2rad(90),
                ]
            ]

            self.graph = {
                0: [4, 10],
                1: [0],
                2: [1],
                3: [2, 6],
                4: [3],
                5: [2, 9],
                6: [5],
                7: [1],
                8: [7],
                9: [8],
                10: [1, 7],
            }
        else:
            raise ValueError(f"No such goal graph for {goal_config}")

        self.goal_reprs = []
        self.edge_reprs = []

        self.current_start_idx = 0
        self.current_goal_idx = 0

        self.goal_threshold = goal_threshold
        self.scale = scale
    
    @staticmethod
    def from_circle(origin : np.ndarray, radius : float, counterclockwise : bool, count : int, scale : float, goal_threshold = GOAL_THRESHOLD):
        """
        Generate a goal graph from a circle.
        """
        goals = []
        start_headings = []
        for i in range(count):
            angle = i * 2 * np.pi / count
            if not counterclockwise:
                angle = -angle
            pos = origin + radius * np.array([np.cos(angle), np.sin(angle)])
            goals.append(pos)
            start_heading = normalize_angle(angle + np.pi / 2 if counterclockwise else angle - np.pi / 2)
            start_headings.append((start_heading - 0.1, start_heading + 0.1))

        graph = {}
        for i in range(count):
            graph[i] = [(i + 1) % count]

        goal_graph = GoalGraph(scale,goal_threshold=goal_threshold)
        goal_graph.goals = goals
        goal_graph.start_headings = start_headings
        goal_graph.graph = graph
        return goal_graph
    

    @property
    def current_start(self):
        return np.array(self.goals[self.current_start_idx]) * self.scale, self.start_headings[self.current_start_idx]

    @property
    def current_goal(self):
        return np.array(self.goals[self.current_goal_idx]) * self.scale

    def is_complete(self):
        return self._ticks_at_current_goal > 0

    def set_goal(self, goal_idx, physics):
        """
        Set a new goal and update the renderables to match.
        """
        for idx, repr in enumerate(self.goal_reprs):
            # opacity = 1.0 if idx == goal_idx else 0.0
            opacity = 1 if idx < 7 else 0.0
            if physics:
                physics.bind(repr).rgba = (*repr.rgba[:3], opacity)
            else:
                repr.rgba = (*repr.rgba[:3], opacity)

        self.current_goal_idx = goal_idx
        self._ticks_at_current_goal = 0

    def tick(self, car_pos, physics):
        """
        Update the goal if the car was at the current goal for at least one tick.
        We need the delay so that the car can get the high reward for reaching
        the goal before the goal changes.
        """
        if self.is_complete():
            self.current_start_idx = self.current_goal_idx
            # self.set_goal(random.choice(self.graph[self.current_start_idx]), physics)
            self.set_goal(self.graph[self.current_start_idx][0], physics)
            return True

        if np.linalg.norm(np.array(car_pos)[:2] - self.current_goal) < self.goal_threshold:
            self._ticks_at_current_goal += 1
        else:
            self._ticks_at_current_goal = 0

        return False

    def reset(self, physics):
        # self.current_start_idx = random.randint(0, len(self.goals) - 1)
        self.current_start_idx = 0
        self.set_goal(random.choice(self.graph[self.current_start_idx]), physics)
        self._ticks_at_current_goal = 0

    def add_renderables(self, mjcf_root, height_lookup, show_edges=False):
        """
        Add renderables to the mjcf root to visualize the goals and (optionally) edges.
        """
        self.clear_renderables()
        RENDER_HEIGHT_OFFSET = 5.0

        self.goal_reprs = [
            mjcf_root.worldbody.add('site',
                                    type="sphere",
                                    size="0.1",
                                    rgba=(0.0, 1.0, 0.0, 0.5),
                                    group=0,
                                    pos=(g[0] * self.scale, g[1] * self.scale, height_lookup((g[0] * self.scale, g[1] * self.scale)) + RENDER_HEIGHT_OFFSET))
            for g in self.goals
        ]

        self.edge_reprs = [
            mjcf_root.worldbody.add('site',
                                    type="cylinder",
                                    size="0.05",
                                    rgba=(1, 1, 1, 0.5),
                                    fromto=(self.goals[s][0] * self.scale, self.goals[s][1] * self.scale, height_lookup((self.goals[s][0] * self.scale, self.goals[s][1] * self.scale)) + RENDER_HEIGHT_OFFSET,
                                            self.goals[g][0] * self.scale, self.goals[g][1] * self.scale, height_lookup((self.goals[g][0] * self.scale, self.goals[g][1] * self.scale)) + RENDER_HEIGHT_OFFSET))
            for s in self.graph for g in self.graph[s]
            if show_edges
        ]
    
    def clear_renderables(self):
        if self.goal_reprs is not None and len(self.goal_reprs) > 0:
            for r in self.goal_reprs:
                r.remove()
            self.goal_reprs = []
        
        if self.edge_reprs is not None and len(self.edge_reprs) > 0:
            for r in self.edge_reprs:
                r.remove()
            self.edge_reprs = []
        

def dfs_longest(edge_graph, start, visited = []):
    if start not in visited:
        visited.append(start)
    
    longest_e = []
    for e in edge_graph[start]:
        if e in visited:
            continue
        visited.append(e)

        c = dfs_longest(edge_graph, e, visited)
        if len(c) > len(longest_e):
            longest_e = c
    
    ret = [start]
    ret.extend(longest_e)
    return ret
    
class GoalRoute():
    def __init__(self, goal_graph : GoalGraph, start_idx : int = 0):
        self.route : typing.List = dfs_longest(goal_graph.graph,start_idx,[])
        self.graph = goal_graph.graph
        self.goal_threshold = goal_graph.goal_threshold
        self.goals = goal_graph.goals
        self.goal_start_headings = goal_graph.start_headings
        self.scale = goal_graph.scale
        self.goal_reprs = None
        self.edge_reprs = None
        self.route_reprs = None
        self.start_idx = 0
        self.current_line_segment_length = np.linalg.norm(self.current_start[0] - self.current_goal)
        self._compute_line_segment_lengths()

    def _compute_line_segment_lengths(self):
        self.line_segment_lengths = []
        for i in range(len(self.route) - 1):
            self.line_segment_lengths.append(
                np.linalg.norm(
                    np.array(self.goals[self.route[i]]) * self.scale - np.array(self.goals[self.route[i+1]]) * self.scale
                )
            )

    def reset(self, physics = None):
        self.set_current_start(0, physics)
    
    def set_current_start(self, start_idx, physics = None):
        if self.goal_reprs is not None:
            for idx, repr in enumerate(self.goal_reprs):
                is_start = idx == self.route[start_idx]
                is_goal = idx == self.route[start_idx + 1] if start_idx < len(self.route) - 1 else False

                if is_start:
                    rgba = (0.0, 0.0, 1.0, 1.0) #blue
                elif is_goal:
                    rgba = (1.0, 0.0, 0.0, 1.0) #red
                else:
                    rgba = (0.0, 1.0, 0.0, 0.5)


                if physics:
                    physics.bind(repr).rgba = rgba
                else:
                    repr.rgba = rgba
        if self.route_reprs is not None:
            for idx, repr in enumerate(self.route_reprs):
                is_current = idx == start_idx

                if is_current:
                    rgba = (1.0, 0.0, 0.0, 1.0) #solid red
                else:
                    rgba = (1.0, 1.0, 0.0, 0.5) #purple


                if physics:
                    physics.bind(repr).rgba = rgba
                else:
                    repr.rgba = rgba

        print("New Goal Route start idx: ", start_idx)
        self.start_idx = start_idx
        if(start_idx < len(self.route) - 1 and start_idx >= 0):
            self.current_line_segment_length = self.line_segment_lengths[start_idx]
            print("current line segment length: ", self.current_line_segment_length)
        else:
            self.current_line_segment_length = 0.0
        

    def add_renderables(self, mjcf_root, height_lookup, show_edges = False, render_height_offset = 5.0):
        self.clear_renderables()
        self.goal_reprs = [
             mjcf_root.worldbody.add('site',
                                    type="sphere",
                                    size="0.1",
                                    rgba=(0.0, 1.0, 0.0, 0.5), #green
                                    group=0,
                                    pos=(g[0] * self.scale, g[1] * self.scale, height_lookup((g[0] * self.scale, g[1] * self.scale)) + render_height_offset))
            for g in self.goals
        ]
        if show_edges:
            self.edge_reprs = []
            self.route_reprs = []
            for s in self.graph:
                try:
                    indexOfS = self.route.index(s)
                except ValueError:
                    indexOfS = -1

                for g in self.graph[s]:
                    
                    try:
                        indexOfG = self.route.index(g)
                    except ValueError:
                        indexOfG = -1
                    
                    if indexOfS == indexOfG - 1 and indexOfS != -1:
                        pass
                        """
                        In-route edges are added here in previous versions, but now they are added later in the code to align with the self.route layout
                        """
                    else:
                        self.edge_reprs.append(mjcf_root.worldbody.add(
                            'site',
                            type="cylinder",
                            size="0.04",
                            rgba=(1, 1, 1, 0.5), #white
                            fromto=(
                                self.goals[s][0] * self.scale, self.goals[s][1] * self.scale, height_lookup((self.goals[s][0] * self.scale, self.goals[s][1] * self.scale)) + RENDER_HEIGHT_OFFSET,
                                self.goals[g][0] * self.scale, self.goals[g][1] * self.scale, height_lookup((self.goals[g][0] * self.scale, self.goals[g][1] * self.scale)) + RENDER_HEIGHT_OFFSET
                            )
                        ))
            for s_idx,g_idx  in zip(self.route[:-1], self.route[1:]):
                s,g = self.goals[s_idx], self.goals[g_idx]
                self.route_reprs.append(mjcf_root.worldbody.add(
                    'site',
                    type="cylinder",
                    size="0.04",
                    rgba=(1, 1, 0, 0.5), #purple
                    fromto=(
                        s[0] * self.scale, s[1] * self.scale, height_lookup((s[0] * self.scale, s[1] * self.scale)) + render_height_offset,
                        g[0] * self.scale, g[1] * self.scale, height_lookup((g[0] * self.scale, g[1] * self.scale)) + render_height_offset
                    )
                ))

        self.set_current_start(self.start_idx)
    
    @property
    def route_length(self):
        return len(self.route) - 1

    def get_closest_point_on_current_segment(self, loc: np.ndarray) -> np.ndarray:
        assert loc.ndim == 1 and loc.shape[0] >= 2

        loc = loc[:2]
        line_normal = self.current_line_normal_vector
        current_start = self.current_start[0]
        current_goal = self.current_goal
        projected_point = loc + np.dot(self.current_start[0] - loc, line_normal) * line_normal
        
        dist_to_start = np.linalg.norm(projected_point - current_start)
        dist_to_goal = np.linalg.norm(projected_point - current_goal)
        line_segment_length = self.current_line_segment_length
        if dist_to_start > line_segment_length or dist_to_goal > line_segment_length:
            return current_start if dist_to_start < dist_to_goal else current_goal
        else:
            return projected_point
    
    @property
    def current_line_normal_vector(self) -> np.ndarray:
        # get the vector perpendicular to the current line segment
        # if self.is_finished:
        #     raise ValueError("Route finished, call reset() to make it valid again")
        
        start_vec = self.current_start[0]
        end_vec = self.current_goal
        unnormalized = np.array([end_vec[1] - start_vec[1], start_vec[0] - end_vec[0]])
        return unnormalized / np.linalg.norm(unnormalized)

    @property
    def is_finished(self) -> bool:
        return self.start_idx + 1 >= len(self.route) or self.start_idx < 0

    def current_progress_on_route(self, pos) -> float:
        if self.is_finished:
            return 1.0

        pos_on_line = self.get_closest_point_on_current_segment(pos)
        dist_to_start = np.linalg.norm(pos_on_line - self.current_start[0])
        total_dist_to_start = np.sum(self.line_segment_lengths[:self.start_idx])
        p = (total_dist_to_start + dist_to_start) / np.sum(self.line_segment_lengths)
        #assert 0 <= p <= 1
        #print(p)
        return p

    def tick(self, car_pos, physics):
        if self.is_finished:
            raise ValueError("Route finished, call reset() to make it valid again")
        if np.linalg.norm(np.array(car_pos)[:2] - self.current_goal) < self.goal_threshold:
            print("Segment " + str(self.start_idx+1) + "/" + str(len(self.route) - 1) + " finished")
            self.set_current_start(self.start_idx + 1, physics)
            if self.is_finished:
                print("Route Finished")
            return True

        return False

    @property
    def current_start(self) -> typing.Tuple[np.ndarray, typing.Tuple[float,float]]:
        return self.get_start_for_idx(self.start_idx), self.get_start_headings_for_idx(self.start_idx)

    @property
    def current_goal(self) -> np.ndarray:
        return self.get_goal_for_idx(self.start_idx)
    
    def get_goal_for_idx(self, idx) -> np.ndarray:
        if idx >= self.route_length:
            return np.array(self.goals[self.route[-1]]) * self.scale
        else:
            return np.array(self.goals[self.route[idx + 1]]) * self.scale

    def get_start_for_idx(self, idx) -> np.ndarray:
        if idx >= self.route_length:
            return np.array(self.goals[self.route[-2]]) * self.scale
        else:
            return np.array(self.goals[self.route[idx]]) * self.scale
    
    def get_start_headings_for_idx(self, idx) -> typing.Tuple[float,float]:
        if idx >= self.route_length:
            return self.get_start_headings_for_idx(self.route_length - 1)
        else:
            return self.goal_start_headings[idx]

    def clear_renderables(self):
        if self.goal_reprs is not None:
            for goal_repr in self.goal_reprs:
                goal_repr.remove()
        
        if self.edge_reprs is not None:
            for edge_repr in self.edge_reprs:
                edge_repr.remove()
        
        if self.route_reprs is not None:
            for route_repr in self.route_reprs:
                route_repr.remove()
rail_mujoco_walker/floors/heightfield_arena/goal_graph.py
realquantumcookie/APRL
from pathlib import Path
import numpy as np
import scipy
from PIL import Image
import os
import struct

from dm_control import composer
from dm_control.mujoco.wrapper.mjbindings import mjlib
from dm_control import mjcf

_FLOOR_ASSETS_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "resources", "floor_assets")
_ARENA_XML_PATH = os.path.join(_FLOOR_ASSETS_PATH, 'heightfield_arena.xml')
_ARENA_HEIGHTMAP_PATH = os.path.join(_FLOOR_ASSETS_PATH, 'textures', 'heightmap.bin')
_ARENA_TEXTURE_PATH = os.path.join(_FLOOR_ASSETS_PATH, 'textures', 'texture_baked_large.png')

class HeightFieldArena(composer.Arena):
    def __init__(self, scale, height_scale=5, name='heightfield_arena'):
        self.scale = scale
        self.height_scale = height_scale
        self.texture_path = _ARENA_TEXTURE_PATH
        self.heightmap_path = _ARENA_HEIGHTMAP_PATH

        # We can read in the heightmap and store it as a numpy array in self.img and self.img_size
        # with open(self.heightmap_path, 'rb') as f:
        #     self.img_size = struct.unpack('2I', f.read(8))
        #     num_pixels = self.img_size[0]*self.img_size[1]
        #     self.img = np.array(struct.unpack(f'{num_pixels}f', f.read(4*num_pixels))).reshape(self.img_size)
        #     self.img = np.flip(self.img, axis=1)
        #     self.img = (self.img - self.img.min()) / (self.img.max() - self.img.min()) * self.scale / 5
        #     print("Height field size", self.img_size)

        super(HeightFieldArena, self).__init__(name=name)

        self._mjcf_root.visual.scale.framelength = 10
        self._mjcf_root.visual.scale.framewidth = 5

    def _build(self, name):
        # Don't call super()._build here because we want our own XML file
        self._mjcf_root = mjcf.from_path(_ARENA_XML_PATH)
        if name:
            self._mjcf_root.model = name

        self._mjcf_root.option.gravity = (0, 0, -9.81)

        # =========================== Init heightfield and texturize ===========================

        self._hfield = self._mjcf_root.asset.add('hfield',
                                                 name='terrain',
                                                 file=f'{self.heightmap_path}',
                                                 size=(self.scale, self.scale, self.scale / self.height_scale, 0.1))

        self._texture = self._mjcf_root.asset.add('texture',
                                                  name='aesthetic_texture',
                                                  file=f'{self.texture_path}',
                                                  type='2d')

        self._material = self._mjcf_root.asset.add('material',
                                                   texrepeat=(1, 1),
                                                   name='aesthetic_material',
                                                   texture=self._texture,
                                                   texuniform='false')

        self._terrain_geom = self._mjcf_root.worldbody.add(
            'geom',
            name='terrain',
            type='hfield',
            rgba=(1.0, 1.0, 1.0, 1.0),
            pos=(0, 0, -0.01),
            hfield='terrain',
            material=self._material)


        # =========================== Add camera ===========================
        # self._mjcf_root.worldbody.add('camera',
        #                               mode="fixed",
        #                               pos="0 0 25",
        #                               euler="0 0 0")


        # =========================== Add walls to limit robot movement ===========================
        self.walls = []

        self.walls.append(
            self._mjcf_root.worldbody.add('geom',
                                          type="box",
                                          contype="1",
                                          conaffinity="1",
                                          size=(0.1, self.scale + 1, 0.7),
                                          rgba=(0.3, 0.3, 0.9, 0.5),
                                          pos=(self.scale, 0, 0)))

        self.walls.append(
            self._mjcf_root.worldbody.add('geom',
                                          type="box",
                                          contype="1",
                                          conaffinity="1",
                                          size=(0.1, self.scale + 1, 0.7),
                                          rgba=(0.3, 0.3, 0.9, 0.5),
                                          pos=(-self.scale, 0, 0)))

        self.walls.append(
            self._mjcf_root.worldbody.add('geom',
                                          type="box",
                                          contype="1",
                                          conaffinity="1",
                                          size=(self.scale + 1, 0.1, 0.7),
                                          rgba=(0.3, 0.3, 0.9, 0.5),
                                          pos=(0, self.scale, 0)))

        self.walls.append(
            self._mjcf_root.worldbody.add('geom',
                                          type="box",
                                          contype="1",
                                          conaffinity="1",
                                          size=(self.scale + 1, 0.1, 0.7),
                                          rgba=(0.3, 0.3, 0.9, 0.5),
                                          pos=(0, -self.scale, 0)))

    def height_lookup(self, pos):
        """Returns the height of the terrain at the given position."""
        return 0.15

    def in_bounds(self, pos):
        eps = 1.0
        return np.abs(pos).max() < self.scale + eps


rail_mujoco_walker/floors/heightfield_arena/heightfield_arena.py
realquantumcookie/APRL
import typing
import numpy as np
import random

HEIGHTFIELD_ARENA_GOALS = {
    "full": {
        "points": [
            (np.array([0.05, 0.05]), (0.0, 2 * np.pi)),
            (np.array([0.33, -0.15]), (0.0, 2 * np.pi)),
            (np.array([0.3, -0.45]), (0.0, 2 * np.pi)),
            (np.array([0.65, -0.1]), (0.0, 2 * np.pi)),
            (np.array([0.525, 0.2]), (0.0, 2 * np.pi)),
            (np.array([0.6, 0.75]), (0.0, 2 * np.pi)),
            (np.array([0.3, 0.6]), (0.0, 2 * np.pi)),
            (np.array([-0.2, 0.6]), (0.0, 2 * np.pi)),
            (np.array([-0.5, 0.57]), (0.0, 2 * np.pi)),
            (np.array([-0.6, 0.05]), (0.0, 2 * np.pi)),
            (np.array([-0.7, -0.6]), (0.0, 2 * np.pi)),
            (np.array([0.0, -0.8]), (0.0, 2 * np.pi)),
        ],
        "connections": {
            0: [1, 6, 7, 9, 10, 11],
            1: [0, 2, 3, 4],
            2: [1, 3, 11],
            3: [1, 2, 4],
            4: [1, 3, 6],
            5: [4, 6],
            6: [0, 4, 5, 7],
            7: [0, 6, 8],
            8: [7, 9],
            9: [0, 8, 10],
            10: [0, 9, 11],
            11: [0, 2, 10],
        }
    },
    "edges": {
        "points": [
            (np.array([0.9, 0.0]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            (np.array([0.0, 0.9]), (np.deg2rad(180) - 0.1, np.deg2rad(180) + 0.1)),
            (np.array([-0.9, 0.0]), (np.deg2rad(-90) - 0.1, np.deg2rad(-90) + 0.1)),
            (np.array([0.0, -0.9]), (np.deg2rad(0) - 0.1, np.deg2rad(0) + 0.1)),
        ],
        "connections": {
            0: [1],
            1: [2],
            2: [3],
            3: [0],
        }
    },
    "corners": {
        "points": [
            (np.array([0.9, 0.9]), (np.deg2rad(135) - 0.1, np.deg2rad(135) + 0.1)),
            (np.array([-0.9, 0.9]), (np.deg2rad(-135) - 0.1, np.deg2rad(-135) + 0.1)),
            (np.array([-0.9, -0.9]), (np.deg2rad(-45) - 0.1, np.deg2rad(-45) + 0.1)),
            (np.array([0.9, -0.9]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)),
        ],
        "connections": {
            0: [1],
            1: [2],
            2: [3],
            3: [0],
        }
    },
    "ring_dense": {
        "points": [
            (np.array([0.9, 0.9]), (np.deg2rad(135) - 0.1, np.deg2rad(135) + 0.1)),
            (np.array([0.0, 0.9]), (np.deg2rad(-180) - 0.1, np.deg2rad(-180) + 0.1)),
            (np.array([-0.9, 0.9]), (np.deg2rad(-135) - 0.1, np.deg2rad(-135) + 0.1)),
            (np.array([-0.9, 0.0]), (np.deg2rad(-90) - 0.1, np.deg2rad(-90) + 0.1)),
            (np.array([-0.9, -0.9]), (np.deg2rad(-45) - 0.1, np.deg2rad(-45) + 0.1)),
            (np.array([0.0, -0.9]), (np.deg2rad(0) - 0.1, np.deg2rad(0) + 0.1)),
            (np.array([0.9, -0.9]), (np.deg2rad(-45) - 0.1, np.deg2rad(-45) + 0.1)),
            (np.array([0.9, 0.0]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
        ],
        "connections": {
            0: [1, 2],
            1: [2, 3],
            2: [3, 4],
            3: [4, 5],
            4: [5, 6],
            5: [6, 7],
            6: [7, 0],
            7: [0, 1],
        }
    },
    "ring_small_inner": {
        "points": [
            (np.array([-0.10888671875, 0.609375]), (np.deg2rad(-45) - 0.1, np.deg2rad(-45) + 0.1)),
            (np.array([0.07421875, 0.072265625]), (np.deg2rad(-110) - 0.1, np.deg2rad(-110) + 0.1)),
            (np.array([0.0, 0.0]), (np.deg2rad(-110) - 0.1, np.deg2rad(-110) + 0.1)),
            (np.array([-.1, -0.25]), (np.deg2rad(-120) - 0.1, np.deg2rad(-120) + 0.1)),
            (np.array([-0.45, 0.1]), (np.deg2rad(110) - 0.1, np.deg2rad(110) + 0.1)),
            (np.array([-0.5849609375, 0.072265625]), (np.deg2rad(110) - 0.1, np.deg2rad(110) + 0.1)),
            (np.array([-0.49951171875, 0.67041015625]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)),
        ],
        "connections": {
            0: [1],
            1: [2],
            2: [3],
            3: [4],
            4: [5],
            5: [6],
            6: [0],
        }
    },
    "small_inner_1fork": {
        "points": [
            (np.array([-0.10888671875, 0.609375]), (np.deg2rad(-45) - 0.1, np.deg2rad(-45) + 0.1)),
            (np.array([0.07421875, 0.072265625]), (np.deg2rad(-110) - 0.1, np.deg2rad(-110) + 0.1)),
            (np.array([-0.072265625, -0.220703125]), (np.deg2rad(-120) - 0.1, np.deg2rad(-120) + 0.1)),
            (np.array([-0.5849609375, 0.072265625]), (np.deg2rad(110) - 0.1, np.deg2rad(110) + 0.1)),
            (np.array([-0.49951171875, 0.67041015625]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)),
            (np.array([-0.2919921875, -0.5380859375]), (np.deg2rad(-160) - 0.1, np.deg2rad(-160) + 0.1)),
            (np.array([-0.6826171875, -0.46484375]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
        ],
        "connections": {
            0: [1],
            1: [2],
            2: [3, 5],
            3: [4],
            4: [0],
            5: [6],
            6: [3],
        }
    },
    "hike": {
        "points": [
            (np.array([-0.10888671875, 0.609375]), (np.deg2rad(-45) - 0.1, np.deg2rad(-45) + 0.1)),
            (np.array([0.07421875, 0.072265625]), (np.deg2rad(-110) - 0.1, np.deg2rad(-110) + 0.1)),
            (np.array([-0.5849609375, 0.072265625]), (np.deg2rad(110) - 0.1, np.deg2rad(110) + 0.1)),
            (np.array([-0.49951171875, 0.67041015625]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)),
            # (np.array([-0.45, -0.35]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            (np.array([-.35, -.05]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            # (np.array([-.3, 0.3]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
        ],
        "connections": {
            0: [1],
            1: [4],
            # 5: [4],
            4: [2],
            2: [3],
            3: [0],

        }
    },
    "dense_inner": {
        "points": [
            # our original ring
            (np.array([-0.1, 0.58]), (np.deg2rad(-45) - 0.1, np.deg2rad(-45) + 0.1)), # 0
            (np.array([0.1, 0.072265625]), (np.deg2rad(-110) - 0.1, np.deg2rad(-110) + 0.1)), # 1
            (np.array([0.0, 0.0]), (np.deg2rad(-110) - 0.1, np.deg2rad(-110) + 0.1)), # 2
            (np.array([-.1, -0.25]), (np.deg2rad(-120) - 0.1, np.deg2rad(-120) + 0.1)), # 3
            (np.array([-0.45, 0.1]), (np.deg2rad(110) - 0.1, np.deg2rad(110) + 0.1)),  # 4  
            (np.array([-0.57, 0.05]), (np.deg2rad(110) - 0.1, np.deg2rad(110) + 0.1)), # 5
            (np.array([-0.52, 0.62]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)), # 6
            (np.array([-0.45, 0.7]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)), # 7
        ],
        "connections": {
            0: [1, 7],
            1: [0, 2],
            2: [1, 3],
            3: [2, 4],
            4: [5, 3],
            5: [6, 4],
            6: [5, 7],
            7: [6, 0],
        }
    },
    "full_dense": {
        "points": [
            # our original ring
            (np.array([-0.1, 0.58]), (np.deg2rad(-45) - 0.1, np.deg2rad(-45) + 0.1)),
            (np.array([0.1, 0.072265625]), (np.deg2rad(-110) - 0.1, np.deg2rad(-110) + 0.1)),
            (np.array([0.0, 0.0]), (np.deg2rad(-110) - 0.1, np.deg2rad(-110) + 0.1)),
            (np.array([-.1, -0.25]), (np.deg2rad(-120) - 0.1, np.deg2rad(-120) + 0.1)),
            (np.array([-0.3, -.05]), (np.deg2rad(110) - 0.1, np.deg2rad(110) + 0.1)), # 4
            (np.array([-0.57, 0.05]), (np.deg2rad(110) - 0.1, np.deg2rad(110) + 0.1)), # 5
            (np.array([-0.52, 0.62]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)), # 6
            # (np.array([-0.45, 0.7]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)), # 
            # now expand the left bottom corner
            (np.array([-0.7, -0.6]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)),
            (np.array([-.3, -.55]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)),
            # from bottom left corner back to center
            (np.array([0.1, -0.85]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            (np.array([0.3, -0.5]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            (np.array([0.325, -0.07]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            (np.array([0.17, -0.05]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            (np.array([0.27, 0.57]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            (np.array([0.1, 0.75]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            (np.array([0.6, 0.05]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            (np.array([0.5, 0]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            (np.array([0.6, -.3]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            (np.array([-0.45, 0.7]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)), # 7
        ],
        "connections": {
            0: [1, 18, 14],
            1: [0, 2, 12],
            2: [1, 3, 4],
            3: [2, 4, 8],
            4: [5, 3, 2, 7],
            5: [6, 7],
            6: [18, 5],
            7: [5, 8],
            8: [3, 7, 9],
            9: [8, 10],
            10: [9, 11, 17],
            11: [10, 12, 16],
            12: [11, 1, 13],
            13: [14, 12, 15],
            14: [13, 0],
            15: [16, 13],
            16: [15, 11, 17],
            17: [16, 10],
            18: [6, 0],
        }
    },
    # "hike": {
    #     "points": [
    #         (np.array([-0.10888671875, 0.609375]), (np.deg2rad(-45) - 0.1, np.deg2rad(-45) + 0.1)),
    #         (np.array([0.07421875, 0.072265625]), (np.deg2rad(-110) - 0.1, np.deg2rad(-110) + 0.1)),
    #         (np.array([-0.5849609375, 0.072265625]), (np.deg2rad(110) - 0.1, np.deg2rad(110) + 0.1)),
    #         (np.array([-0.49951171875, 0.67041015625]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)),
    #         (np.array([-0.45, -0.35]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
    #         (np.array([-.35, -.05]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
    #         # (np.array([-.3, 0.3]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
    #     ],
    #     "connections": {
    #         0: [1],
    #         1: [5],
    #         5: [4],
    #         4: [2],
    #         2: [3],
    #         3: [0],

    #     }
    # },
    "small_inner_graph": {
        "points": [
            (np.array([-0.10888671875, 0.609375]), (np.deg2rad(-45) - 0.1, np.deg2rad(-45) + 0.1)),
            (np.array([0.0498046875, 0.1943359375]), (np.deg2rad(-90) - 0.1, np.deg2rad(-90) + 0.1)),
            (np.array([-0.072265625, -0.220703125]), (np.deg2rad(-120) - 0.1, np.deg2rad(-120) + 0.1)),
            (np.array([-0.5849609375, 0.072265625]), (np.deg2rad(110) - 0.1, np.deg2rad(110) + 0.1)),
            (np.array([-0.49951171875, 0.67041015625]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)),
            (np.array([-0.2919921875, -0.5380859375]), (np.deg2rad(-160) - 0.1, np.deg2rad(-160) + 0.1)),
            (np.array([-0.6826171875, -0.46484375]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            (np.array([0.2939453125, -0.0498046875]), (np.deg2rad(-30) - 0.1, np.deg2rad(-30) + 0.1)),
            (np.array([0.2939453125, -0.4892578125]), (np.deg2rad(-110) - 0.1, np.deg2rad(-110) + 0.1)),
            (np.array([0.025390625, -0.4892578125]), (np.deg2rad(120) - 0.1, np.deg2rad(120) + 0.1)),
            (np.array([0.28, 0.53]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
        ],
        "connections": {
            0: [1],
            1: [2, 7, 10],
            2: [3, 5],
            3: [4],
            4: [0],
            5: [6],
            6: [3],
            7: [8, 10],
            8: [9],
            9: [5],
            10: [0],
        },
    },
    "small_inner_graph_reversed": {
        "points": [
            (np.array([-0.10888671875, 0.609375]), (np.deg2rad(-45) - 0.1, np.deg2rad(-45) + 0.1)),
            (np.array([0.0498046875, 0.1943359375]), (np.deg2rad(-90) - 0.1, np.deg2rad(-90) + 0.1)),
            (np.array([-0.072265625, -0.220703125]), (np.deg2rad(-120) - 0.1, np.deg2rad(-120) + 0.1)),
            (np.array([-0.5849609375, 0.072265625]), (np.deg2rad(110) - 0.1, np.deg2rad(110) + 0.1)),
            (np.array([-0.49951171875, 0.67041015625]), (np.deg2rad(45) - 0.1, np.deg2rad(45) + 0.1)),
            (np.array([-0.2919921875, -0.5380859375]), (np.deg2rad(-160) - 0.1, np.deg2rad(-160) + 0.1)),
            (np.array([-0.6826171875, -0.46484375]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
            (np.array([0.2939453125, -0.0498046875]), (np.deg2rad(-30) - 0.1, np.deg2rad(-30) + 0.1)),
            (np.array([0.2939453125, -0.4892578125]), (np.deg2rad(-110) - 0.1, np.deg2rad(-110) + 0.1)),
            (np.array([0.025390625, -0.4892578125]), (np.deg2rad(120) - 0.1, np.deg2rad(120) + 0.1)),
            (np.array([0.28, 0.53]), (np.deg2rad(90) - 0.1, np.deg2rad(90) + 0.1)),
        ],
        "connections": {
            0: [4, 10],
            1: [0],
            2: [1],
            3: [2, 6],
            4: [3],
            5: [2, 9],
            6: [5],
            7: [1],
            8: [7],
            9: [8],
            10: [1, 7],
        }
    }
}


rail_mujoco_walker/floors/heightfield_arena/heightfield_arena_goals.py
realquantumcookie/APRL
from .sim_robot import DMWalkerForRailSimWalkerObservables, RailSimWalkerDMControl, DMWalkerForRailSimWalker
from .sim_a1 import A1SimWalker
from .sim_go1 import Go1SimWalker
rail_mujoco_walker/robots/__init__.py
realquantumcookie/APRL
import os

import numpy as np
from .sim_robot import RailSimWalkerDMControl, SIM_ASSET_DIR
from functools import cached_property

_A1_XML_PATH = os.path.join(SIM_ASSET_DIR, 'robot_assets', 'a1', 'xml', 'a1.xml')

"""
A1 simulation robot using DM_Control suite
Joint limits are based on https://github.com/unitreerobotics/unitree_legged_sdk/blob/v3.3.1/include/unitree_legged_sdk/a1_const.h
They are also cross-checked with XML definition file for a1
"""
class A1SimWalker(RailSimWalkerDMControl):
    # _INIT_QPOS = np.asarray([0.0, 0.9, -1.8] * 4)
    # _QPOS_OFFSET = np.asarray([0.2, 0.4, 0.4] * 4)
    # _INIT_QPOS = np.asarray([-0.1, 0.5, -1.4] * 4)

    def __init__(self, Kp=60, Kd=6, action_interpolation : bool = True, limit_action_range : float = 1.0, *args, **kwargs):
        RailSimWalkerDMControl.__init__(
            self,
            XML_FILE = _A1_XML_PATH,
            Kp = Kp,
            Kd = Kd,
            action_interpolation = action_interpolation,
            limit_action_range=limit_action_range,
            *args,
            **kwargs
        )

    @cached_property
    def joint_qpos_init(self) -> np.ndarray:
        return np.asarray([0.0 / 180 * np.pi, 45.0/180*np.pi, -90.0 / 180 * np.pi] * 4)
    
    @cached_property
    def joint_qpos_sitting(self) -> np.ndarray:
        return np.array([0.0 / 180 * np.pi, 70.0/180*np.pi, -150.0 / 180 * np.pi] * 4)

    @cached_property
    def joint_qpos_offset(self) -> np.ndarray:
        return np.array([0.2, 0.4, 0.4] * 4)

    @cached_property
    def joint_qpos_mins(self) -> np.ndarray:
        return np.asarray([-0.802, -1.05, -2.7] * 4)

    @cached_property
    def joint_qpos_maxs(self) -> np.ndarray:
        return np.asarray([0.802, 4.19, -0.916] * 4)
rail_mujoco_walker/robots/sim_a1.py
realquantumcookie/APRL
import os

import numpy as np
from .sim_robot import RailSimWalkerDMControl, SIM_ASSET_DIR
from functools import cached_property

_Go1_XML_PATH = os.path.join(SIM_ASSET_DIR, 'robot_assets', 'go1', 'xml', 'go1.xml')

"""
Go1 simulation robot using DM_Control suite
Joint limits are based on https://github.com/unitreerobotics/unitree_legged_sdk/blob/go1/include/unitree_legged_sdk/go1_const.h
They are also cross-checked with XML definition file for go1 (Credit to Kevin for exporting the XML file from Unitree_ROS)
The XML file in unitree_mujoco is actually wrong, it is a copy of a1.xml
"""
class Go1SimWalker(RailSimWalkerDMControl):
    # _INIT_QPOS = np.asarray([0.0, 0.9, -1.8] * 4)
    # _QPOS_OFFSET = np.asarray([0.2, 0.4, 0.4] * 4)

    def __init__(self, Kp=60, Kd=6, action_interpolation : bool = True, limit_action_range : float = 1.0, *args, **kwargs):
        RailSimWalkerDMControl.__init__(
            self,
            XML_FILE = _Go1_XML_PATH,
            Kp = Kp,
            Kd = Kd,
            action_interpolation = action_interpolation,
            limit_action_range = limit_action_range,
            *args,
            **kwargs
        )
        RailSimWalkerDMControl.control_timestep = 0.05
        RailSimWalkerDMControl.control_subtimestep = 0.002

    @cached_property
    def joint_qpos_init(self) -> np.ndarray:
        # shift_delta = np.array([-0.1, -0.1, 0.1])
        # real_shift = shift_delta.repeat(4).reshape((3,4)).T.flatten()
        # real_shift[3::6] = -real_shift[3::6]
        real_shift = np.array([
            -0.1, -0.2, 0.0,
            0.1, -0.2, 0.0,
            -0.1, 0.0, 0.0,
            0.1, 0.0, 0.0
        ])
        return np.array([0.0 / 180 * np.pi, 45.0/180*np.pi, -90.0 / 180 * np.pi] * 4) + real_shift

    @cached_property
    def joint_qpos_sitting(self) -> np.ndarray:
        return np.array([0.0 / 180 * np.pi, 70.0/180*np.pi, -150.0 / 180 * np.pi] * 4)
    
    @cached_property
    def joint_qpos_offset(self) -> np.ndarray:
        return np.array([0.2, 0.4, 0.4] * 4)

    @cached_property
    def joint_qpos_mins(self) -> np.ndarray:
        return np.asarray([-1.047, -0.663, -2.721] * 4)

    @cached_property
    def joint_qpos_maxs(self) -> np.ndarray:
        return np.asarray([1.047, 2.966, -0.837] * 4)


rail_mujoco_walker/robots/sim_go1.py
realquantumcookie/APRL
from functools import cached_property
import typing

import numpy as np
from dm_control import composer, mjcf, mujoco
from dm_control.composer.observation import observable
from dm_control.locomotion.walkers import base
from dm_control.utils.transformations import quat_to_euler
from dm_env import specs
from rail_walker_interface import BaseWalkerInSim, BaseWalker, BaseWalkerWithFootContact
from dm_control.mjcf import Physics, Element
import transforms3d as tr3d
import os
from ..utils import find_dm_control_non_contacting_height, settle_physics
from collections import deque

SIM_ASSET_DIR = os.path.join(os.path.dirname(__file__), '../resources')

"""
BaseRobotObservables is a class that defines the observables of the robot.
Note that the observables defined in the inherited class will be included in the observables of the robot.
"""
class DMWalkerForRailSimWalkerObservables(composer.Observables): #(base.WalkerObservables):
    def __init__(self, entity : "DMWalkerForRailSimWalker", foot_contact_threshold : np.ndarray, delayed_substep : int = 20):
        self._delayed_substep = delayed_substep
        super().__init__(entity)
        self._entity = entity
        self._foot_contact_threshold = foot_contact_threshold
    
    @composer.observable
    def joints_pos(self):
        #return observable.MJCFFeature('qpos', self._entity.observable_joints)
        return self._entity.get_delayed_observable(
            "joints_pos",
            delay_substeps=self._delayed_substep,
            lambda_fn=lambda physics: physics.bind(self._entity.observable_joints).qpos
        )

    @composer.observable
    def joints_vel(self):
        #return observable.MJCFFeature('qvel', self._entity.observable_joints)
        return self._entity.get_delayed_observable(
            "joints_vel",
            delay_substeps=self._delayed_substep,
            lambda_fn=lambda physics: physics.bind(self._entity.observable_joints).qvel
        )

    @composer.observable
    def sensors_gyro(self):
        #return observable.MJCFFeature('sensordata',self._entity.mjcf_model.sensor.gyro)
        return self._entity.get_delayed_observable(
            "sensors_gyro",
            delay_substeps=self._delayed_substep,
            lambda_fn=lambda physics: physics.bind(self._entity.mjcf_model.sensor.gyro).sensordata
        )

    @composer.observable
    def sensors_framequat(self):
        # return observable.MJCFFeature('sensordata',self._entity.mjcf_model.sensor.framequat)
        return self._entity.get_delayed_observable(
            "sensors_framequat",
            delay_substeps=self._delayed_substep,
            lambda_fn=lambda physics: physics.bind(self._entity.mjcf_model.sensor.framequat).sensordata
        )

    @composer.observable 
    def torques(self):
        # return observable.MJCFFeature('force', self._entity.actuators)
        return self._entity.get_delayed_observable(
            "torques",
            delay_substeps=self._delayed_substep,
            lambda_fn=lambda physics: physics.bind(self._entity.actuators).force
        )
    
    @composer.observable
    def foot_forces(self):
        #return observable.MJCFFeature('sensordata', self._entity.mjcf_model.sensor.touch)
        return self._entity.get_delayed_observable(
            "foot_forces",
            delay_substeps=self._delayed_substep,
            lambda_fn=lambda physics: physics.bind(self._entity.mjcf_model.sensor.touch).sensordata
        )

    @composer.observable
    def foot_forces_normalized(self):
        # return observable.Generic(
        #     lambda physics: (physics.bind(self._entity.mjcf_model.sensor.touch).sensordata / foot_force_max_values).astype(np.float32)
        # )
        return self._entity.get_delayed_observable(
            "foot_forces_normalized",
            delay_substeps=self._delayed_substep,
            lambda_fn=lambda physics: (physics.bind(self._entity.mjcf_model.sensor.touch).sensordata / 50.0).astype(np.float32)
        )

    @composer.observable
    def foot_forces_normalized_masked(self):
        foot_force_max_values = np.array([100.0, 100.0, 60.0, 60.0])
        def lambda_func(physics):
                foot_forces = physics.bind(self._entity.mjcf_model.sensor.touch).sensordata
                normalized = np.array(foot_forces / foot_force_max_values, dtype=np.float32)
                normalized[-1] = 0.0
                return normalized
        # return observable.Generic(lambda_func)
        return self._entity.get_delayed_observable(
            "foot_forces_normalized_masked",
            delay_substeps=self._delayed_substep,
            lambda_fn=lambda_func
        )

    @composer.observable
    def foot_contacts(self):
        # return observable.Generic(
        #     lambda physics: (physics.bind(self._entity.mjcf_model.sensor.touch).sensordata >= self._foot_contact_threshold).astype(np.float32)
        # )
        return self._entity.get_delayed_observable(
            "foot_contacts",
            delay_substeps=self._delayed_substep,
            lambda_fn=lambda physics: (physics.bind(self._entity.mjcf_model.sensor.touch).sensordata >= self._foot_contact_threshold).astype(np.float32)
        )

    @composer.observable
    def sensors_local_velocimeter(self):
        # return observable.Generic(
        #     lambda physics: self._entity.get_velocity(physics)
        # )
        return self._entity.get_delayed_observable(
            "sensors_local_velocimeter",
            delay_substeps=self._delayed_substep,
            lambda_fn=lambda physics: self._entity.get_velocity(physics)
        )

    @composer.observable
    def sensors_local_velocimeter_x(self):
        return self._entity.get_delayed_observable(
            "sensors_local_velocimeter_x",
            delay_substeps=self._delayed_substep,
            lambda_fn=lambda physics: self._entity.get_velocity(physics)[0:1]
        )

    @composer.observable
    def imu(self):
        # return observable.Generic(lambda physics: self._entity.get_imu(physics))
        return self._entity.get_delayed_observable(
            "imu",
            delay_substeps=self._delayed_substep,
            lambda_fn=lambda physics: self._entity.get_imu(physics)
        )
    
    @composer.observable
    def sensors_accelerometer(self):
        # return observable.MJCFFeature('sensordata', self._entity.mjcf_model.sensor.accelerometer)
        return self._entity.get_delayed_observable(
            "sensors_accelerometer",
            delay_substeps=self._delayed_substep,
            lambda_fn=lambda physics: physics.bind(self._entity.mjcf_model.sensor.accelerometer).sensordata
        )

    # @composer.observable
    # def body_height(self):
    #     return observable.MJCFFeature('xpos', self._entity.root_body)[2]
    
    # @composer.observable
    # def body_position(self):
    #     return observable.MJCFFeature('xpos', self._entity.root_body)

    # Semantic groupings of Walker observables.
    def _collect_from_attachments(self, attribute_name):
        out = []
        for entity in self._entity.iter_entities(exclude_self=True):
            out.extend(getattr(entity.observables, attribute_name, []))
        return out

    @property
    def kinematic_sensors(self):
        return (
            [
                self.sensors_gyro,
                self.sensors_accelerometer,
                self.sensors_local_velocimeter,
                self.sensors_framequat
            ] +
            self._collect_from_attachments('kinematic_sensors')
        )

    @property
    def dynamic_sensors(self):
        return self._collect_from_attachments('dynamic_sensors')

    @property
    def proprioception(self):
        return ([self.joints_pos, self.joints_vel, self.torques] +
                self._collect_from_attachments('proprioception'))


class DMWalkerForRailSimWalker(base.Walker):
    # For some reason mujoco failed to read the <default/motor> flag in XML file, hard-coding this into the code.
    DEFAULT_CTRL_LIMIT_MIN = -33.5
    DEFAULT_CTRL_LIMIT_MAX = 33.5

    def __init__(self, XML_FILE : str, action_interpolation : bool, power_protect_factor : float = 0.1, *args, **kwargs):
        assert power_protect_factor > 0.0 and power_protect_factor <= 1.0
        self._XML_FILE = XML_FILE
        self.target_action : typing.Optional[np.ndarray] = None
        self.prev_action : typing.Optional[np.ndarray] = None
        self.last_obs : typing.Dict[str,typing.Any] = None
        self._control_timestep = 0.05
        self._control_subtimestep = 0.002
        self._substep_count = 0
        self.action_interpolation = action_interpolation
        self.power_protect_factor = power_protect_factor
        self._last_physics : typing.Optional[Physics] = None
        super().__init__(*args,**kwargs)

    def _build(
            self,
            name: typing.Optional[str] = "robot",
            kp: float = 60,
            kd: float = 6,
            foot_contact_threshold = np.array([10,10,10,10]),
            delay_substeps : int = 10,
        ):
        self._mjcf_root = mjcf.from_path(self._XML_FILE)
        if name:
            self._mjcf_root.model = name
        # Find MJCF elements that will be exposed as attributes.
        self._root_body = self._mjcf_root.find('body', 'trunk')
        self._root_body.pos[-1] = 0.125

        self._foot_sites = [
            self._mjcf_root.find('site', f'{joint_name}_foot_site')
            for joint_name in ["FR", "FL", "RR", "RL"]
        ]

        self._joints : list[Element] = [] #self.mjc_model.find_all('joint')

        self._actuators : list[Element] = self.mjcf_model.find_all('actuator')

        for actuator in self._actuators:
            self._joints.append(actuator.joint)

        self.feet_geom_names = ["RL", "RR", "FL", "FR"]
        # self.feet_geoms = [self.root_body.find("geom", name) for name in feet_geom_names]
        self._foot_friction = 0.8

        self.kp = kp
        self.kd = kd

        self.last_obs = {
            "3d_location": np.zeros(3),
            "3d_linear_velocity": np.zeros(3),
            "3d_local_velocity": np.zeros(3),
            "angular_velocity": np.zeros(3),
            "framequat_wijk": tr3d.euler.euler2quat(0,0,0),
            "roll_pitch_yaw": np.zeros(3),
            "3d_acceleration_local": np.zeros(3),
            "joint_pos": np.zeros(len(self._joints)),
            "joint_vel": np.zeros(len(self._joints)),
            "joint_acc" : np.zeros(len(self._joints)),
            "foot_force": np.zeros(4),
            "torques": np.zeros(len(self._actuators)),
        }
        self._delayed_observable_lambdas : typing.Dict[str, typing.Callable[[Physics], typing.Any]] = {}
        self._delayed_observable_queues : typing.Dict[str, deque] = {}
        self._last_step_position = None
        self._built_observable = DMWalkerForRailSimWalkerObservables(self, foot_contact_threshold, delay_substeps)
        self._delay_substeps = delay_substeps
        delayed_joint_acc_observable = self.get_delayed_observable(
            'joint_acc',
            delay_substeps,
            lambda physics : physics.bind(self.observable_joints).qacc
        )

    def get_delayed_observable(self, name : str, delay_substeps : int, lambda_fn : typing.Callable[[Physics], typing.Any]) -> observable.Generic:
        self._delayed_observable_lambdas[name] = lambda_fn
        self._delayed_observable_queues[name] = deque(maxlen=max(delay_substeps,1))
        def get_obs(physics : Physics):
            if len(self._delayed_observable_queues[name]) == 0:
                return self._delayed_observable_lambdas[name](physics)
            else:
                return self._delayed_observable_queues[name][0]
        return observable.Generic(get_obs)
    
    def read_delayed_observables(self, name: str):
        if name not in self._delayed_observable_queues:
            raise ValueError(f"Observable {name} not found.")
        if len(self._delayed_observable_queues[name]) == 0:
            return self._delayed_observable_lambdas[name](self._last_physics)
        else:
            return self._delayed_observable_queues[name][0]

    def action_spec(self, physics : Physics):
        minimum = []
        maximum = []
        for actuator in self.actuators:
            joint = actuator.joint

            joint_range = physics.bind(joint).range
            minimum.append(joint_range[0])
            maximum.append(joint_range[1])

        return specs.BoundedArray(
            shape=(len(minimum),),
            dtype=np.float32,
            minimum=minimum,
            maximum=maximum
            #name='\t'.join([actuator.name for actuator in self.actuators])
        )

    @cached_property
    def ctrllimits(self):
        minimum = []
        maximum = []
        for actuator in self.actuators:
            if actuator.ctrllimited == "true" or actuator.ctrllimited:
                minimum.append(actuator.ctrlrange[0])
                maximum.append(actuator.ctrlrange[1])
            elif actuator.ctrllimited is None:
                minimum.append(self.DEFAULT_CTRL_LIMIT_MIN)
                maximum.append(self.DEFAULT_CTRL_LIMIT_MAX)
            else:
                # no limit
                minimum.append(-np.inf)
                maximum.append(np.inf)

        return minimum, maximum

    def apply_action(self, physics : Physics, desired_qpos : np.ndarray, random_state : np.random.RandomState):
        joints_bind = physics.bind(self.joints)
        qpos = joints_bind.qpos
        qvel = joints_bind.qvel

        action = self.kp * (desired_qpos - qpos) - self.kd * qvel
        minimum, maximum = self.ctrllimits
        minimum = np.asarray(minimum)
        maximum = np.asarray(maximum)

        minimum *= self.power_protect_factor
        maximum *= self.power_protect_factor
        action = np.clip(action, minimum, maximum)
        physics.bind(self.actuators).ctrl = action

    def set_target_action(self, target_action : np.ndarray):
        # print(f"New action {target_action} set, last substep count: ", self._substep_count)
        self.prev_action = self.target_action
        self.target_action = target_action
        self._substep_count = 0

    def _build_observables(self):
        return self._built_observable
        
    def get_imu(self, physics):
        quat = physics.bind(self.mjcf_model.sensor.framequat).sensordata
        roll, pitch, yaw = quat_to_euler(quat)

        gyro = physics.bind(self.mjcf_model.sensor.gyro).sensordata
        dr, dp, dy = gyro

        return np.array([roll, pitch, dr, dp])

    @property
    def root_body(self):
        return self._root_body

    @property
    def foot_friction(self):
        friction_vals = []
        for geom_name in self.feet_geom_names:
            friction_vals.append(self._last_physics.named.model.geom_friction[f"robot/{geom_name}"][0])
        return np.mean(friction_vals)
        # return self._foot_friction
    
    @foot_friction.setter
    def foot_friction(self, value : float):
        print("here")
        for geom_name in self.feet_geom_names:
            self._last_physics.model.geom('robot/'+geom_name).friction[0] = value
        # for geom in self.feet_geoms:
        #     geom.friction = (value, 0.02, 0.01)
        self._foot_friction = value

    @property
    def joints(self):
        """List of joint elements belonging to the arm."""
        return self._joints

    @property
    def observable_joints(self):
        return self._joints

    @property
    def actuators(self):
        """List of actuator elements belonging to the arm."""
        return self._actuators

    @property
    def mjcf_model(self):
        """Returns the `mjcf.RootElement` object corresponding to this robot."""
        return self._mjcf_root

    @composer.cached_property
    def root_body_linvel_sensor(self):
        return self._mjcf_root.find("sensor", "trunk_linvel")

    def get_roll_pitch_yaw_rad(self, physics : Physics):
        quat = self.get_framequat(physics) #physics.bind(self.mjcf_model.sensor.framequat).sensordata
        return quat_to_euler(quat)

    def get_velocity(self, physics : Physics):
        velocimeter = physics.bind(self.mjcf_model.sensor.velocimeter)
        return velocimeter.sensordata + np.random.normal(0.0, 0.05, size=(3,))
    
    def get_position(self, physics : Physics):
        return physics.bind(self.root_body).xpos
    
    def set_position(self, physics : Physics, position : np.ndarray):
        physics.bind(self.root_body).xpos = position
    
    def get_framequat(self, physics : Physics) -> np.ndarray:
        return physics.bind(self.root_body).xquat
    
    def set_framequat(self, physics : Physics, quat : np.ndarray):
        physics.bind(self.root_body).xquat = quat
    
    def get_foot_positions(self, physics : Physics):
        if self._foot_sites[0] is None:
            return None
        return [physics.bind(foot).xpos for foot in self._foot_sites]

    def initialize_episode_mjcf(self, random_state):
        """Callback executed when the MJCF model is modified between episodes."""
        pass

    def after_compile(self, physics, random_state):
        """Callback executed after the Mujoco Physics is recompiled."""
        self._last_physics = physics

    def initialize_episode(self, physics, random_state):
        """Callback executed during episode initialization."""
        self._last_physics = physics
        self._last_step_position = self.get_position(physics)
        self.refresh_observation(physics)

    def before_step(self, physics, random_state):
        """Callback executed before an agent control step."""
        self._last_physics = physics

    def before_substep(self, physics, random_state):
        """Callback executed before a simulation step."""
        self._last_physics = physics

        # === Apply action ===
        if self.target_action is not None:
            if self.action_interpolation and self.prev_action is not None:
                current_qpos = physics.bind(self.observable_joints).qpos
                target_action = self.target_action
                substep_progress = min((self._substep_count + 1) * self._control_subtimestep / self._control_timestep, 1.0)
                current_target_action = (target_action - self.prev_action) * substep_progress + self.prev_action
            else:
                current_target_action = self.target_action
            self.apply_action(physics, current_target_action, random_state)
        self._substep_count += 1
        

    def after_substep(self, physics, random_state):
        """A callback which is executed after a simulation step."""
        self._last_physics = physics
        # Required for the linvel sensor.
        mujoco.mj_subtreeVel(physics.model.ptr, physics.data.ptr)
        # Update delayed observables
        for obs_name, obs_lambda in self._delayed_observable_lambdas.items():
            self._delayed_observable_queues[obs_name].append(obs_lambda(physics))

    def after_step(self, physics, random_state):
        """Callback executed after an agent control step."""
        self._last_physics = physics
    
    def refresh_observation(self,physics):
        # === Update observations read by the robot class ===
        location, quaternion = self.get_position(physics).copy(), self.get_framequat(physics).copy()
        if self._last_step_position is None:
            global_velocity = np.zeros(3)
        else:
            global_velocity = (location - self._last_step_position) / self._control_timestep + np.random.normal(0.0, 0.05, size=(3,))
        
        local_velocity = tr3d.quaternions.rotate_vector(
            global_velocity, 
            tr3d.quaternions.qinverse(quaternion)
        )

        delayed_quaternion = self.read_delayed_observables("sensors_framequat").copy()
        delayed_velocity_local = self.read_delayed_observables("sensors_local_velocimeter").copy()
        delayed_velocity_global = tr3d.quaternions.rotate_vector(delayed_velocity_local, quaternion)
        
        self.last_obs = {
            "3d_location": location, # No need to delay this since this is only used to reset the robot
            "3d_linear_velocity": delayed_velocity_global, # global_velocity,
            "3d_local_velocity": delayed_velocity_local, #local_velocity, 
            "angular_velocity": self.read_delayed_observables("sensors_gyro").copy(), # physics.bind(self.mjcf_model.sensor.gyro).sensordata.copy(),
            "framequat_wijk": delayed_quaternion, #quaternion, 
            "roll_pitch_yaw": quat_to_euler(delayed_quaternion), #quat_to_euler(quaternion), 
            "3d_acceleration_local": self.read_delayed_observables("sensors_accelerometer").copy(), # physics.bind(self.mjcf_model.sensor.accelerometer).sensordata.copy(),
            "joint_pos": self.read_delayed_observables("joints_pos").copy(), #physics.bind(self.observable_joints).qpos.copy(),
            "joint_vel": self.read_delayed_observables("joints_vel").copy(), #physics.bind(self.observable_joints).qvel.copy(),
            "joint_acc" : self.read_delayed_observables("joint_acc").copy(), #physics.bind(self.observable_joints).qacc.copy(), 
            "foot_force": self.read_delayed_observables("foot_forces").copy(), #physics.bind(self.mjcf_model.sensor.touch).sensordata.copy(),
            "torques": self.read_delayed_observables("torques").copy() #physics.bind(self.actuators).force
        }
        self._last_step_position = location
    
    def reset(self, physics):
        self.refresh_observation(physics)
        self.last_obs["3d_linear_velocity"][:] = 0
        self.last_obs["3d_local_velocity"][:] = 0
        self.last_obs["angular_velocity"][:] = 0
        self.last_obs["3d_acceleration_local"][:] = 0
        self.last_obs["joint_vel"][:] = 0
        self.last_obs["joint_acc"][:] = 0
        self.last_obs["torques"][:] = 0


"""
In compliance with `rail_walker_interface`, this class is a wrapper of the robot class.
"""
class RailSimWalkerDMControl(BaseWalker[typing.Dict[str,np.ndarray]],BaseWalkerInSim, BaseWalkerWithFootContact):
    def __init__(
        self, 
        XML_FILE : str, 
        Kp=60, 
        Kd=6, 
        action_interpolation : bool = True, 
        limit_action_range : float = 1.0, 
        power_protect_factor : float = 0.1, 
        foot_contact_threshold : np.ndarray = np.array([20.0,20.0,20.0,20.0]),
        *args, 
        **kwargs
    ):
        self.mujoco_walker = DMWalkerForRailSimWalker(XML_FILE, action_interpolation=action_interpolation, power_protect_factor=power_protect_factor, kp=Kp, kd=Kd, foot_contact_threshold=foot_contact_threshold, *args, **kwargs)
        BaseWalker.__init__(
            self,
            name="robot",
            Kp = Kp,
            Kd = Kd,
            force_real_control_timestep=False,
            limit_action_range=limit_action_range,
            power_protect_factor=power_protect_factor,
        )
        BaseWalkerInSim.__init__(self)
        BaseWalkerWithFootContact.__init__(self)


    @property
    def is_real_robot(self) -> bool:
        return False

    @property
    def action_interpolation(self):
        return self.mujoco_walker.action_interpolation
    
    @action_interpolation.setter
    def action_interpolation(self, value : bool):
        self.mujoco_walker.action_interpolation = value

    @property
    def power_protect_factor(self):
        return self.mujoco_walker.power_protect_factor
    
    @power_protect_factor.setter
    def power_protect_factor(self, value : float):
        assert value > 0 and value <= 1.0
        self.mujoco_walker.power_protect_factor = value
    
    @property
    def foot_friction(self):
        return self.mujoco_walker.foot_friction
    
    @foot_friction.setter
    def foot_friction(self, value : float):
        self.mujoco_walker.foot_friction = value

    @property
    def foot_contact_threshold(self):
        return self.mujoco_walker._built_observable._foot_contact_threshold

    @foot_contact_threshold.setter
    def foot_contact_threshold(self, value : np.ndarray):
        assert value.shape == (4,)
        self.mujoco_walker._built_observable._foot_contact_threshold = value.astype(np.float32)

    @property
    def joint_qpos_init(self) -> np.ndarray:
        pass

    @property
    def joint_qpos_offset(self) -> np.ndarray:
        return np.array([0.2, 0.4, 0.4] * 4)

    @property
    def joint_qpos_mins(self) -> np.ndarray:
        pass

    @property
    def joint_qpos_maxs(self) -> np.ndarray:
        pass

    def get_3d_location(self) -> np.ndarray:
        return self.mujoco_walker.last_obs["3d_location"]

    def set_3d_location(self, target_location : np.ndarray) -> np.ndarray:
        self.mujoco_walker.set_position(self.mujoco_walker._last_physics, target_location)
        self.mujoco_walker.refresh_observation(self.mujoco_walker._last_physics)
    
    def reset_2d_location(self, target_location : np.ndarray, target_quaternion : np.ndarray | None = None, target_qpos : np.ndarray | None = None) -> np.ndarray:
        self.reset_2d_location_with_qpos(
            target_location=target_location,
            target_quaternion=target_quaternion,
            target_qpos=target_qpos
        )

    def reset_2d_location_with_qpos(
        self,
        target_location: np.ndarray,
        target_quaternion: np.ndarray | None = None,
        target_qpos: np.ndarray | None = None,
    ):
        if target_qpos is None:
            target_qpos = self.joint_qpos_init
        find_dm_control_non_contacting_height(self.mujoco_walker._last_physics, self.mujoco_walker, target_location[0], target_location[1], qpos = target_qpos, quat = target_quaternion)
        self.mujoco_walker.target_action = target_qpos
        self.mujoco_walker.refresh_observation(self.mujoco_walker._last_physics)

    # Reset task utilities
    def reset_dropped(self, target_joints : np.ndarray, target_location : np.ndarray, target_quaternion : np.ndarray | None = None, settle : bool = False) -> np.ndarray:
        # Set robot to target joints and location and quaternion, wait for physics to settle
        physics = self.mujoco_walker._last_physics
        physics.bind(self.mujoco_walker.joints).qpos[:] = target_joints
        self.mujoco_walker.set_pose(physics, target_location, target_quaternion)
        if settle:
            # physics.bind(self.mujoco_walker.joints).qpos[:] = target_joints
            # self.mujoco_walker.set_pose(physics, target_location, target_quaternion)
            settle_physics(
                physics,
                self.mujoco_walker
            ) 
        # else:
            # find_dm_control_non_contacting_height(physics, self.mujoco_walker, target_location[0], target_location[1], qpos = target_joints, quat = target_quaternion)
            # print height of the robot
            # height = physics.bind(self.mujoco_walker.root_body).xpos[-1]
            # print(height)
        self.mujoco_walker.refresh_observation(self.mujoco_walker._last_physics)
        self.mujoco_walker.target_action = target_joints

    def set_framequat_wijk(self, framequat_wijk: np.ndarray) -> None:
        self.mujoco_walker.set_framequat(self.mujoco_walker._last_physics, framequat_wijk)
        self.mujoco_walker.refresh_observation(self.mujoco_walker._last_physics)

    def set_roll_pitch_yaw(self, roll: float, pitch: float, yaw: float) -> None:
        self.set_framequat_wijk(
            tr3d.euler.euler2quat(roll, pitch, yaw)
        )

    @property
    def control_timestep(self) -> float:
        return self.mujoco_walker._control_timestep
    
    @control_timestep.setter
    def control_timestep(self, value : float) -> None:
        self.mujoco_walker._control_timestep = value
    
    @property
    def control_subtimestep(self) -> float:
        return self.mujoco_walker._control_subtimestep

    @control_subtimestep.setter
    def control_subtimestep(self, value : float) -> None:
        self.mujoco_walker._control_subtimestep = value

    @property
    def Kp(self) -> float:
        return self.mujoco_walker.kp
    
    @Kp.setter
    def Kp(self, value : float) -> None:
        self.mujoco_walker.kp = value
    
    @property
    def Kd(self) -> float:
        return self.mujoco_walker.kd
    
    @Kd.setter
    def Kd(self, value : float) -> None:
        self.mujoco_walker.kd = value

    def receive_observation(self) -> bool:
        if self.mujoco_walker._last_physics is None:
            return False
        
        self.mujoco_walker.refresh_observation(self.mujoco_walker._last_physics)
        return self.mujoco_walker.last_obs is not None

    def reset(self) -> None:
        self.mujoco_walker.reset(self.mujoco_walker._last_physics)

    def get_3d_linear_velocity(self) -> np.ndarray:
        return self.mujoco_walker.last_obs["3d_linear_velocity"]

    def get_3d_local_velocity(self) -> np.ndarray:
        return self.mujoco_walker.last_obs["3d_local_velocity"]

    def get_3d_angular_velocity(self) -> np.ndarray:
        return self.mujoco_walker.last_obs["angular_velocity"]

    def get_framequat_wijk(self) -> np.ndarray:
        return self.mujoco_walker.last_obs["framequat_wijk"]

    def get_roll_pitch_yaw(self) -> tuple[float, float, float]:
        return self.mujoco_walker.last_obs["roll_pitch_yaw"]

    def get_last_observation(self) -> typing.Optional[typing.Dict[str,np.ndarray]]:
        return self.mujoco_walker.last_obs

    def get_3d_acceleration_local(self) -> np.ndarray:
        return self.mujoco_walker.last_obs["3d_acceleration_local"]

    def get_joint_qpos(self) -> np.ndarray:
        return self.mujoco_walker.last_obs["joint_pos"]

    def get_joint_qvel(self) -> np.ndarray:
        return self.mujoco_walker.last_obs["joint_vel"]
    
    def get_joint_qacc(self) -> np.ndarray:
        return self.mujoco_walker.last_obs["joint_acc"]

    def get_joint_torques(self) -> np.ndarray:
        return self.mujoco_walker.last_obs["torques"]

    def get_foot_force(self) -> np.ndarray:
        return self.mujoco_walker.last_obs["foot_force"]
    
    def get_foot_force_norm(self) -> np.ndarray:
        return self.get_foot_force() / 50.0
    
    def get_foot_contact(self) -> np.ndarray:
        return self.get_foot_force() >= self.foot_contact_threshold

    def _apply_action(self, action: np.ndarray) -> bool:
        self.mujoco_walker.set_target_action(action)
        return True

    def close(self) -> None:
        return
rail_mujoco_walker/robots/sim_robot.py
realquantumcookie/APRL
from .joystick_task import JoystickPolicyProviderWithDMControl, JoystickPolicyDMControlTask
rail_mujoco_walker/tasks/__init__.py
realquantumcookie/APRL
from typing import Dict, Optional, Tuple, Callable, Any, List

import dm_control.utils.transformations as tr
import numpy as np
from dm_control import composer
from dm_control.composer.observation import observable
from dm_control.utils import rewards
from dm_control.mjcf.physics import Physics
from dm_control.mujoco.engine import Physics as EnginePhysics
from mujoco import MjvScene
# from dm_control.utils import transformations
from dm_control.locomotion import arenas
from dm_control import composer
from dm_control.locomotion.walkers import base
from dm_env import specs
from ..robots.sim_robot import RailSimWalkerDMControl
from rail_walker_interface import JoystickPolicy, BaseWalker, BaseWalkerWithFootContact
from functools import cached_property
from ..utils import add_arrow_to_mjv_scene, add_sphere_to_mjv_scene
from collections import OrderedDict

DEFAULT_CONTROL_TIMESTEP = 0.05
DEFAULT_PHYSICS_TIMESTEP = 0.005

class JoystickPolicyProviderWithDMControl:
    def render_scene_callback(self, task : "JoystickPolicyDMControlTask", physics : EnginePhysics, scene : MjvScene) -> None:
        pass

class JoystickPolicyDMControlTask(composer.Task):
    def __init__(
        self,
        joystick_policy: JoystickPolicy,
        floor : composer.Arena,
        enabled_cameras : List[str] = ["top_camera", "side_camera"]
    ):
        assert isinstance(joystick_policy.robot.unwrapped(),RailSimWalkerDMControl) 
        self.joystick_policy = joystick_policy

        self._floor = floor
        # self._floor.mjcf_model.size.nconmax = 400
        # self._floor.mjcf_model.size.njmax = 2000
        self._floor.add_free_entity(self.robot.mujoco_walker)

        self.set_timesteps(joystick_policy.control_timestep, joystick_policy.control_subtimestep)

        # Add cameras
        if "side_camera" in enabled_cameras:
            self.robot.mujoco_walker.mjcf_model.worldbody.add(
                'camera',
                name='side_camera',
                pos=[0, -1, 0.5],
                xyaxes=[1, 0, 0, 0, 0.342, 0.940],
                mode="trackcom",
                fovy=60
            )

        if "top_camera" in enabled_cameras:
            self.robot.mujoco_walker.mjcf_model.worldbody.add(
                'camera',
                name='top_camera',
                pos=[0, 0, 3],
                xyaxes=[1, 0, 0, 0, 1, 0],
                mode="trackcom",
                fovy=60.0
            )

        for obs in self.robot.mujoco_walker.observables._observables.keys():
            getattr(self.robot.mujoco_walker.observables, obs).enabled = True

        self._arrow_root = None
        self._arrow_target_yaw = None
        self._arrow_velocity = None
        self._arrow_walker_heading = None
        self._sphere_goal = None
        
        self._perturb = True

    # def reload_joystick_policy(self):
        # enabled_observables = self.joystick_policy.enabled_observables
        # for obs in self.robot.mujoco_walker.observables._observables.keys():
        #     if obs not in enabled_observables:
        #         getattr(self.robot.mujoco_walker.observables, obs).enabled = False
        #     else:
        #         getattr(self.robot.mujoco_walker.observables, obs).enabled = True

        # if hasattr(self, "_target_obs") and self._target_obs is not None:
        #     self._target_obs.enabled = False
        #     self._target_obs = None

        # if self.joystick_policy.target_observable is not None:
        #     self._target_obs = observable.Generic(lambda physics: self.joystick_policy.target_observable.get_observation())
        #     self._target_obs.enabled = True
        # else:
        #     self._target_obs = None
        

    @property
    def root_entity(self):
        return self._floor
    
    @property
    def robot(self) -> RailSimWalkerDMControl:
        return self.joystick_policy.robot
    
    @property
    def observables(self):
        obs = super().observables
        # Enforce Observable Order
        return OrderedDict(sorted(obs.items(), key=lambda t: t[0]))

    @property
    def task_observables(self):
        task_observables = super().task_observables
        
        # if self._target_obs is not None:
        #     self._target_obs.enabled = True
        #     task_observables['target_obs'] = self._target_obs
        
        return task_observables
    
    @property
    def control_timestep(self) -> float:
        """Returns the agent's control timestep for this task (in seconds)."""
        return super().control_timestep

    @control_timestep.setter
    def control_timestep(self, value : float):
        """Sets the agent's control timestep for this task (in seconds)."""
        self.joystick_policy.control_timestep = value
        super().control_timestep = value

    @property
    def physics_timestep(self) -> float:
        """Returns the physics timestep for this task (in seconds)."""
        return super().physics_timestep
    
    @physics_timestep.setter
    def physics_timestep(self, value : float):
        """Sets the physics timestep for this task (in seconds)."""
        self.joystick_policy.control_subtimestep = value
        super().physics_timestep = value

    def set_timesteps(self, control_timestep : float, physics_timestep : float):
        self.joystick_policy.control_timestep = control_timestep
        self.joystick_policy.control_subtimestep = physics_timestep
        super().set_timesteps(control_timestep, physics_timestep)

    def get_reward(self, physics):
        rew = self.joystick_policy.get_reward()
        return rew

    def _randomize_foot_friction(
        self, physics: Physics
    ) -> None:
        friction = 1.5
        feet_geom_names = ["RL", "RR", "FL", "FR"]
        for geom_name in feet_geom_names:
            physics.model.geom('robot/'+geom_name).friction[0] = friction

        # logging.warning(f”Randomized foot friction: {friction}“)
        # friction = 0
        # physics.named.model.geom_friction[consts.FOOT_GEOM_NAMES, 0] = friction
        
    def initialize_episode(self, physics: Physics, random_state: np.random.RandomState):
        super().initialize_episode(physics, random_state)
        
        # Measure body weight in Newtons.
        body_mass = physics.named.model.body_subtreemass['robot/trunk']
        self._body_weight = -physics.model.opt.gravity[2] * body_mass

        self._perturb_steps: int = 0
        self._frc: np.ndarray = np.zeros(3)
        # self._randomize_foot_friction(physics)

        self.robot.mujoco_walker._last_physics = physics
        self._floor.initialize_episode(physics, random_state)
        self.joystick_policy.reset(random_state)

    def before_step(self, physics : Physics, action : np.ndarray, random_state : np.random.RandomState):
        self.robot.mujoco_walker._last_physics = physics
        self.joystick_policy.before_step(action, random_state)
        self._perturb_root_body(physics, random_state)

    def before_substep(self, physics, action, random_state):
        self.robot.mujoco_walker._last_physics = physics

    def action_spec(self, physics):
        # return specs.BoundedArray(
        #     shape=(len(self.robot.joint_qpos_mins),),
        #     dtype=np.float32,
        #     minimum=self.robot.action_qpos_mins,
        #     maximum=self.robot.action_qpos_maxs
        #     #name='\t'.join([actuator.name for actuator in self.actuators])
        # )
        return self.robot.mujoco_walker.action_spec(physics)

    def after_step(self, physics, random_state):
        self.robot.mujoco_walker._last_physics = physics
        self.joystick_policy.after_step(random_state)

        # Update render arrows
        self._arrow_root = self.robot.get_3d_location()
        self._arrow_walker_heading = tr.quat_rotate(
            self.robot.get_framequat_wijk(),
            np.array([1.0, 0.0, 0.0])
        ).astype(np.float32)
        self._arrow_velocity = self.robot.get_3d_linear_velocity().astype(np.float32)
        self._arrow_target_yaw = np.array([*self.joystick_policy.target_goal_world_delta_unit, 0],dtype=np.float32)
        if hasattr(self.joystick_policy.reward_provider, "target_linear_velocity"):
            self._arrow_target_yaw *= self.joystick_policy.reward_provider.target_linear_velocity
        
        self._sphere_goal = np.array([*self.joystick_policy.target_goal_world_delta,1.0]) + self._arrow_root
    
    def should_terminate_episode(self, physics):
        self.robot.mujoco_walker._last_physics = physics
        return False

    def get_discount(self, physics):
        return 1.0
    
    def _perturb_root_body(
        self, physics: Physics, random_state: np.random.RandomState
    ) -> None:
        if not self._perturb:
            return

        if self._perturb_steps > 20:
            magnitude = 0.05 * self._body_weight
            angle = random_state.uniform(0, 2 * np.pi)
            vector = np.asarray([np.cos(angle), np.sin(angle), 0])
            self._frc = magnitude * vector
            self._perturb_steps = 0

        physics.named.data.xfrc_applied["robot/trunk", :3] = self._frc
        self._perturb_steps += 1
        
    @cached_property
    def render_scene_callback(self):
        def render_cb(physics : EnginePhysics, scene : MjvScene):
            if isinstance(self.joystick_policy.target_yaw_provider, JoystickPolicyProviderWithDMControl):
                # Call the render_scene_callback of the target_yaw_provider
                self.joystick_policy.target_yaw_provider.render_scene_callback(self, physics, scene)
            
            if isinstance(self.joystick_policy.reward_provider, JoystickPolicyProviderWithDMControl):
                # Call the render_scene_callback of the reward_provider
                self.joystick_policy.reward_provider.render_scene_callback(self, physics, scene)
            
            for termination_provider in self.joystick_policy.termination_providers:
                if isinstance(termination_provider, JoystickPolicyProviderWithDMControl):
                    # Call the render_scene_callback of the termination_provider
                    termination_provider.render_scene_callback(self, physics, scene)
            
            for reset_provider in self.joystick_policy.resetters:
                if isinstance(reset_provider, JoystickPolicyProviderWithDMControl):
                    # Call the render_scene_callback of the reset_provider
                    reset_provider.render_scene_callback(self, physics, scene)

            if self._arrow_root is not None:
                # Add velocity arrow => greenish yellow
                add_arrow_to_mjv_scene(scene, self._arrow_root, self._arrow_root + self._arrow_velocity, radius=0.01, rgba=np.array([0.5, 1.0, 0.0, 1.0], dtype=np.float32))
                # Add target yaw arrow => yellow
                target_height_offset = np.array([0.0,0.0,0.25], dtype=np.float32)
                add_arrow_to_mjv_scene(scene, self._arrow_root + target_height_offset, self._arrow_root + target_height_offset + self._arrow_target_yaw, radius=0.01, rgba=np.array([1.0,1.0,0.0,1.0], dtype=np.float32))
                # Add target location sphere => yellow
                add_sphere_to_mjv_scene(scene,self._sphere_goal, radius=0.2, rgba=np.array([1.0,1.0,0.0,1.0], dtype=np.float32))
                # Add walker heading arrow => green
                add_arrow_to_mjv_scene(scene, self._arrow_root, self._arrow_root + self._arrow_walker_heading, radius=0.01, rgba=np.array([0.0,1.0,0.0,1.0], dtype=np.float32))

            if isinstance(self.robot.unwrapped(), BaseWalkerWithFootContact):
                force_max_length = 0.3
                foot_forces = self.robot.get_foot_force()
                foot_contacts = self.robot.get_foot_contact()
                max_force = np.max(np.abs(foot_forces))
                foot_positions = self.robot.mujoco_walker.get_foot_positions(self.robot.mujoco_walker._last_physics)
                for i, foot_pos in enumerate(foot_positions):
                    if foot_contacts[i]:
                        foot_force = foot_forces[i]
                        # Add foot force arrow => red
                        add_arrow_to_mjv_scene(
                            scene, 
                            foot_pos, 
                            foot_pos + np.array([0.0,0.0,foot_force/max_force * force_max_length]), 
                            radius=0.01, 
                            rgba=np.array([0.0,0.0,1.0,0.5], dtype=np.float32) # blue
                        )


        return render_cb
rail_mujoco_walker/tasks/joystick_task.py
realquantumcookie/APRL
from .render_scene_util import add_arrow_to_mjv_scene, add_sphere_to_mjv_scene
from .resetter_util import find_dm_control_non_contacting_height, find_dm_control_non_contacting_qpos, settle_physics
rail_mujoco_walker/utils/__init__.py
realquantumcookie/APRL
import mujoco
import numpy as np

def add_arrow_to_mjv_scene(scene : mujoco.MjvScene, point1 : np.ndarray, point2 : np.ndarray, radius : float, rgba : np.ndarray):
    if scene.ngeom >= scene.maxgeom:
        return

    scene.ngeom += 1  # increment ngeom

    mujoco.mjv_initGeom(
        scene.geoms[scene.ngeom-1], # geom
        mujoco.mjtGeom.mjGEOM_ARROW, # type
        np.zeros(3), # size
        np.zeros(3), # pos
        np.zeros(9), # mat
        rgba.astype(np.float32), # mat
    )
    
    mujoco.mjv_makeConnector(
        scene.geoms[scene.ngeom-1], # geom
        mujoco.mjtGeom.mjGEOM_ARROW, # type
        radius, # width
        *point1[:3], # a0 - a2
        *point2[:3], # b0 - b2
    )

def add_sphere_to_mjv_scene(scene : mujoco.MjvScene, point : np.ndarray, radius : float, rgba : np.ndarray):
    if scene.ngeom >= scene.maxgeom:
        return
    
    scene.ngeom += 1  # increment ngeom

    mujoco.mjv_initGeom(
        scene.geoms[scene.ngeom-1], # geom
        mujoco.mjtGeom.mjGEOM_SPHERE, # type
        np.array([radius]*3), # size
        point[:3], # pos
        np.zeros(9), # mat
        rgba.astype(np.float32), # mat
    )
rail_mujoco_walker/utils/render_scene_util.py
realquantumcookie/APRL
from dm_control.rl import control
import mujoco
import numpy as np
from dm_control.mjcf.physics import Physics
from typing import Optional


_MAX_SETTLE_PHYSICS_ATTEMPTS = 1
_MIN_SETTLE_PHYSICS_TIME = 0.0
_MAX_SETTLE_PHYSICS_TIME = 2.0

_SETTLE_QVEL_TOL = 1e-3

def settle_physics(
    physics: Physics,
    robot,
    qvel_tol=_SETTLE_QVEL_TOL,
    max_attempts=_MAX_SETTLE_PHYSICS_ATTEMPTS,
    min_time=_MIN_SETTLE_PHYSICS_TIME,
    max_time=_MAX_SETTLE_PHYSICS_TIME,
) -> bool:
    """Steps the physics until the robot root body is at rest."""
    sensor_binding = physics.bind(robot.root_body_linvel_sensor)
    for _ in range(max_attempts):
        original_time = physics.data.time
        while physics.data.time - original_time < max_time:
            physics.step()
            max_qvel = np.max(np.abs(sensor_binding.sensordata))
            if (max_qvel < qvel_tol) and (physics.data.time - original_time) > min_time:
                return True
        physics.data.time = original_time
    return False

def find_dm_control_non_contacting_height(
    physics : Physics,
    walker,
    x_pos : float = 0.0,
    y_pos : float = 0.0,
    qpos : Optional[np.ndarray]=None,
    quat : Optional[np.ndarray]=None,
    maxiter : int=1000
):
    z_pos = 0.0  # Start embedded in the floor.
    num_contacts = 1
    count = 1
    # Move up in 1cm increments until no contacts.
    while num_contacts > 0:
        try:
            with physics.reset_context():
                walker.set_pose(physics, [x_pos, y_pos, z_pos], quaternion=quat)
                if qpos is not None:
                    physics.bind(walker.joints).qpos[:] = qpos.copy()
                
        except control.PhysicsError:
            # We may encounter a PhysicsError here due to filling the contact
            # buffer, in which case we simply increment the height and continue.
            pass
        num_contacts = physics.data.ncon

        z_pos += 0.01
        count += 1
        if count > maxiter:
            raise ValueError('maxiter reached: possibly contacts in null pose of body.')

def find_dm_control_non_contacting_qpos(
    physics : Physics, 
    random_state : np.random.RandomState, 
    walker, 
    maxiter : int =1000
):
    num_contacts = 1
    count = 1
    # Move up in 1cm increments until no contacts.
    while num_contacts > 0:
        try:
            with physics.reset_context():
                walker.set_pose(physics, [0, 0, 100])
                joints_range = physics.bind(walker.joints).range
                qpos = random_state.uniform(joints_range[:, 0],
                                            joints_range[:, 1])
                walker.configure_joints(physics, qpos)
        except control.PhysicsError:
            # We may encounter a PhysicsError here due to filling the contact
            # buffer, in which case we simply increment the height and continue.
            pass
        num_contacts = physics.data.ncon

        count += 1

        if count > maxiter:
            raise ValueError(
                'maxiter reached: possibly contacts in null pose of body.')

    walker.set_pose(physics, [0, 0, 0])

    if num_contacts == 0:
        return qpos
    else:
        return None

rail_mujoco_walker/utils/resetter_util.py
realquantumcookie/APRL
from .velocity_estimators import *
rail_real_walker/__init__.py
realquantumcookie/APRL
from .go1 import Go1RealWalker
from .go1_remote import Go1RealWalkerRemote
rail_real_walker/robots/__init__.py
realquantumcookie/APRL
from rail_walker_interface import BaseWalker, BaseWalkerWithFootContact, BaseWalkerWithJointTemperatureSensor, BaseWalkerWithJoystick, Walker3DVelocityEstimator
from typing import Generic, TypeVar, Optional
import numpy as np
import unitree_go1_wrapper as go1_wrapper
from functools import cached_property
import transforms3d as tr
import multiprocessing
import threading
import time
from ctypes import c_float, c_bool

class Go1RobotThreadRunner:
    REAL_CONTROL_TIMESTEP = 0.002
    MOMENTUM_BETA = 0.95
    def __init__(
        self,
        Kp: float,
        Kd: float,
        power_protect_factor: float,
        action_interpolation: bool,
        torque_cutoff_hip: float,
        torque_cutoff_thigh: float,
        torque_cutoff_calf: float,
    ):
        assert power_protect_factor > 0 and power_protect_factor <= 1

        self.run_thread : Optional[multiprocessing.Process] = None
        self.event_req_observation : Optional[multiprocessing.Event] = None
        self.event_quit : Optional[multiprocessing.Event] = None
        self.pipe_comm : Optional[multiprocessing.Pipe] = None

        self.Kd = multiprocessing.Value(c_float, Kd)
        self.Kp = multiprocessing.Value(c_float, Kp)
        self.torque_cutoff_hip = multiprocessing.Value(c_float, torque_cutoff_hip)
        self.torque_cutoff_thigh = multiprocessing.Value(c_float, torque_cutoff_thigh)
        self.torque_cutoff_calf = multiprocessing.Value(c_float, torque_cutoff_calf)
        self.action_interpolation = multiprocessing.Value(c_bool, action_interpolation)
        self.power_protect_factor = multiprocessing.Value(c_float, power_protect_factor)
        
    def set_up_run_thread(
        self, 
        robot: "Go1RealWalker",
        control_timestep: float = 0.02,
    ):
        self.clean_up_run_thread()
        self.event_req_observation = multiprocessing.Event()
        self.event_quit = multiprocessing.Event()
        self.pipe_comm, child_pipe = multiprocessing.Pipe(duplex=True)
        self.run_thread = multiprocessing.Process(
            target=__class__.thread_run,
            args=(
                robot,
                self,
                robot.robot_interface,
                robot.robot_safety_interface,
                self.event_req_observation,
                self.event_quit,
                child_pipe,
                control_timestep,
            )
        )
        self.run_thread.start()

    @staticmethod
    def cutoff_action(
        current_action : np.ndarray, 
        current_qpos : np.ndarray, 
        current_dqpos : np.ndarray,
        max_torque : np.ndarray,
        Kp : float,
        Kd : float,
    ) -> np.ndarray:
        # Cutoff action to prevent robot from exerting too much torque on joints
        delta_qpos = current_action - current_qpos
        current_d_torque = Kd * current_dqpos
        current_d_torque = np.clip(current_d_torque, -max_torque, max_torque)
        min_delta_qpos = (current_d_torque - max_torque) / Kp
        min_delta_qpos = np.minimum(min_delta_qpos, 0) # Clip so that 0 is always a valid value
        
        max_delta_qpos = (current_d_torque + max_torque) / Kp
        max_delta_qpos = np.maximum(max_delta_qpos, 0) # Clip so that 0 is always a valid value

        clipped_delta_qpos = np.clip(delta_qpos, min_delta_qpos, max_delta_qpos)
        clipped_action = current_qpos + clipped_delta_qpos
        return clipped_action
    
    def clean_up_run_thread(self):
        if self.event_quit is not None and self.run_thread is not None:
            self.event_quit.set()
            self.run_thread.join(5.0)
            if self.run_thread.is_alive():
                self.run_thread.terminate()
            self.run_thread.close()
            self.run_thread = None

        if self.run_thread is not None:
            if self.run_thread.is_alive():
                self.run_thread.terminate()
            self.run_thread.close()
            self.run_thread = None
        
        if self.pipe_comm is not None:
            self.pipe_comm.close()
        
        self.event_quit = None
        self.event_req_observation = None
        self.pipe_comm = None
    
    def async_request_observation(self) -> None:
        if self.event_req_observation is not None:
            self.event_req_observation.set()
        else:
            raise RuntimeError("No active thread to request observation from.")
    
    def async_read_observation(self) -> Optional[tuple[go1_wrapper.LowState, np.ndarray]]:
        if self.pipe_comm is not None and self.pipe_comm.poll():
            return self.pipe_comm.recv()
        else:
            return None

    def read_observation(self) -> tuple[go1_wrapper.LowState, np.ndarray]:
        if self.event_req_observation:
            self.event_req_observation.set()
            return self.pipe_comm.recv()
        else:
            raise RuntimeError("No active thread to read observation from.")

    def async_apply_action(self, action : np.ndarray) -> None:
        if self.pipe_comm is not None:
            self.pipe_comm.send(action)
        else:
            raise RuntimeError("No active thread to apply action to.")   

    @staticmethod
    def thread_run(
        robot: BaseWalker[go1_wrapper.LowState],
        thread_runner: "Go1RobotThreadRunner",
        robot_interface: go1_wrapper.UDP,
        robot_safety_interface: go1_wrapper.Safety,
        event_req_observation: multiprocessing.Event,
        event_quit: multiprocessing.Event,
        pipe_comm: multiprocessing.Pipe,
        control_timestep: float = 0.02
    ):
        def build_command(action : np.ndarray, command : go1_wrapper.LowCmd, Kp : float, Kd : float) -> go1_wrapper.LowCmd:
            for i in range(12):
                command.motorCmd[i].q = action[i]
                command.motorCmd[i].dq = 0.0
                command.motorCmd[i].Kp = Kp
                command.motorCmd[i].Kd = Kd
                command.motorCmd[i].tau = 0.0
            return command

        thread_id = threading.get_ident()
        print(f"Go1 Robot Control Thread Started (id={thread_id}, CONTROL_TIMESTEP={__class__.REAL_CONTROL_TIMESTEP}).")
        
        try:
            robot_interface.recv()
            last_control_motor_states = [go1_wrapper.MotorState.fromNative(robot_interface._low_state_native.motorState[i]) for i in range(12)]
            last_control_qpos = np.array([last_control_motor_states[i].q for i in range(12)])
            last_control_dq = np.array([last_control_motor_states[i].dq for i in range(12)])

            last_state= None
            last_action = None
            curr_action = None
            last_control_time = 0.0
            control_command = robot_interface.initLowCmdData()
            last_action_command_t = 0.0

            power_protect_factor = thread_runner.power_protect_factor.value
            max_torque = np.array([thread_runner.torque_cutoff_hip.value, thread_runner.torque_cutoff_thigh.value, thread_runner.torque_cutoff_calf.value] * 4)
            max_torque *= power_protect_factor
            Kp = thread_runner.Kp.value
            Kd = thread_runner.Kd.value
            action_interpolation = thread_runner.action_interpolation.value

            while True:
                t = time.time()

                # Make sure we are not sending commands too fast (at REAL_CONTROL_TIMESTEP)
                action_command_dt = t - last_action_command_t
                dt = t - last_control_time
                if dt < __class__.REAL_CONTROL_TIMESTEP:
                    time.sleep(__class__.REAL_CONTROL_TIMESTEP - dt)
                    continue
                
                # check if we need to quit
                if event_quit.is_set():
                    break

                t = time.time()
                
                # Update observation / state
                robot_interface.recv()
                last_control_motor_states = [go1_wrapper.MotorState.fromNative(robot_interface._low_state_native.motorState[i]) for i in range(12)]
                last_control_qpos = np.array([last_control_motor_states[i].q for i in range(12)])
                last_control_dq = np.array([last_control_motor_states[i].dq for i in range(12)])

                # Send out observation if requested
                if event_req_observation.is_set():
                    event_req_observation.clear()
                    last_state = robot_interface.getRecvLow()
                    obs_to_send = last_state
                    pipe_comm.send(obs_to_send)

                # set last_action and last_action_command_t if we have a new action
                if pipe_comm.poll():
                    last_action = curr_action
                    curr_action = pipe_comm.recv()
                    last_action_command_t = t
                    action_command_dt = 0.0

                    # Refresh Kp, Kd, max_torque, action_interpolation
                    power_protect_factor = thread_runner.power_protect_factor.value
                    max_torque = np.array([thread_runner.torque_cutoff_hip.value, thread_runner.torque_cutoff_thigh.value, thread_runner.torque_cutoff_calf.value] * 4)
                    max_torque *= power_protect_factor
                    Kp = thread_runner.Kp.value
                    Kd = thread_runner.Kd.value
                    action_interpolation = thread_runner.action_interpolation.value

                    last_state = robot_interface.getRecvLow()
                    
                # Send out action (control)
                if curr_action is not None:
                    action_to_send = curr_action

                    if last_action is not None and action_interpolation:
                        control_progress = max(min(action_command_dt / control_timestep,1.0),0.0)
                        action_to_send = (1.0 - control_progress) * last_action + control_progress * curr_action

                    action_to_send = __class__.cutoff_action(action_to_send, last_control_qpos, last_control_dq, max_torque, Kp, Kd)
                    # print(f"Sending action: {action_to_send}")
                    control_command = build_command(action_to_send, control_command, Kp, Kd)
                    
                    # Send command & Power Protect
                    robot_interface.setToSendLow(control_command)
                    # Power protect
                    robot_safety_interface.powerProtect(10)
                    robot_safety_interface.positionLimit()
                    robot_interface.send()

        finally:
            pipe_comm.close()
            print(f"Go1 Robot Control Thread(id={thread_id}) Ended.")
    
    def __del__(self):
        self.clean_up_run_thread()


class Go1RealWalker(BaseWalker[go1_wrapper.LowState], BaseWalkerWithFootContact, BaseWalkerWithJoystick, BaseWalkerWithJointTemperatureSensor):
    def __init__(
        self, 
        velocity_estimator: Walker3DVelocityEstimator,
        power_protect_factor : float = 0.5,
        foot_contact_threshold: np.ndarray = np.array([100, 100, 100, 100]),
        action_interpolation: bool = True,
        name: Optional[str] = "robot", 
        Kp: float = 5, 
        Kd: float = 1,
        control_timestep : float = 0.05,
        force_real_control_timestep : bool = True,
        limit_action_range : float = 1.0,
        torque_cutoff_hip : float = 23.7,
        torque_cutoff_thigh : float = 23.7,
        torque_cutoff_calf : float = 35.55
    ):  
        self.thread_runner = Go1RobotThreadRunner(
            Kp=Kp, 
            Kd=Kd, 
            power_protect_factor=power_protect_factor,
            action_interpolation=action_interpolation, 
            torque_cutoff_hip=torque_cutoff_hip,
            torque_cutoff_thigh=torque_cutoff_thigh,
            torque_cutoff_calf=torque_cutoff_calf
        )
        assert control_timestep >= 0.02
        BaseWalker.__init__(self,name, Kp, Kd, force_real_control_timestep, limit_action_range, power_protect_factor)
        BaseWalkerWithFootContact.__init__(self)
        BaseWalkerWithJointTemperatureSensor.__init__(self)
        BaseWalkerWithJoystick.__init__(self)
        
        self.velocity_estimator = velocity_estimator
        self._last_state : Optional[go1_wrapper.LowState] = None
        self._last_velocity : np.ndarray = np.zeros(3)
        # self.foot_contact_threshold = foot_contact_threshold
        self.foot_contact_no_contact_threshold = np.array([36.0, 120.0, 90.0, 130.0])
        self.foot_contact_has_contact_threshold = np.array([120.0, 230.0, 340.0, 380.0])
        self.foot_contact_threshold = self.foot_contact_no_contact_threshold * 0.8 + self.foot_contact_has_contact_threshold * 0.2
        #self.foot_contact_threshold = np.array([90, 140, 150, 150])
        self.foot_contact_calibrated = True

        self._control_timestep = control_timestep
        
        self.robot_interface = go1_wrapper.UDP(go1_wrapper.ControlLevel.LOW)
        self.robot_interface.setToSendLow(self.robot_interface.initCommunicationLowCmdData())
        self.robot_interface.send()
        time.sleep(0.05) # Wait for the robot to send back the first state
        self.robot_interface.recv()

        self._last_state = self.robot_interface.getRecvLow() # Fill in _last_state just in case velocity estimator requires it
        self._last_velocity = np.zeros(3)
        self._last_velocity_estimator_t = time.time()
        self._last_joystick_values : list[np.ndarray] = [
            np.zeros(2),
            np.zeros(2)
        ]

        self.robot_safety_interface = go1_wrapper.Safety(self.robot_interface)
        
        self.thread_runner.set_up_run_thread(
            self,
            control_timestep
        )
        self.velocity_estimator.reset(self, self.get_framequat_wijk())
    
    def __copy__(self):
        raise NotImplementedError()

    def __deepcopy__(self, memo):
        raise NotImplementedError()
    
    @property
    def is_real_robot(self) -> bool:
        return True

    @property
    def control_timestep(self) -> float:
        return self._control_timestep

    @property
    def control_subtimestep(self) -> float:
        return Go1RobotThreadRunner.REAL_CONTROL_TIMESTEP
    
    @property
    def power_protect_factor(self) -> float:
        return self.thread_runner.power_protect_factor.value
    
    @power_protect_factor.setter
    def power_protect_factor(self, value: float) -> None:
        assert value > 0.0 and value <= 1.0
        self.thread_runner.power_protect_factor.value = value

    @property
    def Kp(self) -> float:
        return self.thread_runner.Kp.value
    
    @Kp.setter
    def Kp(self, value: float) -> None:
        self.thread_runner.Kp.value = value
    
    @property
    def Kd(self) -> float:
        return self.thread_runner.Kd.value
    
    @Kd.setter
    def Kd(self, value: float) -> None:
        self.thread_runner.Kd.value = value

    @property
    def torque_cutoff_hip(self) -> float:
        return self.thread_runner.torque_cutoff_hip.value
    
    @torque_cutoff_hip.setter
    def torque_cutoff_hip(self, value: float) -> None:
        self.thread_runner.torque_cutoff_hip.value = value
    
    @property
    def torque_cutoff_thigh(self) -> float:
        return self.thread_runner.torque_cutoff_thigh.value
    
    @torque_cutoff_thigh.setter
    def torque_cutoff_thigh(self, value: float) -> None:
        self.thread_runner.torque_cutoff_thigh.value = value

    @property
    def torque_cutoff_calf(self) -> float:
        return self.thread_runner.torque_cutoff_calf.value

    @torque_cutoff_calf.setter
    def torque_cutoff_calf(self, value: float) -> None:
        self.thread_runner.torque_cutoff_calf.value = value

    @property
    def action_interpolation(self) -> bool:
        return self.thread_runner.action_interpolation.value
    
    @action_interpolation.setter
    def action_interpolation(self, value: bool) -> None:
        self.thread_runner.action_interpolation.value = value

    def receive_observation(self) -> bool:
        t = time.time()
        
        # Step robot interface
        new_state = self.thread_runner.read_observation()
        is_new_state_valid = new_state.isValid()
        if not is_new_state_valid:
            if self._last_state is None:
                self._last_state = new_state
        else:
            self._last_state = new_state

        raw_joystick = go1_wrapper.XRockerBtnDataStruct.fromNative(self._last_state.wirelessRemote)
        self._last_joystick_values = [
            np.array([raw_joystick.lx, raw_joystick.ly]),
            np.array([raw_joystick.rx, raw_joystick.ry])
        ]
        
        # Step velocity estimator
        self.velocity_estimator.step(
            self,
            self.get_framequat_wijk(),
            t - self._last_velocity_estimator_t
        )
        self._last_velocity_estimator_t = t
        self._last_velocity = self.velocity_estimator.get_3d_linear_velocity()

        return is_new_state_valid

    @cached_property
    def joint_qpos_init(self) -> np.ndarray:
        real_shift = np.array([
            -0.1, -0.2, 0.0,
            0.1, -0.2, 0.0,
            -0.1, 0.0, 0.0,
            0.1, 0.0, 0.0
        ])
        # shift_delta=np.array([-0.1, -0.1, 0.1])
        # real_shift = shift_delta.repeat(4).reshape((3,4)).T.flatten()
        # real_shift[3::6] = -real_shift[3::6]
        return np.array([go1_wrapper.GO1_HIP_INIT, go1_wrapper.GO1_THIGH_INIT, go1_wrapper.GO1_CALF_INIT] * 4) + real_shift

    @cached_property
    def joint_qpos_sitting(self) -> np.ndarray:
        return np.array([0.0 / 180 * np.pi, 70.0/180*np.pi, -150.0 / 180 * np.pi] * 4)

    @cached_property
    def joint_qpos_offset(self) -> np.ndarray:
        return np.array([0.2, 0.4, 0.4] * 4)

    @cached_property
    def joint_qpos_mins(self) -> np.ndarray:
        return np.array([go1_wrapper.GO1_HIP_MIN, go1_wrapper.GO1_THIGH_MIN, go1_wrapper.GO1_CALF_MIN] * 4)

    @cached_property
    def joint_qpos_maxs(self) -> np.ndarray:
        return np.array([go1_wrapper.GO1_HIP_MAX, go1_wrapper.GO1_THIGH_MAX, go1_wrapper.GO1_CALF_MAX] * 4)

    def reset(self) -> None:
        t = time.time()
        self._last_velocity = np.zeros(3)
        self._last_velocity_estimator_t = t
        self.velocity_estimator.reset(self, self.get_framequat_wijk())
        super().reset()

    def get_3d_linear_velocity(self) -> np.ndarray:
        return self._last_velocity
    
    def get_3d_local_velocity(self) -> np.ndarray:
        try:
            if hasattr(self.velocity_estimator,"get_3d_local_velocity"):
                return self.velocity_estimator.get_3d_local_velocity()
            else:
                return tr.quaternions.rotate_vector(
                    self.get_3d_linear_velocity(), 
                    tr.quaternions.qinverse(self.get_framequat_wijk())
                )
        except:
            print("Error in get_3d_local_velocity, continuing anyway")
            import traceback
            traceback.print_exc()
            return np.zeros(3)

    def get_3d_angular_velocity(self) -> np.ndarray:
        return self._last_state.imu.gyroscope
    
    def get_framequat_wijk(self) -> np.ndarray:
        return self._last_state.imu.quaternion
    
    def get_roll_pitch_yaw(self) -> np.ndarray:
        return self._last_state.imu.rpy

    def get_last_observation(self) -> Optional[go1_wrapper.LowState]:
        return self._last_state

    def get_3d_acceleration_local(self) -> np.ndarray:
        return self._last_state.imu.accelerometer

    def get_joint_qpos(self) -> np.ndarray:
        return np.array([ms.q for ms in self._last_state.motorState[:12]])

    def get_joint_qvel(self) -> np.ndarray:
        return np.array([ms.dq for ms in self._last_state.motorState[:12]])
    
    def get_joint_qacc(self) -> np.ndarray:
        return np.array([ms.ddq for ms in self._last_state.motorState[:12]])

    def get_joint_torques(self) -> np.ndarray:
        return np.array([ms.tauEst for ms in self._last_state.motorState[:12]])
    
    def get_joint_temperature_celsius(self) -> np.ndarray:
        return np.array([ms.temperature for ms in self._last_state.motorState[:12]], dtype=np.float32)
    
    def hibernate(self) -> None:
        if self.thread_runner.run_thread is not None:
            print("Hibernating...")
            temp_action_range = self.limit_action_range # Save action range
            self.limit_action_range = 1.0 # Set action range to 1.0 temporarily
            start_qpos = self.get_joint_qpos()
            delta_qpos = self.joint_qpos_sitting - start_qpos
            for i in range(20):
                self.apply_action(start_qpos + delta_qpos * ((i+1) / 20.0))
            self.limit_action_range = temp_action_range # Restore action range

        self.thread_runner.clean_up_run_thread()
    
    def cancel_hibernate(self) -> None:
        self.robot_interface.setToSendLow(self.robot_interface.initCommunicationLowCmdData())
        self.robot_interface.send()
        time.sleep(0.05) # Wait for the robot to send back the first state
        self.robot_interface.recv()

        self._last_state = self.robot_interface.getRecvLow() # Fill in _last_state just in case velocity estimator requires it
        self._last_velocity = np.zeros(3)
        self._last_velocity_estimator_t = time.time()
        self.thread_runner.set_up_run_thread(self,self.control_timestep)

    def _apply_action(self, action: np.ndarray) -> bool:
        # if np.allclose(action, self.joint_qpos_init):
        #     print("Applying Standing Pose")
        # else:
        #     print("Applying Action", action)
        
        self.thread_runner.async_apply_action(action)
        return True
    
    def get_foot_contact(self) -> np.ndarray:
        return self.get_foot_force() >= self.foot_contact_threshold
    
    def get_foot_force(self) -> np.ndarray:
        return self._last_state.footForce
    
    def get_foot_force_norm(self) -> np.ndarray:
        return (self.get_foot_force() - self.foot_contact_no_contact_threshold) / (self.foot_contact_has_contact_threshold - self.foot_contact_no_contact_threshold)
    
    def get_joystick_values(self) -> list[np.ndarray]:
        return self._last_joystick_values
    
    def close(self) -> None:
        self.thread_runner.clean_up_run_thread()
        self.velocity_estimator.close()
rail_real_walker/robots/go1.py
realquantumcookie/APRL
from .intel_t265 import IntelRealsenseT265Estimator
from .intel_t265_position_est import IntelRealsenseT265PositionEstimator
from .dummy_velocity_estimator import DummyVelocityEstimator
from .remote_t265 import IntelRealsenseT265EstimatorRemote
from .foot_contact_estimator import Go1ForwardKinematicsVelocityProvider
from .kalman_filter_fused_estimator import KalmanFilterFusedVelocityEstimator
rail_real_walker/velocity_estimators/__init__.py
realquantumcookie/APRL
import pyrealsense2 as rs
from rail_walker_interface import BaseWalker, Walker3DVelocityEstimator
from typing import Optional
import multiprocessing
import numpy as np
import transforms3d as tr3d
import time

class DummyVelocityEstimator(Walker3DVelocityEstimator):
    def __init__(
        self
    ):
        super().__init__()
        self.velocity = np.zeros(3)

    def reset(self, robot: BaseWalker, frame_quat: Optional[np.ndarray]):
        self.velocity[:] = 0.0

    def step(self, robot: BaseWalker, frame_quat: Optional[np.ndarray], dt: float):
        self.velocity[:] = 0.0

    def get_3d_linear_velocity(self) -> np.ndarray:
        return self.velocity

    def close(self) -> None:
        pass
rail_real_walker/velocity_estimators/dummy_velocity_estimator.py
realquantumcookie/APRL
import numpy as np
from rail_walker_interface import Walker3DVelocityEstimatorLocal, BaseWalker, BaseWalkerWithFootContact
from typing import Optional, List
import time
import numba

@numba.jit(nopython=True)
def analytical_leg_jacobian(leg_angles : np.ndarray, l_hip_sign : int):
    """
  Computes the analytical Jacobian.
  Args:
  ` leg_angles: a list of 3 numbers for current abduction, hip and knee angle.
    l_hip_sign: whether it's a RIGHT (0) or LEFT(1) leg.
  """
    l_up = 0.2
    l_low = 0.2
    l_hip = 0.08505 * (-1)**(l_hip_sign + 1)

    t1, t2, t3 = leg_angles[0], leg_angles[1], leg_angles[2]
    l_eff = np.sqrt(l_up**2 + l_low**2 + 2 * l_up * l_low * np.cos(t3))

    t_eff = t2 + t3 / 2
    st1 = np.sin(t1)
    ct1 = np.cos(t1)
    st_eff = np.sin(t_eff)
    ct_eff = np.cos(t_eff)
    st3 = np.sin(t3)

    J = np.zeros((3, 3))
    J[0, 0] = 0
    J[0, 1] = -l_eff * ct_eff
    J[0, 2] = l_low * l_up * st3 * st_eff / l_eff - l_eff * ct_eff / 2
    J[1, 0] = -l_hip * st1 + l_eff * ct1 * ct_eff
    J[1, 1] = -l_eff * st1 * st_eff
    J[1, 2] = -l_low * l_up * st1 * st3 * ct_eff / l_eff - l_eff * st1 * st_eff / 2
    J[2, 0] = l_hip * ct1 + l_eff * st1 * ct_eff
    J[2, 1] = l_eff * st_eff * ct1
    J[2, 2] = l_low * l_up * st3 * ct1 * ct_eff / l_eff + l_eff * st_eff * ct1 / 2
    return J

class Go1ForwardKinematicsVelocityProvider(
    Walker3DVelocityEstimatorLocal[BaseWalker]
):
    def __init__(self):
        super().__init__()
        self._local_velocity = np.zeros(3)
    
    def _recalculate_local_velocity(self, robot: BaseWalker, joints_qvel : np.ndarray):
        foot_contacts = robot.get_foot_contact()
        joints_qpos = robot.get_joint_qpos()
        joints_qvel = robot.get_joint_qvel()

        foot_velocities : List[float] = []
        for leg_idx in range(4):
            if foot_contacts[leg_idx]:
                leg_jacob = analytical_leg_jacobian(
                    joints_qpos[leg_idx * 3 : leg_idx * 3 + 3], 
                    leg_idx
                )
                leg_vel = leg_jacob @ joints_qvel[leg_idx * 3 : leg_idx * 3 + 3]
                foot_velocities.append(leg_vel)
        
        if len(foot_velocities) == 0:
            return
        
        foot_velocities = np.vstack(foot_velocities)
        observed_v = np.mean(foot_velocities, axis=0)
        self._local_velocity = -observed_v
        

    def _reset(self, robot: BaseWalker, frame_quat: Optional[np.ndarray]): 
        self._local_velocity = np.zeros(3)
        self._last_qpos = robot.get_joint_qpos().copy()

    def _step(self, robot: BaseWalker, frame_quat: np.ndarray, dt: float):
        joints_qvel = (robot.get_joint_qpos() - self._last_qpos) / dt
        self._recalculate_local_velocity(robot, joints_qvel)
        self._last_qpos = robot.get_joint_qpos().copy()
    
    def get_3d_local_velocity(self) -> np.ndarray:
        return self._local_velocity
rail_real_walker/velocity_estimators/foot_contact_estimator.py
realquantumcookie/APRL
import pyrealsense2 as rs
from rail_walker_interface import BaseWalker, Walker3DVelocityEstimatorLocal, Walker3DFrameQuatEstimator
from typing import Optional
import multiprocessing
import numpy as np
import transforms3d as tr3d
import time

"""
Intel Realsense T265 Velocity/Frame Quaternion Estimator
Please see axis definition here, you need to define how those axis are installed on your robot through the constructor

https://github.com/IntelRealSense/librealsense/blob/master/doc/t265.md
"""
class IntelRealsenseT265Estimator(Walker3DVelocityEstimatorLocal, Walker3DFrameQuatEstimator):
    def __init__(
        self,
        x_axis_on_robot: np.ndarray = np.array([1, 0, 0]),
        y_axis_on_robot: np.ndarray = np.array([0, 1, 0]),
        z_axis_on_robot: np.ndarray = np.array([0, 0, 1]),
    ):
        Walker3DVelocityEstimatorLocal.__init__(self)
        Walker3DFrameQuatEstimator.__init__(self)
        self.comm_pipe : Optional[multiprocessing.Pipe] = None
        self.running_thread : Optional[multiprocessing.Process] = None
        self._enabled = True

        self._local_velocity = np.zeros(3)
        self._last_frame_quat = np.zeros(4)
        self.setup_installation_quat_inv(x_axis_on_robot, y_axis_on_robot, z_axis_on_robot)
        
        
        self.set_up_thread_runner()
    
    @property
    def enabled(self) -> bool:
        return self._enabled
    
    @enabled.setter
    def enabled(self, value: bool) -> None:
        self._enabled = value

    def setup_installation_quat_inv(
        self, 
        x_axis_on_robot: np.ndarray,
        y_axis_on_robot: np.ndarray,
        z_axis_on_robot: np.ndarray
    ):
        assert np.any(x_axis_on_robot != 0) and np.any(y_axis_on_robot != 0) and np.any(z_axis_on_robot != 0)
        assert np.inner(x_axis_on_robot, y_axis_on_robot) == 0 and np.inner(x_axis_on_robot, z_axis_on_robot) == 0 and np.inner(y_axis_on_robot, z_axis_on_robot) == 0
        
        rotation_quat = tr3d.quaternions.mat2quat(
            np.hstack([
                x_axis_on_robot / np.linalg.norm(x_axis_on_robot), 
                y_axis_on_robot / np.linalg.norm(y_axis_on_robot), 
                z_axis_on_robot / np.linalg.norm(z_axis_on_robot),
            ])
        )
        self.installation_quat_inv = tr3d.quaternions.qinverse(rotation_quat)

    def step(self, robot: BaseWalker, frame_quat: Optional[np.ndarray], dt: float):
        self._step(robot, frame_quat, dt)
        local_velocity = self.get_3d_local_velocity()
        self._global_velocity = tr3d.quaternions.rotate_vector(local_velocity, self._last_frame_quat)

    def __copy__(self):
        raise NotImplementedError()
    
    def __deepcopy__(self, memo):
        raise NotImplementedError()

    def _reset(self, robot: BaseWalker, frame_quat: Optional[np.ndarray]): 
        self._local_velocity = np.zeros(3)
        self._last_frame_quat = np.zeros(4)
        self.receive_position_and_framequat()

    def receive_position_and_framequat(self) -> tuple[np.ndarray, np.ndarray]:
        if self.comm_pipe is not None:
            self.comm_pipe.send(True)
            position, framequat = self.comm_pipe.recv()
            return position, framequat
        else:
            print("Realsense t265 not initialized")
            return np.zeros(3), np.zeros(4)

    def _step(self, robot: BaseWalker, frame_quat: Optional[np.ndarray], dt: float):
        tracker_velocity, received_quat = self.receive_position_and_framequat()
        if np.any(np.isnan(tracker_velocity)):
            print("Nan in realsense t265 step")
        
        tracker_velocity = np.nan_to_num(tracker_velocity, nan=0.0)
        #print(tracker_velocity,received_quat)
        if self.enabled:
            self._local_velocity = tr3d.quaternions.rotate_vector(
                tracker_velocity,
                self.installation_quat_inv
            )
        else:
            self._local_velocity = np.zeros(3)
        
        if frame_quat is None:
            #print(self.installation_quat_inv, received_quat)
            self._last_frame_quat = tr3d.quaternions.qmult(self.installation_quat_inv, received_quat)
        else:
            self._last_frame_quat = frame_quat
    
    def get_3d_local_velocity(self) -> np.ndarray:
        return self._local_velocity

    def get_framequat_wijk(self) -> np.ndarray:
        return self._last_frame_quat
    
    def set_up_thread_runner(self):
        self.cleanup_thread_runner()
        self.comm_pipe, child_pipe = multiprocessing.Pipe()
        self.running_thread = multiprocessing.Process(
            target=__class__.thread_runer, 
            args=(child_pipe,),
            daemon=True
        )
        self.running_thread.start()

    def close(self):
        self.cleanup_thread_runner()

    def cleanup_thread_runner(self):
        if self.running_thread is not None:
            self.running_thread.terminate()
            self.running_thread = None
        
        if self.comm_pipe is not None:
            self.comm_pipe.close()
            self.comm_pipe = None

    @staticmethod
    def init_rs_stream():
        pipe = rs.pipeline()
        cfg = rs.config()
        cfg.enable_stream(rs.stream.pose)
        pipe.start(cfg)
        return pipe

    def thread_runer(
        comm_pipe: multiprocessing.Pipe
    ):
        print("Intel T265 thread started")
        try:
            rs_pipe = __class__.init_rs_stream()
        except:
            print("Failed to initialize realsense t265")
            comm_pipe.close()
            return
        print("Successfully initialized realsense t265")
        last_position = np.zeros(3)
        last_frame_quat = np.zeros(4)
        last_time = 0

        last_step_position = np.zeros(3)
        last_step_framequat = np.zeros(4)
        last_step_time = 0
        last_step_velocity = np.zeros(3)
        
        try:
            while True:
                t = time.time()

                if comm_pipe.poll():
                    comm_pipe.recv()
                    if last_step_time > 0 and last_time > 0:
                        step_dt = last_time - last_step_time
                        if step_dt > 0:
                            global_velocity = (last_position - last_step_position) / step_dt
                            try:
                                if tr3d.quaternions.qnorm(last_step_framequat) < 1e-5:
                                    print("Realsense t265 frame quat is zero, using zero velocity")
                                    local_velocity = np.zeros(3)
                                else:
                                    local_velocity = tr3d.quaternions.rotate_vector(
                                        global_velocity,
                                        tr3d.quaternions.qinverse(last_step_framequat)
                                    )
                            except:
                                print("Error in realsense t265 tr3d.quaternions.rotate_vector()")
                                local_velocity = np.zeros(3)
                        else:
                            global_velocity = np.zeros(3)
                            local_velocity = last_step_velocity
                        
                        if np.any(np.isnan(local_velocity)):
                            print("Nan in realsense t265 local_velocity")
                            print(global_velocity, last_step_framequat)
                            local_velocity = last_step_velocity

                    else:
                        local_velocity = np.zeros(3)
                    
                    local_velocity = np.clip(local_velocity, -1.5, 1.5)
                    comm_pipe.send((local_velocity, last_frame_quat))
                    last_step_position = last_position
                    last_step_framequat = last_frame_quat
                    last_step_time = t
                    last_step_velocity = local_velocity
                
                try:
                    frames = rs_pipe.wait_for_frames(20) # 20 ms timeout
                    pose = frames.get_pose_frame()
                    if pose:
                        data = pose.get_pose_data()
                        new_position = np.array([data.translation.x, data.translation.y, data.translation.z])
                        new_frame_quat = np.array([data.rotation.w, data.rotation.x, data.rotation.y, data.rotation.z])
                        if tr3d.quaternions.qnorm(new_frame_quat) < 1e-5:
                            print("Realsense t265 frame quat is zero!")
                            continue
                        last_position = new_position
                        last_frame_quat = new_frame_quat
                        last_time = t
                except RuntimeError:
                    print("Error in realsense t265 rs_pipe.wait_for_frames()")
                except KeyboardInterrupt:
                    print("Keyboard interrupt in realsense t265 thread")
                    import traceback; traceback.print_exc()
                    break

        finally:
            rs_pipe.stop()
            comm_pipe.close()
            print("Intel T265 thread ended")
rail_real_walker/velocity_estimators/intel_t265.py
realquantumcookie/APRL
import pyrealsense2 as rs
from rail_walker_interface import BaseWalker, Walker3DLocationEstimator, Walker3DFrameQuatEstimator
from typing import Optional
import multiprocessing
import numpy as np
import transforms3d as tr3d
import time

def normalize_rad(angle: float) -> float:
    return (angle + np.pi) % (2 * np.pi) - np.pi

"""
Intel Realsense T265 Velocity/Frame Quaternion Estimator
Please see axis definition here, you need to define how those axis are installed on your robot through the constructor

https://github.com/IntelRealSense/librealsense/blob/master/doc/t265.md
"""
class IntelRealsenseT265PositionEstimator(Walker3DLocationEstimator, Walker3DFrameQuatEstimator):
    def __init__(
        self,
        x_axis_on_robot: np.ndarray = np.array([1, 0, 0]),
        y_axis_on_robot: np.ndarray = np.array([0, 1, 0]),
        z_axis_on_robot: np.ndarray = np.array([0, 0, 1]),
    ):
        Walker3DLocationEstimator.__init__(self)
        Walker3DFrameQuatEstimator.__init__(self)
        self.comm_pipe : Optional[multiprocessing.Pipe] = None
        self.running_thread : Optional[multiprocessing.Process] = None
        self._enabled = True
        
        self._last_position = np.zeros(3)
        self._last_frame_quat = np.zeros(4)
        self._rs_quat_estimate = np.zeros(4)
        self.setup_installation_quat_inv(x_axis_on_robot, y_axis_on_robot, z_axis_on_robot)
        self.set_up_thread_runner()

    @property
    def enabled(self) -> bool:
        return self._enabled
    
    @enabled.setter
    def enabled(self, value: bool) -> None:
        self._enabled = value
    
    def setup_installation_quat_inv(
        self, 
        x_axis_on_robot: np.ndarray,
        y_axis_on_robot: np.ndarray,
        z_axis_on_robot: np.ndarray
    ):
        assert np.any(x_axis_on_robot != 0) and np.any(y_axis_on_robot != 0) and np.any(z_axis_on_robot != 0)
        assert np.inner(x_axis_on_robot, y_axis_on_robot) == 0 and np.inner(x_axis_on_robot, z_axis_on_robot) == 0 and np.inner(y_axis_on_robot, z_axis_on_robot) == 0
        
        rotation_quat = tr3d.quaternions.mat2quat(
            np.hstack([
                x_axis_on_robot / np.linalg.norm(x_axis_on_robot), 
                y_axis_on_robot / np.linalg.norm(y_axis_on_robot), 
                z_axis_on_robot / np.linalg.norm(z_axis_on_robot),
            ])
        )
        self.installation_quat_inv = tr3d.quaternions.qinverse(rotation_quat)

    def __copy__(self):
        raise NotImplementedError()
    
    def __deepcopy__(self, memo):
        raise NotImplementedError()

    def update_position_estimate(self):
        local_position_rs_shift, rs_framequat = self.receive_position_shift_and_framequat()
        if not self.enabled:
            return

        rs_robot_framequat = tr3d.quaternions.qmult(
            self.installation_quat_inv,
            rs_framequat
        )
        self._rs_quat_estimate = rs_robot_framequat

        framequat_used = self._last_frame_quat if self._last_frame_quat is not None else rs_robot_framequat
        
        # quat_from_rs_to_robot_global = tr3d.quaternions.qmult(
        #     self.installation_quat_inv,
        #     framequat_used
        # )
        # global_position_shift = tr3d.quaternions.rotate_vector(
        #     local_position_rs_shift,
        #     quat_from_rs_to_robot_global
        # )
        
        local_position_shift = tr3d.quaternions.rotate_vector(
            local_position_rs_shift,
            self.installation_quat_inv
        )
        global_position_shift = tr3d.quaternions.rotate_vector(
            local_position_shift,
            framequat_used
        )
        self._last_position += global_position_shift

    def reset(self, robot: BaseWalker, frame_quat: Optional[np.ndarray]):
        if not np.all(self._last_frame_quat == 0): # If already initialized
            self.update_position_estimate()
        else:
            self.receive_position_shift_and_framequat() # Let RS thread clear up the accumulated position shift

        self._last_frame_quat = frame_quat
        Walker3DLocationEstimator.reset(
            self,
            robot,
            frame_quat
        )
    
    def step(self, robot: BaseWalker, frame_quat: Optional[np.ndarray], dt: float):
        self.update_position_estimate()
        self._last_frame_quat = frame_quat
        Walker3DLocationEstimator.step(
            self,
            robot,
            frame_quat,
            dt
        )
    
    def get_framequat_wijk(self) -> np.ndarray:
        return self._rs_quat_estimate
    
    def _estimate_location(self, robot: BaseWalker, frame_quat: Optional[np.ndarray]) -> np.ndarray:
        return self._last_position

    def receive_position_shift_and_framequat(self) -> tuple[np.ndarray, np.ndarray]:
        if self.comm_pipe is not None:
            self.comm_pipe.send(True)
            position_shift_local, framequat = self.comm_pipe.recv()
            return position_shift_local, framequat
        else:
            raise RuntimeError("Intel T265 thread is not running")
    
    def set_up_thread_runner(self):
        self.cleanup_thread_runner()
        self.comm_pipe, child_pipe = multiprocessing.Pipe()
        self.running_thread = multiprocessing.Process(
            target=__class__.thread_runer, 
            args=(child_pipe,),
            daemon=True
        )
        self.running_thread.start()

    def close(self):
        self.cleanup_thread_runner()

    def cleanup_thread_runner(self):
        if self.running_thread is not None:
            self.running_thread.terminate()
            self.running_thread = None
        
        if self.comm_pipe is not None:
            self.comm_pipe.close()
            self.comm_pipe = None

    @staticmethod
    def init_rs_stream():
        pipe = rs.pipeline()
        cfg = rs.config()
        cfg.enable_stream(rs.stream.pose)
        pipe.start(cfg)
        return pipe

    def thread_runer(
        comm_pipe: multiprocessing.Pipe
    ):
        print("Intel T265 thread started")
        try:
            rs_pipe = __class__.init_rs_stream()
        except:
            print("Failed to initialize realsense t265")
            comm_pipe.close()
            return
        print("Successfully initialized realsense t265")
        last_position = np.zeros(3)
        last_frame_quat = np.zeros(4)
        last_time = 0

        last_step_position = np.zeros(3)
        last_step_framequat = np.zeros(4)
        last_step_time = 0
        
        try:
            while True:
                t = time.time()

                if comm_pipe.poll():
                    comm_pipe.recv()
                    if last_step_time > 0 and last_time > 0:
                        step_dt = last_time - last_step_time
                        if step_dt > 0:
                            global_pos_shift = last_position - last_step_position
                        else:
                            global_pos_shift = np.zeros(3)
                        try:
                            if tr3d.quaternions.qnorm(last_step_framequat) < 1e-5:
                                print("Realsense t265 frame quat is zero, using zero pos shift")
                                local_pos_shift = np.zeros(3)
                            else:
                                local_pos_shift = tr3d.quaternions.rotate_vector(
                                    global_pos_shift,
                                    tr3d.quaternions.qinverse(last_step_framequat)
                                )
                        except:
                            print("Error in realsense t265 tr3d.quaternions.rotate_vector()")
                            local_pos_shift = np.zeros(3)
                    else:
                        local_pos_shift = np.zeros(3)
                    
                    comm_pipe.send((local_pos_shift, last_frame_quat))
                    last_step_position = last_position
                    last_step_framequat = last_frame_quat
                    last_step_time = t
                
                try:
                    frames = rs_pipe.wait_for_frames(20) # 20 ms timeout
                    pose = frames.get_pose_frame()
                    if pose:
                        data = pose.get_pose_data()
                        last_position = np.array([data.translation.x, data.translation.y, data.translation.z])
                        last_frame_quat = np.array([data.rotation.w, data.rotation.x, data.rotation.y, data.rotation.z])
                        last_time = t
                except:
                    print("Error in realsense t265 rs_pipe.wait_for_frames()")


        finally:
            rs_pipe.stop()
            comm_pipe.close()
rail_real_walker/velocity_estimators/intel_t265_position_est.py
realquantumcookie/APRL
from typing import Optional
from .foot_contact_estimator import Go1ForwardKinematicsVelocityProvider
from rail_walker_interface import Walker3DVelocityEstimatorLocal, BaseWalker
from filterpy.kalman import KalmanFilter
import numpy as np
from typing import Optional
import transforms3d as tr3d

DEFAULT_ACCELEROMETER_NOISE = 0.0306
DEFAULT_FK_NOISE = 0.00621
GRAVITY_VECTOR = np.array([0, 0, 9.81])

class KalmanFilterFusedVelocityEstimator(Walker3DVelocityEstimatorLocal):
    def __init__(
        self, 
        accelerometer_variance: float = DEFAULT_ACCELEROMETER_NOISE,
        sensor_variance: float = DEFAULT_FK_NOISE,
        sensor : Walker3DVelocityEstimatorLocal = Go1ForwardKinematicsVelocityProvider(),
        initial_variance: float = 0.1
    ):
        super().__init__()
        self.sensor = sensor

        self._initial_variance = initial_variance
        self.filter = KalmanFilter(dim_x=3, dim_z=3, dim_u=3)
        
        self.filter.x = np.zeros(3)
        self.filter.P = np.eye(3) * self._initial_variance
        self.filter.Q = np.eye(3) * accelerometer_variance
        self.filter.R = np.eye(3) * sensor_variance

        self.filter.H = np.eye(3) # measurement function, y = Hx
        self.filter.F = np.eye(3) # state transition matrix
        self.filter.B = np.eye(3)

        self._estimated_local_velocity = np.zeros(3)
    
    def _update_velocity_estimate(self, robot: BaseWalker, dt: float):

        raw_acceleration = robot.get_3d_acceleration_local()
        gravity_local = tr3d.quaternions.rotate_vector(
            GRAVITY_VECTOR,
            tr3d.quaternions.qinverse(robot.get_framequat_wijk())
        )
        calibrated_local_acc = raw_acceleration - gravity_local
        self.filter.predict(u=calibrated_local_acc * dt)
        self.filter.update(self.sensor.get_3d_local_velocity())
        self._estimated_local_velocity = self.filter.x.copy()

    def _reset(self, robot: BaseWalker, frame_quat: Optional[np.ndarray]):
        self.sensor._reset(robot, frame_quat)
        
        self.filter.x = np.zeros(3)
        # self.filter.P = np.eye(3) * self._initial_variance

        self._estimated_local_velocity = np.zeros(3)

    def _step(self, robot: BaseWalker, frame_quat: np.ndarray, dt: float):
        self.sensor._step(robot, frame_quat, dt)
        self._update_velocity_estimate(robot, dt)
    
    def get_3d_local_velocity(self) -> np.ndarray:
        return self._estimated_local_velocity
rail_real_walker/velocity_estimators/kalman_filter_fused_estimator.py
realquantumcookie/APRL
from rail_walker_interface import BaseWalker, Walker3DVelocityEstimatorLocal, Walker3DFrameQuatEstimator
from .remote_t265_runner import DataPack, RemoteT265Data
from typing import Optional, Any
import numpy as np
import transforms3d as tr3d
import time
import socket
from serde.msgpack import from_msgpack
import errno

class IntelRealsenseT265EstimatorRemote(Walker3DVelocityEstimatorLocal):
    def __init__(
        self,
        x_axis_on_robot: np.ndarray = np.array([1, 0, 0]),
        y_axis_on_robot: np.ndarray = np.array([0, 1, 0]),
        z_axis_on_robot: np.ndarray = np.array([0, 0, 1]),
        server_addr : tuple[str, int] = ("192.168.123.161",6000)
    ):
        Walker3DVelocityEstimatorLocal.__init__(self)
        Walker3DFrameQuatEstimator.__init__(self)
        self.socket_client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket_client.setblocking(True)
        self.socket_client.settimeout(0.5)
        self.socket_client.connect(server_addr)
        self.socket_client.setblocking(False)
        
        self._server_addr = server_addr
        self.data_pack = DataPack(self.deal_with_data)

        self._enabled = True

        self._last_data : Optional[RemoteT265Data] = None
        self._last_data_t : float = 0.0

        self._local_velocity = np.zeros(3)
        self.setup_installation_quat_inv(x_axis_on_robot, y_axis_on_robot, z_axis_on_robot)
    
    def deal_with_data(self, command : bytes, data : bytes, is_last_data : bool, custom : Any) -> Any:
        if command == b"o" and is_last_data:
            try:
                new_data = from_msgpack(RemoteT265Data, data)
            except KeyboardInterrupt as e:
                raise e
            except BaseException as e:
                print("Error in remote T265 Data Decode", e)
                return

            new_t = time.time()
            if np.linalg.norm(new_data.rotation_wxyz) < 1e-6:
                return

            # if self._last_data_t > 0 and new_t - self._last_data_t < 1.0:
            #     dt = new_t - self._last_data_t
            #     tracker_velocity = (new_data.translation - self._last_data.translation) / dt
            #     tracker_local_velocity = tr3d.quaternions.rotate_vector(
            #         tracker_velocity,
            #         tr3d.quaternions.qinverse(self._last_data.rotation_wxyz)
            #     )
            # else:
            tracker_velocity = np.asarray(new_data.velocity)
            tracker_local_velocity = tr3d.quaternions.rotate_vector(
                tracker_velocity,
                tr3d.quaternions.qinverse(np.asarray(new_data.rotation_wxyz))
            )
            local_velocity = tr3d.quaternions.rotate_vector(
                tracker_local_velocity,
                self.installation_quat_inv
            )
            self._last_data = new_data
            self._last_data_t = new_t
            self._local_velocity = local_velocity

    def send_read_obs_data(self) -> None:
        pass
        # self.socket_client.send(self.data_pack.encode(b"o",b""))
    
    def try_receive_data(self) -> None:
        while True:
            try:
                data = self.socket_client.recv(8192)
                if not data:
                    raise RuntimeError("No data received from remote Go1 Runner! The TCP connection might be broken!")
                self.data_pack.feed_data(data, time.time())
                if len(data) < 8192:
                    return
            except socket.error as e:
                if e.errno == errno.EAGAIN or e.errno == errno.EWOULDBLOCK:
                    return
                else:
                    raise e

    @property
    def server_addr(self) -> tuple[str, int]:
        return self._server_addr

    @property
    def enabled(self) -> bool:
        return self._enabled
    
    @enabled.setter
    def enabled(self, value: bool) -> None:
        self._enabled = value

    def setup_installation_quat_inv(
        self, 
        x_axis_on_robot: np.ndarray,
        y_axis_on_robot: np.ndarray,
        z_axis_on_robot: np.ndarray
    ):
        assert np.any(x_axis_on_robot != 0) and np.any(y_axis_on_robot != 0) and np.any(z_axis_on_robot != 0)
        assert np.inner(x_axis_on_robot, y_axis_on_robot) == 0 and np.inner(x_axis_on_robot, z_axis_on_robot) == 0 and np.inner(y_axis_on_robot, z_axis_on_robot) == 0
        
        rotation_quat = tr3d.quaternions.mat2quat(
            np.hstack([
                x_axis_on_robot / np.linalg.norm(x_axis_on_robot), 
                y_axis_on_robot / np.linalg.norm(y_axis_on_robot), 
                z_axis_on_robot / np.linalg.norm(z_axis_on_robot),
            ])
        )
        self.installation_quat_inv = tr3d.quaternions.qinverse(rotation_quat)

    def _step(self, robot: BaseWalker, frame_quat: Optional[np.ndarray], dt: float):
        self.send_read_obs_data()
        self.try_receive_data()

    def __copy__(self):
        raise NotImplementedError()
    
    def __deepcopy__(self, memo):
        raise NotImplementedError()

    def _reset(self, robot: BaseWalker, frame_quat: Optional[np.ndarray]): 
        self.send_read_obs_data()
        self.try_receive_data()
        self._local_velocity = np.zeros(3)
    
    def get_3d_local_velocity(self) -> np.ndarray:
        return self._local_velocity
    
    def close(self) -> None:
        self.socket_client.close()

rail_real_walker/velocity_estimators/remote_t265.py
realquantumcookie/APRL
import socket
from serde import serde
from serde.msgpack import to_msgpack
from dataclasses import dataclass
import time
import errno
import pyrealsense2 as rs
import numpy as np
from enum import Enum
import select
from typing import Any
import os
import psutil
from typing import Tuple

class DataPack:
    def __init__(self, deal_command_callback):
        self.remaining_data = b""
        self.deal_command_callback = deal_command_callback
    
    def feed_data(self, dat: bytes, custom : Any) -> Any:
        self.remaining_data += dat
        return self.stream_decode(custom)
    
    def clear_data(self) -> None:
        self.remaining_data = b""

    def stream_decode(self, custom : Any) -> Any:
        decoded = True
        ret = None
        while decoded:
            decoded, ret = self.try_decode(custom)
        return ret

    def try_decode(self, custom: Any) -> Tuple[bool, Any]:
        if len(self.remaining_data) < 3:
            return False, None
        command = self.remaining_data[0:1]
        length = int.from_bytes(self.remaining_data[1:3], 'big')

        if len(self.remaining_data) < length + 3:
            return False, None
        
        if len(self.remaining_data) < length + 3 + 3:
            last_data = True
        else:
            next_length = int.from_bytes(self.remaining_data[length+3+1:length+3+3], 'big')
            last_data = not len(self.remaining_data) >= (length+3)+3+next_length

        data = self.remaining_data[3:length + 3]
        ret = self.deal_command_callback(command, data, last_data, custom)
        self.remaining_data = self.remaining_data[length + 3:]
        return True, ret
    
    @staticmethod
    def encode(command : bytes, data : bytes) -> bytes:
        length = len(data)
        enc = command + length.to_bytes(2, 'big') + data
        return enc

class RemoteT265Confidence(Enum):
    FAILED = 0x0
    LOW = 0x1
    MEDIUM = 0x2
    HIGH = 0x3

@serde
@dataclass
class RemoteT265Data:
    # acceleration: np.ndarray
    # angular_acceleration: np.ndarray
    # angular_velocity: np.ndarray
    # mapper_confidence: int
    tracker_confidence: int
    rotation_wxyz: Tuple[float, float, float, float]
    translation: Tuple[float, float, float]
    velocity: Tuple[float, float, float]

def empty_t265_data() -> RemoteT265Data:
    return RemoteT265Data(
        # acceleration=np.zeros(3),
        # angular_acceleration=np.zeros(3),
        # angular_velocity=np.zeros(3),
        # mapper_confidence=RemoteT265Confidence.FAILED.value,
        tracker_confidence=RemoteT265Confidence.FAILED.value,
        rotation_wxyz=(1.0, 0.0, 0.0, 0.0),
        translation=(0.0, 0.0, 0.0),
        velocity=(0.0, 0.0, 0.0)
    )

class RemoteT265Updater:
    def __init__(
        self,
        wait_ms : int
    ):
        self.wait_ms = wait_ms
        self.last_data : RemoteT265Data = empty_t265_data()
        self.last_time = 0
        self.rs_pipe = None

    @staticmethod
    def create_rs_stream():
        pipe = rs.pipeline()
        cfg = rs.config()
        cfg.enable_stream(rs.stream.pose)
        pipe.start(cfg)
        return pipe

    def try_init_rs(self) -> bool:
        if self.rs_pipe is None:
            print("Initializing realsense t265")
            try:
                self.rs_pipe = __class__.create_rs_stream()
            except RuntimeError as e:
                return False
        return True
    
    @property
    def is_rs_inited(self) -> bool:
        return self.rs_pipe is not None
    
    @staticmethod
    def convert_rs_pose_to_t265_data(pose : rs.pose) -> RemoteT265Data:
        return RemoteT265Data(
            # acceleration=np.array([pose.acceleration.x, pose.acceleration.y, pose.acceleration.z]),
            # angular_acceleration=np.array([pose.angular_acceleration.x, pose.angular_acceleration.y, pose.angular_acceleration.z]),
            # angular_velocity=np.array([pose.angular_velocity.x, pose.angular_velocity.y, pose.angular_velocity.z]),
            # mapper_confidence=int(pose.mapper_confidence),
            tracker_confidence=int(pose.tracker_confidence),
            rotation_wxyz=(pose.rotation.w, pose.rotation.x, pose.rotation.y, pose.rotation.z),
            translation=(pose.translation.x, pose.translation.y, pose.translation.z),
            velocity=(pose.velocity.x, pose.velocity.y, pose.velocity.z)
        )

    def try_update(self) -> bool:
        if self.rs_pipe is None:
            return False
        try:
            frames = self.rs_pipe.wait_for_frames(self.wait_ms) # variable timeout
        except RuntimeError:
            return False
        t = time.time()
        pose = frames.get_pose_frame()
        if pose:
            data = pose.get_pose_data()
            self.last_data = __class__.convert_rs_pose_to_t265_data(data)
            self.last_time = t
            return True
        return False

    def close(self) -> None:
        if self.rs_pipe is not None:
            self.rs_pipe.stop()
            self.rs_pipe = None

    def __del__(self):
        self.close()


class RemoteT265VelocityEstimatorRunner:
    def __init__(
        self,
        bind_address : Tuple[str, int] = ("", 6000),
        wait_ms : int = 20
    ):
        self.tcp_server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.tcp_server.bind(bind_address)
        self.tcp_server.listen(1)
        self.tcp_server.setblocking(False)
        self.tcp_clients : list[socket.socket] = []

        self._bind_address = bind_address
        self.updater = RemoteT265Updater(wait_ms=wait_ms)
        self.data_pack = DataPack(self.deal_command)
        self._next_to_send = self.data_pack.encode(
            b"o",
            to_msgpack(empty_t265_data())
        )
        self._updater_fail_count : int = 0
    
    def deal_command(self, command : bytes, data : bytes, is_last_data : bool, custom : socket.socket) -> Any:
        if command == b"o":
            custom.send(self._next_to_send)

    @property
    def bind_address(self) -> Tuple[str, int]:
        return self._bind_address

    def broadcast_data(self, data : bytes) -> None:
        for client in self.tcp_clients:
            try:
                num_sent = client.send(data)
            except:
                self.disconnect(client)
                continue
            if num_sent <= 0:
                self.disconnect(client)

    def disconnect(self, socket : socket.socket) -> None:
        print("Disconnecting Socket", socket)
        try:
            self.tcp_clients.remove(socket)
        except ValueError:
            pass
        socket.close()
        if len(self.tcp_clients) <= 0:
            self.data_pack.clear_data()

    def init_rs_until_success(self) -> None:
        while not self.updater.try_init_rs():
            time.sleep(5.0)
        print("Realsense t265 initialized")

    def try_update_t265(self) -> None:
        if self._updater_fail_count > 50:
            print("Reinitializing realsense t265 due to too many failures")
            self.init_rs_until_success()
            self._updater_fail_count = 0
        if self.updater.try_update():
            self._updater_fail_count = 0
            if len(self.tcp_clients) > 0:
                to_send = self.data_pack.encode(
                    b"o",
                    to_msgpack(self.updater.last_data)
                )
                self._next_to_send = to_send
                self.broadcast_data(to_send)
        else:
            self._updater_fail_count += 1

    def run(self):
        pid = os.getpid()
        process = psutil.Process(pid)
        try:
            process.nice(-5)
        except psutil.AccessDenied:
            print("Cannot set process priority, please run as root!")

        self.init_rs_until_success()
        try:
            while True:
                self.try_update_t265()
                if len(self.tcp_clients) <= 0:
                    time.sleep(0.1)
                
                readable, writable, errored = select.select([self.tcp_server] + self.tcp_clients, [], [], 0.0)
                for s in readable:
                    if s is self.tcp_server:
                        client_socket, address = self.tcp_server.accept()
                        self.tcp_clients.append(client_socket)
                        print ("Connection from", address)
                    else:
                        try:
                            data = s.recv(1024)
                        except socket.error as e:
                            if e.errno == errno.ECONNRESET:
                                self.disconnect(s)
                                break
                            elif e.errno != errno.EAGAIN and e.errno != errno.EWOULDBLOCK:
                                raise e
                            continue
                        if data:
                            self.data_pack.feed_data(data, s)
                        else:
                            self.disconnect(s)
                            break
        finally:
            process.nice(5)
    
    def close(self):
        self.updater.close()

    def __del__(self):
        self.close()

rail_real_walker/velocity_estimators/remote_t265_runner.py
realquantumcookie/APRL
from .envs import *
from .joystick_policy import *
rail_walker_gym/__init__.py
realquantumcookie/APRL
CONFIG_KP = 60.0
CONFIG_KD = 5.0
CONFIG_ACTION_INTERPOLATION = False
CONFIG_ACTION_RANGE = 0.3 # similar to 2.0x offset in previous versions, limit joint action range to 30% of joint range
CONFIG_TARGET_VELOCITY = 0.5
CONFIG_USE_ENERGY_PENALTY = True
CONFIG_REAL_POWER_PROTECT_FACTOR = 0.4
CONFIG_FLOOR_SCALE = 20.0
rail_walker_gym/envs/register_config.py
realquantumcookie/APRL
import typing
import gym
import gym.wrappers
from gym.envs.registration import register
from typing import Callable
from .wrappers import *
from rail_walker_interface import BaseWalker, JoystickPolicy
from jaxrl5.wrappers import wrap_gym
# from jaxrl5.wrappers.single_precision import SinglePrecision
# from jaxrl5.wrappers.universal_seed import UniversalSeed

# def wrap_gym(env: gym.Env, rescale_actions: bool = True) -> gym.Env:
#     env = SinglePrecision(env)
#     env = UniversalSeed(env)
#     if rescale_actions:
#         env = gym.wrappers.RescaleAction(env, -1, 1)

#     env = gym.wrappers.ClipAction(env)

#     return env

def common_post_process_wrapper(env : gym.Env, action_range: float, filter_actions : int = 0, frame_stack : int = 3, action_history : int = 3, center_init_action : bool = False, log_print = False):
    assert hasattr(env, "joystick_policy"), "The environment must have a joystick policy"
    joystick_policy : JoystickPolicy = env.joystick_policy
    robot = joystick_policy.robot

    robot.limit_action_range = action_range # Set the action range for the robot
    if log_print:
        print("Action Space before applying bound", env.action_space)
    if frame_stack > 1:
        env = FrameStackWrapper(env, frame_stack, exclude_keys=["target_obs", "target_custom", "target_vel"])
    env = ClipActionToRange(env, robot.action_qpos_mins, robot.action_qpos_maxs) # Clip the action space to the range of the robot
    # env = gym.wrappers.ClipAction(env)

    if filter_actions > 0:
        env = ActionFilterWrapper(env, robot.control_timestep, robot, int(filter_actions))
    if action_history > 0:
        env = AddPreviousActions(env, action_history)
    
    if center_init_action:
        env = RescaleActionAsymmetric(env, -1.0, 1.0, robot.joint_qpos_init)
    else:
        env = gym.wrappers.RescaleAction(env, -1.0, 1.0)

    obs_space_before_flatten = env.observation_space
    if log_print:
        print("Obs Space before flattening", env.observation_space)
    env = gym.wrappers.FlattenObservation(env)
    env = wrap_gym(env, True)
    

    return env, obs_space_before_flatten

def get_formatted_register_env_entry_point(
    entry_point : Callable,
):
    def to_ret(
        make_env_kwargs: dict,
        make_env_kwargs_override: list[Callable[[dict], dict]],
        make_env_kwargs_callback:typing.Optional[Callable[[dict],dict]] = None,
    ) -> gym.Env:
        env_kargs = make_env_kwargs.copy()
        if make_env_kwargs_callback is not None:
            env_kargs = make_env_kwargs_callback(env_kargs)
        for override in make_env_kwargs_override:
            env_kargs = override(env_kargs)

        return entry_point(
            **env_kargs
        )
    return to_ret
    

def formatted_register_env(
    id_format:str,
    format_args:list,
    make_entry_point:Callable,
    register_kwargs_override:typing.Optional[dict],
    make_env_kargs: dict,
    make_env_kargs_override: list[Callable[[dict], dict]],
    make_env_kwargs_callback:typing.Optional[Callable[[dict],dict]] = None,
):
    call_kwargs = dict(
        id=id_format.format(*format_args),
        entry_point=get_formatted_register_env_entry_point(
            make_entry_point
        ),
        kwargs=dict(
            make_env_kwargs=make_env_kargs,
            make_env_kwargs_override=make_env_kargs_override.copy(),
            make_env_kwargs_callback=make_env_kwargs_callback
        )
    )
    if register_kwargs_override is not None:
        call_kwargs.update(register_kwargs_override)
    register(**call_kwargs)

def iter_formatted_register_env(
    id_format:str,
    make_entry_point:Callable,
    format_args_list:typing.List[typing.List[typing.Tuple[
        typing.Union[str, int, float],
        typing.Optional[dict],
        typing.Optional[Callable[[dict],dict]]
    ]]],
    base_register_kwargs:dict = {},
    base_make_env_kwargs:dict = {},
    base_make_env_kwargs_callback:typing.Optional[Callable[[dict],dict]] = None,
):
    def _iter_formatted_register(step, format_args : list[str | int | float] = [], register_kargs_override = {}, make_env_kargs_override : list[Callable[[dict],dict]] = []):
        if step >= len(format_args_list):
            register_kargs = base_register_kwargs.copy()
            register_kargs.update(register_kargs_override)
            formatted_register_env(id_format, format_args, make_entry_point, register_kargs, base_make_env_kwargs, make_env_kargs_override, base_make_env_kwargs_callback)
        else:
            for arg, register_kwargs_override, make_env_kwargs_override_fn in format_args_list[step]:
                
                format_args.append(arg)
                make_env_kargs_override.append(make_env_kwargs_override_fn)
                
                if register_kwargs_override is not None:
                    new_register_kwargs_override = register_kargs_override.copy()
                    new_register_kwargs_override.update(register_kwargs_override)
                else:
                    new_register_kwargs_override = register_kargs_override
                
                _iter_formatted_register(step+1, format_args, new_register_kwargs_override, make_env_kargs_override)
                format_args.pop()
                make_env_kargs_override.pop()
    _iter_formatted_register(0, [], {}, [])

rail_walker_gym/envs/register_helper.py
realquantumcookie/APRL
import gym
from dm_control.locomotion import arenas
from .register_helper import iter_formatted_register_env
from rail_mujoco_walker import RailSimWalkerDMControl, Go1SimWalker, A1SimWalker, JoystickPolicyDMControlTask, HeightFieldArena, HEIGHTFIELD_ARENA_GOALS, CrouchingHeightfieldArena
from rail_walker_interface import JoystickPolicy, BaseWalker, WalkerVelocitySmoother
from typing import Any, Optional
from .wrappers import *
from ..joystick_policy import *
from ..joystick_policy_mujoco import *
import numpy as np
from .register_config import *
from mujoco_utils import composer_utils
from dm_control import composer
# from dm_env_wrappers import ActionNoiseWrapper

def make_sim_env(
    robot: RailSimWalkerDMControl,
    joystick_policy_parameters : dict[str,Any],
    floor_override : str = "floor",
    render_height : int = 256,
    render_width : int = 256,
    height_scale : float = 5.0,
) -> gym.Env:
    print("Making sim env")
    if isinstance(robot.unwrapped(), Go1SimWalker):
        robot.control_timestep = 0.05
        robot.control_subtimestep = 0.005
    elif isinstance(robot, A1SimWalker):
        robot.control_timestep = 0.05
        robot.control_subtimestep = 0.002
    
    joystick_policy = JoystickPolicy(
        robot = robot,
        **joystick_policy_parameters
    )

    print("Joystick Policy Parameters", joystick_policy_parameters)
    
    # Initialize floor regardless of floor_override
    if floor_override == "floor":
        floor = arenas.Floor(size=(CONFIG_FLOOR_SCALE,CONFIG_FLOOR_SCALE))
        floor._top_camera.remove()
    elif floor_override == "heightfield_arena":
        floor = HeightFieldArena(CONFIG_FLOOR_SCALE, height_scale=height_scale)
    elif floor_override == "crouching_heightfield_arena":
        floor = CrouchingHeightfieldArena(CONFIG_FLOOR_SCALE)
    else:
        raise ValueError(f"Invalid floor_override {floor_override}, must be one of ['floor', 'heightfield_arena', 'crouching_heightfield_arena']")
    
    task = JoystickPolicyDMControlTask(
        joystick_policy=joystick_policy,
        floor=floor
    )

    env = composer_utils.Environment(
        task=task,
        strip_singleton_obs_buffer_dim=True,
        recompile_physics=False
    )
    # env = composer.Environment(
    #     task=task,
    #     strip_singleton_obs_buffer_dim=True,
    # )
    # env = ActionNoiseWrapper(env, scale=0.05)

    env = DMControlMultiCameraRenderWrapper(
        env,
        task.render_scene_callback,
        render_height,
        render_width
    )
    env = RailWalkerMujocoComplianceWrapper(env)
    return env

iter_formatted_register_env(
    id_format = "{}ResetMujoco-v0",
    make_entry_point = make_sim_env,
    format_args_list = [
        [
            ("Go1", None, lambda kwargs: {
                **kwargs,
                "robot": Go1SimWalker(
                    Kp = CONFIG_KP,
                    Kd = CONFIG_KD,
                    action_interpolation = CONFIG_ACTION_INTERPOLATION,
                    limit_action_range = CONFIG_ACTION_RANGE
                )
            }),
            ("A1", None, lambda kwargs: {
                **kwargs,
                "robot": A1SimWalker(
                    Kp = CONFIG_KP,
                    Kd = CONFIG_KD,
                    action_interpolation = CONFIG_ACTION_INTERPOLATION,
                    limit_action_range=CONFIG_ACTION_RANGE
                )
            })
        ],
    ],
    base_register_kwargs = dict(
        
    ),
    base_make_env_kwargs = dict(
        
    ),
    base_make_env_kwargs_callback = lambda kwargs: {
        **kwargs,
        "floor_override": "floor",
        "joystick_policy_parameters": dict(
            reward_provider=ResetRewardProvider(CONFIG_USE_ENERGY_PENALTY),
            target_yaw_provider=JoystickPolicyForwardOnlyTargetProvider(),
            termination_providers=[],
            truncation_providers=[],
            resetters=[ResetPolicyResetter(init_dist=np.array([1.0, 0.0, 0.0]))],
            target_observable = None,
            enabled_observables = ['joints_pos', 'imu']
        )
    }
)

iter_formatted_register_env(
    id_format = "{}EasyResetMujoco-v0",
    make_entry_point = make_sim_env,
    format_args_list = [
        [
            ("Go1", None, lambda kwargs: {
                **kwargs,
                "robot": Go1SimWalker(
                    Kp = CONFIG_KP,
                    Kd = CONFIG_KD,
                    action_interpolation = CONFIG_ACTION_INTERPOLATION,
                    limit_action_range=CONFIG_ACTION_RANGE
                )
            }),
            ("A1", None, lambda kwargs: {
                **kwargs,
                "robot": A1SimWalker(
                    Kp = CONFIG_KP,
                    Kd = CONFIG_KD,
                    action_interpolation = CONFIG_ACTION_INTERPOLATION,
                    limit_action_range=CONFIG_ACTION_RANGE
                )
            })
        ],
    ],
    base_register_kwargs = dict(
        
    ),
    base_make_env_kwargs = dict(
        
    ),
    base_make_env_kwargs_callback = lambda kwargs: {
        **kwargs,
        "floor_override": "floor",
        "joystick_policy_parameters": dict(
            reward_provider=ResetRewardProvider(CONFIG_USE_ENERGY_PENALTY),
            target_yaw_provider=JoystickPolicyForwardOnlyTargetProvider(),
            termination_providers=[],
            truncation_providers=[],
            resetters=[ResetPolicyResetter(init_dist=np.array([0.5, 0.5, 0.0]))],
            target_observable = None,
            enabled_observables = ['joints_pos', 'imu']
        )
    }
)

iter_formatted_register_env(
    id_format = "{}GatedResetMujoco-v0",
    make_entry_point = make_sim_env,
    format_args_list = [
        [
            ("Go1", None, lambda kwargs: {
                **kwargs,
                "robot": Go1SimWalker(
                    Kp = CONFIG_KP,
                    Kd = CONFIG_KD,
                    action_interpolation = CONFIG_ACTION_INTERPOLATION,
                    limit_action_range=CONFIG_ACTION_RANGE
                )
            }),
            ("A1", None, lambda kwargs: {
                **kwargs,
                "robot": A1SimWalker(
                    Kp = CONFIG_KP,
                    Kd = CONFIG_KD,
                    action_interpolation = CONFIG_ACTION_INTERPOLATION,
                    limit_action_range=CONFIG_ACTION_RANGE
                )
            })
        ],
    ],
    base_register_kwargs = dict(
        
    ),
    base_make_env_kwargs = dict(
        
    ),
    base_make_env_kwargs_callback = lambda kwargs: {
        **kwargs,
        "floor_override": "floor",
        "joystick_policy_parameters": dict(
            reward_provider=GatedResetRewardProvider(CONFIG_USE_ENERGY_PENALTY),
            target_yaw_provider=JoystickPolicyForwardOnlyTargetProvider(),
            termination_providers=[],
            truncation_providers=[],
            resetters=[ResetPolicyResetter()],
            target_observable = None,
            enabled_observables = ['joints_pos', 'imu']
        )
    }
)

iter_formatted_register_env(
    id_format = "{}GatedEasyResetMujoco-v0",
    make_entry_point = make_sim_env,
    format_args_list = [
        [
            ("Go1", None, lambda kwargs: {
                **kwargs,
                "robot": Go1SimWalker(
                    Kp = CONFIG_KP,
                    Kd = CONFIG_KD,
                    action_interpolation = CONFIG_ACTION_INTERPOLATION,
                    limit_action_range=CONFIG_ACTION_RANGE
                )
            }),
            ("A1", None, lambda kwargs: {
                **kwargs,
                "robot": A1SimWalker(
                    Kp = CONFIG_KP,
                    Kd = CONFIG_KD,
                    action_interpolation = CONFIG_ACTION_INTERPOLATION,
                    limit_action_range=CONFIG_ACTION_RANGE
                )
            })
        ],
    ],
    base_register_kwargs = dict(
        
    ),
    base_make_env_kwargs = dict(
        
    ),
    base_make_env_kwargs_callback = lambda kwargs: {
        **kwargs,
        "floor_override": "floor",
        "joystick_policy_parameters": dict(
            reward_provider=GatedResetRewardProvider(CONFIG_USE_ENERGY_PENALTY),
            target_yaw_provider=JoystickPolicyForwardOnlyTargetProvider(),
            termination_providers=[],
            truncation_providers=[],
            resetters=[ResetPolicyResetter(init_dist=np.array([0.5, 0.5, 0.0]))],
            target_observable = None,
            enabled_observables = ['joints_pos', 'imu']
        )
    }
)

def make_sanity_lambda_sim(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_yaw_provider'] = JoystickPolicyForwardOnlyTargetProvider()
    kwargs['joystick_policy_parameters']['resetters'] = [JoystickPolicyPointInSimResetter(np.zeros(2),0)]
    return kwargs

def make_autojoystick_lambda_sim(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_yaw_provider'] = JoystickPolicyAutoJoystickTargetProvider()
    kwargs['joystick_policy_parameters']['resetters'] = [JoystickPolicyPointInSimResetter(np.zeros(2),0)]
    kwargs['robot'] = WalkerVelocitySmoother(
        kwargs['robot'], 
        np.array([0.05, 0.05, 0.1, 0.1, 0.2, 0.5])
    )
    return kwargs

def make_autojoystick_simple_lambda_sim(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_yaw_provider'] = JoystickPolicyAutoJoystickSimpleTargetProvider()
    kwargs['joystick_policy_parameters']['resetters'] = [JoystickPolicyPointInSimResetter(np.zeros(2),0)]
    kwargs['robot'] = WalkerVelocitySmoother(
        kwargs['robot'], 
        np.array([0.05, 0.05, 0.1, 0.1, 0.2, 0.5])
    )
    return kwargs

def make_autojoystick_simple_plus_lambda_sim(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_yaw_provider'] = JoystickPolicyAutoJoystickSimplePlusTargetProvider()
    kwargs['joystick_policy_parameters']['resetters'] = [JoystickPolicyPointInSimResetter(np.zeros(2),0)]
    kwargs['robot'] = WalkerVelocitySmoother(
        kwargs['robot'], 
        np.array([0.05, 0.05, 0.1, 0.1, 0.2, 0.5])
    )
    return kwargs

def make_smoothed_target_lambda(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_yaw_provider'] = JoystickPolicySmoothedTargetProvider(
        kwargs['joystick_policy_parameters']['target_yaw_provider'],
        max_rate=5.0/180.0*np.pi,
    )
    return kwargs

def make_limited_target_lambda(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_yaw_provider'] = JoystickPolicyLimitedTargetProvider(
        kwargs['joystick_policy_parameters']['target_yaw_provider'],
        max_delta_angle=30.0/180.0*np.pi,
    )
    return kwargs


def make_target_delta_yaw_obs_lambda(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_observable'] = JoystickPolicyTargetDeltaYawObservable()
    return kwargs

def make_cos_sin_target_delta_yaw_obs_lambda(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_observable'] = JoystickPolicyCosSinTargetDeltaYawObservable()
    return kwargs

def make_target_delta_yaw_and_dist_obs_lambda(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_observable'] = JoystickPolicyTargetDeltaYawAndDistanceObservable()
    return kwargs

def make_empty_obs_lambda(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_observable'] = None
    return kwargs

def make_strict_reward_lambda(
    kwargs: dict[str,Any]
) -> dict[str, Any]:
    kwargs['joystick_policy_parameters']['reward_provider'] = JoystickPolicyStrictRewardProvider()
    return kwargs

def make_sep_reward_lambda(
    kwargs: dict[str,Any]
) -> dict[str, Any]:
    kwargs['joystick_policy_parameters']['reward_provider'] = JoystickPolicySeperateRewardProvider()
    return kwargs

def make_eth_reward_lambda(
    kwargs: dict[str, Any]
) -> dict[str, Any]:
    kwargs["joystick_policy_parameters"]["reward_provider"] = JoystickPolicyETHRewardProvider()
    return kwargs

def make_awi_reward_lambda(
    kwargs: dict[str, Any]
) -> dict[str, Any]:
    kwargs["joystick_policy_parameters"]["reward_provider"] = WalkInTheParkRewardProvider()
    return kwargs

iter_formatted_register_env(
    id_format = "{}{}Mujoco-{}{}-{}-v0",
    make_entry_point = make_sim_env,
    format_args_list = [
        [
            ("Go1", None, lambda kwargs: {
                **kwargs,
                "robot": Go1SimWalker(
                    Kp = CONFIG_KP,
                    Kd = CONFIG_KD,
                    action_interpolation = CONFIG_ACTION_INTERPOLATION,
                    limit_action_range=CONFIG_ACTION_RANGE,
                )
            }),
            ("A1", None, lambda kwargs: {
                **kwargs,
                "robot": A1SimWalker(
                    Kp = CONFIG_KP,
                    Kd = CONFIG_KD,
                    action_interpolation = CONFIG_ACTION_INTERPOLATION,
                    limit_action_range=CONFIG_ACTION_RANGE
                )
            }),
        ],
        [
            ("AutoJoystick", None, make_autojoystick_lambda_sim),
            ("AutoJoystickSimple", None, make_autojoystick_simple_lambda_sim),
            ("AutoJoystickSimplePlus", None, make_autojoystick_simple_plus_lambda_sim),
            ("Sanity", None, make_sanity_lambda_sim),
        ],
        [
            ("Smoothed", None, make_smoothed_target_lambda),
            ("Limited", None, make_limited_target_lambda),
            ("", None, lambda kwargs: kwargs)
        ],
        [
            ("TDY", None, make_target_delta_yaw_obs_lambda),
            ("CosSinTDY", None, make_cos_sin_target_delta_yaw_obs_lambda),
            ("TDYDist", None, make_target_delta_yaw_and_dist_obs_lambda),
            ("Empty", None, make_empty_obs_lambda)
        ],
        [
            ("StrictRew", None, make_strict_reward_lambda),
            ("SepRew", None, make_sep_reward_lambda),
            ("ETHRew", None, make_eth_reward_lambda),
            ("AWIRew", None, make_awi_reward_lambda)
        ]
    ],
    base_register_kwargs = dict(
        
    ),
    base_make_env_kwargs = dict(
        
    ),
    base_make_env_kwargs_callback = lambda kwargs: {
        **kwargs,
        # "floor_override": "floor",
        "joystick_policy_parameters": dict(
            termination_providers=[JoystickPolicyRollPitchTerminationConditionProvider(30/180*np.pi, -25/180*np.pi, 20/180*np.pi)],
            truncation_providers=[],
            enabled_observables=[
                "joints_pos",
                "joints_vel",
                "imu",
                "sensors_local_velocimeter",
                "torques",
                "foot_forces_normalized",
            ],
        )
    }
)

rail_walker_gym/envs/register_mujoco.py
realquantumcookie/APRL
import gym
from .register_helper import iter_formatted_register_env
from rail_walker_interface import JoystickPolicy, BaseWalker, WalkerVelocitySmoother, JoystickEnvImpl
from rail_real_walker import DummyVelocityEstimator, IntelRealsenseT265Estimator, IntelRealsenseT265EstimatorRemote, Go1ForwardKinematicsVelocityProvider, KalmanFilterFusedVelocityEstimator
from rail_real_walker.robots import Go1RealWalker, Go1RealWalkerRemote
from typing import Any
from .wrappers import *
from ..joystick_policy import *
from ..joystick_policy_real import *
import numpy as np
from .register_config import *
import copy

def make_real_env(
    robot: BaseWalker,
    joystick_policy_parameters : dict[str,Any]
) -> gym.Env:
    print("Making real env")
    
    joystick_policy = JoystickPolicy(
        robot = robot,
        **joystick_policy_parameters
    )

    print("Joystick Policy Parameters", joystick_policy_parameters)
    
    env = JoystickEnvImpl(joystick_policy)
    return env

def make_sanity_lambda_real(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_yaw_provider'] = JoystickPolicyForwardOnlyTargetProvider()
    return kwargs

def make_autojoystick_lambda_real(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_yaw_provider'] = JoystickPolicyAutoJoystickTargetProvider()
    return kwargs

def make_realjoystick_lambda_real(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_yaw_provider'] = JoystickPolicyGo1JoystickTargetProvider()
    kwargs['robot'] = WalkerVelocitySmoother(
        kwargs['robot'], 
        np.array([0.05,0.05,0.1,0.1,0.2,0.5])
    )
    return kwargs

def make_realjoystick_simple_lambda_real(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_yaw_provider'] = JoystickPolicyGo1JoystickSimpleTargetProvider()
    kwargs['robot'] = WalkerVelocitySmoother(
        kwargs['robot'], 
        np.array([0.05,0.05,0.1,0.1,0.2,0.5])
    )
    return kwargs

def make_smoothed_target_lambda(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_yaw_provider'] = JoystickPolicySmoothedTargetProvider(
        kwargs['joystick_policy_parameters']['target_yaw_provider'],
        max_rate=5.0/180.0*np.pi,
    )
    return kwargs

def make_limited_target_lambda(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_yaw_provider'] = JoystickPolicyLimitedTargetProvider(
        kwargs['joystick_policy_parameters']['target_yaw_provider'],
        max_delta_angle=30.0/180.0*np.pi,
    )
    return kwargs

def make_target_delta_yaw_obs_lambda(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_observable'] = JoystickPolicyTargetDeltaYawObservable()
    return kwargs

def make_cos_sin_target_delta_yaw_obs_lambda(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_observable'] = JoystickPolicyCosSinTargetDeltaYawObservable()
    return kwargs

def make_target_delta_yaw_and_dist_obs_lambda(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_observable'] = JoystickPolicyTargetDeltaYawAndDistanceObservable()
    return kwargs

def make_empty_obs_lambda(
    kwargs: dict[str,Any]
) -> dict[str,Any]:
    kwargs['joystick_policy_parameters']['target_observable'] = None
    return kwargs

def make_strict_reward_lambda(
    kwargs: dict[str,Any]
) -> dict[str, Any]:
    kwargs['joystick_policy_parameters']['reward_provider'] = JoystickPolicyStrictRewardProvider()
    return kwargs

def make_sep_reward_lambda(
    kwargs: dict[str,Any]
) -> dict[str, Any]:
    kwargs['joystick_policy_parameters']['reward_provider'] = JoystickPolicySeperateRewardProvider()
    return kwargs

iter_formatted_register_env(
    id_format = "{}{}Real-{}{}-{}-v0",
    make_entry_point = make_real_env,
    format_args_list = [
        [
            ("Go1", None, lambda kwargs: {
                **kwargs,
                "robot": Go1RealWalker(
                    velocity_estimator=IntelRealsenseT265Estimator(
                        x_axis_on_robot=np.array([-1, 0, 0]),
                        y_axis_on_robot=np.array([0, 0, 1]),
                        z_axis_on_robot=np.array([0, 1, 0]),
                    ),
                    power_protect_factor=CONFIG_REAL_POWER_PROTECT_FACTOR,
                    Kp=CONFIG_KP,
                    Kd=CONFIG_KD,
                    limit_action_range=CONFIG_ACTION_RANGE,
                    action_interpolation=CONFIG_ACTION_INTERPOLATION
                )
            })
        ],
        [
            ("Sanity", None, make_sanity_lambda_real),
            ("AutoJoystick", None, make_autojoystick_lambda_real),
            ("RealJoystick", None, make_realjoystick_lambda_real),
            ("RealJoystickSimple", None, make_realjoystick_simple_lambda_real)
        ],
        [
            ("Smoothed", None, make_smoothed_target_lambda),
            ("Limited", None, make_limited_target_lambda),
            ("", None, lambda kwargs: kwargs)
        ],
        [
            ("TDY", None, make_target_delta_yaw_obs_lambda),
            ("CosSinTDY", None, make_cos_sin_target_delta_yaw_obs_lambda),
            ("TDYDist", None, make_target_delta_yaw_and_dist_obs_lambda),
            ("Empty", None, make_empty_obs_lambda)
        ],
        [
            ("StrictRew", None, make_strict_reward_lambda),
            ("SepRew", None, make_sep_reward_lambda)
        ]
    ],
    base_register_kwargs = dict(
        
    ),
    base_make_env_kwargs = dict(
        
    ),
    base_make_env_kwargs_callback = lambda kwargs: {
        **kwargs,
        "joystick_policy_parameters": dict(
            termination_providers=[JoystickPolicyRollPitchTerminationConditionProvider(30/180*np.pi, -30/180*np.pi, 25/180 * np.pi)],
            truncation_providers=[],
            resetters=[JoystickPolicyManualResetter()],
            enabled_observables = [
                "joints_pos",
                "joints_vel",
                "imu",
                "sensors_local_velocimeter",
                "torques", 
                "foot_forces_normalized"
            ]
        )
    }
)

# For debugging real world reset policy

def lambda_sitting_real_env(
    kwargs
):
    new_kwargs = copy.copy(kwargs)
    new_kwargs["joystick_policy_parameters"]["resetters"] = [JoystickPolicyManualResetter(target_pose="sitting")]
    return new_kwargs

def lambda_standing_real_env(
    kwargs
):
    new_kwargs = copy.copy(kwargs)
    new_kwargs["joystick_policy_parameters"]["resetters"] = [JoystickPolicyManualResetter(target_pose="standing")]
    return new_kwargs

iter_formatted_register_env(
    id_format = "{}{}ResetReal-v0",
    make_entry_point = make_real_env,
    format_args_list = [
        [
            ("Go1", None, lambda kwargs: {
                **kwargs,
                "robot": Go1RealWalker(
                    velocity_estimator=DummyVelocityEstimator(),
                    power_protect_factor=CONFIG_REAL_POWER_PROTECT_FACTOR,
                    Kp=CONFIG_KP,
                    Kd=CONFIG_KD,
                    limit_action_range=CONFIG_ACTION_RANGE,
                    action_interpolation=CONFIG_ACTION_INTERPOLATION
                )
            }),
        ],
        [
            ("Sitting", None, lambda_sitting_real_env),
            ("Standing", None, lambda_standing_real_env)
        ]
    ],
    base_register_kwargs = dict(
        
    ),
    base_make_env_kwargs = dict(
        
    ),
    base_make_env_kwargs_callback = lambda kwargs: {
        **kwargs,
        "joystick_policy_parameters": dict(
            reward_provider=JoystickPolicyStrictRewardProvider(),
            target_yaw_provider=JoystickPolicyForwardOnlyTargetProvider(),
            termination_providers=[],
            truncation_providers=[],
            resetters=[],
            target_observable = None,
            enabled_observables = ["joints_pos", "imu"]
        )
    }
)
rail_walker_gym/envs/register_real.py
realquantumcookie/APRL
from .obs_wrappers import AddPreviousActions
from .action_filter import ActionFilterWrapper
from .action_wrappers import ClipActionToRange
from .render_wrappers import WanDBVideoWrapper, RenderViewerWrapper
from .dm_to_gym import DMControlMultiCameraRenderWrapper, RailWalkerMujocoComplianceWrapper
from .rollout_collect import RolloutCollector, Rollout
from .resetter_policy_supported import ResetterPolicySupportedEnvironment
from .frame_stack import FrameStackWrapper
from .joystick_policy_truncation import JoystickPolicyTruncationWrapper
from .view_target import JoystickTargetViewer
from .relabel_wrapper import RelabelAggregateWrapper, RelabelTargetProvider, SymmetricRelabelWrapper
from .action_rescale import RescaleActionAsymmetric
rail_walker_gym/envs/wrappers/__init__.py
realquantumcookie/APRL
# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Two types of filters which can be applied to policy output sequences.

1. Simple exponential filter
2. Butterworth filter - lowpass or bandpass

The implementation of the butterworth filter follows scipy's lfilter
https://github.com/scipy/scipy/blob/v1.2.1/scipy/signal/signaltools.py

We re-implement the logic in order to explicitly manage the y states

The filter implements::
       a[0]*y[n] = b[0]*x[n] + b[1]*x[n-1] + ... + b[M]*x[n-M]
                             - a[1]*y[n-1] - ... - a[N]*y[n-N]

We assume M == N.
"""

import collections

import gym
import numpy as np
from absl import logging
from scipy.signal import butter
from rail_walker_interface import BaseWalker
from typing import Any

ACTION_FILTER_ORDER = 2
ACTION_FILTER_LOW_CUT = 0.0
ACTION_FILTER_HIGH_CUT = 4.0

class ActionFilter(object):
    """Implements a generic lowpass or bandpass action filter."""

    def __init__(self, a, b, order, num_joints, ftype='lowpass'):
        """Initializes filter.

    Either one per joint or same for all joints.

    Args:
      a: filter output history coefficients
      b: filter input coefficients
      order: filter order
      num_joints: robot DOF
      ftype: filter type. 'lowpass' or 'bandpass'
    """
        self.num_joints = num_joints
        if isinstance(a, list):
            self.a = a
            self.b = b
        else:
            self.a = [a]
            self.b = [b]

        # Either a set of parameters per joint must be specified as a list
        # Or one filter is applied to every joint
        if not ((len(self.a) == len(self.b) == num_joints) or
                (len(self.a) == len(self.b) == 1)):
            raise ValueError('Incorrect number of filter values specified')

        # Normalize by a[0]
        for i in range(len(self.a)):
            self.b[i] /= self.a[i][0]
            self.a[i] /= self.a[i][0]

        # Convert single filter to same format as filter per joint
        if len(self.a) == 1:
            self.a *= num_joints
            self.b *= num_joints
        self.a = np.stack(self.a)
        self.b = np.stack(self.b)

        if ftype == 'bandpass':
            assert len(self.b[0]) == len(self.a[0]) == 2 * order + 1
            self.hist_len = 2 * order
        elif ftype == 'lowpass':
            assert len(self.b[0]) == len(self.a[0]) == order + 1
            self.hist_len = order
        else:
            raise ValueError('%s filter type not supported' % (ftype))

        logging.info('Filter shapes: a: %s, b: %s', self.a.shape, self.b.shape)
        logging.info('Filter type:%s', ftype)

        self.yhist = collections.deque(maxlen=self.hist_len)
        self.xhist = collections.deque(maxlen=self.hist_len)
        self.reset()

    def reset(self):
        """Resets the history buffers to 0."""
        self.yhist.clear()
        self.xhist.clear()
        for _ in range(self.hist_len):
            self.yhist.appendleft(np.zeros((self.num_joints, 1)))
            self.xhist.appendleft(np.zeros((self.num_joints, 1)))

    def filter(self, x):
        """Returns filtered x."""
        xs = np.concatenate(list(self.xhist), axis=-1)
        ys = np.concatenate(list(self.yhist), axis=-1)
        y = np.multiply(x, self.b[:, 0]) + np.sum(
            np.multiply(xs, self.b[:, 1:]), axis=-1) - np.sum(
                np.multiply(ys, self.a[:, 1:]), axis=-1)
        self.xhist.appendleft(x.reshape((self.num_joints, 1)).copy())
        self.yhist.appendleft(y.reshape((self.num_joints, 1)).copy())
        return y

    def init_history(self, x):
        x = np.expand_dims(x, axis=-1)
        for i in range(self.hist_len):
            self.xhist[i] = x
            self.yhist[i] = x


class ActionFilterButter(ActionFilter):
    """Butterworth filter."""

    def __init__(self,
                 lowcut=None,
                 highcut=None,
                 sampling_rate=None,
                 order=ACTION_FILTER_ORDER,
                 num_joints=None):
        """Initializes a butterworth filter.

    Either one per joint or same for all joints.

    Args:
      lowcut: list of strings defining the low cutoff frequencies.
        The list must contain either 1 element (same filter for all joints)
        or num_joints elements
        0 for lowpass, > 0 for bandpass. Either all values must be 0
        or all > 0
      highcut: list of strings defining the high cutoff frequencies.
        The list must contain either 1 element (same filter for all joints)
        or num_joints elements
        All must be > 0
      sampling_rate: frequency of samples in Hz
      order: filter order
      num_joints: robot DOF
    """
        self.lowcut = ([float(x) for x in lowcut]
                       if lowcut is not None else [ACTION_FILTER_LOW_CUT])
        self.highcut = ([float(x) for x in highcut]
                        if highcut is not None else [ACTION_FILTER_HIGH_CUT])
        if len(self.lowcut) != len(self.highcut):
            raise ValueError(
                'Number of lowcut and highcut filter values should '
                'be the same')
        # import ipdb; ipdb.set_trace()
        if sampling_rate is None:
            raise ValueError('sampling_rate should be provided.')

        if num_joints is None:
            raise ValueError('num_joints should be provided.')

        if np.any(self.lowcut):
            if not np.all(self.lowcut):
                raise ValueError('All the filters must be of the same type: '
                                 'lowpass or bandpass')
            self.ftype = 'bandpass'
        else:
            self.ftype = 'lowpass'

        a_coeffs = []
        b_coeffs = []
        for i, (l, h) in enumerate(zip(self.lowcut, self.highcut)):
            if h <= 0.0:
                raise ValueError('Highcut must be > 0')

            b, a = self.butter_filter(l, h, sampling_rate, order)
            logging.info(
                'Butterworth filter: joint: %d, lowcut: %f, highcut: %f, '
                'sampling rate: %d, order: %d, num joints: %d', i, l, h,
                sampling_rate, order, num_joints)
            b_coeffs.append(b)
            a_coeffs.append(a)

        super(ActionFilterButter, self).__init__(a_coeffs, b_coeffs, order,
                                                 num_joints, self.ftype)

    def butter_filter(self, lowcut, highcut, fs, order=5):
        """Returns the coefficients of a butterworth filter.

    If lowcut = 0, the function returns the coefficients of a low pass filter.
    Otherwise, the coefficients of a band pass filter are returned.
    Highcut should be > 0

    Args:
      lowcut: low cutoff frequency
      highcut: high cutoff frequency
      fs: sampling rate
      order: filter order
    Return:
      b, a: parameters of a butterworth filter
    """
        # import ipdb; ipdb.set_trace()
        nyq = 0.5 * fs
        low = lowcut / nyq
        high = highcut / nyq
        if low:
            b, a = butter(order, [low, high], btype='band')
        else:
            b, a = butter(order, [high], btype='low')
        return b, a


class ActionFilterExp(ActionFilter):
    """Filter by way of simple exponential smoothing.

  y = alpha * x + (1 - alpha) * previous_y
  """

    def __init__(self, alpha, num_joints):
        """Initialize the filter.

    Args:
      alpha: list of strings defining the alphas.
        The list must contain either 1 element (same filter for all joints)
        or num_joints elements
        0 < alpha <= 1
      num_joints: robot DOF
    """
        self.alphas = [float(x) for x in alpha]
        logging.info('Exponential filter: alpha: %d', self.alphas)

        a_coeffs = []
        b_coeffs = []
        for a in self.alphas:
            a_coeffs.append(np.asarray([1., a - 1.]))
            b_coeffs.append(np.asarray([a, 0]))

        order = 1
        self.ftype = 'lowpass'

        super(ActionFilterExp, self).__init__(a_coeffs, b_coeffs, order,
                                              num_joints, self.ftype)


class ActionFilterHelper(object):

    def __init__(self,
                 control_timestep: float,
                 num_joints: int,
                 init_qpos: np.ndarray,
                 highcut: int = -1):
        self._init_qpos = init_qpos

        # Initialize the filter.
        sampling_rate = int(
            1.0 / control_timestep
        )  # need to be able to get control frequency from env
        # TODO: Probably add an assertion that it's an int.

        if highcut > 0:
            self._action_filter = ActionFilterButter(
                sampling_rate=sampling_rate,
                num_joints=num_joints,
                highcut=[highcut])
        else:
            self._action_filter = ActionFilterButter(
                sampling_rate=sampling_rate, num_joints=num_joints)

    def reset(self):
        self._action_filter.init_history(self._init_qpos)

    def filter(self, action):
        # Filter the action
        action_dtype = action.dtype
        action = self._action_filter.filter(action)
        action = action.astype(action_dtype)

        return action

class ActionFilterWrapper(gym.ActionWrapper):
    """Filter the action."""

    def __init__(self, env : gym.Env, control_timestep : float, walker : BaseWalker, highcut: int = 4):
        super().__init__(env)
        self._action_filter_helper = ActionFilterHelper(control_timestep, walker.joint_nums, walker.joint_qpos_init, highcut)

    def reset(self, *args, seed: int | None = None, options: dict[str, Any] | None = None, **kwargs):
        self._action_filter_helper.reset()
        return self.env.reset(*args, seed=seed, options=options, **kwargs)

    def action(self, action):
        action = self._action_filter_helper.filter(action)
        return action
rail_walker_gym/envs/wrappers/action_filter.py
realquantumcookie/APRL
import gym
import gym.spaces
from typing import Optional, Union
import numpy as np

class RescaleActionAsymmetric(gym.ActionWrapper):
    def __init__(
        self, 
        env : gym.Env, 
        low : Union[float, np.ndarray],
        high : Union[float, np.ndarray],
        center_action : Optional[np.ndarray] = None
    ):
        super().__init__(env)
        self._center_action = center_action
        self.low = low
        self.high = high
    
    @property
    def center_action(self) -> np.ndarray:
        return self._center_action if self._center_action is not None else (
            (self.env.action_space.high + self.env.action_space.low) / 2.0
        )
    
    @center_action.setter
    def center_action(self, value : Optional[np.ndarray]):
        self._center_action = value
    
    @property
    def action_space(self) -> gym.spaces.Box:
        return gym.spaces.Box(
            low = self.low,
            high = self.high,
            shape=self.env.action_space.shape,
            dtype=self.env.action_space.dtype
        )

    def action(self, action : np.ndarray):
        new_center = (self.high + self.low) / 2.0
        
        new_delta_action = action - new_center
        below_center_idx = new_delta_action < 0
        other_idx = np.logical_not(below_center_idx)

        new_delta_high = self.high - new_center
        if not isinstance(new_delta_high, float):
            new_delta_high = new_delta_high[other_idx]
        new_delta_low = new_center - self.low
        if not isinstance(new_delta_low, float):
            new_delta_low = new_delta_low[below_center_idx]

        center_action = self.center_action

        delta_center = new_delta_action.copy()
        delta_center[below_center_idx] *= (center_action[below_center_idx] - self.env.action_space.low[below_center_idx]) / new_delta_low
        delta_center[other_idx] *= (self.env.action_space.high[other_idx] - center_action[other_idx]) / new_delta_high
        ret = (center_action + delta_center).astype(self.env.action_space.dtype)
        # print(action, "=>", ret)
        # print(self.env.action_space)
        return ret

    def reverse_action(self, action : np.ndarray):
        delta_center = action - self.center_action
        below_center_idx = delta_center < 0
        other_idx = np.logical_not(below_center_idx)

        new_center = (self.high + self.low) / 2.0
        new_delta_high = self.high - new_center
        if not isinstance(new_delta_high, float):
            new_delta_high = new_delta_high[other_idx]
        new_delta_low = new_center - self.low
        if not isinstance(new_delta_low, float):
            new_delta_low = new_delta_low[below_center_idx]

        center_action = self.center_action

        new_delta_center = delta_center.copy()
        new_delta_center[below_center_idx] *= new_delta_low / (center_action[below_center_idx] - self.env.action_space.low[below_center_idx])
        new_delta_center[other_idx] *= new_delta_high / (self.env.action_space.high[other_idx] - center_action[other_idx])
        return new_center + new_delta_center

rail_walker_gym/envs/wrappers/action_rescale.py
realquantumcookie/APRL
import gym
import numpy as np

class ClipActionToRange(gym.ActionWrapper):

    def __init__(self, env : gym.Env, min_action, max_action):
        super().__init__(env)

        min_action = np.asarray(min_action)
        max_action = np.asarray(max_action)

        min_action = min_action + np.zeros(env.action_space.shape,
                                           dtype=env.action_space.dtype)

        max_action = max_action + np.zeros(env.action_space.shape,
                                           dtype=env.action_space.dtype)

        min_action = np.maximum(min_action, env.action_space.low)
        max_action = np.minimum(max_action, env.action_space.high)

        self.action_space = gym.spaces.Box(
            low=min_action,
            high=max_action,
            shape=env.action_space.shape,
            dtype=env.action_space.dtype,
        )
    
    def action(self, action):
        return np.clip(action, self.action_space.low, self.action_space.high)

rail_walker_gym/envs/wrappers/action_wrappers.py
realquantumcookie/APRL
from typing import Callable, Optional, Any
import numpy as np
import math
import gym
import gym.spaces
import dmcgym
import dmcgym.env
from dm_control import composer
from dm_control.rl import control
from dm_control.mujoco.engine import Physics as EnginePhysics
from rail_walker_interface import JoystickEnvironment, WalkerEnvironment, BaseWalker, JoystickPolicy
from rail_mujoco_walker import JoystickPolicyDMControlTask
import mujoco
import dm_env
from collections import OrderedDict

class RailWalkerMujocoComplianceWrapper(gym.ObservationWrapper, JoystickEnvironment, WalkerEnvironment):
    def __init__(self, env: gym.Env):
        assert isinstance(env.unwrapped, dmcgym.DMCGYM)
        assert isinstance(env.unwrapped._env.task, JoystickPolicyDMControlTask)
        gym.Wrapper.__init__(
            self, env
        )
        JoystickEnvironment.__init__(self)
        WalkerEnvironment.__init__(self)
    
    @property
    def observation_space(self) -> gym.Space:
        robot_space = self.env.observation_space
        real_obs_space = {}
        for key, space in robot_space.items():
            if key.startswith("robot/") and key[len("robot/"):] in self.joystick_policy.enabled_observables:
                real_obs_space[key] = space        
        
        if self.joystick_policy.target_observable is not None:
            real_obs_space["target_obs"] = self.joystick_policy.target_observable.get_observation_spec()

        target_custom_data_spec = self.joystick_policy.target_yaw_provider.get_target_custom_data_observable_spec()
        if self.joystick_policy.enable_target_custom_obs and target_custom_data_spec is not None:
            real_obs_space["target_custom"] = target_custom_data_spec
        
        if not self.joystick_policy.target_yaw_provider.is_target_velocity_fixed():
            real_obs_space["target_vel"] = gym.spaces.Box(
                low = 0.0,
                high = np.inf,
                shape=(1,)
            )
        
        # Enforce order
        real_obs_space = OrderedDict(sorted(real_obs_space.items(), key=lambda t: t[0]))
        obs_space = gym.spaces.Dict(real_obs_space)
        return obs_space

    @property
    def action_space(self) -> gym.Space:
        return gym.spaces.Box(
            low = self.robot.action_qpos_mins,
            high = self.robot.action_qpos_maxs,
            dtype = np.float32
        )

    @property
    def robot(self) -> BaseWalker:
        task : JoystickPolicyDMControlTask = self.env.unwrapped._env.task
        return task.robot

    @property
    def joystick_policy(self) -> JoystickPolicy:
        task : JoystickPolicyDMControlTask = self.env.unwrapped._env.task
        return task.joystick_policy
    
    def set_joystick_policy(self, joystick_policy: JoystickPolicy):
        task : JoystickPolicyDMControlTask = self.env.unwrapped._env.task
        task.joystick_policy = joystick_policy
    
    @property
    def is_resetter_policy(self) -> bool:
        return False

    def step(self, action):
        obs, reward, done, info = self.env.step(action)
        info.update(self.joystick_policy.last_info)
        done = done or self.joystick_policy.should_terminate() or self.joystick_policy.should_truncate()
        if self.joystick_policy.should_truncate():
            info["TimeLimit.truncated"] = True
        return self.observation(obs), reward, done, info

    def observation(self, observation):
        obs = {}
        for key, value in observation.items():
            if key.startswith("robot/") and key[len("robot/"):] in self.joystick_policy.enabled_observables:
                obs[key] = value
        if self.joystick_policy.target_observable is not None:
            obs["target_obs"] = self.joystick_policy.target_observable.get_observation()
        
        if not self.joystick_policy.target_yaw_provider.is_target_velocity_fixed():
            target_vel = self.joystick_policy._target_velocity
            obs["target_vel"] = np.array([target_vel])
        
        target_custom_data_obs = self.joystick_policy.target_yaw_provider.get_target_custom_data_observable()
        if self.joystick_policy.enable_target_custom_obs and target_custom_data_obs is not None:
            obs["target_custom"] = target_custom_data_obs
        return obs

class DMControlMultiCameraRenderWrapper(dmcgym.DMCGYM):    
    def __init__(
        self,
        env: composer.Environment | control.Environment | dm_env.Environment,
        scene_callback : Optional[Callable[[EnginePhysics, mujoco.MjvScene], None]] = None,
        render_height: int = 84,
        render_width: int = 84,
    ):
        self.metadata["render_modes"].append("multi_camera")
        super().__init__(
            env
        )
        self.render_height = render_height
        self.render_width = render_width
        self.scene_callback = scene_callback

    def render(
        self,
        *args,
        **kwargs
    ) -> np.ndarray | None:
        return self._internal_render(
            height=self.render_height,
            width=self.render_width,
            scene_callback=self.scene_callback
        )
    
    def reset(self, seed : int | None = None, options: dict[str, Any] | None = None, **kwargs):
        if seed is not None:
            self.seed(seed)
        obs = super().reset()
        if kwargs.get("return_info",False):
            return obs, {}
        else:
            return obs

    def _internal_render(
        self,
        height: int = 84,
        width: int = 84,
        scene_callback = None
    ):
        
        physics = self._env.physics
        num_cameras = physics.model.ncam
        num_columns = int(math.ceil(math.sqrt(num_cameras)))
        num_rows = int(math.ceil(float(num_cameras) / num_columns))
        frame = np.zeros((num_rows * height, num_columns * width, 3), dtype=np.uint8)
        for col in range(num_columns):
            for row in range(num_rows):
                camera_id = row * num_columns + col
                if camera_id >= num_cameras:
                    break
                subframe = physics.render(
                    camera_id=camera_id, height=height, width=width, scene_callback=scene_callback
                )
                frame[
                    row * height : (row + 1) * height, col * width : (col + 1) * width
                ] = subframe
        return frame


rail_walker_gym/envs/wrappers/dm_to_gym.py
realquantumcookie/APRL
import gym
import gym.spaces
import collections
import copy
import numpy as np
from typing import List

class FrameStackWrapper(gym.ObservationWrapper):
    """Wrapper that stacks observations along a new final axis."""
    @staticmethod
    def stack_box_spec(spec : gym.spaces.Box, num_frames : int):
        low = np.expand_dims(spec.low, axis=-1).repeat(num_frames, axis=-1)
        high = np.expand_dims(spec.high, axis=-1).repeat(num_frames, axis=-1)
        new_shape = spec.shape + (num_frames,)
        return gym.spaces.Box(low, high, shape=new_shape, dtype=spec.dtype)

    @staticmethod
    def stack_multibinary_spec(spec: gym.spaces.MultiBinary, num_frames: int):
        new_spec = copy.copy(spec)
        if isinstance(spec.n, int):
            new_spec.n = (spec.n,num_frames)
        else:
            new_spec.n = spec.n + (num_frames,)
        return new_spec

    @staticmethod
    def stack_dict_spec(spec: gym.spaces.Dict, exclude_keys: List[str], num_frames: int):
        new_spec = copy.copy(spec)
        for key, value in spec.spaces.items():
            if key in exclude_keys:
                new_spec[key] = value
            elif isinstance(value, gym.spaces.Box):
                new_spec[key] = FrameStackWrapper.stack_box_spec(value, num_frames)
            elif isinstance(value, gym.spaces.Dict):
                new_spec[key] = FrameStackWrapper.stack_dict_spec(value, num_frames)
            elif isinstance(value, gym.spaces.MultiBinary):
                new_spec[key] = FrameStackWrapper.stack_multibinary_spec(value, num_frames)
            else:
                raise NotImplementedError
        return new_spec
    
    @staticmethod
    def stack_box_obs(all_obs : list[np.ndarray], num_frames : int):
        if len(all_obs) < num_frames:
            all_obs = [np.zeros_like(all_obs[0])]*(num_frames - len(all_obs)) + all_obs
        return np.stack(all_obs, axis=-1)
    
    @staticmethod
    def stack_multibinary_obs(all_obs : list[np.ndarray], num_frames : int):
        return __class__.stack_box_obs(all_obs, num_frames)

    @staticmethod
    def stack_dict_obs(all_obs : list[dict[str,np.ndarray]], exclude_keys: List[str], num_frames : int):
        new_obs = {}
        for key, value in all_obs[-1].items():
            if key in exclude_keys:
                new_obs[key] = value
            elif isinstance(value,dict) or isinstance(value, collections.OrderedDict):
                new_obs[key] = FrameStackWrapper.stack_dict_obs(
                    [obs[key] for obs in all_obs],
                    num_frames
                )
            elif isinstance(value, np.ndarray):
                if value.dtype == np.int8:
                    new_obs[key] = FrameStackWrapper.stack_multibinary_obs(
                        [obs[key] for obs in all_obs],
                        num_frames
                    )
                else:
                    new_obs[key] = FrameStackWrapper.stack_box_obs(
                        [obs[key] for obs in all_obs],
                        num_frames
                    )
            else:
                raise NotImplementedError
        return new_obs

    def __init__(
        self,
        env: gym.Env,
        num_frames: int = 4,
        exclude_keys: List[str] = [],
    ):
        super().__init__(env)
        assert isinstance(env.observation_space, gym.spaces.Dict) or isinstance(env.observation_space, gym.spaces.Box) or isinstance(env.observation_space, gym.spaces.MultiBinary), "FrameStackWrapper only works with Dict, Box, or MultiBinary observation spaces"
        assert isinstance(env.observation_space, gym.spaces.Dict) or len(exclude_keys) == 0, "FrameStackWrapper does not support excluding keys from Dict observation spaces, {}".format(env.observation_space)

        self.previous_queue = collections.deque([],maxlen=num_frames)
        if isinstance(env.observation_space, gym.spaces.Dict):
            self.observation_space = __class__.stack_dict_spec(env.observation_space, exclude_keys, num_frames)
        elif isinstance(env.observation_space, gym.spaces.Box):
            self.observation_space = __class__.stack_box_spec(env.observation_space, num_frames)
        elif isinstance(env.observation_space, gym.spaces.MultiBinary):
            self.observation_space = __class__.stack_multibinary_spec(env.observation_space, num_frames)
        else:
            raise NotImplementedError
        self._last_joystick_truncation = False
        self.exclude_keys = exclude_keys

    def reset(self, *args, **kwargs):
        resetret = self.env.reset(*args, **kwargs)
        if self._last_joystick_truncation:
            self.previous_queue.pop()
        else:
            self.previous_queue.clear()
        self._last_joystick_truncation = False
        if kwargs.get("return_info", False):
            obs, info = resetret
            self.previous_queue.append(obs)
            return self.observation(obs), info
        else:
            obs = resetret
            self.previous_queue.append(obs)
            return self.observation(obs)

    def step(self, action):
        obs, rew, done, info  = self.env.step(action)
        if done and "TimeLimit.joystick_target_change" in info and info["TimeLimit.joystick_target_change"]:
            self._last_joystick_truncation = True
        else:
            self._last_joystick_truncation = False
        
        self.previous_queue.append(obs)
        return self.observation(obs), rew, done, info

    @property
    def num_frames(self):
        return self.previous_queue.maxlen

    def observation(self, observation):
        if isinstance(observation, dict) or isinstance(observation, collections.OrderedDict):
            new_obs = __class__.stack_dict_obs(self.previous_queue, self.exclude_keys, self.num_frames)
        elif isinstance(observation, np.ndarray):
            if observation.dtype == np.int8:
                new_obs = __class__.stack_multibinary_obs(self.previous_queue, self.num_frames)
            else:
                new_obs = __class__.stack_box_obs(list(self.previous_queue), self.num_frames)
        return new_obs
rail_walker_gym/envs/wrappers/frame_stack.py
realquantumcookie/APRL
import gym
from rail_mujoco_walker.robots.sim_robot import RailSimWalkerDMControl
from rail_walker_interface import BaseWalker, JoystickPolicy, JoystickEnvironment, JoystickPolicyTargetProvider, JoystickPolicyTerminationConditionProvider
from rail_mujoco_walker import DMWalkerForRailSimWalker, RailSimWalkerDMControl
from typing import List, Optional, Tuple, Dict, Any
import numpy as np

class JoystickPolicyRollout:
    def __init__(self) -> None:
        self.joint_qpos : List[np.ndarray] = []
        self.joint_qvel : List[np.ndarray] = []
        self.joint_torques : List[np.ndarray] = []
        self.framequats : List[np.ndarray] = []
        self.linear_3d_velocities : List[np.ndarray] = []
        self.angular_3d_velocities : List[np.ndarray] = []
        self.target_world_deltas : List[np.ndarray] = []
        self.target_linear_velocities : List[float] = []
        self.terminations : List[bool] = []
        self.truncations : List[bool] = []

    def append(
        self,
        joystick_policy: JoystickPolicy
    ):
        self.joint_qpos.append(joystick_policy.robot.get_joint_qpos())
        self.joint_qvel.append(joystick_policy.robot.get_joint_qvel())
        self.joint_torques.append(joystick_policy.robot.get_joint_torques())
        self.framequats.append(joystick_policy.robot.get_framequat_wijk())
        self.linear_3d_velocities.append(joystick_policy.robot.get_3d_linear_velocity())
        self.angular_3d_velocities.append(joystick_policy.robot.get_3d_angular_velocity())
        self.target_world_deltas.append(joystick_policy.target_goal_world_delta())
        self.target_linear_velocities.append(joystick_policy._target_velocity)
        self.terminations.append(joystick_policy.should_terminate())
        self.truncations.append(joystick_policy.should_truncate())
    
    @property
    def dones(self) -> List[bool]:
        return [truncation or termination for truncation, termination in zip(self.truncations, self.terminations)]

    def current_episode(self):
        dones = self.dones
        last_done = len(self.dones) - 1
        last_start = last_done
        while last_start > 0 and not dones[last_start - 1]:
            last_start -= 1
        
        ret = __class__()
        ret.joint_qpos = self.joint_qpos[last_start:last_done+1]
        ret.joint_qvel = self.joint_qvel[last_start:last_done+1]
        ret.joint_torques = self.joint_torques[last_start:last_done+1]
        ret.framequats = self.framequats[last_start:last_done+1]
        ret.linear_3d_velocities = self.linear_3d_velocities[last_start:last_done+1]
        ret.angular_3d_velocities = self.angular_3d_velocities[last_start:last_done+1]
        ret.target_world_deltas = self.target_world_deltas[last_start:last_done+1]
        ret.target_linear_velocities = self.target_linear_velocities[last_start:last_done+1]
        ret.terminations = self.terminations[last_start:last_done+1]
        ret.truncations = self.truncations[last_start:last_done+1]
        return ret

    def last_episode(self):
        # Find the last done
        last_done = len(self.dones) - 1
        while last_done > 0 and not self.dones[last_done]:
            last_done -= 1

        last_start = last_done
        while last_start > 0 and not self.dones[last_start - 1]:
            last_start -= 1
        
        # Return the last episode
        ret = __class__()
        ret.joint_qpos = self.joint_qpos[last_start:last_done+1]
        ret.joint_qvel = self.joint_qvel[last_start:last_done+1]
        ret.joint_torques = self.joint_torques[last_start:last_done+1]
        ret.framequats = self.framequats[last_start:last_done+1]
        ret.linear_3d_velocities = self.linear_3d_velocities[last_start:last_done+1]
        ret.angular_3d_velocities = self.angular_3d_velocities[last_start:last_done+1]
        ret.target_world_deltas = self.target_world_deltas[last_start:last_done+1]
        ret.target_linear_velocities = self.target_linear_velocities[last_start:last_done+1]
        ret.terminations = self.terminations[last_start:last_done+1]
        ret.truncations = self.truncations[last_start:last_done+1]
        return ret
    
    def export_npz(self, path : str) -> None:
        np.savez_compressed(
            path, 
            joint_qpos=self.joint_qpos,
            joint_qvel=self.joint_qvel,
            joint_torques=self.joint_torques,
            framequats=self.framequats,
            linear_3d_velocities=self.linear_3d_velocities,
            angular_3d_velocities=self.angular_3d_velocities,
            target_world_deltas=self.target_world_deltas,
            target_linear_velocities=self.target_linear_velocities,
            terminations=self.terminations,
            truncations=self.truncations
        )
    
    def is_current_episode_empty(self) -> bool:
        dones = self.dones
        return len(dones) == 0 or dones[-1]

    def clear(self) -> None:
        self.joint_qpos.clear()
        self.joint_qvel.clear()
        self.joint_torques.clear()
        self.framequats.clear()
        self.linear_3d_velocities.clear()
        self.angular_3d_velocities.clear()
        self.target_world_deltas.clear()
        self.target_linear_velocities.clear()
        self.terminations.clear()
        self.truncations.clear()

    @staticmethod
    def import_npz(path : str):
        roll = __class__()
        data = np.load(path)
        roll.joint_qpos = list(data["joint_qpos"])
        roll.joint_qvel = list(data["joint_qvel"])
        roll.joint_torques = list(data["joint_torques"])
        roll.framequats = list(data["framequats"])
        roll.linear_3d_velocities = list(data["linear_3d_velocities"])
        roll.angular_3d_velocities = list(data["angular_3d_velocities"])
        roll.target_world_deltas = list(data["target_world_deltas"])
        roll.target_linear_velocities = list(data["target_linear_velocities"])
        roll.terminations = list(data["terminations"])
        roll.truncations = list(data["truncations"])
        return roll
    
    def __len__(self) -> int:
        return len(self.dones)

class JoystickPolicyRecorderWrapper(gym.Wrapper, JoystickEnvironment):
    def __init__(self, env: gym.Env):
        assert hasattr(env, "joystick_policy"), "env must have a joystick_policy attribute"
        assert isinstance(env.joystick_policy, JoystickPolicy), "env.joystick_policy must be a JoystickPolicy"
        gym.Wrapper.__init__(self, env)
        JoystickEnvironment.__init__(self)
        self.recorded_joystick_policy_rollout = JoystickPolicyRollout()

    @property
    def joystick_policy(self) -> JoystickPolicy:
        return self.env.joystick_policy
    
    def set_joystick_policy(self, joystick_policy: JoystickPolicy):
        self.env.set_joystick_policy(joystick_policy)

    def step(self, action: Any) -> Tuple[Any, float, bool, dict]:
        ret = self.env.step(action)
        self.recorded_joystick_policy_rollout.append(
            self.joystick_policy
        )
        return ret
    
    def reset(self, *args, **kwargs) -> Any:
        ret = self.env.reset(*args, **kwargs)
        self.recorded_joystick_policy_rollout.append(
            self.joystick_policy
        )
        return ret
    
class JoystickPolicyReplayTargetProvider(JoystickPolicyTargetProvider[RailSimWalkerDMControl], JoystickPolicyTerminationConditionProvider[RailSimWalkerDMControl]):
    def __init__(
        self,
        joystick_policy_rollout: JoystickPolicyRollout
    ):
        JoystickPolicyTargetProvider.__init__(self)
        JoystickPolicyTerminationConditionProvider.__init__(self)

        self.joystick_policy_rollout = joystick_policy_rollout

        self._is_target_velocity_fixed = True
        for target_velocity in self.joystick_policy_rollout.target_linear_velocities:
            if target_velocity != self.joystick_policy_rollout.target_linear_velocities[0]:
                self._is_target_velocity_fixed = False
                break
        
        self._current_step = 0
    
    def get_target_goal_world_delta(self, Robot: RailSimWalkerDMControl) -> np.ndarray:
        return self.joystick_policy_rollout.target_world_deltas[self._current_step % len(self.joystick_policy_rollout)]

    def get_target_velocity(self, Robot: RailSimWalkerDMControl) -> float:
        return self.joystick_policy_rollout.target_linear_velocities[self._current_step % len(self.joystick_policy_rollout)]
    
    def is_target_velocity_fixed(self) -> bool:
        return self._is_target_velocity_fixed

    def step_target(
        self, 
        Robot: RailSimWalkerDMControl, 
        info_dict: dict[str, Any], 
        randomState: np.random.RandomState
    ) -> None:
        self._current_step += 1
        current_position = Robot.get_3d_location()
        recorded_velocity = self.joystick_policy_rollout.linear_3d_velocities[self._current_step % len(self.joystick_policy_rollout)]
        new_2d_location = current_position[:2] + recorded_velocity[:2] * Robot.control_timestep
        Robot.reset_2d_location_with_qpos(
            new_2d_location,
            self.joystick_policy_rollout.framequats[self._current_step % len(self.joystick_policy_rollout)],
            self.joystick_policy_rollout.joint_qpos[self._current_step % len(self.joystick_policy_rollout)]
        )


    def reset_target(
        self, 
        Robot: RailSimWalkerDMControl, 
        info_dict: dict[str, Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider, 
        randomState: np.random.RandomState
    ) -> None:
        self._current_step += 1 if self._current_step == 0 else 0
        Robot.reset_2d_location_with_qpos(
            np.zeros(2),
            self.joystick_policy_rollout.framequats[self._current_step % len(self.joystick_policy_rollout)],
            self.joystick_policy_rollout.joint_qpos[self._current_step % len(self.joystick_policy_rollout)]
        )

    def reset_termination_condition(
        self, 
        Robot: RailSimWalkerDMControl, 
        info_dict: dict[str, Any], 
        termination_provider_triggered, 
        randomState: np.random.RandomState
    ) -> None:
        return

    def step_termination_condition(
        self, 
        Robot: RailSimWalkerDMControl, 
        target_goal_world_delta: np.ndarray, 
        target_goal_local: np.ndarray, 
        target_yaw: float, 
        target_delta_yaw: float, 
        target_velocity: float, 
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw: float, 
        target_custom_data: Optional[Any], 
        enable_target_custom_obs: bool, 
        info_dict: dict[str, Any], 
        randomState: np.random.RandomState
    ) -> None:
        return
    
    def should_terminate(self) -> bool:
        return self.joystick_policy_rollout.dones[self._current_step % len(self.joystick_policy_rollout)]

rail_walker_gym/envs/wrappers/joystick_policy_replay.py
realquantumcookie/APRL
import gym
from gym.core import Env
import numpy as np
from rail_walker_interface import JoystickEnvironment, JoystickPolicy
from rail_walker_interface.joystick_policy.joystick_policy import JoystickPolicy

class JoystickPolicyTruncationWrapper(gym.Wrapper, JoystickEnvironment):
    def __init__(self, env: Env):
        assert hasattr(env, "joystick_policy") and isinstance(env.joystick_policy, JoystickPolicy)
        gym.Wrapper.__init__(
            self, env
        )
        JoystickEnvironment.__init__(self)
        self._last_obs = None
        self.random_state = np.random.RandomState()
    
    @property
    def joystick_policy(self) -> JoystickPolicy:
        return self.env.joystick_policy
    
    def set_joystick_policy(self, joystick_policy: JoystickPolicy):
        return self.env.set_joystick_policy(joystick_policy)
    
    def step(self, action):
        obs, reward, done, info = self.env.step(action)
        if not done and (
            self.joystick_policy.lock_target and self.joystick_policy.has_after_after_step
        ):
            done = True
            info['TimeLimit.truncated'] = True
            info['TimeLimit.joystick_target_change'] = True
            self._last_obs = obs
        else:
            self._last_obs = None
        return obs, reward, done, info
    
    def seed(self, seed=None):
        self.random_state = np.random.RandomState(seed)
        return self.env.seed(seed)
    
    def reset(self, **kwargs):
        if self._last_obs is not None:
            assert self.joystick_policy.has_after_after_step
            self.joystick_policy.after_after_step(self.random_state)
            if self.joystick_policy.target_observable is not None:
                self._last_obs["target_obs"] = self.joystick_policy.target_observable.get_observation()
            target_obs = self.joystick_policy.target_yaw_provider.get_target_custom_data_observable()
            if self.joystick_policy.enable_target_custom_obs and target_obs is not None:
                self._last_obs["target_custom"] = target_obs
            obs = self._last_obs
            self._last_obs = None
            if kwargs.get("return_info", False):
                return obs, self.joystick_policy.last_info
            else:
                return obs
            
        else:
            return self.env.reset(**kwargs)
rail_walker_gym/envs/wrappers/joystick_policy_truncation.py
realquantumcookie/APRL
import copy
from collections import deque
import gym
import numpy as np

class AddPreviousActions(gym.ObservationWrapper):
    def __init__(self, env : gym.Env, action_history: int = 1, *args, **kwargs):
        self.actions = deque(maxlen=action_history)
        super().__init__(env, *args, **kwargs)

        assert isinstance(env.observation_space, gym.spaces.Dict)
        assert 'actions' not in env.observation_space.spaces

        # Add actions to observation space
        new_obs = copy.copy(env.observation_space.spaces)
        low = np.repeat(env.action_space.low, repeats=action_history, axis=0)
        high = np.repeat(env.action_space.high, repeats=action_history, axis=0)
        action_space = gym.spaces.Box(low, high)
        new_obs['actions'] = action_space
        self.observation_space = gym.spaces.Dict(new_obs)
        self._last_joystick_truncation = False

    def reset(self, *args, **kwargs):
        if not self._last_joystick_truncation:
            for _ in range(self.actions.maxlen):
                self.actions.append(np.zeros_like(self.action_space.low))
        self._last_joystick_truncation = False
        return super().reset(*args, **kwargs)

    def step(self, action):
        obs, rew, done, info = super().step(action)
        if done and "TimeLimit.joystick_target_change" in info and info["TimeLimit.joystick_target_change"]:
            self._last_joystick_truncation = True
        else:
            self._last_joystick_truncation = False
        
        self.actions.append(copy.deepcopy(action))
        return obs, rew, done, info

    def observation(self, observation):
        observation = copy.copy(observation)
        observation['actions'] = np.concatenate(self.actions)
        return observation
rail_walker_gym/envs/wrappers/obs_wrappers.py
realquantumcookie/APRL
from typing import Union
import gym
from rail_walker_interface import JoystickPolicy, JoystickEnvImpl
from rail_real_walker.robots import *
from rail_mujoco_walker import JoystickPolicyDMControlTask, Go1SimWalker, A1SimWalker
from ...joystick_policy import *
from dm_control.locomotion import arenas
from mujoco_utils import composer_utils
from dm_control import composer
from .dm_to_gym import DMControlMultiCameraRenderWrapper, RailWalkerMujocoComplianceWrapper
import numpy as np

class RealRenderWrapperWithSim(gym.Wrapper):
    def __init__(
        self,
        env : gym.Env
    ):
        assert isinstance(env.unwrapped, JoystickEnvImpl)
        super().__init__(env)
        self.init_sim_environment(256, 256)
    
    @property
    def real_env(self) -> JoystickEnvImpl:
        return self.env
    
    def init_sim_environment(self, render_height, render_width):
        assert isinstance(self.real_env.robot, Go1RealWalker)
        sim_robot = Go1SimWalker(
            Kp = self.real_env.robot.Kp,
            Kd = self.real_env.robot.Kd,
            action_interpolation=self.real_env.robot.action_interpolation,
            limit_action_range=self.real_env.robot.limit_action_range,
        )
        sim_robot.control_timestep = self.real_env.robot.control_timestep
        sim_robot.control_subtimestep = self.real_env.robot.control_subtimestep
        floor = arenas.Floor(
            size=(20.0, 20.0)
        )
        floor._top_camera.remove()
            
        sim_task = JoystickPolicyDMControlTask(
            joystick_policy=JoystickPolicy(
                robot=sim_robot,
                reward_provider=JoystickPolicyStrictRewardProvider(),
                target_yaw_provider=JoystickPolicyForwardOnlyTargetProvider(),
                termination_providers=[],
                truncation_providers=[],
                resetters=[],
                initializers=[],
                target_observable=None,
                enabled_observables=[]
            ),
            floor=floor
        )
        dm_env = composer_utils.Environment(
            task=sim_task,
            strip_singleton_obs_buffer_dim=True,
            recompile_physics=False,
        )
        # dm_env = composer.Environment(
        #     task=sim_task,
        #     strip_singleton_obs_buffer_dim=True
        # )
        sim_env = DMControlMultiCameraRenderWrapper(
            env=dm_env,
            scene_callback=sim_task.render_scene_callback,
            render_height=render_height,
            render_width=render_width,
        )
        sim_env = RailWalkerMujocoComplianceWrapper(
            sim_env,
        )
        self.sim_env = sim_env
    
    def align_sim_to_real(self):
        sim_robot : Go1SimWalker = self.sim_env.robot
        sim_policy : JoystickPolicy = self.sim_env.joystick_policy
        sim_task : JoystickPolicyDMControlTask = self.sim_env.task
        sim_robot.Kp = self.real_env.robot.Kp
        sim_robot.Kd = self.real_env.robot.Kd
        sim_robot.action_interpolation = self.real_env.robot.action_interpolation
        sim_robot.limit_action_range = self.real_env.robot.limit_action_range
        sim_robot.mujoco_walker._last_physics.bind(sim_robot.mujoco_walker.joints).qpos[:] = self.real_env.robot.get_joint_qpos()
        sim_robot.mujoco_walker.set_framequat(sim_robot.mujoco_walker._last_physics, self.real_env.robot.get_framequat_wijk())
        sim_policy._target_delta_yaw = self.real_env.joystick_policy.target_delta_yaw
        sim_policy._target_yaw = self.real_env.joystick_policy.target_yaw
        sim_policy._target_goal_world_delta = self.real_env.joystick_policy.target_goal_world_delta
        sim_policy._target_goal_local = self.real_env.joystick_policy.target_goal_local
        sim_robot.mujoco_walker.refresh_observation(sim_robot.mujoco_walker._last_physics)

    def step(self, action):
        ret = self.env.step(action)
        self.align_sim_to_real()
        return ret

    def reset(self, *args, **kwargs):
        ret = super().reset(**kwargs)
        self.sim_env.reset()
        self.align_sim_to_real()
        return ret
    
    def render(self, mode="human", **kwargs):
        return self.sim_env.render(mode=mode, **kwargs)

    @property
    def control_timestep(self):
        return self.real_env.robot.control_timestep
rail_walker_gym/envs/wrappers/real_render.py
realquantumcookie/APRL
import gym
import gym.spaces
from gym.core import Env
from jaxrl5.data.replay_buffer import ReplayBuffer
from rail_walker_interface import BaseWalker, JoystickPolicy, JoystickPolicyTargetProvider, JoystickPolicyTerminationConditionProvider, JoystickPolicyRewardProvider, JoystickEnvironment
import numpy as np
from typing import Any, Optional, List

from rail_walker_interface.joystick_policy.joystick_policy import JoystickPolicy

class RelabelTargetProvider(JoystickPolicyTargetProvider):
    def __init__(self, target_provider : JoystickPolicyTargetProvider, steps_every_new_target : int = 5):
        self.target_provider = target_provider
        self.steps_every_new_target = steps_every_new_target
        self._target_velocity = 0.0
        self._new_target_velocities : List[float] = []
        self._steps_until_new_target = 0

    def get_target_goal_world_delta(self, Robot: BaseWalker) -> np.ndarray:
        return self.target_provider.get_target_goal_world_delta(Robot)
    
    def get_target_velocity(self, Robot: BaseWalker) -> float:
        return self._target_velocity
    
    def is_target_velocity_fixed(self) -> bool:
        return False
    
    def has_target_changed(self) -> bool:
        return self._steps_until_new_target <= 0 or self.target_provider.has_target_changed()

    def step(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str,Any], 
        randomState : np.random.RandomState
    ) -> None:
        self._new_target_velocities.append(Robot.get_3d_local_velocity()[0])

    def after_step(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str,Any], 
        randomState : np.random.RandomState
    ) -> None:
        self._target_velocity = np.mean(self._new_target_velocities)
        self._steps_until_new_target = self.steps_every_new_target
        self._new_target_velocities.clear()

    def reset(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str,Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        self._steps_until_new_target = self.steps_every_new_target
        self._new_target_velocities.clear()

class RelabelAggregateWrapper(gym.Wrapper):
    def __init__(self, env: gym.Env, max_steps : int = 100_000, other_envs : list[gym.Env] = []):
        super().__init__(env)
        self.other_envs = other_envs
        for other_env in other_envs:
            assert hasattr(other_env, "joystick_policy"), "Other envs must have a joystick policy"
            other_joystick : JoystickPolicy = other_env.joystick_policy
            other_joystick.termination_providers.clear()

        self.prev_obs_list = []
        self._replay_buffer = ReplayBuffer(
            env.observation_space,
            env.action_space,
            capacity=max_steps,
        )
        self._enabled = True

    def enable_relabel(self):
        self._enabled = True

    def disable_relabel(self):
        self._enabled = False

    def get_relabel_replay_buffer(self) -> ReplayBuffer:
        return self._replay_buffer
    
    def set_relabel_replay_buffer(self, buffer : ReplayBuffer):
        self._replay_buffer = buffer
    
    def step(self, action):
        ret = super().step(action)
        
        if self._enabled:
            new_obs_list = []
            for idx, env in enumerate(self.other_envs):
                prev_obs = self.prev_obs_list[idx]
                obs, rew, done, info = env.step(action)
                done = done or ret[2]
                truncated = "TimeLimit.truncated" in info and info['TimeLimit.truncated']
                if (not done) or truncated:
                    mask = 1.0
                else:
                    mask = 0.0
                self._replay_buffer.insert(dict(
                    observations=prev_obs,
                    actions=action,
                    rewards=rew,
                    masks=mask,
                    dones=done,
                    next_observations=obs
                ))
                if 'TimeLimit.joystick_target_change' in info and info['TimeLimit.joystick_target_change']:
                    obs, info = env.reset(return_info=True)
                    done = False
            
                new_obs_list.append(obs)
            self.prev_obs_list = new_obs_list
        return ret
    
    def reset(self, *args, **kwargs):
        ret = super().reset(*args, **kwargs)
        self.prev_obs_list = [env.reset() for env in self.other_envs]
        return ret

class SymmetricRelabelWrapper(gym.Wrapper, JoystickEnvironment):
    def __init__(self, env: gym.Env, obs_space_before_flatten : gym.Space, max_steps : int = 100_000):
        super().__init__(env)
        self._replay_buffer = ReplayBuffer(
            env.observation_space,
            env.action_space,
            capacity=max_steps,
        )
        self.obs_space_before_flatten = obs_space_before_flatten
        self.prev_obs_relabeled = None
        self._enabled = True

    @property
    def joystick_policy(self) -> JoystickPolicy:
        return self.env.joystick_policy
    
    def set_joystick_policy(self, joystick_policy: JoystickPolicy):
        return self.env.set_joystrick_policy(joystick_policy)
    
    @property
    def is_resetter_policy(self) -> bool:
        return self.env.is_resetter_policy
    
    def relabel_action(self, action : np.ndarray):
        assert action.shape == (12,), "Action must be of shape (12,)"
        ret_action = np.zeros_like(action)
        ret_action[[0, 6]] = -action[[3, 9]]
        ret_action[[3, 9]] = -action[[0, 6]]
        ret_action[[1, 2, 7, 8]] = action[[4, 5, 10, 11]]
        ret_action[[4, 5, 10, 11]] = action[[1, 2, 7, 8]]
        return ret_action

    def relabel_obs(self, obs : np.ndarray) -> np.ndarray:
        ret_obs_unflattened = gym.spaces.unflatten(self.obs_space_before_flatten, obs)
        if "target_obs" in ret_obs_unflattened and ret_obs_unflattened["target_obs"] is not None:
            ret_obs_unflattened["target_obs"] = -ret_obs_unflattened["target_obs"] # Only supports relabeling TDY for now
        if "robot/joints_pos" in ret_obs_unflattened and ret_obs_unflattened["robot/joints_pos"] is not None:
            ret_obs_unflattened["robot/joints_pos"] = self.relabel_joint_potentially_stacked(ret_obs_unflattened["robot/joints_pos"])
        if "robot/joints_vel" in ret_obs_unflattened and ret_obs_unflattened["robot/joints_vel"] is not None:
            ret_obs_unflattened["robot/joints_vel"] = self.relabel_joint_potentially_stacked(ret_obs_unflattened["robot/joints_vel"])
        if "robot/torques" in ret_obs_unflattened and ret_obs_unflattened["robot/torques"] is not None:
            ret_obs_unflattened["robot/torques"] = self.relabel_joint_potentially_stacked(ret_obs_unflattened["robot/torques"])
        if "robot/foot_contacts" in ret_obs_unflattened and ret_obs_unflattened["robot/foot_contacts"] is not None:
            ret_obs_unflattened["robot/foot_contacts"] = self.relabel_feet_potentially_stacked(ret_obs_unflattened["robot/foot_contacts"])
        if "robot/foot_forces" in ret_obs_unflattened and ret_obs_unflattened["robot/foot_forces"] is not None:
            ret_obs_unflattened["robot/foot_forces"] = self.relabel_feet_potentially_stacked(ret_obs_unflattened["robot/foot_forces"])
        if "robot/foot_forces_normalized" in ret_obs_unflattened and ret_obs_unflattened["robot/foot_forces_normalized"] is not None:
            ret_obs_unflattened["robot/foot_forces_normalized"] = self.relabel_feet_potentially_stacked(ret_obs_unflattened["robot/foot_forces_normalized"])
        return gym.spaces.flatten(self.obs_space_before_flatten, ret_obs_unflattened)

    def relabel_feet_potentially_stacked(self, feet_obs : np.ndarray) -> np.ndarray:
        if feet_obs.ndim > 1:
            assert feet_obs.shape[-1] == 4
            return np.stack([
                self.relabel_feet_potentially_stacked(feet_obs[i]) for i in range(feet_obs.shape[0])
            ], axis=0)
        else:
            assert feet_obs.shape == (4,)
            ret_obs = np.zeros_like(feet_obs)
            ret_obs[[0, 2]] = feet_obs[[1, 3]]
            ret_obs[[1, 3]] = feet_obs[[0, 2]]
            return ret_obs

    def relabel_joint_potentially_stacked(self, joint_obs : np.ndarray) -> np.ndarray:
        if joint_obs.ndim > 1:
            assert joint_obs.shape[-1] == 12
            return np.stack([
                self.relabel_joint_potentially_stacked(joint_obs[i]) for i in range(joint_obs.shape[0])
            ], axis=0)
        else:
            return self.relabel_action(joint_obs)

    def enable_relabel(self):
        self._enabled = True

    def disable_relabel(self):
        self._enabled = False

    def get_relabel_replay_buffer(self) -> ReplayBuffer:
        return self._replay_buffer
    
    def set_relabel_replay_buffer(self, buffer : ReplayBuffer):
        self._replay_buffer = buffer
    
    def step(self, action):
        ret = super().step(action)
        
        if self._enabled:
            relabeled_action = self.relabel_action(action)
            relabeled_obs = self.relabel_obs(ret[0])
            truncated = "TimeLimit.truncated" in ret[3] and ret[3]['TimeLimit.truncated']
            self._replay_buffer.insert(dict(
                observations=self.prev_obs_relabeled,
                actions=relabeled_action,
                rewards=ret[1],
                masks=1.0 if not ret[2] or truncated else 0.0,
                dones=ret[3],
                next_observations=relabeled_obs
            ))
            self.prev_obs_relabeled = relabeled_obs
        return ret
    
    def reset(self, *args, **kwargs):
        ret = super().reset(*args, **kwargs)
        if kwargs.get("return_info", False):
            obs = ret[0]
        else:
            obs = ret
        self.prev_obs_relabeled = self.relabel_obs(obs)
        return ret
rail_walker_gym/envs/wrappers/relabel_wrapper.py
realquantumcookie/APRL
import gym
from typing import Optional
import wandb
import numpy as np
import cv2 as cv
from typing import Any
import pygame
from PIL import Image

class WanDBVideoWrapper(gym.Wrapper):
    def __init__(
        self, 
        env : gym.Env, 
        video_format : str = "mp4", # "mp4" or "gif" or "webm" or "ogg"
        log_name = "training/video",
        record_every_n_steps = -1,
        frame_rate : Optional[int] = None,
        video_length_limit : int = 0, # Change to > 0 to limit video length
    ):
        super().__init__(env)
        if frame_rate is None:
            if hasattr(env,"joystick_policy"):
                control_timestep = env.joystick_policy.robot.control_timestep
            elif hasattr(env.unwrapped._env, "control_timestep"):
                control_timestep = getattr(env.unwrapped._env, "control_timestep")()
            else:
                raise AttributeError(
                "Environment must have a control_timestep() method."
                )

            frame_rate = int(1.0 / control_timestep)
            print("WanDBVideoWrapper: control_timestep = {}, framerate = {}".format(control_timestep, frame_rate))
                
        
        self._frame_rate = frame_rate
        self._frames = []
        self.log_name = log_name
        self.wandb_step = 0
        self._last_recorded_iter = 0
        self._current_recording_iter = 0
        self.video_format = video_format
        self.enableWandbVideo = True
        self.record_every_n_steps = record_every_n_steps
        self.video_length_limit = video_length_limit
        self._last_joystick_truncation = False
    
    @property
    def is_recording(self):
        return self.enableWandbVideo and (
            self._last_recorded_iter < self._current_recording_iter
        )

    @property
    def should_start_recording(self):
        return self.enableWandbVideo and not(self.is_recording) and (
            self.wandb_step - self._last_recorded_iter >= self.record_every_n_steps or
            self.record_every_n_steps < 0
        )

    def set_wandb_step(self, step : int) -> None:
        self.wandb_step = step

    def step(self, action):
        obs, rew, done, info = self.env.step(action) #ret = (obs, reward, done, info)
        self._try_append_record()
        if "TimeLimit.joystick_target_change" in info and info["TimeLimit.joystick_target_change"]:
            self._last_joystick_truncation = True
        elif done:
            self._terminate_record()
            self._last_joystick_truncation = False
        self.wandb_step += 1
        return obs, rew, done, info

    def reset(self, *args, seed: int | None = None, options: dict[str, Any] | None = None, **kwargs):
        if not self._last_joystick_truncation:
            self._terminate_record()
            if self.should_start_recording:
                setattr(self.env.unwrapped, "_is_video_rec", True)
            else:
                setattr(self.env.unwrapped, "_is_video_rec", False)
        ret = self.env.reset(*args, seed=seed, options=options, **kwargs)
        if not self._last_joystick_truncation:
            self._try_start_record()
        self._last_joystick_truncation = False
        return ret

    # Helper methods.

    def _try_append_record(self):
        if self.is_recording:
            self._frames.append(self.env.render())
        
    def _try_start_record(self):
        if self.should_start_recording:
            if hasattr(self.env,"prepended_frames") and self.env.prepended_frames is not None:
                self._frames = self.env.prepended_frames
            else:
                self._frames = []
            
            self._frames.append(self.env.render())
            self._current_recording_iter = self.wandb_step

    def _terminate_record(self) -> None:
        self._last_recorded_iter = self._current_recording_iter
        if len(self._frames) <= 0:
            return

        if self.video_length_limit > 0:
            self._frames = self._frames[-self.video_length_limit:]
        
        print("Writing {} to wandb (Step {})".format(self.log_name, self.wandb_step))
        wandb.log(
            {
                # RGBA => BGRA
                self.log_name: wandb.Video(
                    np.stack(self._frames).transpose(0, 3, 1, 2), 
                    fps=self._frame_rate, 
                    format=self.video_format
                )
            }, 
            step=self.wandb_step #self._current_recording_iter
        )
        self._frames = []

class _RenderWrapperViewer:
    def __init__(
        self
    ):
        self.screen = None
        self.clock = None
        self.fps = 60
    
    def init_pygame(self, x, y, window_name : str = "Enviroment Viewer") -> None:
        pygame.init()
        self.screen = pygame.display.set_mode((x, y), pygame.HWSURFACE | pygame.DOUBLEBUF)
        pygame.display.set_caption(window_name)
        self.clock = pygame.time.Clock()

    @property
    def has_inited(self) -> bool:
        return self.screen is not None

    def render(self, image : Image) -> bool:
        if self.screen is None:
            self.init_pygame(image.width, image.height)
        
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return False
        
        image_surface = pygame.image.fromstring(image.tobytes(), image.size, image.mode).convert()
        self.screen.fill((0,0,0))
        self.screen.blit(image_surface, (0, 0))
        pygame.display.flip()
        self.clock.tick(self.fps)
        return True


class RenderViewerWrapper(gym.Wrapper):
    def __init__(self, env, window_name = "Environment Render", control_timestep = 0.05):
        super().__init__(env)
        self.window_name = window_name
        self._viewer = _RenderWrapperViewer()
        self._viewer.fps = int(1.0 / control_timestep)

    def step(self, action):
        ret = self.env.step(action) #ret = (obs, reward, done, info)

        if ret[2]:
            self._terminate_display()
        else:
            render = self.env.render()
            if render is not None:
                self._display(render)
            else:
                self._terminate_display()
        
        return ret

    def reset(self, *args, seed: int | None = None, options: dict[str, Any] | None = None, **kwargs):
        self._terminate_display()
        ret = self.env.reset(*args,seed=seed, options=options, **kwargs)
        return ret

    def _display(self, frame):
        frame = Image.fromarray(frame,mode="RGB")
        if not self._viewer.has_inited:
            self._viewer.init_pygame(frame.width, frame.height, window_name=self.window_name)
        self._viewer.render(frame)
    
    def _terminate_display(self):
        pass

rail_walker_gym/envs/wrappers/render_wrappers.py
realquantumcookie/APRL
import gym
import gym.spaces
from rail_walker_interface import JoystickEnvironment, JoystickPolicy, JoystickPolicyResetter
from rail_mujoco_walker import JoystickPolicyDMControlTask
from jaxrl5.agents.agent import Agent as JaxAgent
import typing
import numpy as np
from dmcgym.env import DMCGYM
from rail_walker_interface import JoystickEnvImpl


def resetter_rescale_action(action, old_min, old_max, new_min, new_max):
    return (action - old_min) / (old_max - old_min) * (new_max - new_min) + new_min

class ResetterPolicySupportedEnvironment(gym.Wrapper, JoystickEnvironment):
    @staticmethod
    def compute_obs_act_space_reset_agent(
        env: typing.Union[gym.Env,JoystickEnvironment],
        resetter_joystick_policy: JoystickPolicy,
        wrap_resetter_env_lambda : typing.Callable[[gym.Env], gym.Env] = lambda env: env,
        resetter_action_scale : typing.Optional[float] = 1.0,
    ):
        old_joy_policy = env.joystick_policy
        old_action_scale = old_joy_policy.robot.limit_action_range

        if resetter_action_scale is not None:
            old_joy_policy.robot.limit_action_range = resetter_action_scale
        env.set_joystick_policy(resetter_joystick_policy)
        
        reset_env = wrap_resetter_env_lambda(env)
        obs_space, act_space = reset_env.observation_space, reset_env.action_space
        
        old_joy_policy.robot.limit_action_range = old_action_scale
        env.set_joystick_policy(old_joy_policy)
        return obs_space, act_space

    def __init__(
        self, 
        env: typing.Union[gym.Env,JoystickEnvironment],
        resetter_joystick_policy: JoystickPolicy,
        resetter_agent: JaxAgent,
        max_seconds : float = 5.0,
        intercept_all_reset_calls = True,
        wrap_resetter_env_lambda : typing.Callable[[gym.Env], gym.Env] = lambda env: env,
        fallback_resetter: typing.Optional[JoystickPolicyResetter] = None,
        max_trials : int = 1,
        resetter_Kp : typing.Optional[float] = None,
        resetter_Kd : typing.Optional[float] = None,
        resetter_action_interpolation : typing.Optional[bool] = None,
        resetter_action_scale : typing.Optional[float] = None,
        resetter_power_protect_factor : typing.Optional[float] = None,
    ):
        assert hasattr(env, "joystick_policy"), "The environment must have a joystick policy"
        assert hasattr(env, "set_joystick_policy"), "The environment must have a set_joystick_policy method"
        
        gym.Wrapper.__init__(
            self, env
        )
        JoystickEnvironment.__init__(self)
        self.resetter_policy_joystick_policy = resetter_joystick_policy
        self.resetter_policy_agent = resetter_agent
        self.resetter_policy_max_seconds = max_seconds
        self.resetter_env = None
        self.resetter_Kp = resetter_Kp
        self.resetter_Kd = resetter_Kd
        self.resetter_action_interpolation = resetter_action_interpolation
        self.resetter_action_scale = resetter_action_scale
        self.resetter_power_protect_factor = resetter_power_protect_factor

        self.first_reset = True
        self.wrap_resetter_env_lambda = wrap_resetter_env_lambda
        self.random_state = np.random.RandomState(0)
        self.intercept_all_reset_calls = intercept_all_reset_calls

        self.max_trials = max_trials
        self.fallback_resetter = fallback_resetter

        self.replace_cache = None
        
        # Hacking the render method
        self.prepended_frames = None

    def replace_resetter_settings(self):
        assert self.replace_cache is None, "The replace cache is already set"
        self.replace_cache = {}
        if self.resetter_Kp is not None:
            self.replace_cache["Kp"] = self.joystick_policy.robot.Kp
            self.joystick_policy.robot.Kp = self.resetter_Kp
        if self.resetter_Kd is not None:
            self.replace_cache["Kd"] = self.joystick_policy.robot.Kd
            self.joystick_policy.robot.Kd = self.resetter_Kd
        if self.resetter_action_interpolation is not None:
            self.replace_cache["action_interpolation"] = self.joystick_policy.robot.action_interpolation
            self.joystick_policy.robot.action_interpolation = self.resetter_action_interpolation
        if self.resetter_action_scale is not None:
            self.replace_cache["action_scale"] = self.joystick_policy.robot.limit_action_range
            self.joystick_policy.robot.limit_action_range = self.resetter_action_scale
        if self.resetter_power_protect_factor is not None:
            self.replace_cache["power_protect_factor"] = self.joystick_policy.robot.power_protect_factor
            self.joystick_policy.robot.power_protect_factor = self.resetter_power_protect_factor
        
        self.replace_cache["joystick_policy"] = self.joystick_policy
        self.set_joystick_policy(self.resetter_policy_joystick_policy)

    def recover_original_settings(self):
        assert self.replace_cache is not None, "The replace cache is not set"
        if "Kp" in self.replace_cache:
            self.joystick_policy.robot.Kp = self.replace_cache["Kp"]
        if "Kd" in self.replace_cache:
            self.joystick_policy.robot.Kd = self.replace_cache["Kd"]
        if "action_interpolation" in self.replace_cache:
            self.joystick_policy.robot.action_interpolation = self.replace_cache["action_interpolation"]
        if "action_scale" in self.replace_cache:
            self.joystick_policy.robot.limit_action_range = self.replace_cache["action_scale"]
        if "power_protect_factor" in self.replace_cache:
            self.joystick_policy.robot.power_protect_factor = self.replace_cache["power_protect_factor"]
        self.set_joystick_policy(self.replace_cache["joystick_policy"])
        self.replace_cache = None

    @property
    def intercept_all_reset_calls(self):
        return self._intercept_all_reset_calls
    
    @intercept_all_reset_calls.setter
    def intercept_all_reset_calls(self, value):
        self._intercept_all_reset_calls = value
        self._inject_mujoco()
    
    def _inject_mujoco(self):
        if isinstance(self.env.unwrapped, DMCGYM) and isinstance(self.env.unwrapped._env.task,JoystickPolicyDMControlTask):
            if self._intercept_all_reset_calls:
                self.env.unwrapped._env.task.should_terminate_episode = lambda *args, **kwargs: False
            else:
                self.env.unwrapped._env.task.should_terminate_episode = JoystickPolicyDMControlTask.should_terminate_episode

    def _init_resetter_env(self):
        self.replace_resetter_settings()
        print("=================== Initializing reset policy environment ===================")
        print("Resetter Policy Environment Action Space Before Wrapping", self.env.action_space)
        tmp_resetter_env = gym.Wrapper(self.env)
        # Intercept the reset call to resetter_env
        def new_reset(*args, **kwargs):
            self.resetter_policy_joystick_policy.reset(
                self.random_state
            )
            if kwargs.get("return_info",False):
                ret = self.receive_observation(tmp_resetter_env)[:2]
            else:
                ret = self.receive_observation(tmp_resetter_env)[0]
            return ret
        tmp_resetter_env.reset = new_reset
        
        resetter_env = self.wrap_resetter_env_lambda(tmp_resetter_env)
        obs = resetter_env.reset(seed=0)
        self.resetter_policy_agent.eval_actions(obs) # Jit the resetter policy
        self.resetter_env = resetter_env
        self.recover_original_settings()
        print("=================== Initialized reset policy environment ===================")

    def receive_observation(self,env:gym.Env):
        # if isinstance(self.env, DMCGYM) and isinstance(self.env._env, dmcenv.Environment):
        #     dm_control_env = self.env._env
        #     dm_control_env._observation_updater.update()
        #     obs = dmc_obs2gym_obs(dm_control_env._observation_updater.get_observation())
        

        if isinstance(self.env,JoystickEnvImpl):
            obs = self.env._get_obs()
            done = self.env.joystick_policy.should_terminate() or self.env.joystick_policy.should_truncate()
            info = self.env.joystick_policy.last_info
            return obs, info, done

        last_action = self.joystick_policy.robot.get_joint_qpos()
        rescaled_action = resetter_rescale_action(
            last_action,
            self.joystick_policy.robot.action_qpos_mins, #env.unwrapped.action_space.low,
            self.joystick_policy.robot.action_qpos_maxs, #env.unwrapped.action_space.high,
            env.action_space.low,
            env.action_space.high
        )
        rescaled_action = np.clip(rescaled_action.astype(env.action_space.dtype), env.action_space.low, env.action_space.high)
        obs, _, done, info = env.step(rescaled_action)
        return obs, info, done

    def perform_resetter_policy(self):
        print("====================== Resetter Policy ======================")
        print("Replacing the joystick policy with resetter policy")
        is_recording = hasattr(self.env.unwrapped, "_is_video_rec") and self.env.unwrapped._is_video_rec
        recorded_frames = []
        self.replace_resetter_settings()
        obs, info = self.resetter_env.reset(return_info=True)
        done = self.resetter_policy_joystick_policy.should_terminate() or self.resetter_policy_joystick_policy.should_truncate()

        self.resetter_policy_joystick_policy.reset(self.random_state)
        if is_recording:
            recorded_frames.append(self.env.render(mode="rgb_array"))

        max_steps = int(self.resetter_policy_max_seconds / self.joystick_policy.control_timestep)
        
        for _ in range(max_steps):
            if done:
                break
            action = self.resetter_policy_agent.eval_actions(obs)
            next_obs, reward, done, info = self.resetter_env.step(action)
            if is_recording:
                recorded_frames.append(self.env.render(mode="rgb_array"))
            obs = next_obs
        
        print("Resetter policy done")
        self.recover_original_settings()
        if is_recording:
            self.prepended_frames = recorded_frames
        print("=============================================================")

    def reset(self, seed : int | None = None, *args, **kwargs):
        # before resetting the environment, we first perform action based on resetter policy
        # and then reset the environment
        if seed is not None:
            self.random_state = np.random.RandomState(seed)
        
        if self.first_reset:
            ret = self.env.reset(*args,seed=seed,**kwargs)
            self._init_resetter_env()
            self.first_reset = False
        else:
            if not self.intercept_all_reset_calls:
                self.perform_resetter_policy()
                ret = self.env.reset(*args,seed=seed,**kwargs)
            else:
                done = True
                
                for _ in range(self.max_trials):
                    if not done:
                        break
                    self.perform_resetter_policy()
                    # When receiving observation, we might be stepping to cause the episode to terminate
                    # In that case, we need to loop resetting the environment until we get a valid observation
                    joystick_policy : JoystickPolicy = self.joystick_policy
                    joystick_policy.reset(self.random_state)
                    obs, info, done = self.receive_observation(self.env)
                
                if done and self.fallback_resetter is not None:
                    self.fallback_resetter.reset(self.joystick_policy.robot, self.joystick_policy._info_dict, self.joystick_policy._termination_reason, self.random_state)
                    joystick_policy.reset(self.random_state)
                    obs, info, done = self.receive_observation(self.env)
                    if done:
                        raise RuntimeError("Fallback resetter failed")

                if kwargs.get("return_info", False):
                    ret = obs, {}
                else:
                    ret = obs
        return ret

    @property
    def observation_space(self) -> gym.Space:
        return super().observation_space
    
    @property
    def action_space(self) -> gym.Space:
        return super().action_space

    @property
    def joystick_policy(self) -> JoystickPolicy:
        return self.env.joystick_policy

    def set_joystick_policy(self, joystick_policy : JoystickPolicy):
        self.env.set_joystick_policy(joystick_policy)

    @property
    def is_resetter_policy(self) -> bool:
        return True
rail_walker_gym/envs/wrappers/resetter_policy_supported.py
realquantumcookie/APRL
import numpy as np
import gym
from typing import Any, Tuple, List

class Rollout:
    def __init__(self) -> None:
        self.observations = []
        self.actions = []
        self.rewards = []
        self.dones = []

    def append(self, observation : np.ndarray, action : np.ndarray, reward : float, done : bool) -> None:
        self.observations.append(observation)
        self.actions.append(action)
        self.rewards.append(reward)
        self.dones.append(done)
    
    def current_episode(self) -> Tuple[np.ndarray, np.ndarray, float, bool]:
        last_done = len(self.dones) - 1
        last_start = last_done
        while last_start > 0 and not self.dones[last_start - 1]:
            last_start -= 1
        
        return np.array(self.observations[last_start:last_done+1]), np.array(self.actions[last_start:last_done+1]), np.array(self.rewards[last_start:last_done+1]), np.array(self.dones[last_start:last_done+1])

    def last_episode(self) -> Tuple[List[np.ndarray], List[np.ndarray], List[float], List[bool]]:
        # Find the last done
        last_done = len(self.dones) - 1
        while last_done > 0 and not self.dones[last_done]:
            last_done -= 1

        last_start = last_done
        while last_start > 0 and not self.dones[last_start - 1]:
            last_start -= 1
        
        # Return the last episode
        return np.array(self.observations[last_start:last_done+1]), np.array(self.actions[last_start:last_done+1]), np.array(self.rewards[last_start:last_done+1]), np.array(self.dones[last_start:last_done+1])

    def export_npz(self, path : str) -> None:
        np.savez_compressed(path, observations=self.observations, actions=self.actions, rewards=self.rewards, dones=self.dones)
    
    def export_last_episode_npz(self, path : str) -> None:
        observations, actions, rewards, dones = self.last_episode()
        np.savez_compressed(path, observations=observations, actions=actions, rewards=rewards, dones=dones)
    
    def export_current_episode_npz(self, path : str) -> None:
        observations, actions, rewards, dones = self.current_episode()
        np.savez_compressed(path, observations=observations, actions=actions, rewards=rewards, dones=dones)

    def is_current_episode_empty(self) -> bool:
        return len(self.dones) == 0 or self.dones[-1]

    def clear(self) -> None:
        self.observations.clear()
        self.actions.clear()
        self.rewards.clear()
        self.dones.clear()

    @staticmethod
    def import_npz(path : str):
        roll = Rollout()
        data = np.load(path)
        roll.observations = list(data["observations"])
        roll.actions = list(data["actions"])
        roll.rewards = list(data["rewards"])
        roll.dones = list(data["dones"])
        return roll
    
    def __len__(self) -> int:
        return len(self.dones)

class RolloutCollector(gym.Wrapper):
    def __init__(self, env : gym.Env) -> None:
        super().__init__(env)
        self.collected_rollouts = Rollout()
        self.next_observation = None

    def step(self, action) -> Tuple[Any, float, bool, dict]:
        ret = self.env.step(action)
        self.collected_rollouts.append(self.next_observation, action, ret[1], ret[2])
        self.next_observation = ret[0]
        return ret
    
    def reset(self, *args, **kwargs):
        ret = self.env.reset(*args,**kwargs)
        if isinstance(ret, tuple) and isinstance(ret[1], dict):
            self.next_observation = ret[0]
        else:
            self.next_observation = ret
        return ret





rail_walker_gym/envs/wrappers/rollout_collect.py
realquantumcookie/APRL
from typing import Union, Any
import gym
from rail_walker_interface import BaseWalker, JoystickPolicy, JoystickEnvironment
import numpy as np
import matplotlib.pyplot as plt
from rail_walker_interface.joystick_policy.joystick_policy import JoystickPolicy

class JoystickTargetViewer(gym.Wrapper, JoystickEnvironment):
    def __init__(self, env: gym.Env):
        assert hasattr(env, 'joystick_policy')
        gym.Wrapper.__init__(self, env)
        JoystickEnvironment.__init__(self)
        plt.ion()
        plt.xlim(-2, 2)
        plt.ylim(-2, 2)
        self.heading_arrow = plt.arrow(0, 0, 0, 1, color='g', width=0.1)
        self.tdy_arrow = plt.arrow(0, 0, 0, 1, color='y', width=0.1)

    @property
    def joystick_policy(self) -> JoystickPolicy:
        return self.env.joystick_policy
    
    def set_joystick_policy(self, joystick_policy: JoystickPolicy):
        return self.env.set_joystick_policy(joystick_policy)
    
    def step(self, action):
        ret = self.env.step(action)
        self.update()
        return ret
    
    def reset(self, **kwargs) -> Any | tuple[Any, dict]:
        ret = super().reset(**kwargs)
        self.update()
        return ret

    def update(self):
        heading_yaw = self.joystick_policy.robot.get_roll_pitch_yaw()[2]
        self.heading_arrow.set_data(
            x=0,
            y=0,
            dx=np.cos(heading_yaw),
            dy=np.sin(heading_yaw)
        )
        if hasattr(self.joystick_policy.reward_provider, 'target_linear_velocity'):
            target_velocity = self.joystick_policy.reward_provider.target_linear_velocity
        else:
            target_velocity = 1.0
        
        ty_with_tv = self.joystick_policy.target_goal_world_delta_unit * target_velocity
        self.tdy_arrow.set_data(
            x=0,
            y=0,
            dx=ty_with_tv[0],
            dy=ty_with_tv[1]
        )
        plt.pause(0.001)
rail_walker_gym/envs/wrappers/view_target.py
realquantumcookie/APRL
from .condition_providers import JoystickPolicyRollPitchTerminationConditionProvider, JoystickPolicyResetterEarlyStopTruncationProvider
from .resetters import JoystickPolicyManualResetter
from .reward_providers import WalkInTheParkRewardProvider, JoystickPolicySeperateRewardProvider, JoystickPolicyStrictRewardProvider, JoystickPolicyETHRewardProvider
from .target_providers import JoystickPolicyAutoJoystickTargetProvider, JoystickPolicyForwardOnlyTargetProvider, JoystickPolicyAutoJoystickSimpleTargetProvider, JoystickPolicyAutoJoystickSimplePlusTargetProvider
from .target_obs_providers import JoystickPolicyCosSinTargetDeltaYawObservable, JoystickPolicyTargetDeltaYawObservable, JoystickPolicyTargetDeltaYawAndDistanceObservable
from .reward_util import calculate_energy, calculate_qvel_penalty
from .target_yaw_wrappers import JoystickPolicyTargetProviderWrapper, JoystickPolicyLimitedTargetProvider, JoystickPolicySmoothedTargetProvider
rail_walker_gym/joystick_policy/__init__.py
realquantumcookie/APRL
import numpy as np
from rail_walker_interface import JoystickPolicyTerminationConditionProvider
from rail_walker_interface import BaseWalker
from typing import Any, Optional


class JoystickPolicyRollPitchTerminationConditionProvider(JoystickPolicyTerminationConditionProvider[BaseWalker]):
    def __init__(self, terminate_roll_rad : float = 60.0 / 180.0 * np.pi, terminate_pitch_min_rad = -60.0 / 180.0 * np.pi, terminate_pitch_max_rad = 60.0/180.0 * np.pi, reset_fall_count : bool = False) -> None:
        super().__init__()
        self.terminate_roll_rad = terminate_roll_rad
        self.terminate_pitch_min_rad = terminate_pitch_min_rad
        self.terminate_pitch_max_rad = terminate_pitch_max_rad
        self._should_term = False
        self.reset_fall_count = reset_fall_count
        self.fall_count = 0
    
    def should_terminate(self) -> bool:
        return self._should_term

    def add_fall_count_info(self, info_dict: dict[str,Any]) -> None:
        self.fall_count += 1
        info_dict["fall_count"] = self.fall_count

    def step_termination_condition(
        self, 
        Robot: BaseWalker, 
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw : float,
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> None:
        roll, pitch, _ = Robot.get_roll_pitch_yaw()
        if abs(roll) > self.terminate_roll_rad or pitch < self.terminate_pitch_min_rad or pitch > self.terminate_pitch_max_rad:
            self._should_term = True
            self.add_fall_count_info(info_dict)
        else:
            self._should_term = False

    def reset_termination_condition(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str,Any], 
        termination_provider_triggered,
        randomState: np.random.RandomState
    ) -> None:
        self._should_term = False
        if self.reset_fall_count:
            self.fall_count = 0
        if "fall_count" not in info_dict:
            info_dict["fall_count"] = self.fall_count

class JoystickPolicyResetterEarlyStopTruncationProvider(JoystickPolicyTerminationConditionProvider[BaseWalker]):
    def __init__(
            self, 
            terminate_roll_pitch_rad : float = 10.0 / 180.0 * np.pi, 
            terminate_qpos_range : float = 0.3,
            terminate_qvel_range : float = 0.02
        ) -> None:
        super().__init__()
        self.terminate_roll_pitch_rad = terminate_roll_pitch_rad
        self.terminate_qpos_range = terminate_qpos_range
        self.terminate_qvel_range = terminate_qvel_range
        self._should_term = False
    
    def should_terminate(self) -> bool:
        return self._should_term

    def step_termination_condition(
        self, 
        Robot: BaseWalker, 
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw : float,
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> None:
        roll, pitch, yaw = Robot.get_roll_pitch_yaw()
        qpos = Robot.get_joint_qpos()
        qvel = Robot.get_joint_qvel()

        dqpos = qpos - Robot.joint_qpos_init
        dqpos_maxrange = Robot.joint_qpos_maxs - Robot.joint_qpos_mins

        rp_terminate = abs(roll) < self.terminate_roll_pitch_rad or abs(pitch) < self.terminate_roll_pitch_rad
        qpos_terminate = np.all(dqpos < dqpos_maxrange * self.terminate_qpos_range)
        qvel_terminate = np.all(abs(qvel) < self.terminate_qvel_range)
        self._should_term = rp_terminate and qpos_terminate and qvel_terminate
    
    def reset_termination_condition(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str,Any], 
        termination_provider_triggered,
        randomState: np.random.RandomState
    ) -> None:
        self._should_term = False

rail_walker_gym/joystick_policy/condition_providers.py
realquantumcookie/APRL
import numpy as np
from rail_walker_interface import JoystickPolicyResetter, JoystickPolicyTerminationConditionProvider
from rail_walker_interface import BaseWalker, BaseWalkerWithFootContact, BaseWalkerWithJoystick
from typing import Any
import time

class JoystickPolicyManualResetter(JoystickPolicyResetter[BaseWalker]):
    def __init__(self, target_pose : str = "standing", max_seconds : float = 1.0, offset_factor : float = 0.02):
        assert target_pose in ['standing', 'sitting']
        assert max_seconds > 0
        assert offset_factor >= 0 and offset_factor <= 1
        
        self.target_pose = target_pose
        self.max_seconds = max_seconds
        self.offset_factor = offset_factor
        super().__init__()
    
    def perform_reset(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str,Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        print()
        print("=================== Manual Reset ===================")
        if isinstance(Robot.unwrapped(), BaseWalkerWithJoystick):
            print("Lift up the robot and press down the right joystick to continue")
            if hasattr(Robot, "foot_contact_no_contact_threshold") and hasattr(Robot, "foot_contact_has_contact_threshold") and not hasattr(Robot, "foot_contact_calibrated"):
                print("Calibrating foot contact thresholds")
            while True:
                Robot.receive_observation()
                if Robot.get_joystick_values()[-1][1] < -0.5:
                    break
                time.sleep(Robot.control_timestep)
            if hasattr(Robot, "foot_contact_no_contact_threshold") and hasattr(Robot, "foot_contact_has_contact_threshold") and not hasattr(Robot, "foot_contact_calibrated"):
                Robot.foot_contact_no_contact_threshold = Robot.get_foot_force()
        else:
            if hasattr(Robot, "foot_contact_no_contact_threshold") and hasattr(Robot, "foot_contact_has_contact_threshold") and not hasattr(Robot, "foot_contact_calibrated"):
                input("Lift up the robot and press Enter to continue")
                Robot.receive_observation()
                Robot.foot_contact_no_contact_threshold = Robot.get_foot_force()
            else:
                input("Lift up the robot or place the robot on the ground and press Enter to continue")
            
        
        print("Resetting robot")
        print()

        target_qpos = Robot.joint_qpos_sitting if self.target_pose == "sitting" else Robot.joint_qpos_init
        Reset_Min_QPos = target_qpos - (target_qpos - Robot.joint_qpos_mins) * self.offset_factor
        Reset_Max_QPos = target_qpos + (Robot.joint_qpos_maxs - target_qpos) * self.offset_factor

        start_t = time.time()
        t = start_t

        Robot.receive_observation()
        start_qpos = Robot.get_joint_qpos()
        delta_qpos = target_qpos - start_qpos
        
        while (t - start_t <= self.max_seconds):
            # print(f"Resetting robot trial {step}")
            progress = (t - start_t) / self.max_seconds * 2.0 # Take half of the time to move to the target pose
            progress = min(progress, 1.0)
            Robot.apply_action(start_qpos + delta_qpos * progress)
            if Robot.receive_observation():
                current_qpos = Robot.get_joint_qpos()
                if np.all(current_qpos >= Reset_Min_QPos) and np.all(current_qpos <= Reset_Max_QPos):
                    print("Reset successful")
                    break
            t = time.time()
        
        if isinstance(Robot.unwrapped(), BaseWalkerWithFootContact):
            print("Finished Foot Force: ", Robot.get_foot_force())
        
        if isinstance(Robot.unwrapped(), BaseWalkerWithJoystick):
            print("Finished, press down the right joystick to continue")
            while True:
                Robot.receive_observation()
                if Robot.get_joystick_values()[-1][1] < -0.5:
                    break
                time.sleep(Robot.control_timestep)
        else:
            input("Finished, press Enter to continue")
        
        if hasattr(Robot, "foot_contact_no_contact_threshold") and hasattr(Robot, "foot_contact_has_contact_threshold") and not hasattr(Robot, "foot_contact_calibrated"):
            Robot.foot_contact_has_contact_threshold = Robot.get_foot_force()
            Robot.foot_contact_calibrated = True
            if hasattr(Robot, "foot_contact_threshold"):
                try:
                    setattr(Robot, "foot_contact_threshold", (Robot.foot_contact_has_contact_threshold + Robot.foot_contact_no_contact_threshold)/2.0)
                except:
                    pass
        
        print("====================================================")



rail_walker_gym/joystick_policy/resetters.py
realquantumcookie/APRL
import numpy as np
from typing import Any,Optional
from rail_walker_interface import JoystickPolicyRewardProvider, JoystickPolicyTerminationConditionProvider
from rail_walker_interface import BaseWalker
from rail_mujoco_walker import RailSimWalkerDMControl
import transforms3d as tr3d
from dm_control.utils import rewards
import typing
from collections import deque

from rail_walker_interface.joystick_policy.joystick_interfaces import JoystickPolicyTerminationConditionProvider
from rail_walker_interface.robot.robot import BaseWalker
from .reward_util import near_quadratic_bound, calculate_gaussian_activation, calculate_torque

JOINT_WEIGHTS = np.array([1.0, 0.75, 0.5] * 4)
CONTACT_DELTA_QPOS_THRESHOLD = -0.2
CONTACT_DELTA_FORCE_THRESHOLD = 0.4

class WalkInTheParkRewardProvider(JoystickPolicyRewardProvider[BaseWalker]):
    def __init__(self) -> None:
        super().__init__()
        self.rew = 0.0
    
    def get_reward(self) -> float:
        return self.rew

    def reset_reward(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str,Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        self.rew = 0.0
    
    def step_reward(
        self, 
        Robot: BaseWalker, 
        action_target_qpos: np.ndarray,
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw : float,
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> None:
        roll, pitch, yaw = Robot.get_roll_pitch_yaw()
        velocity_local = Robot.get_3d_local_velocity()
        projected_x_velocity = np.cos(pitch) * velocity_local[0]
        reward_v = rewards.tolerance(
            projected_x_velocity,
            bounds=(target_velocity,2*target_velocity),
            margin=2*target_velocity,
            value_at_margin=0,
            sigmoid='linear'
        ) * (target_velocity / 0.5) * 10.0
        penalty_drpy = 1.0 * np.abs(Robot.get_3d_angular_velocity()[-1])
        reward_perstep = reward_v - penalty_drpy
        info_dict["reward_v"] = reward_v
        info_dict["penalty_drpy"] = penalty_drpy
        reward_perstep *= max(Robot.get_foot_contact()) if hasattr(Robot, "get_foot_contact") else 1.0
        self.rew = reward_perstep
        

class JoystickPolicyStrictRewardProvider(JoystickPolicyRewardProvider[BaseWalker]):
    def __init__(
            self,
            energy_penalty_weight: float = 0.0,
            smooth_torque_penalty_weight: float = 0.0,
            joint_diagonal_penalty_weight: float = 0.0,
            joint_shoulder_penalty_weight : float = 0.0,
            joint_acc_penalty_weight : float = 0.0,
            thigh_torque_penalty_weight: float = 0.0,
            joint_vel_penalty_weight : float = 0.0,
            pitch_rate_penalty_factor: float = 0.0,
            roll_rate_penalty_factor: float = 0.0,
            qpos_penalty_weight: float = 0.0,
            contact_reward_weight : float = 0.0
        ) -> None:
        self.energy_penalty_weight = energy_penalty_weight
        self.smooth_torque_penalty_weight = smooth_torque_penalty_weight
        self.joint_diagonal_penalty_weight = joint_diagonal_penalty_weight
        self.joint_shoulder_penalty_weight = joint_shoulder_penalty_weight
        self.joint_acc_penalty_weight = joint_acc_penalty_weight
        self.joint_vel_penalty_weight = joint_vel_penalty_weight
        self.pitch_rate_penalty_factor = pitch_rate_penalty_factor
        self.roll_rate_penalty_factor = roll_rate_penalty_factor
        self.qpos_penalty_weight = qpos_penalty_weight
        self.thigh_torque_penalty_weight = thigh_torque_penalty_weight
        self.contact_reward_weight = contact_reward_weight
        self.rew = 0.0
        self._last_torque = None
    
    def get_reward(self) -> float:
        return self.rew

    def reset_reward(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str,Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        self.rew = 0.0
        self._last_torque = Robot.get_joint_torques().copy()
        self._last_joint_qpos = Robot.get_joint_qpos().copy()
        self._last_contacts = Robot.get_foot_contact().copy() if hasattr(Robot, "get_foot_contact") else None
        self._last_foot_force_norm = Robot.get_foot_force_norm().copy() if hasattr(Robot, "get_foot_force_norm") else None
        
    def calculate_velocity_reward_norm(
        self, 
        Robot: BaseWalker, 
        action_target_qpos: np.ndarray,
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw : float,
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> float:
        velocity_local = Robot.get_3d_local_velocity()
        roll, pitch, yaw = Robot.get_roll_pitch_yaw()

        # reward_v = rewards.tolerance(
        #     (np.cos(pitch) * velocity_local[0]),
        #     bounds=(target_velocity,
        #             target_velocity + 0.1),
        #     margin=target_velocity,
        #     value_at_margin=0,
        #     sigmoid='linear'
        # ) * target_velocity
        projected_x_velocity = np.cos(pitch) * velocity_local[0]

        reward_v = near_quadratic_bound(
            projected_x_velocity, 
            target_velocity, 
            target_velocity, 
            target_velocity,
            "linear", #"gaussian", #"linear",
            1.6,
            0.0
        )
        
        reward_v *= (1 + np.cos(target_delta_yaw)) / 2
        return reward_v

    def calculate_contact_reward_unnorm(
        self,
        Robot: BaseWalker
    ) -> float:
        if hasattr(Robot, "get_foot_contact"):
            joint_qpos = Robot.get_joint_qpos()
            diff_joint_qpos = joint_qpos - self._last_joint_qpos
            
            prev_contacts = self._last_contacts
            prev_forces = self._last_foot_force_norm
            contacts : np.ndarray = Robot.get_foot_contact()
            forces : np.ndarray = Robot.get_foot_force_norm()

            rew_norms = []
            qpos_threshold = CONTACT_DELTA_QPOS_THRESHOLD * Robot.control_timestep
            force_threshold = CONTACT_DELTA_FORCE_THRESHOLD * Robot.control_timestep
            for i, contact in enumerate(prev_contacts):
                joint_index_start = 3 * i
                delta_force = forces[i] - prev_forces[i]
                delta_qpos = diff_joint_qpos[joint_index_start+1:joint_index_start+3]
                if contact:
                    condition_rew = np.sum(delta_qpos) >= qpos_threshold or (delta_force < -force_threshold)
                else:
                    condition_rew = np.sum(delta_qpos) <= -qpos_threshold or (delta_force > force_threshold)
                rew_norms.append(np.any(condition_rew))
            rew_norm = np.mean(rew_norms)
            return rew_norm
        else:
            return 0.0

    def step_reward(
        self, 
        Robot: BaseWalker, 
        action_target_qpos: np.ndarray,
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw : float,
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> None:
        velocity_norm = np.linalg.norm(Robot.get_3d_linear_velocity())

        if target_velocity >= 0.05:
            reward_v = self.calculate_velocity_reward_norm(
                Robot, action_target_qpos, target_goal_world_delta, target_goal_local, target_yaw,target_delta_yaw, target_velocity, velocity_to_goal,change_in_abs_target_delta_yaw, target_custom_data, enable_target_custom_obs, info_dict,randomState
            )
            reward_v *= max(target_velocity, 1/5) # Scale so that the reward is more consistent with same velocity but different target velocities
            reward_v *= 20.0
            step_energy_coefficeint = self.energy_penalty_weight
        else:
            velocity_global = Robot.get_3d_linear_velocity()
            reward_close_to_0 = 1 - np.linalg.norm(velocity_global) * 5 # Encourage velocity to be close to 0
            reward_v = reward_close_to_0 * 3
            step_energy_coefficeint = 0.0
        
        info_dict["reward_v"] = reward_v
        
        # Calculate upright coefficient
        up = tr3d.quaternions.quat2mat(Robot.get_framequat_wijk())[-1,-1] # See if the z-axis is pointing up

        # upright_coefficient = rewards.tolerance(up,
        #                             bounds=(0.9, float('inf')),
        #                             sigmoid='quadratic',
        #                             margin=0.9,
        #                             value_at_margin=0)
        upright_coefficient = (0.5*up + 0.5)**2 
        info_dict["reward_upright_coefficient"] = upright_coefficient

        reward_v_constrained = reward_v * upright_coefficient
        info_dict["reward_v_constrained"] = reward_v_constrained
        
        # Calculate Qpos Reward
        if enable_target_custom_obs and target_custom_data is not None and isinstance(target_custom_data, dict) and "should_crouch" in target_custom_data.keys() and target_custom_data["should_crouch"]:
            tar_pose = Robot.joint_qpos_crouch
        else:
            tar_pose = Robot.joint_qpos_init
        
        offset = Robot.joint_qpos_offset
        bounds = (tar_pose - offset, tar_pose + offset)

        qpos = Robot.get_joint_qpos()

        qpos_penalty_norm = 1.0
        for i in range(len(qpos)):
            single_qpos_rew = rewards.tolerance(
                qpos[i],
                bounds=(bounds[0][i], bounds[1][i]),
                margin=offset[i],
                value_at_margin=0.6
            )
            qpos_penalty_norm *= single_qpos_rew
        #####################################
        info_dict["penalty_qpos_normalized"] = qpos_penalty_norm
        qpos_penalty = (0.6 - qpos_penalty_norm) * self.qpos_penalty_weight
        info_dict["penalty_qpos"] = qpos_penalty


        # shifted_qpos = np.concatenate((
        #     qpos[0:3], [-qpos[3]], qpos[4:6], qpos[6:9], [-qpos[9]], qpos[10:]
        # ))
        # qpos_by_fr = np.stack(
        #     (
        #         np.stack((shifted_qpos[0:3], shifted_qpos[3:6]), axis=0),
        #         np.stack((shifted_qpos[6:9], shifted_qpos[9:12]), axis=0)
        #     ),
        #     axis=0
        # )
        # avg_qpos_by_fr = np.mean(qpos_by_fr, axis=1).flatten()
        # real_bounds = []
        # for bound in bounds:
        #     real_bound = np.concatenate((
        #         bound[0:3], bound[6:9]
        #     ))
        #     real_bounds.append(real_bound)
        # real_offset = np.concatenate((
        #     offset[0:3], offset[6:9]
        # ))
        # qpos_penalty_norm = 1.0
        # for i in range(len(avg_qpos_by_fr)):
        #     single_qpos_rew = rewards.tolerance(
        #         avg_qpos_by_fr[i],
        #         bounds=(real_bounds[0][i], real_bounds[1][i]),
        #         margin=real_offset[i],
        #         value_at_margin=0.1
        #     )
        #     qpos_penalty_norm *= single_qpos_rew
        # #####################################
        # info_dict["penalty_qpos_normalized"] = qpos_penalty_norm
        # qpos_penalty = (0.6-qpos_penalty_norm) * self.qpos_penalty_weight
        # info_dict["penalty_qpos"] = qpos_penalty

         # Angular velocity penalty
        angular_velocity = Robot.get_3d_angular_velocity()
        pitch_rate_penalty = self.pitch_rate_penalty_factor * (np.abs(angular_velocity[1]) ** 1.4) - 0.5
        roll_rate_penalty = self.roll_rate_penalty_factor * (np.abs(angular_velocity[0]) ** 1.4) - 0.5
        info_dict["pitch_rate_penalty"] = pitch_rate_penalty
        info_dict["roll_rate_penalty"] = roll_rate_penalty

        # Calculate Energy / Qvel Penalty
        joint_qvel = Robot.get_joint_qvel()
        joint_torques = Robot.get_joint_torques()
        energy = np.sum(np.abs(joint_qvel * joint_torques))
        # energy_sq = np.mean(np.abs(qvel) * (np.abs(torque) ** 1.5))

        if step_energy_coefficeint > 0.0:
            #energy_penalty = -0.5 * np.exp(-0.1*energy)
            energy_penalty = step_energy_coefficeint * (energy - 50 * (2 ** (Robot.limit_action_range / 0.3 - 1)))
            # energy_penalty = step_energy_coefficeint * (np.mean(np.abs(Robot.get_joint_torques()) ** 1.5) - 15)
            # energy_penalty = 0.04 * energy_sq - 1
            # energy_penalty = 0.02 * np.mean((np.abs(torque) ** 1.3)) - 1
        else:
            energy_penalty = 0 #-np.exp(-0.1*np.linalg.norm(qvel))
        
        info_dict['energy'] = energy
        #info_dict["energy_sq"] = energy_sq
        info_dict['penalty_energy'] = energy_penalty
        
        # Calculate Smooth Torque Penalty
        # pd_joint_torques = calculate_torque(
        #     Robot.get_joint_qpos(),
        #     Robot.get_joint_qvel(),
        #     action_target_qpos,
        #     Robot.Kp,
        #     Robot.Kd
        # )
        diff_torque = joint_torques - self._last_torque #pd_joint_torques - self._last_torque
        diff_torque_norm = np.linalg.norm(diff_torque) ** 1.5
        
        smooth_torque_penalty = self.smooth_torque_penalty_weight * (diff_torque_norm - 150) if self.smooth_torque_penalty_weight > 0 else 0.0
        info_dict["smooth_torque_penalty"] = smooth_torque_penalty
        info_dict["diff_torque_norm"] = diff_torque_norm

        # Calculate joint diagonal / shoulder symmetry penalty
        FR_qvel = joint_qvel[0:3]
        FL_qvel = joint_qvel[3:6]
        FL_qvel[0] *= -1
        RR_qvel = joint_qvel[6:9]
        RL_qvel = joint_qvel[9:12]
        RL_qvel[0] *= -1
        # diagonal_difference = (FR_qvel[1] - RL_qvel[1])**2.0 + (FL_qvel[1] - RR_qvel[1])**2.0
        
        if hasattr(Robot, "get_foot_contact"):
            contacts = Robot.get_foot_contact()
            shoulder_difference = np.linalg.norm(
                FR_qvel[1:] - RL_qvel[1:]
            ) * (np.sum(contacts[0:2]) == 1.0) + np.linalg.norm(
                FL_qvel[1:] - RR_qvel[1:]
            ) * (np.sum(contacts[2:]) == 1.0) * 2.0
        else:
            shoulder_difference = np.linalg.norm(FR_qvel[1:] - FL_qvel[1:]) + 2 * np.linalg.norm(RR_qvel[1:] - RL_qvel[1:]) # penalize harder on the back legs
        
        diagonal_difference = np.linalg.norm(FR_qvel[1:] - RL_qvel[1:]) + np.linalg.norm(FL_qvel[1:] - RR_qvel[1:])

        # shoulder_center = np.linalg.norm(FR_qvel[1:] + FL_qvel[1:]) + 2 * np.linalg.norm(RR_qvel[1:] + RL_qvel[1:])
        diagonal_difference_penalty = self.joint_diagonal_penalty_weight * diagonal_difference
        diagonal_difference_penalty *= (0.5 + velocity_norm) / (1.0 + 4.0 * np.abs(target_delta_yaw) / np.pi)

        shoulder_difference_penalty = self.joint_shoulder_penalty_weight * -(shoulder_difference)
        shoulder_difference_penalty *= (0.5 + velocity_norm) / (1.0 + 0.5 * np.abs(target_delta_yaw) / np.pi)

        info_dict["diagonal_difference_penalty"] = diagonal_difference_penalty
        info_dict["shoulder_difference_penalty"] = shoulder_difference_penalty

        # Calculate thigh torque penalty
        thigh_torques = joint_torques[1::3]
        thigh_torque_norm = np.linalg.norm(thigh_torques)
        thigh_torque_penalty = self.thigh_torque_penalty_weight * (thigh_torque_norm - 15)
        info_dict["thigh_torque_norm"] = thigh_torque_norm
        info_dict["thigh_torque_penalty"] = thigh_torque_penalty

        # Calculate joint acceleration penalty
        joint_qacc = Robot.get_joint_qacc()
        joint_qacc_penalty = self.joint_acc_penalty_weight * np.linalg.norm(joint_qacc, ord=1)
        info_dict["joint_qacc_penalty"] = joint_qacc_penalty

        # Calculate joint velocity penalty
        joint_qvel_penalty = self.joint_vel_penalty_weight * np.linalg.norm(joint_qvel)
        info_dict["joint_qvel_penalty"] = joint_qvel_penalty

        # Calculate contact reward
        contact_reward = self.calculate_contact_reward_unnorm(Robot) * self.contact_reward_weight
        info_dict["contact_reward"] = contact_reward

        self._last_torque = joint_torques.copy() #pd_joint_torques.copy()
        self._last_joint_qpos = qpos.copy()
        self._last_contacts = Robot.get_foot_contact().copy() if hasattr(Robot, "get_foot_contact") else None
        self._last_foot_force_norm = Robot.get_foot_force_norm().copy() if hasattr(Robot, "get_foot_force_norm") else None

        rew_perstep = reward_v_constrained - qpos_penalty - energy_penalty - smooth_torque_penalty - roll_rate_penalty - pitch_rate_penalty - diagonal_difference_penalty - shoulder_difference_penalty - thigh_torque_penalty - joint_qacc_penalty - joint_qvel_penalty + contact_reward
        self.rew = rew_perstep

class JoystickPolicySeperateRewardProvider(JoystickPolicyStrictRewardProvider):
    def __init__(self, smooth_tdy_steps : int = 5, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.queue_tdy = deque(maxlen=smooth_tdy_steps)
    
    @property
    def smooth_tdy_steps(self) -> int:
        return self.queue_tdy.maxlen
    
    @smooth_tdy_steps.setter
    def smooth_tdy_steps(self, value : int) -> None:
        self.queue_tdy = deque(maxlen=value)

    def calculate_velocity_reward_norm(
        self, 
        Robot: BaseWalker, 
        action_target_qpos: np.ndarray,
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw : float,
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> float:
        reward_v = super().calculate_velocity_reward_norm(
            Robot,
            action_target_qpos,
            target_goal_world_delta,
            target_goal_local,
            target_yaw,
            target_delta_yaw,
            target_velocity,
            velocity_to_goal,
            change_in_abs_target_delta_yaw,
            target_custom_data,
            enable_target_custom_obs,
            info_dict,
            randomState
        )
        reward_v /= (1 + np.cos(target_delta_yaw)) / 2 # Revert the cosine reward
        return reward_v
    
    def step_reward(
        self, 
        Robot: BaseWalker, 
        action_target_qpos: np.ndarray,
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw : float,
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> None:
        super().step_reward(
            Robot, 
            action_target_qpos,
            target_goal_world_delta, 
            target_goal_local, 
            target_yaw, 
            target_delta_yaw, 
            target_velocity,
            velocity_to_goal, 
            change_in_abs_target_delta_yaw, 
            target_custom_data, 
            enable_target_custom_obs,
            info_dict, 
            randomState
        )
        self.queue_tdy.append(change_in_abs_target_delta_yaw)
        if target_velocity >= 0.05:
            rew_step = self.rew

            smoothed_change_in_tdy = np.mean(self.queue_tdy)

            #if target_delta_yaw != 0.0:
            target_change_tdy = 45.0 / 180.0 * np.pi * Robot.control_timestep * (target_delta_yaw / np.pi)
            target_change_tdy = -np.abs(target_change_tdy)
            
            # else:
            #     target_change_tdy = 0.0
            
            info_dict["target_change_tdy"] = target_change_tdy

            target_change_tdy_left_slack = 45.0 / 180.0 * np.pi * Robot.control_timestep
            target_change_tdy_right_slack = np.clip(
                -target_change_tdy + 20.0 / 180.0 * np.pi * Robot.control_timestep, target_change_tdy_left_slack, None
            )
            
            rew_change_in_abs_tdy = near_quadratic_bound(
                smoothed_change_in_tdy, 
                target_change_tdy, 
                target_change_tdy_left_slack,
                target_change_tdy_right_slack,
                "gaussian",
                1.6,
                0.1
            )
            
            # rew_change_in_abs_tdy *= max(np.abs(target_delta_yaw) / np.pi, 1/5)
            velocity_norm = np.linalg.norm(Robot.get_3d_linear_velocity())
            rew_change_in_abs_tdy *= 4.0 * np.sqrt(self.smooth_tdy_steps) * (0.1 + velocity_norm)
            #rew_change_in_abs_tdy *= 1.0 * np.sqrt(self.smooth_tdy_steps)
            info_dict["rew_change_in_abs_TDY"] = rew_change_in_abs_tdy
            rew_step += rew_change_in_abs_tdy
            self.rew = rew_step

    def reset_reward(self, Robot: BaseWalker, info_dict: dict[str, Any], termination_provider_triggered: JoystickPolicyTerminationConditionProvider, randomState: np.random.RandomState) -> None:
        super().reset_reward(Robot, info_dict, termination_provider_triggered, randomState)
        self.queue_tdy.clear()
        info_dict["rew_change_in_abs_TDY"] = 0.0

class JoystickPolicyETHRewardProvider(JoystickPolicyRewardProvider[BaseWalker]):
    def __init__(self) -> None:
        super().__init__()
        self.rew = 0.0

    def get_reward(self) -> float:
        return self.rew

    def reset_reward(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str, Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider, 
        randomState: np.random.RandomState
    ) -> None:
        self.rew = 0.0
    
    def step_reward(
        self, 
        Robot: BaseWalker, 
        action_target_qpos: np.ndarray,
        target_goal_world_delta: np.ndarray, 
        target_goal_local: np.ndarray, 
        target_yaw: float, 
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw: float, 
        target_custom_data: Any | None, 
        enable_target_custom_obs: bool, 
        info_dict: dict[str, Any], 
        randomState: np.random.RandomState
    ) -> None:
        target_angular_velocity = 45.0 / 180.0 * np.pi * (target_delta_yaw / np.pi)

        local_velocity = Robot.get_3d_local_velocity()
        linear_velocity = Robot.get_3d_linear_velocity()
        angular_velocity = Robot.get_3d_angular_velocity()
        roll, pitch, yaw = Robot.get_roll_pitch_yaw()

        target_linear_velocity_2d = np.array([
            np.cos(yaw) * target_velocity,
            np.sin(yaw) * target_velocity,
        ])

        velocity_term = 1.0 * Robot.control_timestep * calculate_gaussian_activation(linear_velocity[:2] - target_linear_velocity_2d)
        angular_term = 0.5 * Robot.control_timestep * calculate_gaussian_activation(angular_velocity[2] - target_angular_velocity)
        linear_velocity_penalty = -4 * Robot.control_timestep * linear_velocity[2] ** 2
        angular_velocity_penalty = -0.05 * Robot.control_timestep * np.linalg.norm(angular_velocity[:2])**2
        joint_motion_penalty = -0.001 * Robot.control_timestep * (np.linalg.norm(Robot.get_joint_qacc())**2 + np.linalg.norm(Robot.get_joint_qvel())**2)
        joint_torque_penalty = -0.00002 * Robot.control_timestep * np.linalg.norm(Robot.get_joint_torques())**2
        
        info_dict["velocity_term"] = velocity_term
        info_dict["angular_term"] = angular_term
        info_dict["linear_velocity_penalty"] = linear_velocity_penalty
        info_dict["angular_velocity_penalty"] = angular_velocity_penalty
        info_dict["joint_motion_penalty"] = joint_motion_penalty
        info_dict["joint_torque_penalty"] = joint_torque_penalty

        self.rew = 2 + velocity_term + angular_term + linear_velocity_penalty + angular_velocity_penalty + joint_motion_penalty + joint_torque_penalty
rail_walker_gym/joystick_policy/reward_providers.py
realquantumcookie/APRL
from rail_walker_interface import BaseWalker
import numpy as np
import typing

def calculate_energy(joint_qvel : np.ndarray, joint_torques : np.ndarray) -> float:
    return np.sum(np.abs(joint_qvel * joint_torques))

def calculate_qvel_penalty(robot: BaseWalker) -> typing.Tuple[float, float]:
    qvel = robot.get_joint_qvel()
    qvel_penalty = -np.exp(-0.1*np.linalg.norm(qvel))

    return qvel_penalty, qvel

def near_quadratic_bound(value, target, left_margin, right_margin, out_of_margin_activation : str | None = "linear", power = 2.0, value_at_margin = 0.0):
    delta = value-target
    fract = delta/right_margin if delta > 0 else delta/left_margin
    
    if out_of_margin_activation is None or out_of_margin_activation != "near_quadratic":
        clipped_fract = np.clip(fract, -1.0, 1.0)
        rew = 1 - (1-value_at_margin) * (np.abs(clipped_fract) ** power)
        oodfract = fract - clipped_fract
        if out_of_margin_activation == "linear":
            rew -= (1-value_at_margin) * np.abs(oodfract)
        elif out_of_margin_activation == "quadratic":
            rew -= (1-value_at_margin) * (oodfract ** 2)
        elif out_of_margin_activation == "gaussian":
            rew += value_at_margin * np.exp(-oodfract**2/0.25)
    elif out_of_margin_activation == "near_quadratic":
        rew = 1 - (1-value_at_margin) * (np.abs(fract) ** power)
    return rew

def calculate_torque(current_qpos : np.ndarray, current_qvel : np.ndarray, target_qpos : np.ndarray, Kp : float, Kd : float):
    return Kp * (target_qpos - current_qpos) - Kd * current_qvel

def calculate_gaussian_activation(x : float | np.ndarray):
    return np.exp(-np.linalg.norm(x)**2/0.25)

rail_walker_gym/joystick_policy/reward_util.py
realquantumcookie/APRL
from rail_walker_interface import JoystickPolicyTargetObservable, BaseWalker, JoystickPolicyTerminationConditionProvider
import gym
import gym.spaces
import numpy as np
from typing import Any, Optional

class JoystickPolicyTargetDeltaYawObservable(JoystickPolicyTargetObservable[BaseWalker]):
    def __init__(self) -> None:
        super().__init__()
        self.obs : float = 0.0

    def get_observation_spec(self) -> gym.Space:
        return gym.spaces.Box(low=-np.pi, high=np.pi, shape=(1,), dtype=np.float32)
    
    def step_target_obs(
        self,
        Robot: BaseWalker, 
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw: float, 
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> None:
        self.obs = target_delta_yaw

    def reset_target_obs(
        self, 
        Robot: BaseWalker, 
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw: float, 
        target_delta_yaw: float, 
        target_velocity: float,
        info_dict: dict[str,Any], 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        self.obs = target_delta_yaw

    def get_observation(self) -> Any:
        return np.array([self.obs], dtype=np.float32)
    
class JoystickPolicyCosSinTargetDeltaYawObservable(JoystickPolicyTargetObservable[BaseWalker]):
    def __init__(self) -> None:
        super().__init__()
        self.obs = np.zeros(2, dtype=np.float32)

    def get_observation_spec(self) -> gym.Space:
        return gym.spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32)
    
    def step_target_obs(
        self,
        Robot: BaseWalker, 
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw: float, 
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> None:
        self.obs = np.array([np.cos(target_delta_yaw), np.sin(target_delta_yaw)], dtype=np.float32)

    def reset_target_obs(
        self, 
        Robot: BaseWalker, 
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw: float, 
        target_delta_yaw: float, 
        target_velocity: float,
        info_dict: dict[str,Any], 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        self.obs = np.array([np.cos(target_delta_yaw), np.sin(target_delta_yaw)], dtype=np.float32)

    def get_observation(self) -> Any:
        return self.obs
    
class JoystickPolicyTargetDeltaYawAndDistanceObservable(JoystickPolicyTargetObservable[BaseWalker]):
    def __init__(self) -> None:
        super().__init__()
        self.obs = np.zeros(2, dtype=np.float32)

    def get_observation_spec(self) -> gym.Space:
        return gym.spaces.Box(low=np.array([-np.pi,-np.inf]), high=np.array([np.pi,np.inf]), shape=(2,), dtype=np.float32)
    
    def step_target_obs(
        self,
        Robot: BaseWalker, 
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw: float, 
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> None:
        self.obs = np.array([target_delta_yaw,np.linalg.norm(target_goal_world_delta)], dtype=np.float32)

    def reset_target_obs(
        self, 
        Robot: BaseWalker, 
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw: float, 
        target_delta_yaw: float, 
        target_velocity: float,
        info_dict: dict[str,Any], 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        self.obs = np.array([target_delta_yaw,np.linalg.norm(target_goal_world_delta)], dtype=np.float32)


    def get_observation(self) -> Any:
        return self.obs
rail_walker_gym/joystick_policy/target_obs_providers.py
realquantumcookie/APRL
import gym
import gym.spaces
import numpy as np
from rail_walker_interface import JoystickPolicyTargetProvider, JoystickPolicyTerminationConditionProvider
from rail_walker_interface import BaseWalker
from typing import Any, Optional
from rail_walker_interface.joystick_policy.joystick_interfaces import JoystickPolicyTerminationConditionProvider

from rail_walker_interface.robot.robot import BaseWalker

def normalize_rad(rad : float) -> float:
    return (rad + np.pi) % (2 * np.pi) - np.pi

class JoystickPolicyForwardOnlyTargetProvider(JoystickPolicyTargetProvider[BaseWalker]):
    def __init__(self, ahead_distance : float = 1.0, target_linear_velocity : float = 1.0) -> None:
        super().__init__()
        self.ahead_distance = ahead_distance
        self.target_linear_velocity = target_linear_velocity
        self._world_delta = np.zeros(2)

    def get_target_goal_world_delta(self, Robot: BaseWalker) -> np.ndarray:
        return self._world_delta
    
    def get_target_velocity(self, Robot: BaseWalker) -> float:
        return self.target_linear_velocity
    
    def is_target_velocity_fixed(self) -> bool:
        return True

    def has_target_changed(self) -> bool:
        return False

    def step_target(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str,Any], 
        randomState : np.random.RandomState
    ) -> None:
        _, _, yaw = Robot.get_roll_pitch_yaw()
        self._update_target(yaw)

    def reset_target(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str,Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        _, _, yaw = Robot.get_roll_pitch_yaw()
        self._update_target(yaw)

    def _update_target(self, yaw : float):
        self._world_delta[0] = np.cos(yaw) * self.ahead_distance
        self._world_delta[1] = np.sin(yaw) * self.ahead_distance

"""
This class mimics a behavior of a joystick, it follows a random walk time series and, 
with a small chance each time, randomly jumps to a sample from a uniform distribution in the range [-pi, pi)
"""
class JoystickPolicyAutoJoystickTargetProvider(JoystickPolicyTargetProvider[BaseWalker]):
    def __init__(
        self,
        random_walk_std_per_step : list[tuple[float, float]]= [
            (0.6, 30 / 180 * np.pi), # with 0.8 probability, sample from a normal distribution with std 30 degrees
            (0.4, 90 / 180 * np.pi), # with 0.2 probability, sample from a normal distribution with std 90 degrees
        ],
        random_jump_prob : float = 0.4,
        ahead_distance : float = 1.0,
        min_target_velocity : float = 0.0,
        max_target_velocity : float = 1.0,
        max_velocity_change_rate : float = 0.3,
        velocity_probability_distribution : Optional[
            list[tuple[float, float]]
        ] = [
            (0.1, 0.05),
            (0.1, 0.3),
            (0.5, 0.7),
            (0.1, 0.9),
            (0.2, 1.0)
        ],
        target_merge_to_zero_interval : float = 10 / 180 * np.pi,
        keep_target_delta_yaw_steps_avg : int = 20 * 3,
        keep_velocity_steps_avg : int = 20 * 3,
    ):
        self.random_walk_std_per_step= random_walk_std_per_step
        self.random_jump_prob= random_jump_prob
        self.ahead_distance = ahead_distance
        self.min_target_velocity = min_target_velocity
        self.max_target_velocity = max_target_velocity
        self.max_velocity_change_rate = max_velocity_change_rate
        self.velocity_probability_distribution = velocity_probability_distribution
        self.target_merge_to_zero_interval = target_merge_to_zero_interval
        self.keep_target_delta_yaw_steps_avg = keep_target_delta_yaw_steps_avg
        self.keep_velocity_steps_avg = keep_velocity_steps_avg

        self._target_delta_yaw = 0.0
        self._target_velocity = (min_target_velocity + max_target_velocity) / 2
        self._steps_until_next_target_velocity_change = 0
        self._steps_until_next_target_delta_yaw_change = 0
        self._should_change_velocity = False
        self._should_change_delta_yaw = False
    
    def get_target_goal_world_delta(self, Robot: BaseWalker) -> np.ndarray:
        target_yaw = self._target_delta_yaw + Robot.get_roll_pitch_yaw()[2]
        return np.array([
            np.cos(target_yaw), np.sin(target_yaw)
        ]) * self.ahead_distance
    
    def get_target_velocity(self, Robot: BaseWalker) -> float:
        return self._target_velocity
    
    def is_target_velocity_fixed(self) -> bool:
        return False

    def has_target_changed(self) -> bool:
        return self._should_change_velocity or self._should_change_delta_yaw

    def step_target(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str,Any], 
        randomState : np.random.RandomState
    ) -> None:
        if self._steps_until_next_target_delta_yaw_change > 1:
            self._steps_until_next_target_delta_yaw_change -= 1
            self._should_change_delta_yaw = False
        else:
            self._should_change_delta_yaw = True
        
        if self._steps_until_next_target_velocity_change > 1:
            self._steps_until_next_target_velocity_change -= 1
            self._should_change_velocity = False
        else:
            self._should_change_velocity = True
    
    def after_step_target(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str,Any], 
        randomState : np.random.RandomState
    ) -> None:
        if self._should_change_delta_yaw:
            self._steps_until_next_target_delta_yaw_change = randomState.poisson(self.keep_target_delta_yaw_steps_avg) if self.keep_target_delta_yaw_steps_avg > 0 else 0
            if randomState.uniform() < self.random_jump_prob:
                self._target_delta_yaw = randomState.uniform(-np.pi, np.pi)
            else:
                rnum = randomState.uniform()
                t_prob = 0
                t_std = 0.5 / 180 * np.pi # default to 0.5 degrees
                for prob, std in self.random_walk_std_per_step:
                    t_prob += prob
                    if rnum < t_prob:
                        t_std = std
                        break
                
                self._target_delta_yaw += randomState.normal(0, t_std)
                self._target_delta_yaw = normalize_rad(self._target_delta_yaw)
            if np.abs(self._target_delta_yaw) < self.target_merge_to_zero_interval / 2:
                self._target_delta_yaw = 0.0
            
            self._should_change_delta_yaw = False
        if self._should_change_velocity:
            self._steps_until_next_target_velocity_change = randomState.poisson(self.keep_velocity_steps_avg) if self.keep_velocity_steps_avg > 0 else 0
            rand_change_in_v_unit = randomState.uniform()
            if self.velocity_probability_distribution is None:
                change_in_velocity = rand_change_in_v_unit * (self.max_velocity_change_rate * 2) - self.max_velocity_change_rate
            else:
                probability_le_current_velocity = 0.0
                for i in range(len(self.velocity_probability_distribution)):
                    p, v = self.velocity_probability_distribution[i]
                    if v < self._target_velocity:
                        probability_le_current_velocity += p
                    else:
                        prev_p, prev_v = self.velocity_probability_distribution[i-1] if i > 0 else (0.0, self.min_target_velocity)
                        probability_le_current_velocity += (self._target_velocity - prev_v) / (v - prev_v) * p
                        break
                if rand_change_in_v_unit <= probability_le_current_velocity:
                    change_in_velocity = (rand_change_in_v_unit / probability_le_current_velocity - 1.0) * self.max_velocity_change_rate
                else:
                    probability_g_current_velocity = 1.0 - probability_le_current_velocity
                    change_in_velocity = (rand_change_in_v_unit - probability_le_current_velocity) / probability_g_current_velocity * self.max_velocity_change_rate
            
            self._target_velocity += change_in_velocity
            self._target_velocity = np.clip(self._target_velocity, self.min_target_velocity, self.max_target_velocity)
            self._should_change_velocity = False

    def reset_target(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str,Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        self._target_delta_yaw = randomState.uniform(-np.pi, np.pi)
        if np.abs(self._target_delta_yaw) < self.target_merge_to_zero_interval / 2:
            self._target_delta_yaw = 0.0
        rand_v = randomState.uniform()
        if self.velocity_probability_distribution is None:
            self._target_velocity = rand_v * (self.max_target_velocity - self.min_target_velocity) + self.min_target_velocity
        else:
            accumulated_probability = 0.0
            for i in range(len(self.velocity_probability_distribution)):
                p,v = self.velocity_probability_distribution[i]
                if rand_v > accumulated_probability + p:
                    accumulated_probability += p
                else:
                    prev_p, prev_v = self.velocity_probability_distribution[i-1] if i > 0 else (0.0, self.min_target_velocity)
                    self._target_velocity = prev_v + (rand_v - accumulated_probability) / p * (v - prev_v)
                    break

        self._steps_until_next_target_velocity_change = randomState.poisson(self.keep_velocity_steps_avg) if self.keep_velocity_steps_avg > 0 else 0
        self._steps_until_next_target_delta_yaw_change = randomState.poisson(self.keep_target_delta_yaw_steps_avg) if self.keep_target_delta_yaw_steps_avg > 0 else 0
        self._should_change_velocity = False
        self._should_change_delta_yaw = False

class JoystickPolicyAutoJoystickSimpleTargetProvider(JoystickPolicyTargetProvider[BaseWalker]):
    def __init__(
        self,
        ahead_distance : float = 1.0,
        avg_num_steps_per_target : int = 20 * 2,
        target_linear_velocity : float = 0.5
    ) -> None:
        super().__init__()
        self.ahead_distance = ahead_distance
        self.avg_num_steps_per_target = avg_num_steps_per_target
        self.target_linear_velocity = target_linear_velocity
        self._target_delta_yaw = 0.0
        self._next_target_delta_yaw = 0.0
        self._steps_until_next_target = 0

    def get_world_delta(self, target_yaw : float) -> np.ndarray:
        return np.array([np.cos(target_yaw), np.sin(target_yaw)]) * self.ahead_distance
    
    def get_target_velocity(self, Robot: BaseWalker) -> float:
        return self.target_linear_velocity
    
    def is_target_velocity_fixed(self) -> bool:
        return True

    def get_target_goal_world_delta(self, Robot: BaseWalker) -> np.ndarray:
        return self.get_world_delta(self._target_delta_yaw + Robot.get_roll_pitch_yaw()[2])
    
    def has_target_changed(self) -> bool:
        return self._target_delta_yaw != self._next_target_delta_yaw
    
    def step_target(self, Robot: BaseWalker, info_dict: dict[str, Any], randomState: np.random.RandomState) -> None:
        if self._steps_until_next_target >= 1:
            self._steps_until_next_target -= 1
        else:
            self._steps_until_next_target = randomState.poisson(self.avg_num_steps_per_target)
            self._next_target_delta_yaw = randomState.choice(
                np.array([-np.pi/2.0, 0.0, np.pi/2.0])
            )
    
    def after_step_target(self, Robot: BaseWalker, info_dict: dict[str, Any], randomState: np.random.RandomState) -> None:
        self._target_delta_yaw = self._next_target_delta_yaw
    
    def reset_target(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str, Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider, 
        randomState: np.random.RandomState
    ) -> None:
        self._next_target_delta_yaw = randomState.choice(
            np.array([-np.pi/2.0, 0.0, np.pi/2.0])
        )
        self._target_delta_yaw = self._next_target_delta_yaw
        self._steps_until_next_target = randomState.poisson(self.avg_num_steps_per_target)

class JoystickPolicyAutoJoystickSimplePlusTargetProvider(JoystickPolicyTargetProvider[BaseWalker]):
    def __init__(
        self,
        ahead_distance : float = 1.0,
        avg_num_steps_per_target : int = 20 * 2,
        min_target_velocity : float = 0.2,
        max_target_velocity : float = 1.0,
    ) -> None:
        super().__init__()
        self.ahead_distance = ahead_distance
        self.avg_num_steps_per_target = avg_num_steps_per_target
        self.min_target_velocity = min_target_velocity
        self.max_target_velocity = max_target_velocity

        self._target_delta_yaw = 0.0
        self._target_velocity = 0.0
        self._next_target_delta_yaw = 0.0
        self._next_target_velocity = 0.0
        self._steps_until_next_target = 0
        
    
    @property
    def current_custom_target_velocity(self) -> float:
        return self._target_velocity
    
    def get_world_delta(self, target_yaw : float) -> np.ndarray:
        return np.array([np.cos(target_yaw), np.sin(target_yaw)]) * self.ahead_distance
    
    def get_target_velocity(self, Robot: BaseWalker) -> float:
        return self._target_velocity
    
    def is_target_velocity_fixed(self) -> bool:
        return False

    def get_target_goal_world_delta(self, Robot: BaseWalker) -> np.ndarray:
        return self.get_world_delta(self._target_delta_yaw + Robot.get_roll_pitch_yaw()[2])
    
    def has_target_changed(self) -> bool:
        return self._target_delta_yaw != self._next_target_delta_yaw or self._target_velocity != self._next_target_velocity
    
    def step_target(self, Robot: BaseWalker, info_dict: dict[str, Any], randomState: np.random.RandomState) -> None:
        if self._steps_until_next_target >= 1:
            self._steps_until_next_target -= 1
        else:
            self._steps_until_next_target = randomState.poisson(self.avg_num_steps_per_target)
            self._next_target_delta_yaw = randomState.choice(
                np.array([-np.pi/2.0, 0.0, np.pi/2.0])
            )
            self._next_target_velocity = randomState.uniform(self.min_target_velocity, self.max_target_velocity)
    
    def after_step_target(self, Robot: BaseWalker, info_dict: dict[str, Any], randomState: np.random.RandomState) -> None:
        self._target_delta_yaw = self._next_target_delta_yaw
        self._target_velocity = self._next_target_velocity
    
    def reset_target(
        self, 
        Robot: BaseWalker, 
        info_dict: dict[str, Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider, 
        randomState: np.random.RandomState
    ) -> None:
        self._next_target_delta_yaw = randomState.choice(
            np.array([-np.pi/2.0, 0.0, np.pi/2.0])
        )
        self._next_target_velocity = randomState.uniform(self.min_target_velocity, self.max_target_velocity)
        self._target_delta_yaw = self._next_target_delta_yaw
        self._target_velocity = self._next_target_velocity
        self._steps_until_next_target = randomState.poisson(self.avg_num_steps_per_target)
    
rail_walker_gym/joystick_policy/target_providers.py
realquantumcookie/APRL
from typing import Any, Optional
from dm_control.mujoco.engine import Physics as EnginePhysics
import gym
from rail_mujoco_walker.tasks.joystick_task import JoystickPolicyDMControlTask
from rail_walker_interface import JoystickPolicy, JoystickPolicyTargetProvider, BaseWalker, JoystickPolicyTerminationConditionProvider
import numpy as np
from dm_control.mjcf.physics import Physics
import transforms3d as tr3d
import typing
from rail_mujoco_walker import add_arrow_to_mjv_scene, JoystickPolicyProviderWithDMControl
from mujoco import MjvScene

def normalize_rad(rad : float) -> float:
    return (rad + np.pi) % (2 * np.pi) - np.pi

_RobotType = typing.TypeVar("_RobotType", bound=BaseWalker)
class JoystickPolicyTargetProviderWrapper(typing.Generic[_RobotType], JoystickPolicyTargetProvider[_RobotType], JoystickPolicyProviderWithDMControl):
    def __init__(
        self,
        provider: JoystickPolicyTargetProvider[_RobotType]
    ):
        JoystickPolicyTargetProvider.__init__(self)
        JoystickPolicyProviderWithDMControl.__init__(self)
        self.provider = provider
        self.original_target : np.ndarray = np.zeros(2)
        self.target : np.ndarray = np.zeros(2)
    
    def get_target_goal_world_delta(self, Robot: _RobotType) -> np.ndarray:
        return self.target
    
    def get_target_velocity(self, Robot: _RobotType) -> float:
        return self.provider.get_target_velocity(Robot)
    
    def is_target_velocity_fixed(self) -> bool:
        return self.provider.is_target_velocity_fixed()

    def get_target_custom_data(self) -> Any | None:
        return self.provider.get_target_custom_data()

    def get_target_custom_data_observable(self) -> Any | None:
        return self.provider.get_target_custom_data_observable()

    def get_target_custom_data_observable_spec(self) -> gym.Space | None:
        return self.provider.get_target_custom_data_observable_spec()

    def calculate_new_target(
        self,
        Robot: _RobotType,
        info_dict: dict[str,typing.Any],
        randomState: np.random.RandomState
    ):
        return self.original_target

    def step(
        self, 
        Robot: _RobotType, 
        info_dict: dict[str,typing.Any], 
        randomState : np.random.RandomState
    ) -> None:
        self.provider.step(
            Robot,
            info_dict,
            randomState
        )
        self.original_target = self.provider.get_target_goal_world_delta(Robot)
        self.target = self.calculate_new_target(Robot, info_dict, randomState)

    def reset(
        self, 
        Robot: _RobotType, 
        info_dict: dict[str,typing.Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        self.provider.reset(
            Robot,
            info_dict,
            termination_provider_triggered,
            randomState
        )
        self.original_target = self.provider.get_target_goal_world_delta(Robot)
        self.target = self.calculate_new_target(Robot, info_dict, randomState)

    def render_scene_callback(self, task: JoystickPolicyDMControlTask, physics: Physics, scene: MjvScene) -> None:
        if hasattr(self.provider, "render_scene_callback"):
            self.provider.render_scene_callback(task, physics, scene)
        
        orig_target_unit_3d = np.array([*self.original_target, 0.0])
        orig_target_unit_3d /= np.linalg.norm(orig_target_unit_3d)

        if not np.isclose(self.target, self.original_target).all():
            root_arrow = task._arrow_root + np.array([0.0,0.0,0.25])
            add_arrow_to_mjv_scene(
                scene,
                root_arrow,
                root_arrow + orig_target_unit_3d * 0.5,
                0.01,
                np.array([0.0, 1.0, 1.0, 1.0]) #cyan for original target
            )


_LimitedRobotType = typing.TypeVar("_LimitedRobotType", bound=BaseWalker)
class JoystickPolicyLimitedTargetProvider(typing.Generic[_LimitedRobotType], JoystickPolicyTargetProviderWrapper[_LimitedRobotType]):
    def __init__(
        self, 
        provider: JoystickPolicyTargetProvider[_LimitedRobotType],
        max_delta_angle: float = 30.0 / 180 * np.pi, # 30.0 degrees max
    ):
        JoystickPolicyTargetProviderWrapper.__init__(self, provider)
        self.max_delta_angle = max_delta_angle
    
    def calculate_new_target(
        self,
        Robot: _RobotType,
        info_dict: dict[str,typing.Any],
        randomState: np.random.RandomState
    ):
        roll, pitch, yaw = Robot.get_roll_pitch_yaw()
        original_target_yaw = np.arctan2(self.original_target[1], self.original_target[0])
        target_delta_yaw = normalize_rad(original_target_yaw - yaw)

        clipped_target_delta_yaw = np.clip(target_delta_yaw, -self.max_delta_angle, self.max_delta_angle)
        
        target = np.array([
            np.cos(yaw + clipped_target_delta_yaw),
            np.sin(yaw + clipped_target_delta_yaw)
        ]) # unit vector pointing in the clipped direction
        target *= np.maximum(np.inner(self.original_target, target), 0.1) # project the original target onto the clipped direction
        return target

_SmoothedRobotType = typing.TypeVar("_SmoothedRobotType", bound=BaseWalker)
class JoystickPolicySmoothedTargetProvider(typing.Generic[_SmoothedRobotType], JoystickPolicyTargetProviderWrapper[_SmoothedRobotType]):
    def __init__(
        self,
        provider: JoystickPolicyTargetProvider,
        max_rate: float = 30.0 / 180 * np.pi, # 30.0 degrees max per step
    ):
        JoystickPolicyTargetProviderWrapper.__init__(self, provider)
        self.max_rate = max_rate

    def calculate_new_target(
        self,
        Robot: _RobotType,
        info_dict: dict[str,typing.Any],
        randomState: np.random.RandomState
    ):
        original_target_yaw = np.arctan2(self.original_target[1], self.original_target[0])
        last_target_yaw = np.arctan2(self.target[1], self.target[0])
        change_in_target_yaw = normalize_rad(original_target_yaw - last_target_yaw)

        clipped_change_in_target_yaw = np.clip(change_in_target_yaw, -self.max_rate, self.max_rate)
        new_target_yaw = normalize_rad(last_target_yaw + clipped_change_in_target_yaw)
        
        target = np.array([
            np.cos(new_target_yaw),
            np.sin(new_target_yaw)
        ]) # unit vector pointing in the clipped direction
        target *= np.maximum(np.inner(self.original_target, target), 0.1) # project the original target onto the clipped direction
        return target
rail_walker_gym/joystick_policy/target_yaw_wrappers.py
realquantumcookie/APRL
from .resetters_sim import JoystickPolicyPointInSimResetter, JoystickPolicyLastPositionAndYawResetter, ResetPolicyResetter
from .reward_providers_sim import ResetRewardProvider, GatedResetRewardProvider
rail_walker_gym/joystick_policy_mujoco/__init__.py
realquantumcookie/APRL
import numpy as np
from rail_walker_interface import JoystickPolicyResetter, JoystickPolicyTerminationConditionProvider, BaseWalkerInSim, BaseWalker
from rail_mujoco_walker import RailSimWalkerDMControl
import transforms3d as tr3d
from typing import Any, Optional
import copy

class JoystickPolicyPointInSimResetter(JoystickPolicyResetter[BaseWalkerInSim]):
    def __init__(
        self,
        respawn_pos: np.ndarray = np.zeros(2),
        respawn_yaw: float = 0.0,
    ):
        super().__init__()
        self.respawn_pos = respawn_pos
        self.respawn_yaw = respawn_yaw
        self._last_should_crouch = False

    def get_respawn_pose(
        self,
        Robot: BaseWalkerInSim,
        randomState: np.random.RandomState
    ) -> tuple[np.ndarray, np.ndarray]:
        return self.respawn_pos, np.array([0, 0, self.respawn_yaw])
    
    def step_resetter(
        self, 
        Robot: BaseWalkerInSim, 
        target_goal_world_delta: np.ndarray, 
        target_goal_local: np.ndarray, 
        target_yaw: float, 
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw: float, 
        target_custom_data: Optional[Any], 
        enable_target_custom_obs : bool,
        info_dict: dict[str, Any], 
        randomState: np.random.RandomState
    ) -> None:
        if target_custom_data is not None and "should_crouch" in target_custom_data and target_custom_data["should_crouch"]:
            self._last_should_crouch = True
        else:
            self._last_should_crouch = False
    
    def perform_reset(
        self, 
        Robot: BaseWalkerInSim, 
        info_dict: dict[str,Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        respawn_pos, respawn_roll_pitch_yaw = self.get_respawn_pose(Robot, randomState)
        target_qpos = Robot.joint_qpos_init if not self._last_should_crouch else Robot.joint_qpos_crouch
        
        if respawn_pos is not None and respawn_roll_pitch_yaw is not None:
            respawn_quat = tr3d.euler.euler2quat(*respawn_roll_pitch_yaw)
            Robot.reset_2d_location(respawn_pos, respawn_quat, target_qpos)


class JoystickPolicyLastPositionAndYawResetter(JoystickPolicyPointInSimResetter):
    def __init__(
        self,
        random_respawn_yaw : bool = False,
        respawn_yaw_override : Optional[float] = None,
        position_if_init : Optional[np.ndarray] = None,
    ):
        super().__init__()
        self.random_respawn_yaw = random_respawn_yaw
        self.respawn_yaw_override = respawn_yaw_override
        self.last_position = position_if_init
        self.last_yaw = respawn_yaw_override
    
    def get_respawn_pose(self, Robot: BaseWalkerInSim, randomState: np.random.RandomState) -> tuple[np.ndarray, np.ndarray]:
        if self.last_yaw is None:
            respawn_yaw = randomState.uniform(-np.pi, np.pi)
        else:
            respawn_yaw = self.last_yaw
        
        return self.last_position, np.array([0, 0, respawn_yaw])
    
    def step_resetter(
        self, 
        Robot: BaseWalkerInSim, 
        target_goal_world_delta: np.ndarray, 
        target_goal_local: np.ndarray, 
        target_yaw: float, 
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw: float, 
        target_custom_data: Optional[Any], 
        enable_target_custom_obs : bool,
        info_dict: dict[str, Any], 
        randomState: np.random.RandomState
    ) -> None:
        if self.respawn_yaw_override is not None:
            self.last_yaw = self.respawn_yaw_override
        elif self.random_respawn_yaw:
            self.last_yaw = None
        else:
            self.last_yaw = Robot.get_roll_pitch_yaw()[2]

        self.last_position = Robot.get_3d_location()[:2]
        super().step_resetter(Robot, target_goal_world_delta, target_goal_local, target_yaw, target_delta_yaw, target_velocity, velocity_to_goal, change_in_abs_target_delta_yaw, target_custom_data, enable_target_custom_obs, info_dict, randomState)

class ResetPolicyResetter(JoystickPolicyResetter[RailSimWalkerDMControl]):
    def __init__(
        self,
        respawn_pos: np.ndarray = np.zeros(2),
        init_dist: np.ndarray = np.array([0.2, 0.2, 0.6]),
    ):
        self.respawn_pos = respawn_pos
        self._init_dist = init_dist
        super().__init__()
    
    def perform_reset(
        self, 
        Robot: RailSimWalkerDMControl, 
        info_dict: dict[str,Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        initialization = randomState.choice(['sitting', 'standing', 'fallen'], p=list(self._init_dist))
        if initialization == 'sitting':
            target_pose = np.asarray([0.0 / 180 * np.pi, 70.0/180*np.pi, -150.0 / 180 * np.pi] * 4)
            target_location = np.asarray([0, 0, 0.02])
            target_quat = None
        elif initialization == 'standing':
            target_pose = Robot.joint_qpos_init
            target_location = np.array([0, 0, 0.21])
            target_quat = None
        elif initialization == 'fallen':
            nominal_pose = copy.deepcopy(Robot.joint_qpos_init)
            target_location = np.array([0, 0, randomState.uniform(low=0.4, high=0.5)])
            root_rot = randomState.uniform(low=[-3 * np.pi / 4, -3 * np.pi / 4, -np.pi],
                                            high=[3 * np.pi / 4, 3 * np.pi / 4, np.pi])
            target_quat = tr3d.euler.euler2quat(root_rot[0], root_rot[1], root_rot[2])

            joint_lim_low = np.asarray([-0.802851455917, -1.0471975512, -2.69653369433] * 4)
            joint_lim_high = np.asarray([0.802851455917, 4.18879020479, -0.916297857297] * 4)
            joint_pose_size = len(joint_lim_low)

            joint_dir = randomState.randint(0, 2, joint_pose_size).astype(np.float32)
            lim_pose = (1.0 - joint_dir) * joint_lim_low + joint_dir * joint_lim_high

            pose_lerp = randomState.uniform(low=0, high=1, size=joint_pose_size)
            pose_lerp = pose_lerp * pose_lerp * pose_lerp
            target_pose = (1.0 - pose_lerp) * nominal_pose + pose_lerp * lim_pose
            

        Robot.reset_dropped(
            target_pose,
            target_location,
            target_quat,
            settle = initialization == 'fallen'
        )

rail_walker_gym/joystick_policy_mujoco/resetters_sim.py
realquantumcookie/APRL
from rail_walker_interface import JoystickPolicyRewardProvider, JoystickPolicyTerminationConditionProvider
from rail_mujoco_walker import JoystickPolicyProviderWithDMControl, add_sphere_to_mjv_scene, add_arrow_to_mjv_scene, JoystickPolicyDMControlTask, RailSimWalkerDMControl, find_dm_control_non_contacting_height
from dm_control.mujoco.engine import Physics as EnginePhysics
import mujoco
import numpy as np
from typing import Any, Tuple, Optional
import transforms3d as tr3d

JOINT_WEIGHTS = np.array([1.0, 0.75, 0.5] * 4)

class ResetRewardProvider(JoystickPolicyRewardProvider[RailSimWalkerDMControl],JoystickPolicyProviderWithDMControl):
    SUCCESS_RATE_REW_THRES = 0.8

    def __init__(
            self,
            use_energy_penalty: bool = False,
        ) -> None:
        JoystickPolicyRewardProvider.__init__(self)
        JoystickPolicyProviderWithDMControl.__init__(self)
        self.use_energy_penalty = use_energy_penalty
        self.rew = 0.0
        self.standup_height : float = None
    
    def get_reward(self) -> float:
        return self.rew

    def reset_reward(
        self, 
        Robot: RailSimWalkerDMControl, 
        info_dict: dict[str,Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        self.rew = 0.0
        info_dict["reward_roll"] = 0.0
        info_dict["reward_height"] = 0.0
        info_dict["reward_pose"] = 0.0
        info_dict["reward_qvel"] = 0.0
        info_dict["reward_stand"] = 0.0
        info_dict["cos_dist"] = 0.0
        info_dict["energy"] = 0.0
        info_dict["success_rate"] = 0.0
        self.standup_height = 0.325

        # # Retrieve the target height
        # physics = Robot.mujoco_walker._last_physics
        # # Remember the previous state of the robot
        # prev_joint = Robot.get_joint_qpos()
        # prev_pos, prev_quat = Robot.mujoco_walker.get_pose(physics)
        # # Set the robot to standup pose
        # find_dm_control_non_contacting_height(physics, Robot.mujoco_walker, prev_pos[0], prev_pos[1], Robot.joint_qpos_init, prev_quat)
        # self.standup_height = Robot.mujoco_walker.get_position(physics)[2] # retrieve the height of the robot
        # # reset the robot back to previous state
        # with physics.reset_context():
        #     Robot.mujoco_walker.set_pose(physics, prev_pos, prev_quat)
        #     physics.bind(Robot.mujoco_walker.joints).qpos[:] = prev_joint
        # #print("Standup height: ", self.standup_height)

    def _log_success_rate(
        self, 
        info_dict: dict[str,Any], 
        rew_roll : float,
        rew_height : float,
        rew_pose : float
    ) -> None:
        if rew_roll > __class__.SUCCESS_RATE_REW_THRES and rew_height > __class__.SUCCESS_RATE_REW_THRES and rew_pose > __class__.SUCCESS_RATE_REW_THRES:
            info_dict["success_rate"] = 1.0
        else:
            info_dict["success_rate"] = 0.0

    def _calc_reward_roll(
        self,
        Robot: RailSimWalkerDMControl
    ) -> Tuple[float, float]:
        cos_dist = tr3d.quaternions.quat2mat(Robot.get_framequat_wijk())[-1,-1] # See if the z-axis is pointing up
        r_roll = (0.5 * cos_dist + 0.5)**2
        return r_roll, cos_dist

    def _calc_reward_stand(
        self,
        Robot: RailSimWalkerDMControl
    ) -> float:
        # height match reward
        tar_h = self.standup_height
        # root_h = Robot.get_root_height()
        root_h = Robot.get_3d_location()[-1]
        # root_h = physics.bind(self._robot.root_body).xpos[-1]
        h_err = tar_h - root_h
        h_err /= tar_h
        h_err = np.clip(h_err, 0.0, 1.0)
        r_height = 1.0 - h_err

        # pose match reward
        joint_pose = Robot.get_joint_qpos()
        tar_pose = Robot.joint_qpos_init
        pose_diff = tar_pose - joint_pose
        pose_diff = JOINT_WEIGHTS * JOINT_WEIGHTS * pose_diff * pose_diff
        pose_err = np.sum(pose_diff)
        r_pose = np.exp(-0.6 * pose_err)

        # pose velocity reward
        tar_vel = 0.0
        joint_vel = Robot.get_joint_qvel()
        vel_diff = tar_vel - joint_vel
        vel_diff = vel_diff * vel_diff
        vel_err = np.sum(vel_diff)
        r_vel = np.exp(-0.02 * vel_err)

        r_stand = 0.2 * r_height + 0.6 * r_pose + 0.2 * r_vel

        return r_stand, r_height, r_pose, r_vel
    
    def render_scene_callback(self, task : JoystickPolicyDMControlTask, physics : EnginePhysics, scene : mujoco.MjvScene) -> None:
        robot_loc = task.robot.get_3d_location()
        sphere_loc = robot_loc.copy()
        sphere_loc[2] = self.standup_height
        #add_sphere_to_mjv_scene(scene, sphere_loc, 0.1, np.array([1.0, 0.0, 0.0, 1.0]))
        add_arrow_to_mjv_scene(scene, robot_loc, sphere_loc, 0.01, np.array([1.0, 0.0, 0.0, 1.0]))

        
    def step_reward(
        self, 
        Robot: RailSimWalkerDMControl, 
        action_target_qpos: np.ndarray,
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw : float,
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> None:
        roll_w = stand_w = 0.5
        roll_threshold = np.cos(0.2 * np.pi)

        r_roll, root_cos_dist = self._calc_reward_roll(Robot)
        r_stand, r_height, r_pose, r_vel = self._calc_reward_stand(Robot)

        info_dict["reward_roll"] = r_roll
        info_dict["reward_height"] = r_height
        info_dict["reward_pose"] = r_pose
        info_dict["reward_qvel"] = r_vel
        info_dict["cos_dist"] = root_cos_dist

        if root_cos_dist > roll_threshold:
            r_stand = r_stand
        else:
            r_stand = 0.0

        info_dict["reward_stand"] = r_stand

        reward = roll_w * r_roll + stand_w * r_stand
        standing_reward = reward

        #  Calculate Energy / Qvel Penalty
        qvel = Robot.get_joint_qvel()
        torque = Robot.get_joint_torques()
        energy = np.sum(np.abs(qvel * torque))
        info_dict["energy"] = energy
        
        if self.use_energy_penalty:
            energy_reward = - 0.02 * energy
        else:
            energy_reward = 0.0
        
        self._log_success_rate(
            info_dict,
            r_roll,
            r_height,
            r_pose
        )

        self.rew = standing_reward + energy_reward


class GatedResetRewardProvider(ResetRewardProvider):
    def step_reward(
        self, 
        Robot: RailSimWalkerDMControl,
        action_target_qpos: np.ndarray, 
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw : float,
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> None:

        r_roll, root_cos_dist = self._calc_reward_roll(Robot)
        info_dict["reward_roll"] = r_roll
        info_dict["cos_dist"] = root_cos_dist

        r_stand, r_height, r_pose, r_vel = self._calc_reward_stand(Robot)

        info_dict["reward_roll"] = r_roll
        info_dict["reward_height"] = r_height
        info_dict["reward_pose"] = r_pose
        info_dict["reward_qvel"] = r_vel
        info_dict["cos_dist"] = root_cos_dist

        #  Calculate Energy / Qvel Penalty
        qvel = Robot.get_joint_qvel()
        torque = Robot.get_joint_torques()
        energy = np.sum(np.abs(qvel * torque))
        info_dict["energy"] = energy
        
        if self.use_energy_penalty:
            energy_reward = - 0.02 * energy
        else:
            energy_reward = 0.0

        self._log_success_rate(
            info_dict,
            r_roll,
            r_height,
            r_pose
        )

        self.rew = (r_roll * (1 + r_stand) + energy_reward) / 2.0

rail_walker_gym/joystick_policy_mujoco/reward_providers_sim.py
realquantumcookie/APRL
from .target_providers import JoystickPolicyGo1JoystickTargetProvider, JoystickPolicyGo1JoystickSimpleTargetProvider
rail_walker_gym/joystick_policy_real/__init__.py
realquantumcookie/APRL
from rail_real_walker.robots.go1 import Go1RealWalker
from rail_walker_interface import JoystickPolicyTargetProvider, JoystickPolicyTerminationConditionProvider
from unitree_go1_wrapper import XRockerBtnDataStruct
import numpy as np
from typing import Any, Optional
import gym
import gym.spaces

def normalize_rad(rad : float) -> float:
    return ((rad + np.pi) % (2 * np.pi)) - np.pi

"""
Real Joystick Target Provider using the Go1's controller
"""
class JoystickPolicyGo1JoystickTargetProvider(JoystickPolicyTargetProvider[Go1RealWalker]):
    def __init__(
        self,
        ahead_distance : float = 1.0,
        min_target_velocity : float = 0.0,
        max_target_velocity : float = 1.0,
        max_velocity_change_rate : float = 0.2,
        max_target_yaw_change_rate : float = 5.0 / 180.0 * np.pi,
        min_hold_steps : int = 20 * 1,
        max_hold_steps : int = 20 * 3,
        target_merge_to_zero_interval : float = 10 / 180 * np.pi,
    ):
        self.ahead_distance = ahead_distance
        self.min_target_velocity = min_target_velocity
        self.max_target_velocity = max_target_velocity
        self.max_velocity_change_rate = max_velocity_change_rate
        self.max_tdy_change_rate = max_target_yaw_change_rate
        self.min_hold_steps = min_hold_steps
        self.max_hold_steps = max_hold_steps
        self.target_merge_to_zero_interval = target_merge_to_zero_interval

        self._tagret_delta_yaw = 0.0
        self._new_target_delta_yaw= 0.0
        self._target_velocity = (min_target_velocity + max_target_velocity) / 2
        self._new_target_velocity = self._target_velocity
        self._remaining_hold_min_steps = min_hold_steps
        self._remaining_hold_max_steps = max_hold_steps
    
    def get_target_goal_world_delta(self, Robot: Go1RealWalker) -> np.ndarray:
        return self.get_world_delta(Robot, self._target_delta_yaw)

    def get_target_velocity(self, Robot: Go1RealWalker) -> float:
        return self._target_velocity
    
    def is_target_velocity_fixed(self) -> bool:
        return False

    def step_target(
        self, 
        Robot: Go1RealWalker, 
        info_dict: dict[str,Any], 
        randomState : np.random.RandomState
    ) -> None:
        roll, pitch, yaw = Robot.get_roll_pitch_yaw()
        real_tdy, real_target_v = self.get_joystick_real_target(Robot, yaw)
        if np.abs(real_tdy) < self.target_merge_to_zero_interval / 2.0:
            real_tdy = 0.0
        self._new_target_delta_yaw = real_tdy
        self._new_target_velocity = real_target_v
        self._remaining_hold_max_steps -= 1
        self._remaining_hold_min_steps -= 1
    
    def get_joystick_real_target(self, Robot : Go1RealWalker, robot_yaw : float) -> tuple[float, float]:
        joystick_raw = Robot.get_joystick_values()[0]
        
        if np.all(joystick_raw <= 0.05):
            return 0.0, 0.0

        local_target_yaw = np.arctan2(-joystick_raw[0], joystick_raw[1])
        real_target_v = np.sqrt(joystick_raw[0]**2 + joystick_raw[1]**2) / np.sqrt(2) * self.max_target_velocity
        return local_target_yaw, real_target_v

    def has_target_changed(self) -> bool:
        return self._remaining_hold_max_steps <= 0 or ((
            np.abs(normalize_rad(self._new_target_delta_yaw - self._target_delta_yaw)) > self.max_tdy_change_rate or
            np.abs(self._new_target_velocity - self._target_velocity) > self.max_velocity_change_rate
        ) and self._remaining_hold_min_steps <= 0)

    def after_step_target(
        self, 
        Robot: Go1RealWalker, 
        info_dict: dict[str,Any], 
        randomState : np.random.RandomState
    ) -> None:
        self._remaining_hold_min_steps = self.min_hold_steps
        self._remaining_hold_max_steps = self.max_hold_steps
        self._target_delta_yaw = self._new_target_delta_yaw
        self._target_velocity = self._new_target_velocity
        return super().after_step(Robot, info_dict, randomState)

    def reset_target(
        self, 
        Robot: Go1RealWalker, 
        info_dict: dict[str,Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        self._remaining_hold_max_steps = self.max_hold_steps
        self._remaining_hold_min_steps = self.min_hold_steps

        # Update joystick target
        self._target_delta_yaw, self._target_velocity = self.get_joystick_real_target(Robot, Robot.get_roll_pitch_yaw()[2])
        if np.abs(self._target_delta_yaw) < self.target_merge_to_zero_interval / 2.0:
            self._target_delta_yaw = 0.0
        
        self._new_target_delta_yaw = self._target_delta_yaw
        self._new_target_velocity = self._target_velocity

    def get_world_delta(self, robot: Go1RealWalker, tdy : float) -> np.ndarray:
        current_yaw = robot.get_roll_pitch_yaw()[2]
        target_yaw = normalize_rad(tdy + current_yaw)
        return np.array([np.cos(target_yaw),np.sin(target_yaw)]) * self.ahead_distance

class JoystickPolicyGo1JoystickSimpleTargetProvider(JoystickPolicyGo1JoystickTargetProvider):
    def __init__(
        self, 
        ahead_distance: float = 1, 
        min_target_velocity: float = 0, 
        max_target_velocity: float = 1, 
        max_velocity_change_rate: float = 0.2, 
        min_hold_steps: int = 20 * 1, 
        max_hold_steps: int = 20 * 3
    ):
        super().__init__(
            ahead_distance, 
            min_target_velocity, 
            max_target_velocity, 
            max_velocity_change_rate, 
            40/180*np.pi, 
            min_hold_steps, 
            max_hold_steps, 
            np.pi / 2
        )
    
    def step_target(self, Robot: Go1RealWalker, info_dict: dict[str, Any], randomState: np.random.RandomState) -> None:
        super().step_target(Robot, info_dict, randomState)
        sign_ntdy = np.sign(self._new_target_delta_yaw)
        self._new_target_delta_yaw = sign_ntdy * np.pi / 2
        self._new_target_velocity = 0.5 if sign_ntdy == 0 else 0.3
        #self._new_target_velocity = 0.8

    def reset_target(
        self, 
        Robot: Go1RealWalker, 
        info_dict: dict[str, Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider, 
        randomState: np.random.RandomState
    ) -> None:
        super().reset_target(Robot, info_dict, termination_provider_triggered, randomState)
        sign_ntdy = np.sign(self._new_target_delta_yaw)
        self._new_target_delta_yaw = sign_ntdy * np.pi / 2
        self._new_target_velocity = 0.5 if sign_ntdy == 0 else 0.3
        #self._new_target_velocity = 0.8
        self._target_delta_yaw = self._new_target_delta_yaw
        self._target_velocity = self._new_target_velocity
rail_walker_gym/joystick_policy_real/target_providers.py
realquantumcookie/APRL
from .robot import *
from .environment import *
from .joystick_policy import *
rail_walker_interface/__init__.py
realquantumcookie/APRL
from .env import WalkerEnvironment, JoystickEnvironment
from .joystick_real import JoystickEnvObservationExtractor, JoystickEnvImpl
rail_walker_interface/environment/__init__.py
realquantumcookie/APRL
from ..robot import BaseWalker
from ..joystick_policy.joystick_policy import JoystickPolicy


class WalkerEnvironment:
    @property
    def robot(self) -> BaseWalker:
        pass

class JoystickEnvironment:
    @property
    def joystick_policy(self) -> JoystickPolicy:
        pass

    def set_joystick_policy(self, joystick_policy: JoystickPolicy):
        pass

    @property
    def is_resetter_policy(self) -> bool:
        return False
rail_walker_interface/environment/env.py
realquantumcookie/APRL
import gym
import gym.spaces
from .env import WalkerEnvironment, JoystickEnvironment
from ..robot import BaseWalker, BaseWalkerWithFootContact
from ..joystick_policy import JoystickPolicy
from functools import cached_property
import numpy as np
from typing import Optional, Any
import copy
from collections import OrderedDict

class JoystickEnvObservationExtractor:
    def __init__(self, env : "JoystickEnvImpl"):
        self.env = env

    @cached_property
    def observation_spec(self) -> gym.spaces.Dict:
        ret_dict = {
            "robot/joints_pos": gym.spaces.Box(
                low=self.env.robot.joint_qpos_mins,
                high=self.env.robot.joint_qpos_maxs,
                shape=(self.env.robot.joint_nums,),
                dtype=np.float32
            ),
            "robot/joints_vel": gym.spaces.Box(
                low=-np.inf,
                high=np.inf,
                shape=(self.env.robot.joint_nums,),
                dtype=np.float32
            ),
            "robot/imu": gym.spaces.Box(
                low=-np.inf,
                high=np.inf,
                shape=(4,),
            ),
            "robot/sensors_gyro": gym.spaces.Box(
                low=-np.inf,
                high=np.inf,
                shape=(3,),
                dtype=np.float32
            ),
            "robot/sensors_framequat": gym.spaces.Box(
                low=-1.0,
                high=1.0,
                shape=(4,),
                dtype=np.float32
            ),
            "robot/torques": gym.spaces.Box(
                low=-np.inf,
                high=np.inf,
                shape=(self.env.robot.joint_nums,),
                dtype=np.float32
            ),
            "robot/sensors_local_velocimeter": gym.spaces.Box(
                low=-np.inf,
                high=np.inf,
                shape=(3,),
                dtype=np.float32
            ),
            "robot/sensors_local_velocimeter_x": gym.spaces.Box(
                low=-np.inf,
                high=np.inf,
                shape=(1,),
                dtype=np.float32
            ),
            "robot/sensors_accelerometer": gym.spaces.Box(
                low=-np.inf,
                high=np.inf,
                shape=(3,),
                dtype=np.float32
            ),
        }
        if isinstance(self.env.robot.unwrapped(), BaseWalkerWithFootContact):
            ret_dict["robot/foot_forces"] = gym.spaces.Box(
                low=-np.inf,
                high=np.inf,
                shape=(4,),
                dtype=np.float32
            )
            ret_dict["robot/foot_forces_normalized"] = gym.spaces.Box(
                low=-np.inf,
                high=np.inf,
                shape=(4,),
                dtype=np.float32
            )
            ret_dict["robot/foot_forces_normalized_masked"] = gym.spaces.Box(
                low=-np.inf,
                high=np.inf,
                shape=(4,),
                dtype=np.float32
            )
            ret_dict["robot/foot_contacts"] = gym.spaces.Box( # should use MultiBinary but flatten() does not support having multibinary / box spaces in a Dict
                low=0,
                high=1,
                shape=(4,),
                dtype=np.float32
            )
        return gym.spaces.Dict(ret_dict)
    
    def extract_observation(self) -> dict[str,Any]:
        roll, pitch, yaw = self.env.robot.get_roll_pitch_yaw()
        dr, dp, dy = self.env.robot.get_3d_angular_velocity()
        imu = np.array([roll, pitch, dr, dp], dtype=np.float32)

        ret_dict = {
            "robot/joints_pos": self.env.robot.get_joint_qpos(),
            "robot/joints_vel": self.env.robot.get_joint_qvel(),
            "robot/imu": imu,
            "robot/sensors_gyro": self.env.robot.get_3d_angular_velocity(),
            "robot/sensors_framequat": self.env.robot.get_framequat_wijk(),
            "robot/torques": self.env.robot.get_joint_torques(),
            "robot/sensors_local_velocimeter": self.env.robot.get_3d_local_velocity(),
            "robot/sensors_local_velocimeter_x": self.env.robot.get_3d_local_velocity()[0:1],
            "robot/sensors_accelerometer": self.env.robot.get_3d_acceleration_local(),
        }
        if isinstance(self.env.robot.unwrapped(), BaseWalkerWithFootContact):
            ret_dict["robot/foot_forces"] = self.env.robot.get_foot_force()
            ret_dict["robot/foot_contacts"] = self.env.robot.get_foot_contact()
            if hasattr(self.env.robot, "foot_contact_no_contact_threshold") and hasattr(self.env.robot, "foot_contact_has_contact_threshold"):
                ret_dict["robot/foot_forces_normalized"] = (ret_dict["robot/foot_forces"] - self.env.robot.foot_contact_no_contact_threshold) / (self.env.robot.foot_contact_has_contact_threshold - self.env.robot.foot_contact_no_contact_threshold)
            else:
                ret_dict["robot/foot_forces_normalized"] = ret_dict["robot/foot_forces"]
            masked_foot_forces = ret_dict["robot/foot_forces_normalized"].copy()
            masked_foot_forces[-1] = 0.0
            ret_dict["robot/foot_forces_normalized_masked"] = masked_foot_forces
        return ret_dict


class JoystickEnvImpl(gym.Env[dict[str,Any],np.ndarray], WalkerEnvironment, JoystickEnvironment):
    metadata = {
        "render_modes": []
    }

    def __init__(
            self, 
            joystick_policy : JoystickPolicy
        ):
        gym.Env.__init__(self)
        WalkerEnvironment.__init__(self)
        JoystickEnvironment.__init__(self)
        # ====================== Store Parameters ======================
        self._joystick_policy = joystick_policy
        self.obs_extractor = JoystickEnvObservationExtractor(self)
        self.random_state = np.random.RandomState()
    
    @property
    def action_space(self) -> gym.spaces.Box:
        return gym.spaces.Box(
            low=self.robot.action_qpos_mins,
            high=self.robot.action_qpos_maxs,
            dtype=np.float32
        )

    @property
    def observation_space(self) -> gym.spaces.Dict:
        robot_space = self.obs_extractor.observation_spec
        
        real_obs_space = {}
        for key, space in robot_space.items():
            if key.startswith("robot/") and key[len("robot/"):] in self.joystick_policy.enabled_observables:
                real_obs_space[key] = space        
        
        if self.joystick_policy.target_observable is not None:
            real_obs_space["target_obs"] = self.joystick_policy.target_observable.get_observation_spec()
        
        if not self.joystick_policy.target_yaw_provider.is_target_velocity_fixed():
            real_obs_space["target_vel"] = gym.spaces.Box(
                low = 0.0,
                high = np.inf,
                shape=(1,)
            )
        
        target_custom_data_spec = self.joystick_policy.target_yaw_provider.get_target_custom_data_observable_spec()
        if self.joystick_policy.enable_target_custom_obs and target_custom_data_spec is not None:
            real_obs_space["target_custom"] = target_custom_data_spec
        
        # Enforce order
        real_obs_space = OrderedDict(sorted(real_obs_space.items(), key=lambda t: t[0]))
        obs_space = gym.spaces.Dict(real_obs_space)
        return obs_space

    @property
    def joystick_policy(self) -> JoystickPolicy:
        return self._joystick_policy

    def set_joystick_policy(self, joystick_policy: JoystickPolicy):
        self._joystick_policy = joystick_policy

    @property
    def is_resetter_policy(self) -> bool:
        return False

    @property
    def robot(self) -> BaseWalker:
        return self.joystick_policy.robot
    
    def _get_obs(self) -> dict[str,np.ndarray]:
        robot_obs = self.obs_extractor.extract_observation()
        
        real_obs = {}
        for key, value in robot_obs.items():
            if key.startswith("robot/") and key[len("robot/"):] in self.joystick_policy.enabled_observables:
                real_obs[key] = value
        
        if self.joystick_policy.target_observable is not None:
            target_observable_obs = self.joystick_policy.target_observable.get_observation()
            real_obs["target_obs"] = target_observable_obs
        
        if not self.joystick_policy.target_yaw_provider.is_target_velocity_fixed():
            target_vel = self.joystick_policy._target_velocity
            real_obs["target_vel"] = np.array([target_vel], dtype=np.float32)

        target_custom_data_obs = self.joystick_policy.target_yaw_provider.get_target_custom_data_observable()
        if self.joystick_policy.enable_target_custom_obs and target_custom_data_obs is not None:
            real_obs["target_custom"] = target_custom_data_obs

        return real_obs

    def step(self, action : np.ndarray) -> tuple[dict[str,Any], float, bool, dict[str,Any]]:
        # ====================== Step Joystick Policy ======================
        self.joystick_policy.before_step(action, self.random_state)
        infos = self.joystick_policy.after_step(self.random_state)
        
        # ====================== Return ======================
        reward = self.joystick_policy.get_reward()
        obs = self._get_obs()
        truncated = self.joystick_policy.should_truncate()
        if truncated:
            infos["TimeLimit.truncated"] = True
        
        done = self.joystick_policy.should_terminate() or truncated
        
        return obs, reward, done, infos

    def reset(self, seed : int | None = None, options: dict[str, Any] | None = None, **kwargs) -> tuple[dict[str,Any], dict[str,Any]]:
        if seed is not None:
            self.random_state = np.random.RandomState(seed)
        
        infos = self.joystick_policy.reset(self.random_state)
        if kwargs.get("return_info", False):
            return self._get_obs(), infos
        else:
            return self._get_obs()

    def seed(self, seed : int | None = None) -> list[int]:
        self.random_state = np.random.RandomState(seed)

    def render(self, *args, **kwargs):
        raise NotImplementedError

    def close(self):
        self.robot.close()
rail_walker_interface/environment/joystick_real.py
realquantumcookie/APRL
from .joystick_interfaces import JoystickPolicyResetter, JoystickPolicyRewardProvider, JoystickPolicyTargetProvider, JoystickPolicyTerminationConditionProvider, JoystickPolicyTargetObservable
from .joystick_policy import JoystickPolicy
rail_walker_interface/joystick_policy/__init__.py
realquantumcookie/APRL
import gym
import numpy as np
from typing import Optional, Any, TypeVar, Generic, Callable
from ..robot.robot import BaseWalker

_RobotClsTerminationT = TypeVar("_RobotClsTerminationT", bound=BaseWalker)
class JoystickPolicyTerminationConditionProvider(Generic[_RobotClsTerminationT]):
    def should_terminate(self) -> bool:
        raise NotImplementedError()
    
    def step_termination_condition(
            self, 
            Robot: _RobotClsTerminationT, 
            target_goal_world_delta: np.ndarray,
            target_goal_local: np.ndarray,
            target_yaw : float,
            target_delta_yaw: float, 
            target_velocity: float,
            velocity_to_goal: float, 
            change_in_abs_target_delta_yaw : float, 
            target_custom_data: Optional[Any],
            enable_target_custom_obs : bool,
            info_dict: dict[str,Any],
            randomState: np.random.RandomState
        ) -> None:
        pass

    def reset_termination_condition(
        self, 
        Robot: _RobotClsTerminationT,
        info_dict: dict[str,Any], 
        termination_provider_triggered,
        randomState: np.random.RandomState
    ) -> None:
        pass

_RobotClsTargetObsT = TypeVar("_RobotClsTargetObsT", bound=BaseWalker)
class JoystickPolicyTargetObservable(Generic[_RobotClsTargetObsT]):
    def get_observation_spec(self) -> gym.Space:
        raise NotImplementedError()
    
    def step_target_obs(
        self,
        Robot: _RobotClsTargetObsT, 
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw : float,
        target_delta_yaw: float,
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> None:
        pass

    def reset_target_obs(
        self, 
        Robot: _RobotClsTargetObsT, 
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw : float,
        target_delta_yaw: float, 
        target_velocity: float,
        info_dict: dict[str,Any], 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        pass

    def get_observation(self) -> Any:
        raise NotImplementedError()

_RobotClsTargetProviderT = TypeVar("_RobotClsTargetProviderT", bound=BaseWalker)
class JoystickPolicyTargetProvider(Generic[_RobotClsTargetProviderT]):
    def get_target_goal_world_delta(self, Robot: _RobotClsTargetProviderT) -> np.ndarray:
        raise NotImplementedError()
    
    def get_target_velocity(self, Robot: _RobotClsTargetProviderT) -> float:
        return 0.5
    
    def is_target_velocity_fixed(self) -> bool:
        return True

    def get_target_custom_data(self) -> Optional[Any]:
        return None
    
    def get_target_custom_data_observable_spec(self) -> Optional[gym.Space]:
        return None
    
    def get_target_custom_data_observable(self) -> Optional[Any]:
        return None
    
    def has_target_changed(self) -> bool:
        return False

    def step_target(
        self, 
        Robot: _RobotClsTargetProviderT, 
        info_dict: dict[str,Any], 
        randomState : np.random.RandomState
    ) -> None:
        pass

    def after_step_target(
        self, 
        Robot: _RobotClsTargetProviderT, 
        info_dict: dict[str,Any], 
        randomState : np.random.RandomState
    ) -> None:
        pass

    def reset_target(
        self, 
        Robot: _RobotClsTargetProviderT, 
        info_dict: dict[str,Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        pass

_RobotClsRewardProviderT = TypeVar("_RobotClsRewardProviderT", bound=BaseWalker)
class JoystickPolicyRewardProvider(Generic[_RobotClsRewardProviderT]):
    def get_reward(self) -> float:
        raise NotImplementedError()

    def step_reward(
        self, 
        Robot: _RobotClsRewardProviderT, 
        action_target_qpos: np.ndarray,
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw : float,
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> None:
        pass

    def reset_reward(
        self, 
        Robot: _RobotClsRewardProviderT, 
        info_dict: dict[str,Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        pass

_RobotClsResetterT = TypeVar("_RobotClsResetterT", bound=BaseWalker)
class JoystickPolicyResetter(Generic[_RobotClsResetterT]):
    def perform_reset(
        self, 
        Robot: _RobotClsResetterT, 
        info_dict: dict[str,Any], 
        termination_provider_triggered: JoystickPolicyTerminationConditionProvider,
        randomState: np.random.RandomState
    ) -> None:
        pass

    def step_resetter(
        self, 
        Robot: _RobotClsTerminationT, 
        target_goal_world_delta: np.ndarray,
        target_goal_local: np.ndarray,
        target_yaw : float,
        target_delta_yaw: float, 
        target_velocity: float,
        velocity_to_goal: float, 
        change_in_abs_target_delta_yaw : float, 
        target_custom_data: Optional[Any],
        enable_target_custom_obs : bool,
        info_dict: dict[str,Any],
        randomState: np.random.RandomState
    ) -> None:
        pass
rail_walker_interface/joystick_policy/joystick_interfaces.py
realquantumcookie/APRL
from ..robot import BaseWalker, BaseWalkerWithFootContact
from .joystick_interfaces import JoystickPolicyResetter, JoystickPolicyRewardProvider, JoystickPolicyTargetProvider, JoystickPolicyTerminationConditionProvider, JoystickPolicyTargetObservable
import numpy as np
from typing import Optional, Any
import transforms3d as tr3d

def normalize_rad(angle: float) -> float:
    # return np.arctan2(np.sin(angle),np.cos(angle))
    return ((angle + np.pi) % (2 * np.pi)) - np.pi

class JoystickPolicy:
    def __init__(
        self,
        robot: BaseWalker,
        reward_provider: JoystickPolicyRewardProvider,
        target_yaw_provider: JoystickPolicyTargetProvider,
        termination_providers: list[JoystickPolicyTerminationConditionProvider],
        truncation_providers: list[JoystickPolicyTerminationConditionProvider],
        resetters: list[JoystickPolicyResetter],
        initializers: list[JoystickPolicyResetter] = [],
        target_observable: Optional[JoystickPolicyTargetObservable] = None,
        enabled_observables : list[str] = [
            "joints_pos",
            "joints_vel",
            "imu",
            "sensors_local_velocimeter",
            "torques",
            "foot_contacts",
        ],
        lock_target: bool = False,
        enable_target_custom_obs = True
    ):
        self.robot = robot
        self.reward_provider = reward_provider
        self.target_yaw_provider = target_yaw_provider
        self.termination_providers = termination_providers
        self.truncation_providers = truncation_providers
        self.resetters = resetters
        self.initializers = initializers
        self.target_observable = target_observable
        self.enabled_observables = enabled_observables
        self.lock_target = lock_target
        self.enable_target_custom_obs = enable_target_custom_obs

        # Temporary Variables
        self._step_target_qpos = self.robot.get_joint_qpos()

        # Set up task-specific variables
        self._target_goal_world_delta = np.zeros(2)
        self._target_goal_local = np.zeros(2)
        self._target_yaw = 0.0
        self._target_delta_yaw = 0.0
        self._target_velocity = 0.0
        self._target_custom_data = None
        self._rew_step = 0.0
        self._info_dict = {}
        self._has_after_after_step = False
        self._termination_reason : Optional[JoystickPolicyTerminationConditionProvider] = None
        self._truncation_reason : Optional[JoystickPolicyTerminationConditionProvider] = None
        self._inited = False
    
    @property
    def has_after_after_step(self) -> bool:
        return self._has_after_after_step

    @property
    def control_timestep(self) -> float:
        return self.robot.control_timestep

    @control_timestep.setter
    def control_timestep(self, value: float) -> None:
        self.robot.control_timestep = value
    
    @property
    def last_info(self) -> dict[str,Any]:
        return self._info_dict.copy()

    @property
    def control_subtimestep(self) -> float:
        return self.robot.control_subtimestep
    
    @control_subtimestep.setter
    def control_subtimestep(self, value: float) -> None:
        self.robot.control_subtimestep = value

    @property
    def target_yaw(self) -> float:
        return self._target_yaw
    
    @property
    def target_delta_yaw(self) -> float:
        return self._target_delta_yaw

    @property
    def target_goal_world_delta(self) -> np.ndarray:
        return self._target_goal_world_delta.copy()
    
    @property
    def target_goal_local(self) -> np.ndarray:
        return self._target_goal_local.copy()
    
    @property
    def target_custom_data(self) -> Optional[Any]:
        return self._target_custom_data
    
    @property
    def target_goal_world_delta_unit(self) -> np.ndarray:
        norm_goal = np.linalg.norm(self._target_goal_world_delta)
        if norm_goal == 0.0:
            return np.zeros(2)
        else:
            return self._target_goal_world_delta / norm_goal
    
    @property
    def target_goal_local_unit(self) -> np.ndarray:
        norm_goal = np.linalg.norm(self._target_goal_local)
        if norm_goal == 0.0:
            return np.zeros(2)
        else:
            return self._target_goal_local / norm_goal
    

    def __update_target(self) -> float:
        new_target_goal_world_delta = self.target_yaw_provider.get_target_goal_world_delta(self.robot)[:2]
        new_target_velocity = self.target_yaw_provider.get_target_velocity(self.robot)
        _, _, yaw = self.robot.get_roll_pitch_yaw()
        inv_rotation_mat = np.array([
            [np.cos(yaw), np.sin(yaw)],
            [-np.sin(yaw), np.cos(yaw)]
        ])
        new_target_goal_local = inv_rotation_mat @ new_target_goal_world_delta

        new_target_yaw = np.arctan2(new_target_goal_world_delta[1], new_target_goal_world_delta[0]) if np.linalg.norm(new_target_goal_world_delta) > 0.0 else 0.0
        new_target_delta_yaw = normalize_rad(new_target_yaw - self.robot.get_roll_pitch_yaw()[2])
        change_in_abs_target_delta_yaw = self.__get_change_in_abs_target_delta_yaw()

        self._info_dict["target_yaw"] = new_target_yaw
        self._info_dict["target_delta_yaw"] = new_target_delta_yaw
        self._info_dict["target_goal_local_x"] = new_target_goal_local[0]
        self._info_dict["target_goal_local_y"] = new_target_goal_local[1]
        self._info_dict["target_goal_world_delta_x"] = new_target_goal_world_delta[0]
        self._info_dict["target_goal_world_delta_y"] = new_target_goal_world_delta[1]
        self._info_dict["change_in_abs_target_delta_yaw"] = change_in_abs_target_delta_yaw
        self._info_dict["abs_target_delta_yaw"] = np.abs(new_target_delta_yaw)
        self._info_dict["target_velocity"] = new_target_velocity

        self._target_yaw = new_target_yaw
        self._target_delta_yaw = new_target_delta_yaw
        self._target_goal_local = new_target_goal_local
        self._target_goal_world_delta = new_target_goal_world_delta
        self._target_custom_data = self.target_yaw_provider.get_target_custom_data()
        self._target_velocity = new_target_velocity
        return change_in_abs_target_delta_yaw
    
    def __get_change_in_abs_target_delta_yaw(self) -> float:
        new_target_delta_yaw = normalize_rad(self.target_yaw - self.robot.get_roll_pitch_yaw()[2])
        change_in_abs_target_delta_yaw = np.abs(new_target_delta_yaw) - np.abs(self._target_delta_yaw)
        return change_in_abs_target_delta_yaw
    
    def before_step(
        self,
        action: np.ndarray,
        random_state : np.random.RandomState
    ):
        self._step_target_qpos = action
        self.robot.apply_action(action)

    def get_reward(
        self
    ):
        return self._rew_step

    def after_step(
        self,
        random_state : np.random.RandomState
    ) -> dict[str,Any]:
        self._info_dict = {}
        self.robot.receive_observation()

        # Update the target yaw
        self.target_yaw_provider.step_target(
            self.robot,
            self._info_dict,
            random_state
        )
        self._has_after_after_step = self.target_yaw_provider.has_target_changed()
        if not self.lock_target and self._has_after_after_step:
            change_in_abs_target_delta_yaw = self.after_after_step(
                random_state
            )
        else:
            change_in_abs_target_delta_yaw = self.__update_target()

        # Gather info about velocity
        robot_v = self.robot.get_3d_linear_velocity()
        robot_v_norm = np.linalg.norm(robot_v)
        robot_v_to_goal = np.dot(
            robot_v[:2], self.target_goal_world_delta_unit
        )
        robot_v_local = self.robot.get_3d_local_velocity()
        robot_rpy = self.robot.get_roll_pitch_yaw()
        self._info_dict["velocity_norm"] = robot_v_norm
        self._info_dict["velocity_to_goal"] = robot_v_to_goal
        self._info_dict["velocity_local_x"] = robot_v_local[0]
        self._info_dict["velocity_local_y"] = robot_v_local[1]
        self._info_dict["velocity_local_z"] = robot_v_local[2]
        self._info_dict["roll"] = robot_rpy[0]
        self._info_dict["pitch"] = robot_rpy[1]
        self._info_dict["yaw"] = robot_rpy[2]
        self._info_dict["joint_torques"] = np.mean(np.abs(self.robot.get_joint_torques()))
        self._info_dict["joint_qvels"] = np.mean(np.abs(self.robot.get_joint_qvel()))
        self._info_dict["joint_qaccs"] = np.mean(np.abs(self.robot.get_joint_qacc()))
        self._info_dict["joint_velocities"] = np.mean(np.abs(self.robot.get_joint_qvel()))
        if hasattr(self.robot, "get_foot_force"):
            foot_force : np.ndarray = self.robot.get_foot_force()
            if foot_force.shape == (4,):
                foot_force_names = ["FR", "FL", "RR", "RL"]
            else:
                foot_force_names = list(range(foot_force.shape[0]))
            for i in range(len(foot_force_names)):
                self._info_dict["foot_force_" + foot_force_names[i]] = foot_force[i]

        self.reward_provider.step_reward(
            self.robot,
            self._step_target_qpos,
            self.target_goal_world_delta,
            self.target_goal_local,
            self.target_yaw,
            self.target_delta_yaw,
            self._target_velocity,
            robot_v_to_goal,
            change_in_abs_target_delta_yaw,
            self._target_custom_data,
            self.enable_target_custom_obs,
            self._info_dict,
            random_state
        )
        reward_perstep = self.reward_provider.get_reward()
        #assert reward_perstep is not None and reward_perstep != np.nan
        self._info_dict["reward_perstep"] = reward_perstep
        self._rew_step = reward_perstep

        # Step the target yaw observable
        if self.target_observable is not None:
            self.target_observable.step_target_obs(
                self.robot,
                self.target_goal_world_delta,
                self.target_goal_local,
                self.target_yaw,
                self.target_delta_yaw,
                self._target_velocity,
                robot_v_to_goal,
                change_in_abs_target_delta_yaw,
                self._target_custom_data,
                self.enable_target_custom_obs,
                self._info_dict,
                random_state
            )
        
        # Step resetters
        for resetter in self.resetters:
            resetter.step_resetter(
                self.robot,
                self.target_goal_world_delta,
                self.target_goal_local,
                self.target_yaw,
                self.target_delta_yaw,
                self._target_velocity,
                robot_v_to_goal,
                change_in_abs_target_delta_yaw,
                self._target_custom_data,
                self.enable_target_custom_obs,
                self._info_dict,
                random_state
            )

        # Step termination providers
        for termination_provider in self.termination_providers:
            termination_provider.step_termination_condition(
                self.robot,
                self.target_goal_world_delta,
                self.target_goal_local,
                self.target_yaw,
                self.target_delta_yaw,
                self._target_velocity,
                robot_v_to_goal,
                change_in_abs_target_delta_yaw,
                self._target_custom_data,
                self.enable_target_custom_obs,
                self._info_dict,
                random_state
            )
            if termination_provider.should_terminate():
                # print("Termination provider", termination_provider, "terminated the episode")
                self._termination_reason = termination_provider
                break
        
        # Step truncaiton providers
        for truncation_provider in self.truncation_providers:
            truncation_provider.step_termination_condition(
                self.robot,
                self.target_goal_world_delta,
                self.target_goal_local,
                self.target_yaw,
                self.target_delta_yaw,
                self._target_velocity,
                robot_v_to_goal,
                change_in_abs_target_delta_yaw,
                self._target_custom_data,
                self.enable_target_custom_obs,
                self._info_dict,
                random_state
            )
            if truncation_provider.should_terminate():
                # print("Truncation provider", truncation_provider, "truncated the episode")
                self._truncation_reason = truncation_provider
                break

        return self._info_dict.copy()

    def after_after_step(
        self,
        random_state : np.random.RandomState
    ):
        if self._has_after_after_step:
            self.target_yaw_provider.after_step_target(
                self.robot,
                self._info_dict,
                random_state
            )
            change_in_abs_target_delta_yaw = self.__update_target()
            robot_v = self.robot.get_3d_linear_velocity()
            robot_v_to_goal = np.dot(
                robot_v[:2], self.target_goal_world_delta_unit
            )
            # Step the target yaw observable
            if self.target_observable is not None:
                self.target_observable.step_target_obs(
                    self.robot,
                    self.target_goal_world_delta,
                    self.target_goal_local,
                    self.target_yaw,
                    self.target_delta_yaw,
                    self._target_velocity,
                    robot_v_to_goal,
                    change_in_abs_target_delta_yaw,
                    self._target_custom_data,
                    self.enable_target_custom_obs,
                    self._info_dict,
                    random_state
                )
            
            # self.reward_provider.step_ex(
            #     self.robot,
            #     self.target_goal_world_delta,
            #     self.target_goal_local,
            #     self.target_yaw,
            #     self.target_delta_yaw,
            #     robot_v_to_goal,
            #     change_in_abs_target_delta_yaw,
            #     self._target_custom_data,
            #     self.enable_target_custom_obs,
            #     self._info_dict,
            #     random_state
            # )
            # reward_perstep = self.reward_provider.get_reward()
            # #assert reward_perstep is not None and reward_perstep != np.nan
            # self._rew_step = reward_perstep

            self._has_after_after_step = False
            return change_in_abs_target_delta_yaw
        else:
            return 0.0
    
    def reset(self, random_state : np.random.RandomState) -> dict[str,Any]:
        self.robot.receive_observation()
        # Reset the info dict
        self._info_dict = {}

        # Reset the task-specific variables
        self._target_yaw = 0.0
        self._target_delta_yaw = 0.0
        self._has_after_after_step = False

        if not self._inited:
            self._inited = True
            for initializer in self.initializers:
                initializer.perform_reset(
                    self.robot,
                    self._info_dict,
                    self._termination_reason,
                    random_state
                )

        # call the resetters
        for resetter in self.resetters:
            resetter.perform_reset(
                self.robot,
                self._info_dict,
                self._termination_reason,
                random_state
            )

        # Reset the target yaw provider
        self.target_yaw_provider.reset_target(
            self.robot,
            self._info_dict,
            self._termination_reason,
            random_state
        )
        self.__update_target()

        # Reset target yaw obs
        if self.target_observable is not None:
            self.target_observable.reset_target_obs(
                self.robot,
                self.target_goal_world_delta,
                self.target_goal_local,
                self.target_yaw,
                self.target_delta_yaw,
                self._target_velocity,
                self._info_dict,
                self._target_custom_data,
                self.enable_target_custom_obs,
                self._termination_reason,
                random_state
            )

        # Reset reward provider
        self.reward_provider.reset_reward(
            self.robot,
            self._info_dict,
            self._termination_reason,
            random_state
        )

        # Reset termination providers
        for termination_provider in self.termination_providers:
            termination_provider.reset_termination_condition(
                self.robot,
                self._info_dict,
                self._termination_reason,
                random_state
            )
        
        # Reset truncation providers
        for truncation_provider in self.truncation_providers:
            truncation_provider.reset_termination_condition(
                self.robot,
                self._info_dict,
                self._termination_reason,
                random_state
            )
        
        self._termination_reason = None
        self._truncation_reason = None
        self._rew_step = 0.0

        # Reset the robot
        self.robot.reset()
        self.robot.receive_observation()

        for resetter in self.resetters:
            if hasattr(resetter, "last_position"):
                resetter.perform_reset(
                    self.robot,
                    self._info_dict,
                    self._termination_reason,
                    random_state
                )

        return self._info_dict.copy()
    
    def should_terminate(self) -> bool:
        return self._termination_reason is not None

    def should_truncate(self) -> bool:
        return self._truncation_reason is not None
rail_walker_interface/joystick_policy/joystick_policy.py
realquantumcookie/APRL
from .robot import BaseWalker, BaseWalkerLocalizable, BaseWalkerInSim, BaseWalkerWithFootContact, BaseWalkerWithJointTemperatureSensor, BaseWalkerWithJoystick, Walker3DVelocityEstimator, Walker3DFrameQuatEstimator, Walker3DVelocityEstimatorLocal, Walker3DLocationEstimator
from .robot_wrapper import WalkerWrapper
from .wrappers import *

rail_walker_interface/robot/__init__.py
realquantumcookie/APRL
from typing import Generic, Optional, TypeVar
import numpy as np
from functools import cached_property
import gym
import gym.spaces
import transforms3d as tr
import time


_ObsT = TypeVar("_ObsT")
class BaseWalker(Generic[_ObsT]):
    def __init__(
        self, 
        name: Optional[str] = "robot", 
        Kp: float = 5,
        Kd: float = 1,
        force_real_control_timestep : bool = False,
        limit_action_range : float = 1.0,
        power_protect_factor : float = 0.1
    ):
        assert limit_action_range > 0 and limit_action_range <= 1.0
        self.name = name
        self.Kp = Kp
        self.Kd = Kd
        self.force_real_control_timestep = force_real_control_timestep
        self._last_control_t = 0.0
        self.limit_action_range = limit_action_range
        self._power_protect_factor = power_protect_factor

    @property
    def is_real_robot(self) -> bool:
        return False

    @property
    def power_protect_factor(self) -> float:
        return self._power_protect_factor
    
    @power_protect_factor.setter
    def power_protect_factor(self, value: float) -> None:
        assert value >= 0 and value <= 1.0
        self._power_protect_factor = value

    """
    The control_timestep is the time interval between two consecutive model control actions.
    """
    @property
    def control_timestep(self) -> float:
        pass
    
    @property
    def action_interpolation(self) -> bool:
        pass

    """
    The control_subtimestep is the time interval between two consecutive internal control actions. It will also be the physics timestep if in simulation.
    """
    @property
    def control_subtimestep(self) -> float:
        pass

    def receive_observation(self) -> bool:
        pass

    @property
    def joint_qpos_init(self) -> np.ndarray:
        pass

    @property
    def joint_qpos_sitting(self) -> np.ndarray:
        pass

    @cached_property
    def joint_qpos_crouch(self) -> np.ndarray:
        return (self.joint_qpos_init + self.joint_qpos_sitting) / 2.0

    """
    This property will be used to determine the standing range of qpos of the robot.
    """
    @property
    def joint_qpos_offset(self) -> np.ndarray:
        pass

    @property
    def joint_qpos_mins(self) -> np.ndarray:
        pass

    @property
    def joint_qpos_maxs(self) -> np.ndarray:
        pass

    def reset(self) -> None:
        pass

    def get_3d_linear_velocity(self) -> np.ndarray:
        pass

    def get_3d_local_velocity(self) -> np.ndarray:
        pass

    def get_3d_angular_velocity(self) -> np.ndarray:
        pass

    def get_framequat_wijk(self) -> np.ndarray:
        pass

    def get_roll_pitch_yaw(self) -> np.ndarray:
        pass

    def get_last_observation(self) -> Optional[_ObsT]:
        pass

    def get_3d_acceleration_local(self) -> np.ndarray:
        pass

    def get_joint_qpos(self) -> np.ndarray:
        pass

    def get_joint_qvel(self) -> np.ndarray:
        pass

    def get_joint_qacc(self) -> np.ndarray:
        pass

    def get_joint_torques(self) -> np.ndarray:
        pass

    def _apply_action(self, action: np.ndarray) -> bool:
        pass

    def close(self) -> None:
        pass

    def __del__(self):
        self.close()
    
    @property
    def action_qpos_mins(self) -> np.ndarray:
        # delta = -np.minimum(np.abs(self.joint_qpos_mins - self.joint_qpos_init), np.abs(self.joint_qpos_maxs - self.joint_qpos_init))
        # return delta * self.limit_action_range + self.joint_qpos_init
        return (self.joint_qpos_mins - self.joint_qpos_init) * self.limit_action_range + self.joint_qpos_init
    
    @property
    def action_qpos_maxs(self) -> np.ndarray:
        # delta = np.minimum(np.abs(self.joint_qpos_mins - self.joint_qpos_init), np.abs(self.joint_qpos_maxs - self.joint_qpos_init))
        # return delta * self.limit_action_range + self.joint_qpos_init
        return (self.joint_qpos_maxs - self.joint_qpos_init) * self.limit_action_range + self.joint_qpos_init

    def apply_action(self, action: np.ndarray) -> bool:
        action = np.clip(action, self.action_qpos_mins, self.action_qpos_maxs)
        
        if not self.force_real_control_timestep:
            return self._apply_action(action)
        else:
            t = time.time()
            dt = t - self._last_control_t
            if dt >= self.control_timestep:
                self._last_control_t = t
                return self._apply_action(action)
            else:
                time_to_sleep = self.control_timestep - dt
                time.sleep(time_to_sleep)
                self._last_control_t = t + time_to_sleep
                return self._apply_action(action)

    def can_apply_action(self) -> bool:
        t = time.time()
        dt = t - self._last_control_t
        if (not self.force_real_control_timestep) or dt >= self.control_timestep:
            return True
        else:
            return False

    def async_apply_action(self, action: np.ndarray) -> bool:
        if self.can_apply_action():
            self._last_control_t = time.time()
            return self._apply_action(action)
        else:
            return False

    @cached_property
    def joint_nums(self) -> int:
        return len(self.joint_qpos_init)
    
    @cached_property
    def action_spec(self) -> gym.spaces.Box:
        return gym.spaces.Box(
            low=self.joint_qpos_mins, 
            high=self.joint_qpos_maxs, 
            shape=(self.joint_nums,),
            dtype=np.float32
        )

    def unwrapped(self):
        return self

class BaseWalkerWithFootContact:
    def get_foot_contact(self) -> np.ndarray:
        pass

    def get_foot_force(self) -> np.ndarray:
        pass

    def get_foot_force_norm(self) -> np.ndarray:
        pass

class BaseWalkerWithJointTemperatureSensor():
    def get_joint_temperature_celsius(self) -> np.ndarray:
        pass

class BaseWalkerWithJoystick:
    def get_joystick_values(self) -> list[np.ndarray]:
        pass
 
class BaseWalkerLocalizable():
    def get_3d_location(self) -> np.ndarray:
        pass

class BaseWalkerInSim(BaseWalkerLocalizable):
    def set_3d_location(self, target_location : np.ndarray) -> np.ndarray:
        pass

    def reset_2d_location(self, target_location : np.ndarray, target_quaternion : Optional[np.ndarray] = None, target_qpos : Optional[np.ndarray] = None) -> np.ndarray:
        pass

    def set_framequat_wijk(self, framequat_wijk: np.ndarray) -> None:
        pass

    def set_roll_pitch_yaw(self, roll: float, pitch: float, yaw: float) -> None:
        pass

_RobotCls = TypeVar("_RobotCls", bound=BaseWalker)
class Walker3DVelocityEstimator(Generic[_RobotCls]):
    def __init__(self):
        self._enabled = True

    def reset(self, robot: _RobotCls, frame_quat: Optional[np.ndarray]):
        pass

    def step(self, robot: _RobotCls, frame_quat: Optional[np.ndarray], dt: float):
        pass

    def get_3d_linear_velocity(self) -> np.ndarray:
        pass

    def close(self) -> None:
        pass

    @property
    def enabled(self) -> bool:
        return self._enabled
    
    @enabled.setter
    def enabled(self, value: bool) -> None:
        self._enabled = value
    
    def __del__(self):
        self.close()

_RobotClsFramequatEst = TypeVar("_RobotClsFramequatEst", bound=BaseWalker)
class Walker3DFrameQuatEstimator(Generic[_RobotClsFramequatEst]):
    def reset(self, robot: _RobotClsFramequatEst, frame_quat: Optional[np.ndarray]):
        pass

    def step(self, robot: _RobotClsFramequatEst, frame_quat: Optional[np.ndarray], dt: float):
        pass

    def get_framequat_wijk(self) -> np.ndarray:
        pass
    
    def close(self) -> None:
        pass

    @property
    def enabled(self) -> bool:
        return True

    def __del__(self):
        self.close()
    

_RobotClsLocalEst = TypeVar("_RobotClsLocalEst", bound=BaseWalker)
class Walker3DVelocityEstimatorLocal(Walker3DVelocityEstimator[_RobotClsLocalEst], Generic[_RobotClsLocalEst]):
    def __init__(self):
        self._global_velocity = np.zeros(3)

    def reset(self, robot: _RobotClsLocalEst, frame_quat: Optional[np.ndarray]):
        self._global_velocity = np.zeros(3)
        self._reset(robot, frame_quat)

    def step(self, robot: _RobotClsLocalEst, frame_quat: Optional[np.ndarray], dt: float):
        self._step(robot, frame_quat, dt)
        local_velocity = self.get_3d_local_velocity()
        self._global_velocity = tr.quaternions.rotate_vector(local_velocity, frame_quat)

    def get_3d_linear_velocity(self) -> np.ndarray:
        return self._global_velocity
    
    def _reset(self, robot: _RobotClsLocalEst, frame_quat: Optional[np.ndarray]): 
        pass

    def _step(self, robot: _RobotClsLocalEst, frame_quat: np.ndarray, dt: float):
        raise NotImplementedError()
    
    def get_3d_local_velocity(self) -> np.ndarray:
        raise NotImplementedError()

_RobotClsPositionEst = TypeVar("_RobotClsPositionEst", bound=BaseWalker)
class Walker3DLocationEstimator(Generic[_RobotClsPositionEst], Walker3DVelocityEstimator[_RobotClsPositionEst]):
    def __init__(self) -> None:
        Walker3DVelocityEstimator.__init__(self)
        self._last_location = None
        self._last_velocity = None

    def reset(self, robot: _RobotCls, frame_quat: Optional[np.ndarray]):
        self._last_location = self._estimate_location(robot,frame_quat)
        self._last_velocity = np.zeros(3)

    def step(self, robot: _RobotCls, frame_quat: Optional[np.ndarray], dt: float):
        location = self._estimate_location(robot, frame_quat)
        self._last_velocity = (location - self._last_location) / dt
        self._last_location = location

    def get_3d_linear_velocity(self) -> np.ndarray:
        return self._last_velocity

    def get_3d_location(self) -> np.ndarray:
        return self._last_location
    
    def _estimate_location(self, robot: _RobotCls, frame_quat : Optional[np.ndarray]) -> np.ndarray:
        raise NotImplementedError()

    def close(self) -> None:
        pass

    def __del__(self):
        self.close()
    
rail_walker_interface/robot/robot.py
realquantumcookie/APRL
from typing import Optional
from .robot import BaseWalker
from typing import TypeVar, Generic, Any
import numpy as np
import gym.spaces

_WalkerCls = TypeVar("_WalkerCls", bound=BaseWalker)
class WalkerWrapper(BaseWalker, Generic[_WalkerCls]):
    def __init__(self, robot : BaseWalker):
        self.robot = robot
    
    def __getattr__(self, name : str):
        if "robot" in self.__dict__ and hasattr(self.robot, name):
            return getattr(self.robot, name)
        elif name in self.__dict__:
            return self.__dict__[name]
        else:
            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
    
    def __setattr__(self, name : str, value : Any):
        if hasattr(self, 'robot') and hasattr(self.robot, name):
            setattr(self.robot, name, value)
        else:
            self.__dict__[name] = value

    @property
    def is_real_robot(self) -> bool:
        return self.robot.is_real_robot

    @property
    def power_protect_factor(self) -> float:
        return self.robot.power_protect_factor
    
    @power_protect_factor.setter
    def power_protect_factor(self, value: float) -> None:
        self.robot.power_protect_factor = value

    @property
    def control_timestep(self) -> float:
        return self.robot.control_timestep
    
    @property
    def action_interpolation(self) -> bool:
        return self.robot.action_interpolation

    @property
    def control_subtimestep(self) -> float:
        return self.robot.control_subtimestep

    def receive_observation(self) -> bool:
        return self.robot.receive_observation()

    @property
    def joint_qpos_init(self) -> np.ndarray:
        return self.robot.joint_qpos_init

    @property
    def joint_qpos_sitting(self) -> np.ndarray:
        return self.robot.joint_qpos_sitting

    @property
    def joint_qpos_crouch(self) -> np.ndarray:
        return self.robot.joint_qpos_crouch

    @property
    def joint_qpos_offset(self) -> np.ndarray:
        return self.robot.joint_qpos_offset

    @property
    def joint_qpos_mins(self) -> np.ndarray:
        return self.robot.joint_qpos_mins

    @property
    def joint_qpos_maxs(self) -> np.ndarray:
        return self.robot.joint_qpos_maxs

    def reset(self) -> None:
        self.robot.reset()

    def get_3d_linear_velocity(self) -> np.ndarray:
        return self.robot.get_3d_linear_velocity()

    def get_3d_local_velocity(self) -> np.ndarray:
        return self.robot.get_3d_local_velocity()

    def get_3d_angular_velocity(self) -> np.ndarray:
        return self.robot.get_3d_angular_velocity()

    def get_framequat_wijk(self) -> np.ndarray:
        return self.robot.get_framequat_wijk()

    def get_roll_pitch_yaw(self) -> np.ndarray:
        return self.robot.get_roll_pitch_yaw()

    def get_last_observation(self) -> Optional[Any]:
        return self.robot.get_last_observation()

    def get_3d_acceleration_local(self) -> np.ndarray:
        return self.robot.get_3d_acceleration_local()

    def get_joint_qpos(self) -> np.ndarray:
        return self.robot.get_joint_qpos()

    def get_joint_qvel(self) -> np.ndarray:
        return self.robot.get_joint_qvel()
    
    def get_joint_qacc(self) -> np.ndarray:
        return self.robot.get_joint_qacc()

    def get_joint_torques(self) -> np.ndarray:
        return self.robot.get_joint_torques()

    def _apply_action(self, action: np.ndarray) -> bool:
        return self.robot._apply_action(action)

    def close(self) -> None:
        self.robot.close()

    @property
    def action_qpos_mins(self) -> np.ndarray:
        return self.robot.action_qpos_mins
    
    @property
    def action_qpos_maxs(self) -> np.ndarray:
        return self.robot.action_qpos_maxs

    def apply_action(self, action: np.ndarray) -> bool:
        return self.robot.apply_action(action)

    def can_apply_action(self) -> bool:
        return self.robot.can_apply_action()

    def async_apply_action(self, action: np.ndarray) -> bool:
        return self.robot.async_apply_action(action)

    @property
    def joint_nums(self) -> int:
        return self.robot.joint_nums
    
    @property
    def action_spec(self) -> gym.spaces.Box:
        return self.robot.action_spec

    def unwrapped(self):
        return self.robot.unwrapped()
rail_walker_interface/robot/robot_wrapper.py
realquantumcookie/APRL
from .velocity_smoother import WalkerVelocitySmoother
from .readonly_robot import ReadOnlyWalkerWrapper
rail_walker_interface/robot/wrappers/__init__.py
realquantumcookie/APRL
import numpy as np
from ..robot_wrapper import WalkerWrapper

class ReadOnlyWalkerWrapper(WalkerWrapper):
    def receive_observation(self) -> bool:
        return True
    
    def apply_action(self, action: np.ndarray) -> bool:
        return True
rail_walker_interface/robot/wrappers/readonly_robot.py
realquantumcookie/APRL
import numpy as np
from ..robot_wrapper import WalkerWrapper
from functools import cached_property
import transforms3d as tr3d

class SideFlippedWalkerWrapper(WalkerWrapper):
    def receive_observation(self) -> bool:
        ret = self.robot.receive_observation()
        self._local_velocity = self.robot.get_3d_local_velocity()
        self._local_velocity[1] *= -1
        self._linear_velocity = tr3d.quaternions.rotate_vector(self._local_velocity, self.robot.get_framequat_wijk())
        self._local_accel = self.robot.get_3d_acceleration_local()
        self._local_accel[1] *= -1
        self._angular_velocity = self.robot.get_3d_angular_velocity()
        self._angular_velocity[2] *= -1
        return ret

    def get_joint_qpos(self) -> np.ndarray:
        return self._transform_joint_readings(self.robot.get_joint_qpos())

    def _transform_joint_readings(self, readings : np.ndarray) -> np.ndarray:
        return_readings = np.zeros_like(readings)
        return_readings[self.joint_indices_right_hips] = -readings[self.joint_indices_left_hips]
        return_readings[self.joint_indices_left_hips] = -readings[self.joint_indices_right_hips]
        return_readings[self.joint_indices_right_thighs] = readings[self.joint_indices_left_thighs]
        return_readings[self.joint_indices_left_thighs] = readings[self.joint_indices_right_thighs]
        return_readings[self.joint_indices_right_calfs] = readings[self.joint_indices_left_calfs]
        return_readings[self.joint_indices_left_calfs] = readings[self.joint_indices_right_calfs]
        return return_readings

    def get_joint_qvel(self) -> np.ndarray:
        return self._transform_joint_readings(self.robot.get_joint_qvel())
    
    def get_joint_qacc(self) -> np.ndarray:
        return self._transform_joint_readings(self.robot.get_joint_qacc())
    
    def get_joint_torques(self) -> np.ndarray:
        return self._transform_joint_readings(self.robot.get_joint_torques())
    
    @cached_property
    def joint_qpos_init(self) -> np.ndarray:
        return self._transform_joint_readings(self.robot.joint_qpos_init)
    
    @cached_property
    def joint_qpos_sitting(self) -> np.ndarray:
        return self._transform_joint_readings(self.robot.joint_qpos_sitting)
    
    @cached_property
    def joint_qpos_offset(self) -> np.ndarray:
        return self._transform_joint_readings(self.robot.joint_qpos_offset)
    
    @cached_property
    def joint_qpos_crouch(self) -> np.ndarray:
        return self._transform_joint_readings(self.robot.joint_qpos_crouch)
    
    @cached_property
    def joint_qpos_mins(self) -> np.ndarray:
        return self._transform_joint_readings(self.robot.joint_qpos_mins)
    
    @cached_property
    def joint_qpos_maxs(self) -> np.ndarray:
        return self._transform_joint_readings(self.robot.joint_qpos_maxs)
    
    @property
    def action_qpos_mins(self) -> np.ndarray:
        return self._transform_joint_readings(self.robot.action_qpos_mins)
    
    @property
    def action_qpos_maxs(self) -> np.ndarray:
        return self._transform_joint_readings(self.robot.action_qpos_maxs)
    
    def get_3d_linear_velocity(self) -> np.ndarray:
        return self._linear_velocity

    def get_3d_local_velocity(self) -> np.ndarray:
        return self._local_velocity
    
    def get_3d_acceleration_local(self) -> np.ndarray:
        return self._local_accel

    def get_3d_angular_velocity(self) -> np.ndarray:
        return self._angular_velocity
    
    @cached_property
    def joint_indices_left_hips(self):
        return np.array([3, 9])
    
    @cached_property
    def joint_indices_left_thighs(self):
        return np.array([4, 10])
    
    @cached_property
    def joint_indices_left_calfs(self):
        return np.array([5, 11])
    
    @cached_property
    def joint_indices_right_hips(self):
        return np.array([0, 6])
    
    @cached_property
    def joint_indices_right_thighs(self):
        return np.array([1, 7])
    
    @cached_property
    def joint_indices_right_calfs(self):
        return np.array([2, 8])
    
    
rail_walker_interface/robot/wrappers/side_flipped_robot.py
realquantumcookie/APRL
from rail_walker_interface.robot.robot import BaseWalker
from ..robot_wrapper import WalkerWrapper
from ..robot import BaseWalker
from collections import deque
import numpy as np
import transforms3d as tr3d

class WalkerVelocitySmoother(WalkerWrapper):
    def __init__(self, robot: BaseWalker, queue_weights : np.ndarray = np.array([0.1,0.2,0.2,0.2,0.3])):
        super().__init__(robot)
        number_of_queues = len(queue_weights)
        self._velocity_queue = deque([],maxlen=number_of_queues)
        self._local_velocity_queue = deque([],maxlen=number_of_queues)
        self._queue_weights = queue_weights
        self._velocity_estimate = None
        self._local_velocity_estimate = None

    def update_velocity_estimate(self):
        if len(self._velocity_queue) == 0:
            self._velocity_estimate = np.zeros(3)
            self._local_velocity_estimate = np.zeros(3)
        else:
            accumulated_weight = 0.0
            self._velocity_estimate = np.zeros(3)
            self._local_velocity_estimate = np.zeros(3)
            length_of_current_queue = len(self._velocity_queue)
            for i in range(length_of_current_queue):
                weight = self._queue_weights[i + self._velocity_queue.maxlen - length_of_current_queue]
                self._velocity_estimate += weight * self._velocity_queue[i]
                self._local_velocity_estimate += weight * self._local_velocity_queue[i]
                accumulated_weight += weight
            self._velocity_estimate /= accumulated_weight
            self._local_velocity_estimate /= accumulated_weight

    def receive_observation(self) -> bool:
        ret = super().receive_observation()
        self._velocity_queue.append(self.robot.get_3d_linear_velocity())
        self._local_velocity_queue.append(self.robot.get_3d_local_velocity())
        self.update_velocity_estimate()
    
    def reset(self) -> None:
        super().reset()
        self._velocity_queue.clear()
        self._local_velocity_queue.clear()
        self._velocity_queue.append(self.robot.get_3d_linear_velocity())
        self._local_velocity_queue.append(self.robot.get_3d_local_velocity())
        self.update_velocity_estimate()
    
    def get_3d_linear_velocity(self) -> np.ndarray:
        return self._velocity_estimate

    def get_3d_local_velocity(self) -> np.ndarray:
        return self._local_velocity_estimate
rail_walker_interface/robot/wrappers/velocity_smoother.py
realquantumcookie/APRL
# setup.py
from setuptools import setup

setup(
    name='rail_unified_walker',
    version='0.0.1',
    packages=['rail_walker_interface', "rail_mujoco_walker", "rail_real_walker", "unitree_go1_wrapper", "rail_walker_gym", "jaxrl5"]
)
setup.py
realquantumcookie/APRL
import unitree_go1_wrapper as go1_wrapper
import numpy as np
from response_curve import y_knee, y_abduction, y_hip #y_knee
import time

enabled_idx = [0,1,2]
#y_abduction = [go1_wrapper.GO1_HIP_INIT] * 5000
#y_hip = [go1_wrapper.GO1_THIGH_INIT] * 5000
#y_knee = [go1_wrapper.GO1_CALF_INIT] * 5000

if __name__ == "__main__":
    udp = go1_wrapper.UDP(go1_wrapper.ControlLevel.LOW)
    low_cmd = udp.initCommunicationLowCmdData()
    udp.setToSendLow(low_cmd)
    udp.send()
    udp.recv()
    state = udp.getRecvLow()
    
    torques_collected = []
    qpos_collected = []
    qvel_collected = []

    for i in range(12):
        low_cmd.motorCmd[i].Kp = 2.0
        low_cmd.motorCmd[i].Kd = 0.5
        low_cmd.motorCmd[i].tau = 0.0
        low_cmd.motorCmd[i].dq = 0.0
    for i in range(0, 12, 3):
        if i not in enabled_idx:
            low_cmd.motorCmd[i].q = 0.0
    for i in range(1, 12, 3):
        if i not in enabled_idx:
            low_cmd.motorCmd[i].q = 70.0/180.0*np.pi
    for i in range(2, 12, 3):
        if i not in enabled_idx:
            low_cmd.motorCmd[i].q = -150.0/180.0*np.pi

    for i in range(len(y_abduction)):
        target_abduction = y_abduction[i]
        target_hip = y_hip[i]
        target_knee = y_knee[i]
        print(target_abduction, target_hip, target_knee)
        for i in [0, 6]:
            if i in enabled_idx:
                low_cmd.motorCmd[i].q = target_abduction
        for i in [3, 9]:
            if i in enabled_idx:
                low_cmd.motorCmd[i].q = -target_abduction
        for i in range(1, 12, 3):
            if i in enabled_idx:
                low_cmd.motorCmd[i].q = target_hip
        for i in range(2, 12, 3):
            if i in enabled_idx:
                low_cmd.motorCmd[i].q = target_knee
        
        udp.setToSendLow(low_cmd)
        udp.send()
        udp.recv()
        state = udp.getRecvLow()
        torques = np.array([
            state.motorState[i].tauEst for i in range(12)
        ])
        qpos = np.array([
            state.motorState[i].q for i in range(12)
        ])
        qvel = np.array([
            state.motorState[i].dq for i in range(12)
        ])
        torques_collected.append(torques)
        qpos_collected.append(qpos)
        qvel_collected.append(qvel)
    np.savez_compressed(
        "data.npz",
        torques=torques_collected,
        qpos=qpos_collected,
        qvel=qvel_collected,
    )

tests/sysid_tweak/data_collect.py
realquantumcookie/APRL
from unitree_go1_wrapper.native_interface import robot_interface
import unitree_go1_wrapper as go1_wrapper
import numpy as np
from response_curve import y_knee #y_abduction, y_hip, y_knee
import time

y_abduction = [go1_wrapper.GO1_HIP_INIT] * 5000
y_hip = [go1_wrapper.GO1_THIGH_INIT] * 5000
#y_knee = [go1_wrapper.GO1_CALF_INIT] * 5000

if __name__ == "__main__":
    udp = robot_interface.UDP(go1_wrapper.ControlLevel.LOW.value, 8080, "192.168.123.10", 8007)
    low_cmd = robot_interface.LowCmd()
    state = robot_interface.LowState()
    udp.InitCmdData(low_cmd)

    udp.SetSend(low_cmd)
    udp.Send
    udp.Recv()
    udp.GetRecv(state)
    
    torques_collected = []
    qpos_collected = []
    qvel_collected = []

    # low_cmd.motorCmd[0].Kp = 5.0
    # low_cmd.motorCmd[1].Kp = 5.0
    low_cmd.motorCmd[5].Kp = 60.0
    # low_cmd.motorCmd[0].Kd = 2.0
    # low_cmd.motorCmd[1].Kd = 2.0
    low_cmd.motorCmd[5].Kd = 50.0
    for i in range(len(y_abduction)):
        udp.Recv()
        udp.GetRecv(state)

        target_abduction = y_abduction[i]
        target_hip = y_hip[i]
        target_knee = y_knee[i]
        print(target_abduction, target_hip, target_knee)
        # low_cmd.motorCmd[0].q = target_abduction
        # low_cmd.motorCmd[1].q = target_hip
        low_cmd.motorCmd[5].q = target_knee
        low_cmd.motorCmd[5].dq = 0.0
        low_cmd.motorCmd[5].tau = 0.0
        low_cmd.motorCmd[5].Kp = 60.0
        low_cmd.motorCmd[5].Kd = 50.0
        udp.SetSend(low_cmd)
        udp.Send()
        torques = np.array([
            state.motorState[i].tauEst for i in range(12)
        ])
        qpos = np.array([
            state.motorState[i].q for i in range(12)
        ])
        qvel = np.array([
            state.motorState[i].dq for i in range(12)
        ])
        torques_collected.append(torques)
        qpos_collected.append(qpos)
        qvel_collected.append(qvel)
        time.sleep(0.002)
    np.savez_compressed(
        "data.npz",
        torques=torques_collected,
        qpos=qpos_collected,
        qvel=qvel_collected,
    )

tests/sysid_tweak/data_collect_native.py
realquantumcookie/APRL
import unitree_go1_wrapper as go1_wrapper
import numpy as np
import time

print_idxes = list(range(3,6))

if __name__ == "__main__":
    udp = go1_wrapper.UDP(go1_wrapper.ControlLevel.LOW)
    low_cmd = udp.initCommunicationLowCmdData()
    udp.setToSendLow(low_cmd)
    udp.send()
    udp.recv()
    state = udp.getRecvLow()
    

    while True:
        udp.setToSendLow(low_cmd)
        udp.send()
        udp.recv()
        state = udp.getRecvLow()


        torques = np.array([
            state.motorState[i].tauEst for i in range(12)
        ])
        qpos = np.array([
            state.motorState[i].q for i in range(12)
        ])
        qvel = np.array([
            state.motorState[i].dq for i in range(12)
        ])
        print("Qpos", [qpos[i] for i in print_idxes])
    

tests/sysid_tweak/joint_qpos.py
realquantumcookie/APRL
import numpy as np
import matplotlib.pyplot as plt

joint_ranges = {
    "abduction": (-1.047, 1.047),
    "hip": (-0.663, 2.966),
    "knee": (-2.721, -0.837),
}

def square_wave(
    duration: float,
    freq: float,
    period: float,
    off_value: float,
    on_value: float,
) -> np.ndarray:
    dt = 1.0 / freq
    on_time = period / 2
    on_steps = int(on_time / dt)
    t = np.linspace(0, duration, int(duration / dt))
    y = np.ones(len(t)) * off_value
    for i in range(len(t)):
        if i % (2 * on_steps) >= on_steps:
            y[i] = on_value
    return y


duration = 5  # seconds
freq = 500 # Hz
period = 2 # seconds

t = np.linspace(0, duration, int(duration * freq))
y_abduction = square_wave(duration, freq, period, joint_ranges["abduction"][0] + 0.1, joint_ranges["abduction"][1] - 0.1)
y_hip = square_wave(duration, freq, period, joint_ranges["hip"][0] + 0.2, joint_ranges["hip"][1] - 0.2)
y_knee = square_wave(duration, freq, period, joint_ranges["knee"][0] + 0.2, joint_ranges["knee"][1] - 0.2)

# plt.plot(t, y_abduction, label="abduction")
# plt.plot(y_hip, label="hip")
# plt.plot(y_knee, label="knee")
# plt.legend()
# plt.show()
tests/sysid_tweak/response_curve.py
realquantumcookie/APRL
from rail_real_walker import *
from rail_real_walker.robots import Go1RealWalkerRemote
from rail_walker_interface import BaseWalker, BaseWalkerLocalizable, Walker3DVelocityEstimator
import numpy as np
import tqdm

TICK_TIMESTEP = 0.03


if __name__ == "__main__":
    estimator_instance = DummyVelocityEstimator()
    robot = Go1RealWalkerRemote(
        estimator_instance,
        power_protect_factor=0.5,
        Kp=40,
        Kd=5
    )
    
    try:
        print("Press Enter to continue")
        input()
        robot.reset()

        for _ in tqdm.trange(0, 100000):
            robot.apply_action(robot.joint_qpos_init)
            robot.receive_observation()
    except KeyboardInterrupt:
        print("KeyboardInterrupt")
        import traceback
        traceback.print_exc()
    finally:
        robot.close()
tests/test_control_freq.py
realquantumcookie/APRL
from rail_real_walker import *
from rail_real_walker.robots import Go1RealWalker
from rail_walker_interface import BaseWalker, BaseWalkerLocalizable, Walker3DVelocityEstimator
import numpy as np
import matplotlib.pyplot as plt

TICK_TIMESTEP = 0.03


if __name__ == "__main__":
    estimator_instance = DummyVelocityEstimator()
    robot = Go1RealWalker(
        estimator_instance,
        power_protect_factor=0.5,
        Kp=30,
        Kd=3
    )
    
    try:
        print("Press Enter to continue")
        input()
        robot.reset()

        while True:
            robot.apply_action(robot.joint_qpos_init)
            robot.receive_observation()
            foot_contact = robot.get_foot_force()
            print("Foot Contact: ", foot_contact)
    except KeyboardInterrupt:
        print("KeyboardInterrupt")
        import traceback
        traceback.print_exc()
    finally:
        robot.close()
tests/test_foot_contact.py
realquantumcookie/APRL
from rail_real_walker import *
from rail_real_walker.robots import Go1RealWalker
from rail_walker_interface import BaseWalker, BaseWalkerLocalizable, Walker3DVelocityEstimator
import numpy as np
import matplotlib.pyplot as plt

TICK_TIMESTEP = 0.03


if __name__ == "__main__":
    estimator_instance = DummyVelocityEstimator()
    robot = Go1RealWalker(
        estimator_instance,
        power_protect_factor=0.5,
        Kp=40,
        Kd=3
    )
    delta_pos = np.zeros(3)
    
    try:
        print("Press Enter to continue")
        input()
        robot.reset()

        while True:
            robot.receive_observation()
            ipt = input("Input delta pos: ")
            delta_pos = np.array([float(x) for x in ipt.split(",")][:3])
            delta_pos_left = delta_pos.copy()
            delta_pos_left[0] = -delta_pos_left[0]
            delta_pos_real = np.asarray([delta_pos, delta_pos_left, delta_pos, delta_pos_left]).reshape((12,))
            target_pos = robot.joint_qpos_init + delta_pos_real
            print(delta_pos, delta_pos_real, target_pos)
            robot.apply_action(target_pos)
            
    except KeyboardInterrupt:
        print("KeyboardInterrupt")
        import traceback
        traceback.print_exc()
    finally:
        robot.close()
tests/test_joint_init.py
realquantumcookie/APRL
from rail_real_walker import *
from rail_real_walker.robots import Go1RealWalker, Go1RealWalkerRemote
from rail_walker_interface import BaseWalker, BaseWalkerLocalizable, Walker3DVelocityEstimator
import numpy as np
import matplotlib.pyplot as plt

TICK_TIMESTEP = 0.03

def init_estimator() -> Walker3DVelocityEstimator:
    return DummyVelocityEstimator()

def interpolate_next_action(robot : Go1RealWalker, step : int):
    init_coefficient = (np.sin(np.pi * 2 * step / 50) + 1) / 2.0
    sit_coefficient = 1.0 - init_coefficient
    return init_coefficient * robot.joint_qpos_init + sit_coefficient * robot.joint_qpos_sitting


if __name__ == "__main__":
    estimator_instance = init_estimator()
    print("Estimator type: ", type(estimator_instance))
    robot = Go1RealWalker(
        estimator_instance,
        power_protect_factor=0.5,
        Kp=40,
        Kd=5
    )
    
    try:
        with plt.ion():
            plt.xlim(-1, 1)
            plt.ylim(-1, 1)
            FR_thigh_arrow = plt.arrow(0, 0, 0, 0, head_width=0.05, width=0.05, color='r')
            FL_thigh_arrow = plt.arrow(0, 0, 0, 0, head_width=0.05, width=0.05, color='b')
            plt.show()
            robot.reset()

            input("Press Enter to continue...")

            step = 0

            while True:
                robot.apply_action(interpolate_next_action(robot, step))
                robot.receive_observation()
                
                joint_torques = robot.get_joint_torques()
                joint_vels = robot.get_joint_qvel()
                FR_thigh_arrow.set_data(dx=joint_torques[1], dy=joint_vels[1])
                FL_thigh_arrow.set_data(dx=joint_torques[4], dy=joint_vels[4])
                
                plt.plot()
                plt.pause(TICK_TIMESTEP)
                step += 1
    except:
        import traceback; traceback.print_exc()
    
    finally:
        print("Closing...")
        robot.close()
        estimator_instance.close()
        plt.close()
        exit(0)
tests/test_joint_torques.py
realquantumcookie/APRL
from rail_real_walker import *
from rail_real_walker.robots import Go1RealWalker, Go1RealWalkerRemote
from rail_walker_interface import BaseWalker, BaseWalkerLocalizable, Walker3DVelocityEstimator
import numpy as np
import matplotlib.pyplot as plt

TICK_TIMESTEP = 0.03

def init_estimator() -> Walker3DVelocityEstimator:
    return IntelRealsenseT265Estimator(
        x_axis_on_robot=np.array([-1, 0, 0]),
        y_axis_on_robot=np.array([0, 0, 1]),
        z_axis_on_robot=np.array([0, 1, 0]),
    )
    # return Go1ForwardKinematicsVelocityProvider()
    # return KalmanFilterFusedVelocityEstimator()

if __name__ == "__main__":
    estimator_instance = init_estimator()
    print("Estimator type: ", type(estimator_instance))
    robot = Go1RealWalker(
        estimator_instance,
        power_protect_factor=0.5,
        Kp=40,
        Kd=5
    )
    
    try:
        with plt.ion():
            plt.xlim(-1, 1)
            plt.ylim(-1, 1)
            local_arrow = plt.arrow(0, 0, 0, 0, head_width=0.05, width=0.05, color='r')
            global_arrow = plt.arrow(0, 0, 0, 0, head_width=0.05, width=0.05, color='b')
            plt.show()
            robot.reset()

            input("Press Enter to continue...")

            while True:
                robot.apply_action(robot.joint_qpos_init)
                robot.receive_observation()
                
                local_velocity = robot.get_3d_local_velocity()
                local_arrow.set_data(dx=local_velocity[0], dy=local_velocity[1])
                
                global_velocity = robot.get_3d_linear_velocity()
                global_arrow.set_data(dx=global_velocity[0], dy=global_velocity[1])
                plt.plot()
                plt.pause(TICK_TIMESTEP)
    except:
        import traceback; traceback.print_exc()
    
    finally:
        print("Closing...")
        robot.close()
        estimator_instance.close()
        plt.close()
        exit(0)
tests/test_velocity_estimator.py
realquantumcookie/APRL
from rail_real_walker import *
from rail_real_walker.robots import Go1RealWalker
from unitree_go1_wrapper import XRockerBtnDataStruct
from rail_walker_interface import BaseWalker, BaseWalkerLocalizable, Walker3DVelocityEstimator
import numpy as np
import matplotlib.pyplot as plt

TICK_TIMESTEP = 0.03

if __name__ == "__main__":
    estimator_instance = DummyVelocityEstimator()
    print("Estimator type: ", type(estimator_instance))
    robot = Go1RealWalker(
        estimator_instance,
        power_protect_factor=0.5,
        Kp=40,
        Kd=5
    )
    
    try:
        with plt.ion():
            plt.xlim(-1, 1)
            plt.ylim(-1, 1)
            lxly_arrow = plt.arrow(0, 0, 0, 0, head_width=0.05, width=0.05, color='r')
            rxry_arrow = plt.arrow(0, 0, 0, 0, head_width=0.05, width=0.05, color='b')
            plt.show()
            robot.reset()

            input("Press Enter to continue...")

            while True:
                robot.receive_observation()
                print(robot._last_state.wirelessRemote)
                wireless_remote = XRockerBtnDataStruct.fromNative(
                    robot._last_state.wirelessRemote
                )
                
                
                lxly_arrow.set_data(dx=wireless_remote.lx, dy=wireless_remote.ly)
                rxry_arrow.set_data(dx=wireless_remote.rx, dy=wireless_remote.ry)
                plt.plot()
                plt.pause(TICK_TIMESTEP)
    except:
        import traceback; traceback.print_exc()
    
    finally:
        print("Closing...")
        robot.close()
        estimator_instance.close()
        plt.close()
        exit(0)
tests/test_wireless_remote.py
realquantumcookie/APRL
from typing import Optional, Tuple

def get_action_curriculum_planner_linear(
    action_curriculum_steps: int,
    action_curriculum_start: float,
    action_curriculum_end: float,
    start_step: int = 0,
):
    def planned_action_low_high(step: int) -> Optional[Tuple[float, float]]:
        if action_curriculum_steps <= 0:
            return None
        if step - start_step < action_curriculum_steps:
            alpha = (step - start_step) / action_curriculum_steps
            current_action = (1 - alpha) * action_curriculum_start + alpha * action_curriculum_end
            return -current_action, current_action
        else:
            current_action = action_curriculum_end
            return -current_action, current_action
    return planned_action_low_high

def get_action_curriculum_planner_quadratic(
    action_curriculum_steps: int,
    action_curriculum_start: float,
    action_curriculum_end: float,
    start_step: int = 0,
):
    def planned_action_low_high(step: int) -> Optional[Tuple[float, float]]:
        if action_curriculum_steps <= 0:
            return None
        if step - start_step < action_curriculum_steps:
            alpha = (step - start_step) / action_curriculum_steps
            alpha_2 = alpha * alpha
            current_action = (1 - alpha_2) * action_curriculum_start + alpha_2 * action_curriculum_end
            return -current_action, current_action
        else:
            current_action = action_curriculum_end
            return -current_action, current_action
    return planned_action_low_high
    
training/action_curriculum_util.py
realquantumcookie/APRL
import os
import typing
from jaxrl5.agents.agent import Agent as JaxRLAgent
from jaxrl5.data import ReplayBuffer
from flax.training import checkpoints
from natsort import natsorted
import pickle
import shutil
from rail_walker_gym.envs.wrappers.rollout_collect import Rollout

def initialize_project_log(project_dir) -> None:
    os.makedirs(project_dir, exist_ok=True)
   
    chkpt_dir = os.path.join(project_dir, 'checkpoints')
    buffer_dir = os.path.join(project_dir, "buffers")
    rollout_dir = os.path.join(project_dir, "rollouts")
    eval_rollout_dir = os.path.join(project_dir, "eval_rollouts")
    
    os.makedirs(chkpt_dir, exist_ok=True)
    os.makedirs(buffer_dir, exist_ok=True)
    os.makedirs(rollout_dir, exist_ok=True)
    os.makedirs(eval_rollout_dir, exist_ok=True)

def list_checkpoint_steps(project_dir) -> typing.List[int]:
    chkpt_dir = os.path.join(project_dir,'checkpoints')
    chkpts = natsorted(os.listdir(chkpt_dir))
    return [int(chkpt.split('_')[-1]) for chkpt in chkpts]

def list_replay_buffer_steps(project_dir) -> typing.List[int]:
    buffer_dir = os.path.join(project_dir, "buffers")
    buffers = natsorted(os.listdir(buffer_dir))
    return [int(buffer.split('_')[-1].split('.')[0]) for buffer in buffers]

def load_checkpoint_at_step(project_dir, step : int, agent : JaxRLAgent, if_failed_return_step : int = 0) -> typing.Tuple[int, JaxRLAgent]:
    chkpt_dir = os.path.join(project_dir,'checkpoints')
    chkpts = os.listdir(chkpt_dir)
    for chkpt in chkpts:
        chkpt_int = int(chkpt.split('_')[-1])
        if chkpt_int == step:
            agent = checkpoints.restore_checkpoint(os.path.join(chkpt_dir, chkpt), agent)
            return chkpt_int, agent
    
    return if_failed_return_step, agent

def load_checkpoint_file(
    filename : str,
    agent : JaxRLAgent
):
    agent = checkpoints.restore_checkpoint(filename, agent)
    return agent

def load_latest_checkpoint(project_dir, agent : JaxRLAgent, if_failled_return_step: int = 0) -> typing.Tuple[int, JaxRLAgent]:
    chkpt_dir = os.path.join(project_dir,'checkpoints')
    chkpts = natsorted(os.listdir(chkpt_dir))
    if len(chkpts) == 0:
        return if_failled_return_step, agent
    else:
        chkpt_int = int(chkpts[-1].split('_')[-1])
        agent = checkpoints.restore_checkpoint(os.path.join(chkpt_dir, chkpts[-1]), agent)
        return chkpt_int, agent

def load_latest_replay_buffer(project_dir) -> typing.Optional[typing.Tuple[int, ReplayBuffer]]:
    buffer_dir = os.path.join(project_dir, "buffers")
    buffers = natsorted(os.listdir(buffer_dir))
    if len(buffers) == 0:
        return None
    else:
        buffer_int = int(buffers[-1].split('_')[-1].split('.')[0])
        with open(os.path.join(buffer_dir, buffers[-1]),'rb') as f:
            replay_buffer = pickle.load(f)
        return buffer_int, replay_buffer

def load_replay_buffer_at_step(project_dir, step: int) -> typing.Optional[ReplayBuffer]:
    buffer_dir = os.path.join(project_dir, "buffers")
    try:
        with open(os.path.join(buffer_dir, f'buffer_{step}.pkl'),'rb') as f:
            replay_buffer = pickle.load(f)
        return replay_buffer
    except:
        return None

def load_latest_additional_replay_buffer(project_dir) -> typing.Optional[ReplayBuffer]:
    buffer_dir = os.path.join(project_dir, "additional_buffers")
    if not os.path.exists(buffer_dir):
        return None
    buffers = natsorted(os.listdir(buffer_dir))
    if len(buffers) <= 1:
        return None
    else:
        with open(os.path.join(buffer_dir, buffers[-2]),'rb') as f:
            replay_buffer = pickle.load(f)
        return replay_buffer

def load_replay_buffer_file(filename : str) -> ReplayBuffer:
    with open(filename,'rb') as f:
        replay_buffer = pickle.load(f)
    return replay_buffer

def save_checkpoint(project_dir, step : int, agent : JaxRLAgent) -> None:
    chkpt_dir = os.path.join(project_dir,'checkpoints')
    try:
        checkpoints.save_checkpoint(
            chkpt_dir,
            agent,
            step=step,
            keep=10,
            overwrite=True
        )
    except:
        pass

def save_replay_buffer(project_dir, step: int, replay_buffer: ReplayBuffer, delete_old_buffers : bool = True) -> None:
    buffer_dir = os.path.join(project_dir, "buffers")
    if delete_old_buffers:
        try:
            shutil.rmtree(buffer_dir)
        except:
            pass
    try:
        os.makedirs(buffer_dir, exist_ok=True)
        with open(os.path.join(buffer_dir, f'buffer_{step}.pkl'),
                    'wb') as f:
            pickle.dump(replay_buffer, f)
    except:
        pass

def save_rollout(project_dir, step: int, is_training : bool, rollout : Rollout):
    rollout_dir = os.path.join(project_dir, "rollouts" if is_training else "eval_rollouts")
    try:
        os.makedirs(rollout_dir, exist_ok=True)
        rollout.export_npz(os.path.join(rollout_dir, f'rollout_{step}.npz'))
    except:
        pass

def save_episode_rollout(project_dir, step: int, is_training : bool, rollout : Rollout):
    rollout_dir = os.path.join(project_dir, "rollouts" if is_training else "eval_rollouts")
    try:
        os.makedirs(rollout_dir, exist_ok=True)
        save_file = os.path.join(rollout_dir, f'rollout_{step}.npz')
        if rollout.is_current_episode_empty():
            rollout.export_last_episode_npz(save_file)
        else:
            rollout.export_current_episode_npz(save_file)
    except:
        pass

def list_rollout_steps(project_dir, is_training : bool) -> typing.List[int]:
    rollout_dir = os.path.join(project_dir, "rollouts" if is_training else "eval_rollouts")
    rollouts = natsorted(os.listdir(rollout_dir))
    return [int(rollout.split('_')[-1].split('.')[0]) for rollout in rollouts]

def load_rollout_at_step(project_dir, step: int, is_training : bool) -> typing.Optional[Rollout]:
    rollout_dir = os.path.join(project_dir, "rollouts" if is_training else "eval_rollouts")
    try:
        rollout = Rollout.import_npz(os.path.join(rollout_dir, f'rollout_{step}.npz'))
        return rollout
    except:
        return None
training/checkpoint_util.py
realquantumcookie/APRL
import ml_collections
from ml_collections.config_dict import config_dict


def get_config():
    config = ml_collections.ConfigDict()

    config.model_cls = 'SACLearnerWithDynamics'

    config.actor_lr = 3e-4
    config.critic_lr = 3e-4
    config.temp_lr = 3e-4

    config.hidden_dims = (256, 256)

    config.discount = 0.99

    config.num_qs = 2

    config.critic_dropout_rate = 0.01
    config.critic_layer_norm = True

    config.tau = 0.005
    config.init_temperature = 0.1
    config.target_entropy = config_dict.placeholder(float)
    
    config.interior_linear_c = 0.0
    config.interior_quadratic_c = 0.0
    config.exterior_linear_c = 0.0
    config.exterior_quadratic_c = 0.0
    
    config.max_gradient_norm = 100.0

    return config
training/configs/droq_config.py
realquantumcookie/APRL
import ml_collections
from ml_collections.config_dict import config_dict


def get_config():
    config = ml_collections.ConfigDict()

    config.model_cls = 'SACLearnerWithDynamics'

    config.actor_lr = 3e-4
    config.critic_lr = 3e-4
    config.temp_lr = 3e-4

    config.hidden_dims = (256, 256)

    config.discount = 0.99

    config.num_qs = 10
    config.num_min_qs = 2

    config.tau = 0.005
    config.init_temperature = 1.0
    config.target_entropy = config_dict.placeholder(float)

    config.backup_entropy = True

    config.critic_layer_norm = False

    config.ctrl_weight = 0.0

    return config
training/configs/redq_config.py
realquantumcookie/APRL
import ml_collections
from ml_collections.config_dict import config_dict


def get_config():
    config = ml_collections.ConfigDict()

    config.model_cls = 'SACLearner'

    config.actor_lr = 3e-4
    config.critic_lr = 3e-4
    config.temp_lr = 3e-4

    config.hidden_dims = (256, 256)

    config.discount = 0.99

    config.num_qs = 2

    config.critic_dropout_rate = 0.01
    config.critic_layer_norm = True

    config.tau = 0.005
    config.init_temperature = 0.1
    config.target_entropy = config_dict.placeholder(float)
    
    config.interior_linear_c = 0.0
    config.interior_quadratic_c = 0.0
    config.exterior_linear_c = 0.0
    config.exterior_quadratic_c = 0.0
    
    config.max_gradient_norm = config_dict.placeholder(float)

    return config
training/configs/reset_config.py
realquantumcookie/APRL
import ml_collections
from ml_collections.config_dict import config_dict


def get_config():
    config = ml_collections.ConfigDict()

    config.model_cls = 'SACLearnerWithDynamics'

    config.actor_lr = 3e-4
    config.critic_lr = 3e-4
    config.temp_lr = 3e-4

    config.hidden_dims = (256, 256)

    config.discount = 0.99

    config.num_qs = 2

    config.tau = 0.005
    config.init_temperature = 1.0
    config.target_entropy = config_dict.placeholder(float)

    #config.sampled_backup = True
    config.backup_entropy = True

    config.critic_layer_norm = False

    config.ctrl_weight = 0.0

    return config
training/configs/sac_config.py
realquantumcookie/APRL
import ml_collections
from ml_collections.config_dict import config_dict


def get_config():
    config = ml_collections.ConfigDict()

    config.model_cls = 'SACLearnerWithDynamics'

    config.actor_lr = 3e-4
    config.critic_lr = 3e-4
    config.temp_lr = 3e-4

    config.hidden_dims = (256, 256)

    config.discount = 0.99

    config.num_qs = 2

    config.critic_layer_norm = True

    config.tau = 0.005
    config.init_temperature = 1.0
    config.target_entropy = config_dict.placeholder(float)

    #config.sampled_backup = True
    config.backup_entropy = True

    config.ctrl_weight = 0.0
    config.max_gradient_norm = 100.0

    return config
training/configs/sac_ln_config.py
realquantumcookie/APRL
#! /usr/bin/env python
# import gym
# from absl import app, flags
# import rail_walker_gym
# import matplotlib.pyplot as plt
import os

import gym
import numpy as np
import tqdm
import wandb
from absl import app, flags
from flax.training import checkpoints
from jaxrl5.agents import SACLearner
from jaxrl5.agents.sac.sac_learner_wdynamics import SACLearnerWithDynamics
from jaxrl5.agents.agent import Agent as JaxRLAgent
from jaxrl5.data import ReplayBuffer
from jaxrl5.wrappers import wrap_gym
from ml_collections import config_flags
import pickle
from dm_control import viewer
from dm_control.mjcf import export_with_assets
import typing

import mujoco
from natsort import natsorted 
import shutil
import rail_walker_gym
import rail_walker_interface

from checkpoint_util import initialize_project_log, load_latest_checkpoint, save_checkpoint, load_latest_replay_buffer, save_replay_buffer, save_rollout
from eval_util import evaluate
from task_config_util import apply_task_configs

import matplotlib.pyplot as plt

FLAGS = flags.FLAGS
flags.DEFINE_string('env_name', 'Go1SanityMujoco-v0', 'Environment name.')
flags.DEFINE_integer('seed', 42, 'Random seed.')
flags.DEFINE_boolean('generate_data', False, 'Generate data.')

def import_register():
    if "Real" in FLAGS.env_name:
        import rail_walker_gym.envs.register_real
    if "Mujoco" in FLAGS.env_name:
        import rail_walker_gym.envs.register_mujoco

def main(_):
    import_register()
    env : gym.Env = gym.make(FLAGS.env_name)
    env = wrap_gym(env, rescale_actions=True)
    env.seed(FLAGS.seed)
    
    # import ipdb; ipdb.set_trace()
    # export_with_assets(
    #     env.task.root_entity.mjcf_model,
    #     out_dir = "./kevin",
    #     out_file_name="scene.xml"
    # )
    # env = gym.wrappers.RescaleAction(env, -1, 1)
    # env = gym.wrappers.ClipAction(env)
    # env = gym.wrappers.RecordEpisodeStatistics(env, deque_size=1)
    # import ipdb; ipdb.set_trace()
    if FLAGS.generate_data:
        actions = []
    else:
        actions = np.load("actions.npy")
    returns = []
    
    env.joystick_policy.robot.foot_friction = 0.2
    for i in range(10):
        traj_actions = []
        curr_return = 0
        env.reset()
        for j in range(100):
            if FLAGS.generate_data:
                action = env.action_space.sample()
            else:
                action = actions[i][j]
            obs, reward, done, info = env.step(action)
            if FLAGS.generate_data:
                traj_actions.append(action)
            curr_return += reward

        returns.append(curr_return)
        if FLAGS.generate_data:
            actions.append(traj_actions)
        print("Return:", curr_return)

    if FLAGS.generate_data:
        np.save("returns.npy", np.array(returns))
        np.save("actions.npy", np.array(actions))

    # save a video of the saved images
    # import imageio
    # images = []
    # for i in range(50):
        # images.append(imageio.imread(f"img_{i}.png"))
    # imageio.mimsave('movie.gif', images, fps=2)

    # delete images
    # import os
    # for i in range(50):
        # os.remove(f"img_{i}.png")

    # from ipdb import set_trace; set_trace()


if __name__ == '__main__':
    app.run(main)

training/debug.py
realquantumcookie/APRL
import gym
import os
from typing import Tuple, Optional
from jaxrl5.agents.agent import Agent as JaxRLAgent
import numpy as np
from rail_walker_gym.envs.wrappers.rollout_collect import Rollout, RolloutCollector
import wandb
from tqdm import trange
import matplotlib.pyplot as plt
import imageio
import time
from skimage.transform import rescale, resize, downscale_local_mean

IMG_PATH = "topdownmap.png"

def log_visitation(
    env: gym.Env,
    locations : np.ndarray,
    train_step: int = 0,
    project_dir: str = None,
    log_tag: str = "training/visitation_map",
    locations_compare : np.ndarray = None,
) :
    traversible = env.joystick_policy.target_yaw_provider.traversible
    img = plt.imread(IMG_PATH)
    # import ipdb; ipdb.set_trace()
    def plot_heatmap(locations_list, start_idx, end_idx, locations_compare_list=None):
        if locations_compare_list is not None:
            assert len(locations_list) == len(locations_compare_list)

        if end_idx > len(locations_list):
            end_idx = len(locations_list)    
        if end_idx - start_idx < 100:
            return None 
        locations = locations_list[start_idx:end_idx]
        locations = np.array(locations).T / traversible.scale
        
        if locations_compare_list is not None:
            locations_compare = locations_compare_list[start_idx:end_idx]
            locations_compare = np.array(locations_compare).T / traversible.scale
        else:
            locations_compare = None
        
        # Plot the heatmap
        fig, ax = plt.subplots(figsize=(10, 10))  # Set the figure size to create a square plot
        ax.imshow(img, extent=[-1, 1, -1, 1])
        ax.scatter(locations[0], locations[1], marker='.', color='turquoise', s=100, alpha=np.linspace(0, 1, end_idx-start_idx)**5)
        if locations_compare is not None:
            ax.scatter(locations_compare[0], locations_compare[1], marker='.', color='hotpink', s=100, alpha=np.linspace(0, 1, end_idx-start_idx)**5)
        # Set the aspect ratio to 'equal'
        ax.set_aspect('equal')

        # Set the axis limits to -1 and 1
        ax.set_xlim(-1, 1)
        ax.set_ylim(-1, 1)

        # # Add a colorbar
        # cbar = fig.colorbar(im)

        # Add labels and title
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_title('Visitation Density')

        if hasattr(traversible, "goals"):
            # Add stars to specific locations
            star_x = [g[0][0] for g in traversible.goals]
            star_y = [g[0][1] for g in traversible.goals]
            ax.scatter(star_x, star_y, marker='*', color='yellow', s=200)
        elif hasattr(traversible, "route_points"):
            star_x = [g[0][0] for g in traversible.route_points]
            star_y = [g[0][1] for g in traversible.route_points]
            ax.scatter(star_x, star_y, marker='*', color='yellow', s=200)
        else:
            pass

        # Save the plot to a file (temporary file or in-memory buffer)
        unique_id = str(int(np.round(time.time() * 1000)))
        image_path = os.path.join(project_dir, 'heatmap'+unique_id+'.png')
        plt.savefig(image_path)
        heatmap_array = imageio.v2.imread(image_path)

        # image_resized = resize(heatmap_array, (heatmap_array.shape[0] // 4, heatmap_array.shape[1] // 4),
        #                     anti_aliasing=True)
        os.remove(image_path)
        plt.close()
        return heatmap_array
        # return image_resized

    # create a gif of the heatmap over time
    heatmap_images = []
    timesteps = 1000
    for i in range(0, len(locations), timesteps//5):
    # for i in range(0, len(locations), timesteps):
        heatmap_image = plot_heatmap(locations, i, i+timesteps, locations_compare)
        if not (heatmap_image is None):
            heatmap_images.append(heatmap_image)

    if heatmap_images:
        # prepend the first image to the list
        heatmap_images = [255+np.zeros_like(heatmap_images[0])] + heatmap_images

        wandb.log({
            log_tag: wandb.Video(
                np.array(heatmap_images).transpose(0, 3, 1, 2), fps=4, format="gif"
                ),
            }, step=train_step)


def evaluate(
    agent : JaxRLAgent, 
    env: gym.Env, 
    num_episodes: int,
    train_step: int = 0,
    log_wandb: bool = False,
    log_video: bool = False,
    record_rollout : bool = False,
    enable_tqdm: bool = False
) -> Tuple[float, float, Optional[np.ndarray], Optional[Rollout]]:
    if hasattr(env, 'enableWandbVideo'):
        prev_enable_wandb_video = env.enableWandbVideo
        env.enableWandbVideo = False
    
    print("===================== Evaluating =====================")
    if record_rollout:
        env = RolloutCollector(env)
    if log_video:
        videos = []
    else:
        videos = None
    
    for episode_i in trange(num_episodes, disable=not enable_tqdm):
        episode_infos = {}
        observation = env.reset()
        done = False
        while not done:
            action = agent.eval_actions(observation)
            observation, rew, done, info = env.step(action)

            if 'TimeLimit.joystick_target_change' in info and info['TimeLimit.joystick_target_change']:
                observation, info = env.reset(return_info=True)
                done = False
            
            for key, value in info.items():
                if key in ['TimeLimit.truncated', 'TimeLimit.joystick_target_change', 'episode']:
                    continue
                value = info[key]
                if key not in episode_infos:
                    episode_infos[key] = [value]
                else:
                    episode_infos[key].append(value)
            
            if log_video:
                videos.append(env.render())
        
        if log_wandb and episode_i == 0:
            for k, v in episode_infos.items():
                if v is None or len(v) <= 0:
                    continue
                if k in ['fall_count','traversible_finished_lap_count']:
                    to_log = v[-1]
                else:
                    to_log = np.mean(v)
                wandb.log({'evaluation/' + str(k): to_log}, step=train_step)

    eval_return = np.mean(env.return_queue)
    eval_length = np.mean(env.length_queue)
    eval_rollouts = env.collected_rollouts if record_rollout else None
    
    if log_wandb:
        wandb.log({
            'evaluation/return': eval_return,
            'evaluation/length': eval_length
        }, step=train_step)
    
    if hasattr(env, 'enableWandbVideo'):
        env.enableWandbVideo = prev_enable_wandb_video
    
    print("===================== Evaluation Done =====================")

    return eval_return, eval_length, videos, eval_rollouts


def combine(one_dict, other_dict):
    combined = {}
    indices = np.random.permutation(len(one_dict['observations']) + len(other_dict['observations']))
    for k, v in one_dict.items():
        if isinstance(v, dict):
            combined[k] = combine(v, other_dict[k])
        else:
            tmp = np.empty((v.shape[0] + other_dict[k].shape[0], *v.shape[1:]),
                           dtype=v.dtype)
            tmp[:v.shape[0]] = v
            tmp[v.shape[0]:] = other_dict[k]
            combined[k] = tmp[indices]

    return combined

def update_with_mixed_buffers(prior_replay_buffer, online_replay_buffer, batch_size, utd_ratio, agent, prior_buffer_ratio=0.5):
    assert 0 < prior_buffer_ratio < 1
    prior_amount = int(batch_size * utd_ratio * prior_buffer_ratio)
    online_amount = batch_size * utd_ratio - prior_amount
    online_batch = online_replay_buffer.sample(online_amount)
    prior_batch = prior_replay_buffer.sample(prior_amount)
    batch = combine(online_batch, prior_batch)

    return agent.update(batch, utd_ratio)

def update_with_delay_with_mixed_buffers(prior_replay_buffer, online_replay_buffer, batch_size : int, utd_ratio : int, agent : JaxRLAgent, prior_buffer_ratio : float = 0.5, actor_delay : int = 20, output_range : Optional[Tuple[float, float]] = None):
    assert 0 < prior_buffer_ratio < 1
    prior_amount = int(batch_size * utd_ratio * prior_buffer_ratio)
    online_amount = batch_size * utd_ratio - prior_amount
    online_batch = online_replay_buffer.sample(online_amount)
    prior_batch = prior_replay_buffer.sample(prior_amount)
    batch = combine(online_batch, prior_batch)
    return agent.update_with_delay(batch, utd_ratio, output_range=output_range, actor_delay=actor_delay)

def compute_priority_weights(x, const, scale):
    indices = np.arange(x)
    time = indices - x
    weights = np.exp(time / scale) + const
    weights = weights / np.sum(weights)
    return weights

def evaluate_route_following(
    agent : JaxRLAgent, 
    env: gym.Env, 
    max_timesteps: int,
    train_step: int = 0,
    log_video: bool = False,
    prior_replay_buffer = None,
    prior_buffer_ratio = 0.5,
    online_replay_buffer = None,
    batch_size=256,
    utd_ratio=20,
    enable_tqdm=False,
    frame_skip=2,
    seed=0,
    use_recency_bias=False,
) -> Tuple[float, float, Optional[np.ndarray], Optional[Rollout]]:

    eval_info = {
        'return': 0.0,
        'falls' : 0,
        'video' : None,
        'tracked_locations' : [],
        'time_to_lap_completion' : [],
    }

    if log_video:
        videos = []
    else:
        videos = None

    # LS HACK: this is a hack to reset the traversible to the very start of the route
    traversible = env.joystick_policy.target_yaw_provider.traversible
    # and to set the random state of the goal generator
    if hasattr(traversible, "goals"):
        traversible.goal_generator_random_state = np.random.RandomState(seed)

    if prior_replay_buffer:
        print(f"===================== Evaluating online training policy at {train_step} =====================")
    else:
        print(f"===================== Evaluating fixed policy at {train_step} =====================")
    # import ipdb; ipdb.set_trace()
    env.joystick_policy.resetters[1]._inited = False
    env.joystick_policy.resetters[0].last_position = None
    observation = env.reset()
    # print the position of the robot at reset
    # import ipdb; ipdb.set_trace()
    threed_loc = env.joystick_policy.robot.get_3d_location()
    print("Location of robot at reset:", np.around(threed_loc, 2))
    done = False  
    video = []

    time_to_lap_completion = []

    for t in trange(max_timesteps, disable=not enable_tqdm):
        action, agent = agent.sample_actions(observation)
        next_observation, rew, done, info = env.step(action)

        if 'TimeLimit.joystick_target_change' in info and info['TimeLimit.joystick_target_change']:
            observation, info = env.reset(return_info=True)
            done = False
        
        eval_info['tracked_locations'].append(info['2d_location'])
        eval_info['return'] += rew

        truncated = "TimeLimit.truncated" in info and info['TimeLimit.truncated']
        if (not done) or truncated:
            mask = 1.0
        else:
            mask = 0.0
            eval_info['falls'] += 1
        
        if traversible.is_complete:
            time_to_lap_completion.append(t)
            eval_info['time_to_lap_completion'].append(t - time_to_lap_completion[-1] if len(time_to_lap_completion) > 0 else t)


        update_info = None
        if prior_replay_buffer:
            if prior_buffer_ratio < 1.0:
                online_replay_buffer.insert(dict(
                    observations=observation,
                    actions=action,
                    rewards=rew,
                    masks=mask,
                    dones=done,
                    next_observations=next_observation
                ))
                if len(online_replay_buffer) >= 1000:
                    agent, update_info = update_with_mixed_buffers(prior_replay_buffer,
                                                               online_replay_buffer,
                                                               batch_size,
                                                               utd_ratio,
                                                               agent,
                                                               prior_buffer_ratio=prior_buffer_ratio)
            else:
                prior_replay_buffer.insert(dict(
                    observations=observation,
                    actions=action,
                    rewards=rew,
                    masks=mask,
                    dones=done,
                    next_observations=next_observation
                ))
                if use_recency_bias:
                    prioritized_indx = prior_replay_buffer.np_random.choice(
                        len(prior_replay_buffer), 
                        size=batch_size * utd_ratio, 
                        p = compute_priority_weights(len(prior_replay_buffer), 0.01, 50)
                    )

                    batch = prior_replay_buffer.sample(batch_size * utd_ratio, indx=prioritized_indx)
                else:
                    batch = prior_replay_buffer.sample(batch_size * utd_ratio)
                agent, update_info = agent.update(batch, utd_ratio)

        # if update_info:
        #     if t % 100 == 0:
        #         for k, v in update_info.items():
        #             wandb.log({f'testtime_training' + f'/{k}': v.item()}, step=t)

        observation = next_observation

        if log_video and t % frame_skip == 0:
            video.append(env.render())
            if len(video) > 60*20:
                videos.append(video)
                video = []

        if done:
            # import ipdb; ipdb.set_trace()
            assert env.joystick_policy.resetters[1]._inited
            observation, info = env.reset(return_info=True)
            done = False
    
    if log_video:
        if len(video) > 0:
            videos.append(video)

    eval_info['videos'] = videos

    return eval_info
training/eval_util.py
realquantumcookie/APRL
#! /usr/bin/env python
import os

import gym
import numpy as np
import tqdm
import wandb
from absl import app, flags
from flax.training import checkpoints
from jaxrl5.agents import SACLearner
from jaxrl5.agents.sac.sac_learner_wdynamics import SACLearnerWithDynamics
from jaxrl5.agents.agent import Agent as JaxRLAgent
from jaxrl5.data import ReplayBuffer
from jaxrl5.wrappers import wrap_gym
from ml_collections import config_flags
import pickle
from dm_control import viewer
import typing

import mujoco
from natsort import natsorted 
import shutil
import rail_walker_gym
import rail_walker_interface

from checkpoint_util import initialize_project_log, load_latest_checkpoint, save_rollout
from eval_util import evaluate
from task_config_util import apply_task_configs

import time

FLAGS = flags.FLAGS

# ==================== Training Flags ====================
flags.DEFINE_string('env_name', 'Go1SanityMujoco-v0', 'Environment name.')
flags.DEFINE_integer('seed', 42, 'Random seed.')
config_flags.DEFINE_config_file(
    'config',
    'configs/droq_config.py',
    'File path to the training hyperparameter configuration.',
    lock_config=False)

# ==================== Eval Flags ====================
flags.DEFINE_integer('eval_episodes', 1,
                     'Number of episodes used for evaluation.')

# ==================== Log / Save Flags ====================
flags.DEFINE_integer('log_interval', 1, 'Logging interval.')
flags.DEFINE_string("save_dir", "./saved", "Directory to save the model checkpoint and replay buffer.")
flags.DEFINE_string('project_name', 'a1-route-laura', 'wandb project name.')
flags.DEFINE_boolean('tqdm', True, 'Use tqdm progress bar.')

flags.DEFINE_boolean('save_eval_videos', False, 'Save videos during evaluation.')
flags.DEFINE_integer("eval_video_length_limit", 0, "Limit the length of evaluation videos.")
flags.DEFINE_integer("eval_video_interval", 1, "Interval to save videos during evaluation.")
flags.DEFINE_boolean('save_eval_rollouts', False, 'Save rollouts during evaluation.')

flags.DEFINE_boolean('launch_viewer', False, "Launch a windowed viewer for the off-screen rendered environment frames.")
flags.DEFINE_boolean('launch_target_viewer', False, "Launch a windowed viewer for joystick heading and target.")
# ========================================================

# ==================== Joystick Task Flags ====================
config_flags.DEFINE_config_file(
    'task_config',
    'task_configs/default.py',
    'File path to the task/control config parameters.',
    lock_config=False)
config_flags.DEFINE_config_file(
    'reset_agent_config',
    'configs/reset_config.py',
    'File path to the training hyperparameter configuration.',
    lock_config=False)
# ========================================================
def import_register():
    if "Real" in FLAGS.env_name:
        import rail_walker_gym.envs.register_real
    if "Mujoco" in FLAGS.env_name:
        import rail_walker_gym.envs.register_mujoco

def main(_):
    import_register()
    exp_name = FLAGS.env_name
    exp_name += f'_s{FLAGS.seed}'
    #exp_name += f'_maxgr{FLAGS.config.max_gradient_norm:.2f}' if FLAGS.config.max_gradient_norm is not None else ''

    # ==================== Setup WandB ====================
    wandb.init(project=FLAGS.project_name, dir=os.getenv('WANDB_LOGDIR'))
    wandb.run.name = "Eval-" + exp_name
    wandb.config.update(FLAGS)
    
    # ==================== Setup Environment ====================
    eval_env : gym.Env = gym.make(FLAGS.env_name)
    
    try:
        eval_joystick_policy : rail_walker_interface.JoystickPolicy = eval_env.joystick_policy
    except:
        eval_joystick_policy = None
        print("Failed to get joystick policy from environment.")

    eval_env_need_render = FLAGS.launch_viewer or FLAGS.save_eval_videos

    if eval_joystick_policy is not None and eval_joystick_policy.robot.is_real_robot and eval_env_need_render:
        from rail_walker_gym.envs.wrappers.real_render import RealRenderWrapperWithSim
        eval_env = RealRenderWrapperWithSim(eval_env)

    if FLAGS.launch_viewer:
        eval_env = rail_walker_gym.envs.wrappers.RenderViewerWrapper(eval_env)
    
    if FLAGS.launch_target_viewer:
        env = rail_walker_gym.envs.wrappers.JoystickTargetViewer(env)

    task_suffix, eval_env = apply_task_configs(eval_env, FLAGS.env_name, 0, FLAGS.task_config, FLAGS.reset_agent_config, True)
    exp_name += task_suffix
    wandb.run.name = "Eval-" + exp_name

    eval_env.seed(FLAGS.seed)

    if FLAGS.save_eval_rollouts:
        eval_env = rail_walker_gym.envs.wrappers.RolloutCollector(eval_env) # wrap environment to automatically collect rollout

    if FLAGS.save_eval_videos:
        eval_env = rail_walker_gym.envs.wrappers.WanDBVideoWrapper(eval_env, record_every_n_steps=FLAGS.eval_video_interval, log_name="evaluation/video", video_length_limit=FLAGS.eval_video_length_limit) # wrap environment to automatically save video to wandb
        eval_env.enableWandbVideo = True
    
    observation, info = eval_env.reset(return_info=True)
    done = False
    # ==================== Setup Checkpointing ====================
    project_dir = os.path.join(FLAGS.save_dir, exp_name)
    initialize_project_log(project_dir)

    # ==================== Setup Learning Agent and Replay Buffer ====================
    agent_kwargs = dict(FLAGS.config)
    model_cls = agent_kwargs.pop('model_cls')

    agent = globals()[model_cls].create(
        FLAGS.seed, 
        eval_env.observation_space,
        eval_env.action_space,
        **agent_kwargs
    )

    agent_loaded_checkpoint_step, agent = load_latest_checkpoint(project_dir, agent, 0)
    if agent_loaded_checkpoint_step > 0:
        print(f"===================== Loaded checkpoint at step {agent_loaded_checkpoint_step} =====================")
    else:
        print(f"===================== No checkpoint found! =====================")
        print("Check directory:", project_dir)
        exit(0)
    # ==================== Start Eval ====================
    accumulated_info_dict = {}
    episode_counter = 0

    try:
        for i in tqdm.trange(0, 500000, initial=0, disable=not FLAGS.tqdm, smoothing=0.1):
            
            # Force steps in environment WanDBVideoWrapper to be the same as our training steps
            # Notice: we need to make sure that the very outside wrapper of env is WandbVideoWrapper, otherwise setting this will not work
            if FLAGS.save_eval_videos:
                eval_env.set_wandb_step(i)

            action = agent.eval_actions(observation)
            
            # Step the environment
            next_observation, reward, done, info = eval_env.step(action)
            truncated = "TimeLimit.truncated" in info and info['TimeLimit.truncated']
            if (not done) or truncated:
                mask = 1.0
            else:
                mask = 0.0

            if eval_joystick_policy is not None and eval_joystick_policy.robot.is_real_robot and (
                isinstance(eval_joystick_policy.robot.unwrapped(), rail_walker_interface.BaseWalkerWithJoystick) and
                eval_joystick_policy.robot.get_joystick_values()[-1][1] > 0.5
            ):
                print("right joystick up detected, pausing evaluation.")
                while (
                    isinstance(eval_joystick_policy.robot.unwrapped(), rail_walker_interface.BaseWalkerWithJoystick) and
                    eval_joystick_policy.robot.get_joystick_values()[-1][1] > 0.5
                ):
                    eval_joystick_policy.robot.receive_observation()
                    time.sleep(eval_joystick_policy.robot.control_timestep)
                print("right joystick up no longer detected, resuming training.")
                next_observation = eval_env.reset(return_info=False)
                done = False

            if 'TimeLimit.joystick_target_change' in info and info['TimeLimit.joystick_target_change']:
                next_observation, info = eval_env.reset(return_info=True)
                done=False

            observation = next_observation

            # Accumulate info to log
            for key in info.keys():
                if key in ['TimeLimit.truncated', 'TimeLimit.joystick_target_change', 'episode']:
                    continue
                value = info[key]
                if key not in accumulated_info_dict:
                    accumulated_info_dict[key] = [value]
                else:
                    accumulated_info_dict[key].append(value)

            if i % FLAGS.log_interval == 0:
                for k in accumulated_info_dict.keys():
                    v = accumulated_info_dict[k]
                    if v is None or len(v) <= 0:
                        continue
                    if k in ['fall_count','traversible_finished_lap_count']:
                        to_log = v[-1]
                    else:
                        to_log = np.mean(v)
                    wandb.log({'evaluation/' + str(k): to_log}, step=i)
                accumulated_info_dict = {}

            if done:
                for k, v in info['episode'].items():
                    decode = {'r': 'return', 'l': 'length', 't': 'time'}
                    wandb.log({f'evaluation/episode_{decode[k]}': v}, step=i)

                observation, info = eval_env.reset(return_info=True)
                done = False
                if FLAGS.save_eval_rollouts:
                    # Since we clear the collected rollouts after saving, we don't need to worry about saving episode rollout
                    save_rollout(project_dir, i, False, eval_env.collected_rollouts)
                    eval_env.collected_rollouts.clear()
                episode_counter += 1
            
            if episode_counter >= FLAGS.eval_episodes:
                break
    except KeyboardInterrupt:
        print("KeyboardInterrupt, please wait for clean up...")
        import traceback
        traceback.print_exc()
    finally:
        print("======================== Cleaning up ========================")
        # ==================== Save Final Checkpoint ====================
        if FLAGS.save_eval_rollouts and len(eval_env.collected_rollouts) > 0:
            save_rollout(project_dir, i, False, eval_env.collected_rollouts)
        if FLAGS.save_eval_videos:
            eval_env._terminate_record()
        eval_env.close()
        print("======================== Done ========================")

if __name__ == '__main__':
    app.run(main)

training/evaluate.py
realquantumcookie/APRL
import gym
import numpy as np
import tqdm
from absl import app, flags
from jaxrl5.wrappers import wrap_gym
import rail_walker_gym
import rail_walker_interface
import wandb
import os
from checkpoint_util import load_rollout_at_step, save_rollout
from ml_collections import config_flags
from task_config_util import apply_task_configs
import time

# ==================== Replay Flags ====================
flags.DEFINE_string("replay_npz", None, "Path to the replay file to load.")
flags.DEFINE_string('env_name', 'Go1SanityMujoco-v0', 'Environment name.')
flags.DEFINE_integer('seed', 10, 'Random seed.')

# ==================== Log/Save Flags ====================
flags.DEFINE_string('project_name', 'a1-route-laura', 'wandb project name.')
flags.DEFINE_bool("save_video", False, "Save video of replay interactions.")
flags.DEFINE_bool("launch_viewer", False, "Launch a windowed viewer for the off-screen rendered environment frames.")
flags.DEFINE_boolean('launch_target_viewer', False, "Launch a windowed viewer for joystick heading and target.")
flags.DEFINE_string("store_npz", None, "Path to the replay file to store.")
flags.DEFINE_integer('log_interval', 1, 'Logging interval.')

# ==================== Joystick Task Flags ====================
config_flags.DEFINE_config_file(
    'task_config',
    'task_configs/default.py',
    'File path to the task/control config parameters.',
    lock_config=False)
config_flags.DEFINE_config_file(
    'reset_agent_config',
    'configs/reset_config.py',
    'File path to the training hyperparameter configuration.',
    lock_config=False)
# ========================================================

FLAGS = flags.FLAGS

def import_register():
    if "Real" in FLAGS.env_name:
        import rail_walker_gym.envs.register_real
    if "Mujoco" in FLAGS.env_name:
        import rail_walker_gym.envs.register_mujoco

def main(_):
    import_register()
    exp_name = FLAGS.env_name
    exp_name += f'_s{FLAGS.seed}'

    # ==================== Setup WandB ====================
    wandb.init(project=FLAGS.project_name, dir=os.getenv('WANDB_LOGDIR'))
    wandb.run.name = "Replay-" + exp_name
    wandb.config.update(FLAGS)
    
    # ==================== Setup Environment ====================
    env : gym.Env = gym.make(FLAGS.env_name)

    try:
        joystick_policy : rail_walker_interface.JoystickPolicy = env.joystick_policy
    except:
        joystick_policy = None
        print("Failed to get joystick policy from environment.")

    env_need_render = FLAGS.launch_viewer or FLAGS.save_video

    if joystick_policy is not None and joystick_policy.robot.is_real_robot and env_need_render:
        from rail_walker_gym.envs.wrappers.real_render import RealRenderWrapperWithSim
        env = RealRenderWrapperWithSim(env)

    if FLAGS.launch_viewer:
        env = rail_walker_gym.envs.wrappers.RenderViewerWrapper(env)
    

    if FLAGS.launch_target_viewer:
        env = rail_walker_gym.envs.wrappers.JoystickTargetViewer(env)

    task_suffix, env = apply_task_configs(env, FLAGS.env_name, 0, FLAGS.task_config, FLAGS.reset_agent_config, True)
    exp_name += task_suffix
    wandb.run.name = "Replay" + exp_name
    
    env.seed(FLAGS.seed)

    if FLAGS.store_npz is not None:
        env = rail_walker_gym.envs.wrappers.RolloutCollector(env) # wrap environment to automatically collect rollout

    if FLAGS.save_video:
        env = rail_walker_gym.envs.wrappers.WanDBVideoWrapper(env, record_every_n_steps=-1) # wrap environment to automatically save video to wandb
        env.enableWandbVideo = True
    
    observation, info = env.reset(return_info=True)
    done = False

    # ==================== Setup Rollout ====================
    to_replay = rail_walker_gym.envs.wrappers.rollout_collect.Rollout.import_npz(
        FLAGS.replay_npz
    )

    input()
    # ==================== Start Training ====================
    accumulated_info_dict = {}
    try:
        for i in tqdm.trange(0, len(to_replay), initial=0, smoothing=0.1):
            # Force steps in environment WanDBVideoWrapper to be the same as our training steps
            # Notice: we need to make sure that the very outside wrapper of env is WandbVideoWrapper, otherwise setting this will not work
            if FLAGS.save_video:
                env.set_wandb_step(i)

            # Get the action from the replay
            action = to_replay.actions[i]
            
            # Step the environment
            next_observation, reward, done, info = env.step(action)
            truncated = "TimeLimit.truncated" in info and info['TimeLimit.truncated']
            if (not done) or truncated:
                mask = 1.0
            else:
                mask = 0.0

            if joystick_policy is not None and joystick_policy.robot.is_real_robot and (
                isinstance(joystick_policy.robot.unwrapped(), rail_walker_interface.BaseWalkerWithJoystick) and
                joystick_policy.robot.get_joystick_values()[-1][1] > 0.5
            ):
                print("right joystick up detected, pausing rollout replay.")
                while (
                    isinstance(joystick_policy.robot.unwrapped(), rail_walker_interface.BaseWalkerWithJoystick) and
                    joystick_policy.robot.get_joystick_values()[-1][1] > 0.5
                ):
                    joystick_policy.robot.receive_observation()
                    time.sleep(joystick_policy.robot.control_timestep)
                print("right joystick up no longer detected, resuming training.")
                next_observation = env.reset(return_info=False)
                done = False

            if 'TimeLimit.joystick_target_change' in info and info['TimeLimit.joystick_target_change']:
                next_observation, info = env.reset(return_info=True)
                done = False
            observation = next_observation

            # Accumulate info to log
            for key in info.keys():
                if key in ['TimeLimit.truncated', 'TimeLimit.joystick_target_change', 'episode']:
                    continue
                value = info[key]
                if key not in accumulated_info_dict:
                    accumulated_info_dict[key] = [value]
                else:
                    accumulated_info_dict[key].append(value)

            if i % FLAGS.log_interval == 0:
                for k in accumulated_info_dict.keys():
                    v = accumulated_info_dict[k]
                    if v is None or len(v) <= 0:
                        continue

                    if k in ['fall_count', 'traversible_finished_lap_count']:
                        to_log = v[-1]
                    else:
                        to_log = np.mean(v)
                    wandb.log({'replay/' + str(k): to_log}, step=i)
                accumulated_info_dict = {}

            if done:
                for k, v in info['episode'].items():
                    decode = {'r': 'return', 'l': 'length', 't': 'time'}
                    wandb.log({f'training/episode_{decode[k]}': v}, step=i)

                observation, info = env.reset(return_info=True)
                done = False
    except KeyboardInterrupt:
        print("KeyboardInterrupt, Please wait for clean up...")
        import traceback
        traceback.print_exc()
    finally:
        print("======================== Cleaning up ========================")
        # ==================== Save Collected Rollout ====================
        if FLAGS.store_npz is not None:
            env.collected_rollouts.export_npz(FLAGS.store_npz)
        if FLAGS.save_video:
            env._terminate_record()
        env.close()
        print("======================== Done ========================")

if __name__ == '__main__':
    app.run(main)
training/replay_rollout.py
realquantumcookie/APRL
from rail_walker_interface import JoystickPolicy, JoystickEnvironment, BaseWalker, BaseWalkerInSim, ReadOnlyWalkerWrapper, JoystickEnvImpl
from typing import Tuple
import gym
import gym.wrappers
from rail_walker_gym.envs.wrappers.relabel_wrapper import RelabelAggregateWrapper, RelabelTargetProvider, SymmetricRelabelWrapper
from rail_walker_gym.joystick_policy.target_providers import JoystickPolicyAutoJoystickTargetProvider
from rail_walker_gym.envs.register_helper import common_post_process_wrapper
import rail_walker_gym.envs.wrappers as envwrappers
from rail_walker_gym.joystick_policy import JoystickPolicyStrictRewardProvider, JoystickPolicyForwardOnlyTargetProvider, JoystickPolicyResetterEarlyStopTruncationProvider, JoystickPolicyManualResetter
from rail_walker_gym.joystick_policy_mujoco import JoystickPolicyPointInSimResetter, JoystickPolicyLastPositionAndYawResetter
from rail_mujoco_walker import RailSimWalkerDMControl
from checkpoint_util import load_checkpoint_file
from jaxrl5.agents import SACLearner
from jaxrl5.agents.sac.sac_learner_wdynamics import SACLearnerWithDynamics
import copy

def _helper_apply_resetter_wrapper(
    env: JoystickEnvironment,
    task_configs,
    reset_agent_config
):
    robot = env.joystick_policy.robot
    joystick_policy = env.joystick_policy
    
    # change fixed point resetter to initializers
    to_move = []
    for resetter in joystick_policy.resetters:
        if isinstance(resetter, JoystickPolicyPointInSimResetter) or isinstance(resetter, JoystickPolicyLastPositionAndYawResetter):
            to_move.append(resetter)
    for resetter in to_move:
        joystick_policy.resetters.remove(resetter)
    joystick_policy.initializers.extend(to_move)
    if isinstance(env.joystick_policy.robot, BaseWalkerInSim):
        joystick_policy.resetters.append(JoystickPolicyLastPositionAndYawResetter())

    def _helper_wrap_reset_env_fn(env):
        env, _ = common_post_process_wrapper(
            env,
            task_configs.reset_policy_action_range,
            task_configs.reset_policy_filter_actions,
            task_configs.reset_policy_frame_stack,
            task_configs.reset_policy_action_history,
            task_configs.reset_policy_center_init_action,
        )
        return env

    reset_joystick_policy = JoystickPolicy(
        robot,
        JoystickPolicyStrictRewardProvider(),
        JoystickPolicyForwardOnlyTargetProvider(),
        [],
        [JoystickPolicyResetterEarlyStopTruncationProvider()],
        [],
        initializers=[],
        target_observable = None,
        enabled_observables=task_configs.reset_policy_joystick_enabled_observables
    )
    resetter_obs_space, resetter_act_space = envwrappers.ResetterPolicySupportedEnvironment.compute_obs_act_space_reset_agent(
        env, reset_joystick_policy, _helper_wrap_reset_env_fn, task_configs.reset_policy_action_range
    )
    reset_agent_config = dict(reset_agent_config)
    reset_agent_model_cls = reset_agent_config.pop('model_cls')
    reset_agent = globals()[reset_agent_model_cls].create(
        0, 
        resetter_obs_space,
        resetter_act_space, 
        **reset_agent_config
    )
    reset_agent = load_checkpoint_file(
        task_configs.reset_agent_checkpoint, 
        reset_agent
    )
    if isinstance(robot.unwrapped(), BaseWalkerInSim):
        fallback_resetter = JoystickPolicyLastPositionAndYawResetter()
        if isinstance(robot.unwrapped(),RailSimWalkerDMControl):
            intercept_reset = True
        else:
            intercept_reset = False
    else:
        fallback_resetter = JoystickPolicyManualResetter()
        intercept_reset = False

    
    env = envwrappers.ResetterPolicySupportedEnvironment(
        env, 
        reset_joystick_policy, 
        reset_agent, 
        max_seconds=task_configs.reset_policy_max_seconds,
        wrap_resetter_env_lambda=_helper_wrap_reset_env_fn,
        fallback_resetter=fallback_resetter,
        resetter_Kp=task_configs.reset_policy_Kp,
        resetter_Kd=task_configs.reset_policy_Kd,
        resetter_action_interpolation=task_configs.reset_policy_action_interpolation,
        resetter_action_scale=task_configs.reset_policy_action_range,
        resetter_power_protect_factor=task_configs.reset_policy_power_protect_factor,
        intercept_all_reset_calls=intercept_reset
    )
    return env
    
def _apply_task_configs(
    env: JoystickEnvironment,
    task_configs,
    reset_agent_config,
    is_relabel_task : bool = False
) -> Tuple[str, gym.Env]:
    task_suffix = ""
    joystick_policy : JoystickPolicy = env.joystick_policy
    joystick_policy.robot.Kp = task_configs.Kp
    joystick_policy.robot.Kd = task_configs.Kd
    
    task_suffix += "_lt" if task_configs.lock_target else "_nlt"
    joystick_policy.lock_target = task_configs.lock_target

    task_suffix += "_PD{:.2f},{:.2f}".format(task_configs.Kp, task_configs.Kd)

    joystick_policy.robot.action_interpolation = task_configs.action_interpolation
    task_suffix += "_ai" if task_configs.action_interpolation else "_nai"
    task_suffix += str(task_configs.filter_actions) if task_configs.filter_actions > 0 else ""

    joystick_policy.robot.limit_action_range = task_configs.action_range
    task_suffix += "_ar{:.2f}".format(task_configs.action_range)

    task_suffix += "_fs{}".format(task_configs.frame_stack)

    task_suffix += "_ciqa" if task_configs.center_init_action else "_nciqa"

    try:
        joystick_policy.robot.power_protect_factor = task_configs.power_protect_factor
        task_suffix += "_ppf{:.2f}".format(task_configs.power_protect_factor)
    except ImportError:
        pass
    
    if hasattr(joystick_policy.target_yaw_provider, "target_linear_velocity"):
        setattr(joystick_policy.target_yaw_provider, "target_linear_velocity", task_configs.rew_target_velocity)

    if hasattr(joystick_policy.reward_provider, "smooth_tdy_steps"):
        setattr(joystick_policy.reward_provider, "smooth_tdy_steps", task_configs.rew_smooth_change_in_tdy_steps)

    if not joystick_policy.target_yaw_provider.is_target_velocity_fixed():
        task_suffix += "_tvdyn"
        if hasattr(joystick_policy.target_yaw_provider,"min_target_velocity") and hasattr(joystick_policy.target_yaw_provider,"max_target_velocity"):
            setattr(joystick_policy.target_yaw_provider, "min_target_velocity", task_configs.min_target_velocity)
            setattr(joystick_policy.target_yaw_provider, "max_target_velocity", task_configs.max_target_velocity)
            task_suffix += "{:.2f},{:.2f}".format(task_configs.min_target_velocity, task_configs.max_target_velocity)
        if hasattr(joystick_policy.target_yaw_provider,"max_velocity_change_rate"):
            setattr(joystick_policy.target_yaw_provider, "max_velocity_change_rate", task_configs.max_velocity_change_rate)
            task_suffix += "-{:.2f}".format(task_configs.max_velocity_change_rate)
    else:
        task_suffix += "_tv{:.2f}".format(joystick_policy.target_yaw_provider.get_target_velocity(joystick_policy.robot))
    
    if hasattr(joystick_policy.target_yaw_provider, "keep_target_delta_yaw_steps_avg") and hasattr(joystick_policy.target_yaw_provider, "keep_velocity_steps_avg"):
        setattr(joystick_policy.target_yaw_provider, "keep_target_delta_yaw_steps_avg", task_configs.joystick_keep_target_delta_yaw_steps_avg)
        setattr(joystick_policy.target_yaw_provider, "keep_velocity_steps_avg", task_configs.joystick_keep_target_velocity_steps_avg)
        task_suffix += "_keep,yaw{:d},tv{:d}".format(task_configs.joystick_keep_target_delta_yaw_steps_avg, task_configs.joystick_keep_target_velocity_steps_avg)

    if (
        hasattr(joystick_policy.reward_provider, "pitch_rate_penalty_factor") and 
        hasattr(joystick_policy.reward_provider, "roll_rate_penalty_factor")
    ):
        setattr(joystick_policy.reward_provider, "pitch_rate_penalty_factor", task_configs.rew_pitch_rate_penalty_factor)
        setattr(joystick_policy.reward_provider, "roll_rate_penalty_factor", task_configs.rew_roll_rate_penalty_factor)
        task_suffix += "_avpf{:.2f},{:.2f}".format(task_configs.rew_pitch_rate_penalty_factor, task_configs.rew_roll_rate_penalty_factor)
    
    if (
        hasattr(joystick_policy.reward_provider, "joint_diagonal_penalty_weight") and
        hasattr(joystick_policy.reward_provider, "joint_shoulder_penalty_weight")
    ):
        setattr(joystick_policy.reward_provider, "joint_diagonal_penalty_weight", task_configs.rew_joint_diagonal_penalty_weight)
        setattr(joystick_policy.reward_provider, "joint_shoulder_penalty_weight", task_configs.rew_joint_shoulder_penalty_weight)
        task_suffix += "_jspw{:.2e},{:.2e}".format(task_configs.rew_joint_diagonal_penalty_weight, task_configs.rew_joint_shoulder_penalty_weight)

    if (
        hasattr(joystick_policy.reward_provider, "joint_acc_penalty_weight") and
        hasattr(joystick_policy.reward_provider, "joint_vel_penalty_weight") and 
        hasattr(joystick_policy.reward_provider, "thigh_torque_penalty_weight")
    ):
        setattr(joystick_policy.reward_provider, "joint_acc_penalty_weight", task_configs.rew_joint_acc_penalty_weight)
        setattr(joystick_policy.reward_provider, "joint_vel_penalty_weight", task_configs.rew_joint_vel_penalty_weight)
        setattr(joystick_policy.reward_provider, "thigh_torque_penalty_weight", task_configs.rew_thigh_torque_penalty_weight)
        task_suffix += "_qpw{:.2e},{:.2e},{:.2e}".format(task_configs.rew_joint_acc_penalty_weight, task_configs.rew_joint_vel_penalty_weight, task_configs.rew_thigh_torque_penalty_weight)


    if hasattr(joystick_policy.reward_provider, "energy_penalty_weight"):
        setattr(joystick_policy.reward_provider, "energy_penalty_weight", task_configs.rew_energy_penalty_weight)
        task_suffix += "_ep{:.3f}".format(task_configs.rew_energy_penalty_weight) if task_configs.rew_energy_penalty_weight > 0.0 else "_nep"
    
    if hasattr(joystick_policy.reward_provider, "qpos_penalty_weight"):
        setattr(joystick_policy.reward_provider, "qpos_penalty_weight", task_configs.rew_qpos_penalty_weight)
        task_suffix += "_qpw{:.3f}".format(task_configs.rew_qpos_penalty_weight)

    if hasattr(joystick_policy.reward_provider, "smooth_torque_penalty_weight"):
        setattr(joystick_policy.reward_provider, "smooth_torque_penalty_weight", task_configs.rew_smooth_torque_penalty_weight)
        task_suffix += "_stpw{:.2e}".format(task_configs.rew_smooth_torque_penalty_weight)

    if hasattr(joystick_policy.reward_provider, "contact_reward_weight"):
        setattr(joystick_policy.reward_provider, "contact_reward_weight", task_configs.rew_contact_reward_weight)
        task_suffix += "_crw{:.2e}".format(task_configs.rew_contact_reward_weight)

    if task_configs.enable_reset_policy and not is_relabel_task:
        env = _helper_apply_resetter_wrapper(env, task_configs, reset_agent_config)
        task_suffix = "_resetter" + task_suffix
    
    if task_configs.limit_episode_length > 0 and not is_relabel_task:
        env = gym.wrappers.TimeLimit(
            env,
            max_episode_steps=task_configs.limit_episode_length
        )

    env = gym.wrappers.RecordEpisodeStatistics(env, deque_size=1)
    env = envwrappers.JoystickPolicyTruncationWrapper(env)
    
    env, obs_space_before_flatten = common_post_process_wrapper(
        env, 
        action_range = task_configs.action_range, 
        filter_actions=task_configs.filter_actions, 
        frame_stack=task_configs.frame_stack, 
        action_history=task_configs.action_history,
        center_init_action=task_configs.center_init_action,
        log_print=True
    )

    return task_suffix, env, obs_space_before_flatten

def apply_task_configs(
    env: JoystickEnvironment,
    env_name: str,
    max_steps: int,
    task_configs,
    reset_agent_config,
    is_eval : bool = False
) -> Tuple[str, gym.Env]:
    suffix, wrapped_env, obs_space_before_flatten = _apply_task_configs(env, task_configs, reset_agent_config, False)
    if task_configs.relabel_count <= 0 or is_eval:
        return suffix, wrapped_env
    
    if not task_configs.symmetric_relabel:
        relabel_envs = []
        if env.joystick_policy.target_yaw_provider.is_target_velocity_fixed():
            fixed_v = env.joystick_policy.target_yaw_provider.get_target_velocity(env.joystick_policy.robot)
            target_velocity_lambda = lambda robot: fixed_v
            target_fixed_velocity_lambda = lambda: True
        else:
            #target_velocity_lambda = lambda robot: env.joystick_policy.target_yaw_provider.get_target_velocity(robot)
            target_velocity_lambda = None
            target_fixed_velocity_lambda = None

        readonly_robot = ReadOnlyWalkerWrapper(
            env.joystick_policy.robot
        )

        for i in range(task_configs.relabel_count):
            relabel_joystick_task = JoystickPolicy(
                readonly_robot,
                copy.copy(env.joystick_policy.reward_provider),
                RelabelTargetProvider(env.joystick_policy.target_yaw_provider),
                [],
                [],
                [],
                [],
                copy.copy(env.joystick_policy.target_observable),
                env.joystick_policy.enabled_observables
            )
            relabel_env = JoystickEnvImpl(relabel_joystick_task)
            if target_velocity_lambda is not None:
                relabel_joystick_task.target_yaw_provider.get_target_velocity = target_velocity_lambda
            if target_fixed_velocity_lambda is not None:
                relabel_joystick_task.target_yaw_provider.is_target_velocity_fixed = target_fixed_velocity_lambda
            relabel_envs.append(_apply_task_configs(relabel_env, task_configs, reset_agent_config, True)[1])
        
        suffix = "_rlab{}".format(task_configs.relabel_count) + suffix

        ret_env = RelabelAggregateWrapper(
            wrapped_env, max_steps, relabel_envs
        )
    else:
        suffix = "_rlabsym" + suffix
        task_configs.relabel_update_fraction = 0.5 # Force to 0.5
        task_configs.relabel_count = 1 # Force to 1
        ret_env = SymmetricRelabelWrapper(
            wrapped_env, obs_space_before_flatten, max_steps
        )

    return suffix, ret_env
    
training/task_config_util.py
realquantumcookie/APRL
import ml_collections
from ml_collections.config_dict import config_dict


def get_config():
    config = ml_collections.ConfigDict()

    # ====================== Shared Control Configs ======================
    config.Kp = 60.0
    config.Kd = 5.0
    config.action_interpolation = False
    config.filter_actions = 0
    config.action_range = 1.0
    config.center_init_action = False
    config.power_protect_factor = 0.6

    # ====================== Environment Configs ======================
    config.frame_stack = 3
    config.action_history = 3
    config.limit_episode_length = 0 # 0 means no limit

    # ====================== Joystick Policy Configs ======================
    config.lock_target = False
    config.rew_target_velocity = 0.5
    config.rew_smooth_change_in_tdy_steps = 5
    config.rew_qpos_penalty_weight = 0.0
    config.rew_energy_penalty_weight = 0.0
    config.rew_smooth_torque_penalty_weight = 0.0
    config.rew_joint_shoulder_penalty_weight = 0.0
    config.rew_joint_diagonal_penalty_weight = 0.0
    config.rew_joint_acc_penalty_weight = 0.0
    config.rew_joint_vel_penalty_weight = 0.0
    config.rew_pitch_rate_penalty_factor = 0.0
    config.rew_roll_rate_penalty_factor = 0.0
    config.rew_thigh_torque_penalty_weight = 0.0
    config.rew_contact_reward_weight = 0.0

    # ====================== Joystick Target Configs ======================
    config.min_target_velocity = 0.3
    config.max_target_velocity = 1.0
    config.joystick_keep_target_delta_yaw_steps_avg = 20 * 15
    config.joystick_keep_target_velocity_steps_avg = 20 * 15
    config.max_velocity_change_rate = 0.3

    # ====================== Reset Wrapper Configs ======================
    config.enable_reset_policy = False
    config.reset_policy_joystick_enabled_observables = ["joints_pos", "imu"]
    config.reset_policy_max_seconds = 5.0
    config.reset_policy_Kp = 40.0
    config.reset_policy_Kd = 5.0
    config.reset_policy_action_interpolation = True
    config.reset_policy_power_protect_factor = 0.7
    config.reset_policy_filter_actions = 0
    config.reset_policy_action_range = 1.0
    config.reset_policy_frame_stack = 3
    config.reset_policy_action_history = 3
    config.reset_policy_center_init_action = False

    config.reset_agent_checkpoint = "./reset_ckpts/Go1_PD40.00,5.00_ai_ar1.00_fs3_ppf0.80_ep_new"

    # ====================== Relabel Wrapper Configs ======================
    config.relabel_count = 0
    config.symmetric_relabel = False
    config.relabel_update_fraction = 0.5

    return config
training/task_configs/default.py
realquantumcookie/APRL
#! /usr/bin/env python
import os

import gym
import numpy as np
import tqdm
import wandb
from absl import app, flags
from flax.training import checkpoints
from jaxrl5.agents import SACLearner
from jaxrl5.agents.sac.sac_learner_wdynamics import SACLearnerWithDynamics
from jaxrl5.agents.agent import Agent as JaxRLAgent
from jaxrl5.data import ReplayBuffer
from ml_collections import config_flags
import pickle
# from dm_control import viewer
import typing
import time
import mujoco
from natsort import natsorted 
import shutil
import rail_walker_gym
import rail_walker_interface
import matplotlib.pyplot as plt

from checkpoint_util import initialize_project_log, load_latest_checkpoint, save_checkpoint, load_latest_replay_buffer, load_latest_additional_replay_buffer, load_replay_buffer_file, save_replay_buffer, save_rollout
from eval_util import evaluate, evaluate_route_following, log_visitation, update_with_delay_with_mixed_buffers
from task_config_util import apply_task_configs
from action_curriculum_util import get_action_curriculum_planner_linear, get_action_curriculum_planner_quadratic


FLAGS = flags.FLAGS

# ==================== Training Flags ====================
flags.DEFINE_string('env_name', 'Go1SanityMujoco-v0', 'Environment name.')
flags.DEFINE_integer('seed', 42, 'Random seed.')
flags.DEFINE_integer('batch_size', 256, 'Mini batch size.')
flags.DEFINE_integer('max_steps', int(5e5), 'Number of training steps.')
flags.DEFINE_integer('start_training', int(1000),
                     'Number of training steps to start training.')
flags.DEFINE_integer('utd_ratio', 20, 'Update to data ratio.')
flags.DEFINE_integer('action_curriculum_steps', -1, 'Number of steps to go to full action range.')
flags.DEFINE_float('action_curriculum_start', 0.1, 'Smallest action range')
flags.DEFINE_float('action_curriculum_end', 1.0, 'Largest action range')
flags.DEFINE_boolean('action_curriculum_linear', True, 'Use linear action curriculum, otherwise quadratic.')
flags.DEFINE_float('action_curriculum_exploration_eps', 0.0, 'Full Action Range Exploration epsilon for action curriculum.')
flags.DEFINE_boolean('load_buffer', False, 'Load replay buffer.')
flags.DEFINE_boolean('variable_friction', False, 'Use different friction for each quadrant of the heightfield.')
config_flags.DEFINE_config_file(
    'config',
    'configs/droq_config.py',
    'File path to the training hyperparameter configuration.',
    lock_config=False)
flags.DEFINE_integer('actor_delay', 20, 'Number of critic updates per actor update.')
flags.DEFINE_string("additional_buffer", None, "Additional replay buffer to load.")
flags.DEFINE_float("additional_buffer_prior", 0.5, "Prior probability of sampling from additional buffer.")
flags.DEFINE_float('terrain_scale', 0.5, 'Scale of generated terrain.')
flags.DEFINE_integer('reset_interval', -1, 'Number of steps between resets. -1 for no resets.')
flags.DEFINE_boolean('reset_agent', False, 'Whether to reset the learning agent.')
flags.DEFINE_boolean('reset_curriculum', True, 'Whether to reset the curriculum.')
flags.DEFINE_string('reset_criterion', 'time', 'Criterion for resetting the curriculum.')
flags.DEFINE_float('threshold', 1.0, 'Threshold to choose when to reset the curriculum (dynamics based only)')

# ==================== Eval Flags ====================
flags.DEFINE_string("eval_env_name", "", "Environment name for evaluation. If empty, use training environment name.")
flags.DEFINE_integer('eval_episodes', 1,
                     'Number of episodes used for evaluation.')
flags.DEFINE_integer('eval_interval', 10000, 'Eval interval.')

# ==================== Log / Save Flags ====================
flags.DEFINE_integer('log_interval', 1000, 'Logging interval.')
flags.DEFINE_integer("save_interval", 10000, "Save interval.")
flags.DEFINE_string("save_dir", "./saved", "Directory to save the model checkpoint and replay buffer.")
flags.DEFINE_boolean('save_buffer', True, 'Save replay buffer for future training.')
flags.DEFINE_boolean('save_old_buffers', False, 'Keep replay buffers in previous steps.')
flags.DEFINE_string('project_name', 'a1-route-laura', 'wandb project name.')
flags.DEFINE_boolean('tqdm', True, 'Use tqdm progress bar.')

flags.DEFINE_boolean('save_eval_videos', False, 'Save videos during evaluation.')
flags.DEFINE_integer("eval_video_length_limit", 0, "Limit the length of evaluation videos.")
flags.DEFINE_boolean('save_eval_rollouts', False, 'Save rollouts during evaluation.')
flags.DEFINE_boolean('save_training_videos', False, 'Save videos during training.')
flags.DEFINE_integer("training_video_length_limit", 0, "Limit the length of training videos.")
flags.DEFINE_integer("training_video_interval", 3000, "Interval to save training videos.")
flags.DEFINE_boolean('save_training_rollouts', False, 'Save rollouts during training.')

flags.DEFINE_boolean('launch_viewer', False, "Launch a windowed viewer for the off-screen rendered environment frames.")
flags.DEFINE_boolean('launch_target_viewer', False, "Launch a windowed viewer for joystick heading and target.")
# ========================================================

# ==================== Joystick Task Flags ====================
config_flags.DEFINE_config_file(
    'task_config',
    'task_configs/default.py',
    'File path to the task/control config parameters.',
    lock_config=False)
config_flags.DEFINE_config_file(
    'reset_agent_config',
    'configs/reset_config.py',
    'File path to the training hyperparameter configuration.',
    lock_config=False)
flags.DEFINE_integer("leg_dropout_step", -1, "Step to start leg dropout.")
# ========================================================

def import_register():
    if "Real" in FLAGS.env_name:
        import rail_walker_gym.envs.register_real
    if "Mujoco" in FLAGS.env_name:
        import rail_walker_gym.envs.register_mujoco

def compute_priority_weights(x, const, scale):
    indices = np.arange(x)
    t = indices - x
    weights = np.exp(t / scale) + const
    weights = weights / np.sum(weights)
    return weights

def initialize_agent(seed, model_cls, obs_space, act_space, **agent_kwargs):
    agent = globals()[model_cls].create(
        seed, 
        obs_space,
        act_space, 
        **agent_kwargs
    )
    return agent

def get_dropout_apply_action(robot : rail_walker_interface.BaseWalker, dropout_leg_idx : typing.List[int]):
    original_fn = robot.apply_action
    def dropout_apply_action(action : np.ndarray) -> bool:
        action = action.copy()
        action[dropout_leg_idx] = robot.joint_qpos_init[dropout_leg_idx]
        return original_fn(action)
    return dropout_apply_action

def main(_):
    import_register()
    exp_name = FLAGS.env_name
    exp_name += f'_s{FLAGS.seed}_maxgr{FLAGS.config.max_gradient_norm:.2f}'
    
    if FLAGS.reset_interval > 0:
        exp_name += f'_reset{FLAGS.reset_interval}'

    # else:
    #     exp_name += '_stepwise'

    if FLAGS.config.exterior_linear_c > 0.0:
        exp_name += f'_extlin{FLAGS.config.exterior_linear_c:.2f}'
    elif FLAGS.config.exterior_quadratic_c > 0.0:
        exp_name += f'_extquad{FLAGS.config.exterior_quadratic_c:.2f}'
        
    if FLAGS.config.interior_quadratic_c > 0.0:
        exp_name += f'_intquad{FLAGS.config.interior_quadratic_c:.2f}'
    
    if FLAGS.variable_friction:
        exp_name += '_varfriction'
        
    if FLAGS.action_curriculum_steps > 0:
        exp_name += f'_ac{FLAGS.action_curriculum_steps}_{FLAGS.action_curriculum_start}to{FLAGS.action_curriculum_end}'
        if FLAGS.action_curriculum_linear:
            exp_name += '_linearly'
        else:
            exp_name += '_quad'

    if "Bumpy" in FLAGS.env_name:
        exp_name += f"_ts{FLAGS.terrain_scale}"
    # ==================== Setup WandB ====================
    wandb.init(project=FLAGS.project_name, dir=os.getenv('WANDB_LOGDIR'))
    wandb.run.name = exp_name
    wandb.config.update(FLAGS)
    
    # ==================== Setup Environment ====================
    env : gym.Env = gym.make(FLAGS.env_name)
    
    if hasattr(env.unwrapped, "task") and hasattr(env.unwrapped.task, "_floor"):
        env.unwrapped.task._floor.terrain_smoothness = 0.5
        env.unwrapped.task._floor.terrain_scale = FLAGS.terrain_scale

    try:
        joystick_policy : rail_walker_interface.JoystickPolicy = env.joystick_policy
    except:
        joystick_policy = None
        print("Failed to get joystick policy from environment.")
    
    env_need_render = FLAGS.launch_viewer or FLAGS.save_training_videos

    if joystick_policy is not None and joystick_policy.robot.is_real_robot and env_need_render:
        from rail_walker_gym.envs.wrappers.real_render import RealRenderWrapperWithSim
        env = RealRenderWrapperWithSim(env, need_render=env_need_render)

    if FLAGS.launch_viewer:
        env = rail_walker_gym.envs.wrappers.RenderViewerWrapper(env)
    
    if FLAGS.launch_target_viewer:
        env = rail_walker_gym.envs.wrappers.JoystickTargetViewer(env)

    if FLAGS.variable_friction:
        friction_list = [0.8, 0.1, 1.2, 0.2, 0.4, 1.0, 0.8, 0.1, 1.2, 0.2, 0.4, 1.0]
        segments = []
        for f in friction_list:
            segments.append(np.ones(FLAGS.action_curriculum_steps)*f)
        friction_schedule = np.concatenate(segments)

    task_suffix, env = apply_task_configs(env, FLAGS.env_name, FLAGS.max_steps, FLAGS.task_config, FLAGS.reset_agent_config, False)
    # exp_name += task_suffix
    wandb.run.name = exp_name

    if FLAGS.eval_interval <= 0:
        eval_env = env
        eval_joystick_policy = joystick_policy
    else:
        eval_env_name = FLAGS.eval_env_name if FLAGS.eval_env_name else FLAGS.env_name
        eval_env : gym.Env = gym.make(eval_env_name)
        try:
            eval_joystick_policy : rail_walker_interface.JoystickPolicy = eval_env.joystick_policy
        except:
            eval_joystick_policy = None
            print("Failed to get joystick policy from eval environment.")
    
        eval_env_need_render = FLAGS.save_eval_videos
        if eval_joystick_policy is not None and eval_joystick_policy.robot.is_real_robot and eval_env_need_render:
            from rail_walker_gym.envs.wrappers.real_render import RealRenderWrapperWithSim
            eval_env = RealRenderWrapperWithSim(eval_env)
        
        _, eval_env = apply_task_configs(eval_env, FLAGS.env_name, FLAGS.max_steps, FLAGS.task_config, FLAGS.reset_agent_config, True)
    
    env.seed(FLAGS.seed)
    eval_env.seed(FLAGS.seed)

    if FLAGS.save_training_rollouts:
        env = rail_walker_gym.envs.wrappers.RolloutCollector(env) # wrap environment to automatically collect rollout

    if FLAGS.save_training_videos:
        env = rail_walker_gym.envs.wrappers.WanDBVideoWrapper(env, record_every_n_steps=FLAGS.training_video_interval, video_length_limit=FLAGS.training_video_length_limit) # wrap environment to automatically save video to wandb
        env.enableWandbVideo = True
    
    observation, info = env.reset(return_info=True)
    # import ipdb; ipdb.set_trace()
    done = False
    # ==================== Setup Checkpointing ====================
    project_dir = os.path.join(FLAGS.save_dir, exp_name)
    initialize_project_log(project_dir)

    # ==================== Setup Learning Agent and Replay Buffer ====================
    agent_kwargs = dict(FLAGS.config)
    model_cls = agent_kwargs.pop('model_cls')

    agent = initialize_agent(FLAGS.seed, model_cls, env.observation_space, env.action_space, **agent_kwargs)
    agent_loaded_checkpoint_step, agent = load_latest_checkpoint(project_dir, agent, 0)

    if agent_loaded_checkpoint_step > 0:
        print(f"===================== Loaded checkpoint at step {agent_loaded_checkpoint_step} =====================")
    else:
        print(f"===================== No checkpoint found. =====================")

    replay_buffer = None
    if FLAGS.load_buffer:
        rb_load = load_latest_replay_buffer(project_dir)
        if rb_load is not None:
            replay_buffer_loaded_step, replay_buffer = rb_load
            print(f"===================== Loaded replay buffer at step {replay_buffer_loaded_step} =====================")

    additional_replay_buffer = None
    if FLAGS.additional_buffer is not None:
        try:
            additional_replay_buffer = load_replay_buffer_file(FLAGS.additional_buffer)
            print(f"===================== Loaded additional replay buffer from {FLAGS.additional_buffer} =====================")
        except:
            additional_replay_buffer = load_latest_additional_replay_buffer(project_dir)
            if additional_replay_buffer is not None:
                print(f"===================== Loaded additional replay buffer from additional_buffers/ =====================")
        
    if replay_buffer is None:
        replay_buffer = ReplayBuffer(env.observation_space, env.action_space,
                                 FLAGS.max_steps)
        replay_buffer_loaded_step = 0
    
    replay_buffer.seed(FLAGS.seed)


    # ==================== Start Training ====================
    if FLAGS.leg_dropout_step > 0 and agent_loaded_checkpoint_step > FLAGS.leg_dropout_step:
        print("Leg dropout step reached, dropping out leg 2.")
        joystick_policy.robot.apply_action = get_dropout_apply_action(
            joystick_policy.robot,
            [2]
        )
    
    accumulated_info_dict = {}
    jax_jitted_actions = False
    jax_jitted_rb = False

    tracked_locations = []

    if FLAGS.action_curriculum_steps <= 0:
        action_curriculum_planner = lambda x: (-FLAGS.action_curriculum_end, FLAGS.action_curriculum_end)
    elif FLAGS.action_curriculum_linear:
        action_curriculum_planner = get_action_curriculum_planner_linear(FLAGS.action_curriculum_steps, FLAGS.action_curriculum_start, FLAGS.action_curriculum_end)
    else:
        action_curriculum_planner = get_action_curriculum_planner_quadratic(FLAGS.action_curriculum_steps, FLAGS.action_curriculum_start, FLAGS.action_curriculum_end)

    actions = []
    values = []
    td_errors = []
    dyn_errors = []
    # for automated curriculum
    ema = None
    dynamics_loss_mean, M2, estimate_count = 0.0, 0.0, 0

    def update_dynerror_statistics(value):
        nonlocal estimate_count, dynamics_loss_mean, M2
        estimate_count += 1
        delta = value - dynamics_loss_mean
        dynamics_loss_mean += delta / estimate_count
        delta2 = value - dynamics_loss_mean
        M2 += delta * delta2

    falls = []
    
    start_time = time.time()
    try:
        for i in tqdm.trange(agent_loaded_checkpoint_step, FLAGS.max_steps + agent_loaded_checkpoint_step, initial=agent_loaded_checkpoint_step, disable=not FLAGS.tqdm, smoothing=0.1):
            if FLAGS.variable_friction:
                env.unwrapped.task.foot_friction = friction_schedule[i]
            # Force steps in environment WanDBVideoWrapper to be the same as our training steps
            # Notice: we need to make sure that the very outside wrapper of env is WandbVideoWrapper, otherwise setting this will not work
            if FLAGS.save_training_videos:
                env.set_wandb_step(i)
            
            # Hack for Jax Jit
            if i >= FLAGS.start_training and not jax_jitted_actions:
                jax_jitted_actions = True
                print("==================== Jitting Jax Actions ====================")
                if joystick_policy is not None:
                    for _ in range(10):
                        joystick_policy.robot.apply_action(joystick_policy.robot.joint_qpos_init)
                
                for _ in range(5):
                    action, agent = agent.sample_actions(env.observation_space.sample())

                print("==================== Done Jitting Jax ====================")
                hack_jit_obs, hack_jit_info = env.reset(return_info=True)
                observation = hack_jit_obs
                if FLAGS.save_training_videos:
                    env.set_wandb_step(i)
                
            if i >= FLAGS.start_training + agent_loaded_checkpoint_step - replay_buffer_loaded_step and not jax_jitted_rb:
                jax_jitted_rb = True
                print("==================== Jitting Jax updates ====================")
                if joystick_policy is not None:
                    for _ in range(20):
                        joystick_policy.robot.apply_action(joystick_policy.robot.joint_qpos_init)
                
                if hasattr(joystick_policy.robot,"hibernate") and hasattr(joystick_policy.robot,"cancel_hibernate"):
                    joystick_policy.robot.hibernate()
                    
                for _ in tqdm.trange(3, disable=not FLAGS.tqdm):
                    if additional_replay_buffer is None:
                        batch = replay_buffer.sample(FLAGS.batch_size * FLAGS.utd_ratio)
                        agent, update_info = agent.update_with_delay(batch, FLAGS.utd_ratio, output_range=action_curriculum_planner(i), actor_delay=FLAGS.actor_delay)
                    else:
                        agent, update_info = update_with_delay_with_mixed_buffers(
                            additional_replay_buffer,
                            replay_buffer,
                            FLAGS.batch_size,
                            FLAGS.utd_ratio,
                            agent,
                            FLAGS.additional_buffer_prior,
                            FLAGS.actor_delay,
                            action_curriculum_planner(i)
                        )

                if hasattr(joystick_policy.robot,"hibernate") and hasattr(joystick_policy.robot,"cancel_hibernate"):
                    joystick_policy.robot.cancel_hibernate()

                print("==================== Done Jitting Jax ====================")
                hack_jit_obs, hack_jit_info = env.reset(return_info=True)
                observation = hack_jit_obs
                if FLAGS.save_training_videos:
                    env.set_wandb_step(i)
            
            if i < FLAGS.start_training:
                action = env.action_space.sample()
                planned_action_range = action_curriculum_planner(i)
                # planned_action_range = curriculum_planners[which_agent](i)
                if planned_action_range is not None:
                    multiplier = planned_action_range[1]
                    if not(FLAGS.action_curriculum_exploration_eps <= 0.0 or np.random.rand() >= FLAGS.action_curriculum_exploration_eps):
                        multiplier = min(1.0, multiplier * 1.5)
                    
                    action = action * multiplier
                        
            else:
                action, agent = agent.sample_actions(observation)
                
            actions.append(action)
            
            # Step the environment
            next_observation, reward, done, info = env.step(action)
            truncated = "TimeLimit.truncated" in info and info['TimeLimit.truncated']
            if (not done) or truncated:
                mask = 1.0
            else:
                mask = 0.0

            if joystick_policy is not None and joystick_policy.robot.is_real_robot and ((
                #instance(joystick_policy.robot.unwrapped(), rail_walker_interface.BaseWalkerWithFootContact) and 
                #np.count_nonzero(joystick_policy.robot.get_foot_contact()) <= 1
            ) or (
                isinstance(joystick_policy.robot.unwrapped(), rail_walker_interface.BaseWalkerWithJoystick) and
                joystick_policy.robot.get_joystick_values()[-1][1] > 0.5
            )):
                print("Lifted foot or right joystick up detected, not inserting into replay buffer.")
                while (
                    isinstance(joystick_policy.robot.unwrapped(), rail_walker_interface.BaseWalkerWithFootContact) and 
                    np.count_nonzero(joystick_policy.robot.get_foot_contact()) <= 1
                ) or (
                    isinstance(joystick_policy.robot.unwrapped(), rail_walker_interface.BaseWalkerWithJoystick) and
                    joystick_policy.robot.get_joystick_values()[-1][1] > 0.5
                ):
                    joystick_policy.robot.receive_observation()
                    time.sleep(joystick_policy.robot.control_timestep)
                print("Foot down or right joystick down detected, resuming training.")
                next_observation = env.reset(return_info=False)
                done = False
            elif (
                joystick_policy is not None and
                isinstance(joystick_policy.robot, rail_walker_interface.BaseWalkerWithJointTemperatureSensor) and
                np.any(joystick_policy.robot.get_joint_temperature_celsius() > 60)
            ):
                print("Overheated joint detected, resetting robot")
                if hasattr(joystick_policy.robot,"hibernate") and hasattr(joystick_policy.robot,"cancel_hibernate"):
                    joystick_policy.robot.hibernate()
                    input("Press enter to continue")
                    joystick_policy.robot.cancel_hibernate()
                else:
                    input("Press enter to continue")
                next_observation, info = env.reset(return_info=True)
                done = False
            else:
                transition_dict = dict(
                    observations=observation,
                    actions=action,
                    rewards=reward,
                    masks=mask,
                    dones=done,
                    next_observations=next_observation
                )
                # falls.append(info['fall_count'])
                # Insert the transition in the replay buffer
                replay_buffer.insert(transition_dict)
            
            if 'TimeLimit.joystick_target_change' in info and info['TimeLimit.joystick_target_change']:
                next_observation, info = env.reset(return_info=True)
                done = False
            
            observation = next_observation

            # Accumulate info to log
            for key in info.keys():
                if key in ['TimeLimit.truncated', 'TimeLimit.joystick_target_change', 'episode']:
                    continue
                value = info[key]
                if key not in accumulated_info_dict:
                    accumulated_info_dict[key] = [value]
                else:
                    accumulated_info_dict[key].append(value)

            # Update the agent
            if i - agent_loaded_checkpoint_step >= FLAGS.start_training - replay_buffer_loaded_step:
                # Relabel batch
                if FLAGS.task_config.relabel_update_fraction > 0.0 and hasattr(env, "get_relabel_replay_buffer"):
                    batch_sample_num = int(round(FLAGS.batch_size * (1.0 - FLAGS.task_config.relabel_update_fraction)))
                    batch = replay_buffer.sample(batch_sample_num)
                    relabel_replay_buffer : ReplayBuffer = env.get_relabel_replay_buffer()
                    batch_relabel = relabel_replay_buffer.sample(FLAGS.batch_size * FLAGS.utd_ratio - batch_sample_num)
                    batch = dict(
                        observations=np.concatenate([batch['observations'], batch_relabel['observations']], axis=0),
                        actions=np.concatenate([batch['actions'], batch_relabel['actions']], axis=0),
                        rewards=np.concatenate([batch['rewards'], batch_relabel['rewards']], axis=0),
                        masks=np.concatenate([batch['masks'], batch_relabel['masks']], axis=0),
                        dones=np.concatenate([batch['dones'], batch_relabel['dones']], axis=0),
                        next_observations=np.concatenate([batch['next_observations'], batch_relabel['next_observations']], axis=0),
                    )
                else:
                    batch = replay_buffer.sample(FLAGS.batch_size * FLAGS.utd_ratio)
                
                
                agent, update_info = agent.update_with_delay(batch, FLAGS.utd_ratio, output_range=action_curriculum_planner(i), actor_delay=FLAGS.actor_delay)
                # update_dynerror_statistics(update_info['dynamics_loss'])
                
                if i % FLAGS.log_interval == 0:
                    wandb.log({f'training/fps': i/(time.time() - start_time)}, step=i)
                    # for all keys in each update_info, add an agent index to the key
                    for k, v in update_info.items():
                        wandb.log({f'training/agent_0/{k}': v.item()}, step=i)
                    # wandb.log({'training/dynamics_surprise': dyn_loss.item()}, step=i)

            
            if i % FLAGS.save_interval == 0 and i > agent_loaded_checkpoint_step:
                save_checkpoint(os.path.join(project_dir), i, agent)

                if i % (FLAGS.log_interval * 10) == 0 and i > agent_loaded_checkpoint_step and FLAGS.save_buffer:
                    save_replay_buffer(project_dir, i, replay_buffer, not FLAGS.save_old_buffers)

            if i % FLAGS.log_interval == 0 and i > agent_loaded_checkpoint_step:
                for k in accumulated_info_dict.keys():
                    v = accumulated_info_dict[k]
                    if v is None or len(v) <= 0:
                        continue
                    
                    if k == '2d_location':
                        # log_visitation(env, v, i, project_dir)
                        tracked_locations.extend(v)
                        tracked_locations = tracked_locations[-FLAGS.save_interval:]
                    else:
                        if k in ['fall_count', 'traversible_finished_lap_count']:
                            to_log = v[-1]
                        else:
                            to_log = np.mean(v)
                        wandb.log({'training/' + str(k): to_log}, step=i)
                accumulated_info_dict = {}

                ################## Skip this for now ###################

                # Compute histogram
                actions = np.array(actions).flatten()
                plt.clf()
                
                hist, bins = np.histogram(actions, bins='auto')
                fig, ax = plt.subplots()
                ax.bar(bins[:-1], hist, width=np.diff(bins), align='edge')
                ax.set_xlabel('Action')
                ax.set_ylabel('Frequency')
                ax.set_title('Histogram of Actions')
                fig.canvas.draw()
                img = np.array(fig.canvas.renderer.buffer_rgba())[:,:,0:3]
                wandb.log({"training/action_hist": wandb.Image(img)}, step=i)
                
                actions = []

            if FLAGS.leg_dropout_step > 0 and i == FLAGS.leg_dropout_step:
                print("Leg dropout step reached, dropping out leg 2.")
                joystick_policy.robot.apply_action = get_dropout_apply_action(
                    joystick_policy.robot,
                    [2]
                )

            if done:
                if 'episode' not in info:
                    import ipdb; ipdb.set_trace()
                for k, v in info['episode'].items():
                    decode = {'r': 'return', 'l': 'length', 't': 'time'}
                    wandb.log({f'training/episode_{decode[k]}': v}, step=i)
                
                if FLAGS.reset_curriculum and FLAGS.reset_criterion == "dynamics_error":
                    recent_batch_indx = np.arange(len(replay_buffer) - FLAGS.batch_size, len(replay_buffer))
                    recent_batch = replay_buffer.sample(FLAGS.batch_size, indx=recent_batch_indx)
                    
                    dyn_errors.append(agent.compute_dynamics_surprise(recent_batch))            

                    wandb.log({f'training/dyn_surprise': np.mean(dyn_errors)}, step=i)
                    
                    if ema is None:
                        ema = np.mean(dyn_errors)
                    else:
                        ema = 0.5 * np.mean(dyn_errors) + .5 * ema
                    dyn_errors = []

                    if ema > FLAGS.threshold:
                        curr_planner = action_curriculum_planner(i)[1]
                        shrink_by_half = curr_planner * 0.5
                        if FLAGS.action_curriculum_steps <= 0:
                            action_curriculum_planner = lambda x: (-FLAGS.action_curriculum_end, FLAGS.action_curriculum_end)
                        elif FLAGS.action_curriculum_linear:
                            start_value = max(FLAGS.action_curriculum_start, shrink_by_half)
                            action_curriculum_planner = get_action_curriculum_planner_linear(FLAGS.action_curriculum_steps, start_value, FLAGS.action_curriculum_end, start_step=i)
                        else:
                            action_curriculum_planner = get_action_curriculum_planner_quadratic(FLAGS.action_curriculum_steps, FLAGS.action_curriculum_start, FLAGS.action_curriculum_end)
                        # # reset the estimates
                        # dynamics_loss_mean, M2, estimate_count = 0.0, 0.0, 0
                        ema = None

                        if FLAGS.reset_agent:
                            agent = initialize_agent(FLAGS.seed, model_cls, env.observation_space, env.action_space, agent_kwargs)
                        
                planned_action_range = action_curriculum_planner(i)
                # planned_action_range = curriculum_planners[agent_i](i)
                if planned_action_range is not None:
                    low = np.around(planned_action_range[0], 2)
                    high = np.around(planned_action_range[1], 2)
                    wandb.log({'training/action_range_low': low}, step=i)
                    wandb.log({'training/action_range_high': high}, step=i)

                observation, info = env.reset(return_info=True)
                done = False
                if FLAGS.save_training_rollouts:
                    # Since we clear the collected rollouts after saving, we don't need to worry about saving episode rollout
                    save_rollout(project_dir, i, True, env.collected_rollouts)
                    env.collected_rollouts.clear()
        
            if FLAGS.reset_interval > 0 and i % FLAGS.reset_interval == 0 and FLAGS.reset_curriculum and i > 0 and FLAGS.reset_criterion == 'time':
                if FLAGS.action_curriculum_steps <= 0:
                    action_curriculum_planner = lambda x: (-FLAGS.action_curriculum_end, FLAGS.action_curriculum_end)
                elif FLAGS.action_curriculum_linear:
                    action_curriculum_planner = get_action_curriculum_planner_linear(FLAGS.action_curriculum_steps, FLAGS.action_curriculum_start, FLAGS.action_curriculum_end, start_step=i)
                else:
                    action_curriculum_planner = get_action_curriculum_planner_quadratic(FLAGS.action_curriculum_steps, FLAGS.action_curriculum_start, FLAGS.action_curriculum_end)
                
                if FLAGS.reset_agent:
                    agent = initialize_agent(FLAGS.seed, model_cls, env.observation_space, env.action_space, agent_kwargs)

    except KeyboardInterrupt:
        print("KeyboardInterrupt, please wait for clean up...")
        import traceback
        traceback.print_exc()
    except:
        import traceback
        traceback.print_exc()
    finally:
        print("======================== Cleaning up ========================")
        # ==================== Save Final Checkpoint ====================
        if i > agent_loaded_checkpoint_step + FLAGS.start_training:
            save_checkpoint(project_dir, i, agent)
            if FLAGS.save_buffer:
                save_replay_buffer(project_dir, i, replay_buffer, not FLAGS.save_old_buffers)
        if FLAGS.save_training_rollouts and len(env.collected_rollouts) > 0:
            save_rollout(project_dir, i, True, env.collected_rollouts)
        if FLAGS.save_training_videos:
            env._terminate_record()
        
        env.close()
        if eval_env is not env:
            eval_env.close()
        print("======================== Done ========================")

if __name__ == '__main__':
    app.run(main)
training/train_online.py
realquantumcookie/APRL
import dmcgym
import gym
from dm_control import viewer

#import rail_walker_gym.envs.register_real
import rail_walker_gym.envs.register_mujoco
from rail_mujoco_walker import JoystickPolicyDMControlTask
from rail_walker_gym import \
    JoystickPolicyRouteFollow2DTraversibleTargetProvider

env = gym.make("Go1GoalGraphFullSquishedMujoco-TDY-StrictRew-v0")
env.reset()

if isinstance(env.unwrapped, dmcgym.DMCGYM):
    task : JoystickPolicyDMControlTask = env.unwrapped._env.task
    if isinstance(task.joystick_policy.target_yaw_provider, JoystickPolicyRouteFollow2DTraversibleTargetProvider):
        traversible = task.joystick_policy.target_yaw_provider.traversible
        if hasattr(traversible, "add_sim_renderables"):
            traversible.add_sim_renderables(
                task._floor._mjcf_root,
                lambda pos: 0.5, #task._floor.height_lookup,
                render_height_offset=0.0,
                show_edges=True
            )

viewer.launch(env.unwrapped._env)
env.close()
training/view_sim.py
realquantumcookie/APRL
from .comm import *
from .common_constants import *
from .go1_constants import *
from .safety import Safety
from .udp import UDP, UDP_CLIENT_PORT, UDP_SERVER_IP_BASIC, UDP_SERVER_IP_SPORT, UDP_SERVER_PORT_BASIC, UDP_SERVER_PORT_SPORT
from .native_interface import robot_interface
from .joystick import XKeySwitchUnion, XKeySwitchUnionComponents, XRockerBtnDataStruct
unitree_go1_wrapper/__init__.py
realquantumcookie/APRL
# https://github.com/unitreerobotics/unitree_legged_sdk/blob/go1/include/unitree_legged_sdk/comm.h
from .native_interface import robot_interface
import typing
import numpy as np

class HasNativeRepr:
    def toNative(self):
        pass

    def toNativeM(self, native):
        pass

    @staticmethod
    def fromNative(native):
        pass

    def __repr__(self) -> str:
        return str(self.__dict__)

class State:
    def isValid(self):
        raise NotImplementedError()

class BmsCmd(HasNativeRepr):
    def __init__(self, off: int = 0xa5):
        self.off = off #off 0xA5
    
    def toNative(self):
        nat = robot_interface.BmsCmd()
        return self.toNativeM(nat)

    def toNativeM(self, native):
        native.off = self.off
        return native
    
    @staticmethod
    def fromNative(native):
        return __class__(native.off)

class BmsState(HasNativeRepr, State):
    def __init__(
        self, 
        version_h : int, #uint8_t
        version_l : int, #uint8_t
		bms_status : int, #uint8_t
		SOC : int, #SOC 0-100%, uint8_t
		current : int, #mA, int32_t
		cycle : int, #uint16_t
		BQ_NTC : typing.Tuple[int,int], #x1 degrees centigrade, std::array<int8_t, 2>
		MCU_NTC: typing.Tuple[int,int], # x1 degrees centigrade, std::array<int8_t, 2>
		cell_vol : typing.Tuple[int,int,int,int,int,int,int,int,int,int] #cell Voltage mV, std::array<uint16_t, 10> 
    ):
        assert len(BQ_NTC) == 2 and len(MCU_NTC) == 2 and len(cell_vol) == 10

        self.version_h = version_h
        self.version_l = version_l
        self.bms_status = bms_status
        self.SOC = SOC
        self.current = current
        self.cycle = cycle
        self.BQ_NTC = BQ_NTC
        self.MCU_NTC = MCU_NTC
        self.cell_vol = cell_vol
    
    def toNative(self):
        assert len(self.BQ_NTC) == 2 and len(self.MCU_NTC) == 2 and len(self.cell_vol) == 10

        nat = robot_interface.BmsState()
        return self.toNativeM(nat)

    def toNativeM(self, native):
        native.version_h = self.version_h
        native.version_l = self.version_l
        native.bms_status = self.bms_status
        native.SOC = self.SOC
        native.current = self.current
        native.cycle = self.cycle
        for i in range(2):
            native.BQ_NTC[i] = self.BQ_NTC[i]
            native.MCU_NTC[i] = self.MCU_NTC[i]
        for i in range(10):
            native.cell_vol[i] = self.cell_vol[i]
        return native
    
    @staticmethod
    def fromNative(native):
        return __class__(
            native.version_h,
            native.version_l,
            native.bms_status,
            native.SOC,
            native.current,
            native.cycle,
            (native.BQ_NTC[0], native.BQ_NTC[1]),
            (native.MCU_NTC[0], native.MCU_NTC[1]),
            (native.cell_vol[0], native.cell_vol[1], native.cell_vol[2], native.cell_vol[3], native.cell_vol[4], native.cell_vol[5], native.cell_vol[6], native.cell_vol[7], native.cell_vol[8], native.cell_vol[9])
        )

    def isValid(self):
        return (
            self.version_h != 0x0 or
            self.version_l != 0x0 or
            self.bms_status != 0x0 or
            self.SOC != 0x0 or
            self.current != 0x0 or
            self.cycle != 0x0 or
            self.BQ_NTC != (0,0) or
            self.MCU_NTC != (0,0) or
            self.cell_vol != (0,0,0,0,0,0,0,0,0,0)
        )

class Cartesian(HasNativeRepr, State):
    def __init__(self, x: float, y: float, z: float):
        self.x = x
        self.y = y
        self.z = z
    
    def toNative(self):
        nat = robot_interface.Cartesian()
        return self.toNativeM(nat)

    def toNativeM(self, native):
        native.x = self.x
        native.y = self.y
        native.z = self.z
        return native
    
    @staticmethod
    def fromNative(native):
        return __class__(native.x, native.y, native.z)

    def isValid(self):
        return self.x != 0 or self.y != 0 or self.z != 0

class IMU(HasNativeRepr, State):
    def __init__(
        self,
        quaternion: np.ndarray,
        gyroscope: np.ndarray,
        accelerometer: np.ndarray,
        rpy: np.ndarray,
        temperature: int
    ) -> None:
        assert quaternion.shape == (4,)
        assert gyroscope.shape == (3,)
        assert accelerometer.shape == (3,)
        assert rpy.shape == (3,)
        self.quaternion = quaternion
        self.gyroscope = gyroscope
        self.accelerometer = accelerometer
        self.rpy = rpy
        self.temperature = temperature
    
    def toNative(self):
        nat = robot_interface.IMU()
        return self.toNativeM(nat)
    
    def toNativeM(self, native):
        assert self.quaternion.shape == (4,)
        assert self.gyroscope.shape == (3,)
        assert self.accelerometer.shape == (3,)
        assert self.rpy.shape == (3,)

        for i in range(4):
            native.quaternion[i] = self.quaternion[i]
        for i in range(3):
            native.gyroscope[i] = self.gyroscope[i]
            native.accelerometer[i] = self.accelerometer[i]
            native.rpy[i] = self.rpy[i]
        native.temperature = self.temperature
        return native

    @staticmethod
    def fromNative(native):
        return __class__(
            np.array([native.quaternion[0], native.quaternion[1], native.quaternion[2], native.quaternion[3]]),
            np.array([native.gyroscope[0], native.gyroscope[1], native.gyroscope[2]]),
            np.array([native.accelerometer[0], native.accelerometer[1], native.accelerometer[2]]),
            np.array([native.rpy[0], native.rpy[1], native.rpy[2]]),
            native.temperature
        )

    def isValid(self):
        return (
            np.any(self.quaternion != 0) or
            np.any(self.gyroscope != 0) or
            np.any(self.accelerometer != 0) or
            np.any(self.rpy != 0) or
            self.temperature != 0
        )

class LED(HasNativeRepr, State):
    def __init__(
        self,
        r: int,
        g: int,
        b: int
    ) -> None:
        assert r >= 0 and r <= 255
        assert g >= 0 and g <= 255
        assert b >= 0 and b <= 255

        self.r = r
        self.g = g
        self.b = b
    
    def toNative(self):
        nat = robot_interface.LED()
        return self.toNativeM(nat)

    def toNativeM(self, native):
        # assert self.r >= 0 and self.r <= 255
        # assert self.g >= 0 and self.g <= 255
        # assert self.b >= 0 and self.b <= 255
        native.r = self.r
        native.g = self.g
        native.b = self.b
        return native
    
    @staticmethod
    def fromNative(native):
        return __class__(native.r, native.g, native.b)
    
    def isValid(self):
        return self.r != 0 or self.g != 0 or self.b != 0
    
class MotorState(HasNativeRepr, State):
    def __init__(
        self,
        mode: int, #motor working mode Servo : 0x0A, Damping : 0x00，Overheat ： 0x08.
        q: float, #current angle(in radian)
        dq: float, #current velocity(in radian/second)
        ddq: float, #current acceleration (in radian/second/second)
        tauEst: float, #current estimated output torque (N*m)
        q_raw: float,
        dq_raw : float,
        ddq_raw : float,
        temperature: int #current temperature (temperature conduction is slow that leads to lag)
    ):
        self.mode = mode
        self.q = q
        self.dq = dq
        self.ddq = ddq
        self.tauEst = tauEst
        self.q_raw = q_raw
        self.dq_raw = dq_raw
        self.ddq_raw = ddq_raw
        self.temperature = temperature
    
    def toNative(self):
        nat = robot_interface.MotorState()
        return self.toNativeM(nat)
    
    def toNativeM(self, native):
        native.mode = self.mode
        native.q = self.q
        native.dq = self.dq
        native.ddq = self.ddq
        native.tauEst = self.tauEst
        native.q_raw = self.q_raw
        native.dq_raw = self.dq_raw
        native.ddq_raw = self.ddq_raw
        native.temperature = self.temperature
        return native
    
    @staticmethod
    def fromNative(native):
        return __class__(
            native.mode,
            native.q,
            native.dq,
            native.ddq,
            native.tauEst,
            native.q_raw,
            native.dq_raw,
            native.ddq_raw,
            native.temperature
        )

    def isValid(self):
        return (
            self.mode != 0 or
            self.q != 0 or
            self.dq != 0 or
            self.ddq != 0 or
            self.tauEst != 0 or
            self.q_raw != 0 or
            self.dq_raw != 0 or
            self.ddq_raw != 0 or
            self.temperature != 0
        )

class MotorCmd(HasNativeRepr):
    def __init__(
        self,
        mode: int, #motor working mode Servo : 0x0A, Damping : 0x00.
        q: float, #target angle(in radian)
        dq: float, #target velocity(in radian/second)
        tau: float, #target output torque (N*m)
        Kp: float, #desired position stiffness (unit: N.m/rad )
        Kd: float #desired velocity stiffness (unit: N.m/(rad/s) 
    ) -> None:
        self.mode = mode
        self.q = q
        self.dq = dq
        self.tau = tau
        self.Kp = Kp
        self.Kd = Kd
    
    def toNative(self):
        nat = robot_interface.MotorCmd()
        return self.toNativeM(nat)
    
    def toNativeM(self, native):
        native.mode = self.mode
        native.q = self.q
        native.dq = self.dq
        native.tau = self.tau
        native.Kp = self.Kp
        native.Kd = self.Kd
        return native
    
    @staticmethod
    def fromNative(native):
        return __class__(
            native.mode,
            native.q,
            native.dq,
            native.tau,
            native.Kp,
            native.Kd
        )

class LowState(HasNativeRepr, State):
    def __init__(
        self,
        head: typing.Tuple[int, int],
        levelFlag: int,
        frameReserve: int,
        SN: typing.Tuple[int, int],
        version: typing.Tuple[int, int],
        bandWidth: int,
        imu: IMU,
        motorState: typing.List[MotorState],
        bms: BmsState,
        footForce: np.ndarray,
        footForceEst: np.ndarray,
        tick: int,
        wirelessRemote: np.ndarray,
        crc: int
    ):
        assert len(motorState) == 20
        assert footForce.shape == (4,)
        assert footForceEst.shape == (4,)
        assert wirelessRemote.shape == (40,) and wirelessRemote.dtype == np.uint8

        self.head = head
        self.levelFlag = levelFlag
        self.frameReserve = frameReserve
        self.SN = SN
        self.version = version
        self.bandWidth = bandWidth
        self.imu = imu
        self.motorState = motorState
        self.bms = bms
        self.footForce = footForce
        self.footForceEst = footForceEst
        self.tick = tick
        self.wirelessRemote = wirelessRemote
        self.crc = crc

    def toNative(self):
        nat = robot_interface.LowState()
        return self.toNativeM(nat)
    
    def toNativeM(self, native):
        assert len(self.motorState) == 20
        assert self.footForce.shape == (4,)
        assert self.footForceEst.shape == (4,)
        assert self.wirelessRemote.shape == (40,)

        native.head[0] = self.head[0]
        native.head[1] = self.head[1]
        native.levelFlag = self.levelFlag
        native.frameReserve = self.frameReserve
        native.SN[0] = self.SN[0]
        native.SN[1] = self.SN[1]
        native.version[0] = self.version[0]
        native.version[1] = self.version[1]
        native.bandWidth = self.bandWidth
        #native.imu = self.imu.toNative()
        self.imu.toNativeM(native.imu)
        for i in range(20):
            #native.motorState[i] = self.motorState[i].toNative()
            self.motorState[i].toNativeM(native.motorState[i])
        #native.bms = self.bms.toNative()
        self.bms.toNativeM(native.bms)
        for i in range(4):
            native.footForce[i] = self.footForce[i]
        for i in range(4):
            native.footForceEst[i] = self.footForceEst[i]
        native.tick = self.tick
        for i in range(40):
            native.wirelessRemote[i] = self.wirelessRemote[i]
        native.crc = self.crc
        return native

    @staticmethod
    def fromNative(native):
        return __class__(
            (native.head[0], native.head[1]),
            native.levelFlag,
            native.frameReserve,
            (native.SN[0], native.SN[1]),
            (native.version[0], native.version[1]),
            native.bandWidth,
            IMU.fromNative(native.imu),
            [MotorState.fromNative(native.motorState[i]) for i in range(20)],
            BmsState.fromNative(native.bms),
            np.asarray(native.footForce),
            np.asarray(native.footForceEst),
            native.tick,
            np.asarray(native.wirelessRemote, dtype=np.uint8),
            native.crc
        )

    def isValid(self):
        return (
            self.head[0] != 0 or
            self.head[1] != 0 or
            self.levelFlag != 0 or
            self.frameReserve != 0 or
            self.SN[0] != 0 or
            self.SN[1] != 0 or
            self.version[0] != 0 or
            self.version[1] != 0 or
            self.bandWidth != 0 or
            self.imu.isValid() or
            any([m.isValid() for m in self.motorState]) or
            self.bms.isValid() or
            np.any(self.footForce != 0) or
            np.any(self.footForceEst != 0) or
            self.tick != 0 or
            np.any(self.wirelessRemote != 0) or
            self.crc != 0
        )
    
class LowCmd(HasNativeRepr):
    def __init__(
        self,
        head: typing.Tuple[int, int],
        version: typing.Tuple[int, int],
        bandWidth: int,
        motorCmd: typing.List[MotorCmd],
        bms: BmsCmd,
        wirelessRemote: np.ndarray,
        crc: int
    ):
        assert len(motorCmd) == 12
        assert wirelessRemote.shape == (40,) and wirelessRemote.dtype == np.uint8

        self.head = head
        self.version = version
        self.bandWidth = bandWidth
        self.motorCmd = motorCmd
        self.bms = bms
        self.wirelessRemote = wirelessRemote
        self.crc = crc
    
    def toNative(self):
        nat = robot_interface.LowCmd()
        return self.toNativeM(nat)
    
    def toNativeM(self, native):
        # assert len(self.motorCmd) == 20
        # assert self.wirelessRemote.shape == (40,)

        native.head[0] = self.head[0]
        native.head[1] = self.head[1]
        native.version[0] = self.version[0]
        native.version[1] = self.version[1]
        native.bandWidth = self.bandWidth
        for i in range(len(self.motorCmd)):
            #native.motorCmd[i] = self.motorCmd[i].toNative()
            self.motorCmd[i].toNativeM(native.motorCmd[i])
        #native.bms = self.bms.toNative()
        self.bms.toNativeM(native.bms)
        for i in range(40):
            native.wirelessRemote[i] = self.wirelessRemote[i]
        native.crc = self.crc
        return native
    
    @staticmethod
    def fromNative(native):
        return __class__(
            (native.head[0], native.head[1]),
            (native.version[0], native.version[1]),
            native.bandWidth,
            [MotorCmd.fromNative(native.motorCmd[i]) for i in range(12)], # 20 in the native struct
            BmsCmd.fromNative(native.bms),
            np.asarray(native.wirelessRemote, dtype=np.uint8),
            native.crc
        )

class HighState(HasNativeRepr, State):
    def __init__(
        self,
        head: typing.Tuple[int, int],
        levelFlag: int,
        frameReserve: int,
        SN: typing.Tuple[int, int],
        version: typing.Tuple[int, int],
        bandWidth: int,
        imu: IMU,
        motorState: typing.List[MotorState],
        bms: BmsState,
        footForce: np.ndarray,
        footForceEst: np.ndarray,
        mode: int,
        progress: int,
        gaitType: int, # 0.idle  1.trot  2.trot running  3.climb stair  4.trot obstacle
        footRaiseHeight: float, #(unit: m, default: 0.08m), foot up height while walking
        position: np.ndarray, #(unit: m), from own odometry in inertial frame, usually drift
        bodyHeight: float, #(unit: m, default: 0.28m)
        velocity: np.ndarray, #(unit: m/s), forwardSpeed, sideSpeed, rotateSpeed in body frame
        yawSpeed: float, #in rad/s,  rotateSpeed in body frame
        rangeObstacle: np.ndarray,
        footPosition2Body: typing.List[Cartesian], #foot position relative to body
        footSpeed2Body: typing.List[Cartesian], #foot speed relative to body
        wirelessRemote: np.ndarray,
        crc: int
    ):
        assert len(motorState) == 20
        assert footForce.shape == (4,)
        assert footForceEst.shape == (4,)
        assert position.shape == (3,)
        assert velocity.shape == (3,)
        assert rangeObstacle.shape == (4,)
        assert len(footPosition2Body) == 4
        assert len(footSpeed2Body) == 4
        assert wirelessRemote.shape == (40,)

        self.head = head
        self.levelFlag = levelFlag
        self.frameReserve = frameReserve
        self.SN = SN
        self.version = version
        self.bandWidth = bandWidth
        self.imu = imu
        self.motorState = motorState
        self.bms = bms
        self.footForce = footForce
        self.footForceEst = footForceEst
        self.mode = mode
        self.progress = progress
        self.gaitType = gaitType
        self.footRaiseHeight = footRaiseHeight
        self.position = position
        self.bodyHeight = bodyHeight
        self.velocity = velocity
        self.yawSpeed = yawSpeed
        self.rangeObstacle = rangeObstacle
        self.footPosition2Body = footPosition2Body
        self.footSpeed2Body = footSpeed2Body
        self.wirelessRemote = wirelessRemote
        self.crc = crc
    
    def toNative(self):
        nat = robot_interface.HighState()
        return self.toNativeM(nat)
    
    def toNativeM(self, native):
        # assert len(self.motorState) == 20
        # assert self.footForce.shape == (4,)
        # assert self.footForceEst.shape == (4,)
        # assert self.position.shape == (3,)
        # assert self.velocity.shape == (3,)
        # assert self.rangeObstacle.shape == (4,)
        # assert len(self.footPosition2Body) == 4
        # assert len(self.footSpeed2Body) == 4
        # assert self.wirelessRemote.shape == (40,)

        native.head[0] = self.head[0]
        native.head[1] = self.head[1]
        native.levelFlag = self.levelFlag
        native.frameReserve = self.frameReserve
        native.SN[0] = self.SN[0]
        native.SN[1] = self.SN[1]
        native.version[0] = self.version[0]
        native.version[1] = self.version[1]
        native.bandWidth = self.bandWidth
        #native.imu = self.imu.toNative()
        self.imu.toNativeM(native.imu)
        for i in range(20):
            #native.motorState[i] = self.motorState[i].toNative()
            self.motorState[i].toNativeM(native.motorState[i])
        #native.bms = self.bms.toNative()
        self.bms.toNativeM(native.bms)
        for i in range(4):
            native.footForce[i] = self.footForce[i]
        for i in range(4):
            native.footForceEst[i] = self.footForceEst[i]
        native.mode = self.mode
        native.progress = self.progress
        native.gaitType = self.gaitType
        native.footRaiseHeight = self.footRaiseHeight
        for i in range(3):
            native.position[i] = self.position[i]
        native.bodyHeight = self.bodyHeight
        for i in range(3):
            native.velocity[i] = self.velocity[i]
        native.yawSpeed = self.yawSpeed
        for i in range(8):
            native.rangeObstacle[i] = self.rangeObstacle[i]
        for i in range(4):
            #native.footPosition2Body[i] = self.footPosition2Body[i].toNative()
            self.footPosition2Body[i].toNativeM(native.footPosition2Body[i])
        for i in range(4):
            #native.footSpeed2Body[i] = self.footSpeed2Body[i].toNative()
            self.footSpeed2Body[i].toNativeM(native.footSpeed2Body[i])
        for i in range(40):
            native.wirelessRemote[i] = self.wirelessRemote[i]
        native.crc = self.crc
        return native
    
    @staticmethod
    def fromNative(native):
        return __class__(
            (native.head[0], native.head[1]),
            native.levelFlag,
            native.frameReserve,
            (native.SN[0], native.SN[1]),
            (native.version[0], native.version[1]),
            native.bandWidth,
            IMU.fromNative(native.imu),
            [MotorState.fromNative(native.motorState[i]) for i in range(20)],
            BmsState.fromNative(native.bms),
            np.array([native.footForce[i] for i in range(4)]),
            np.array([native.footForceEst[i] for i in range(4)]),
            native.mode,
            native.progress,
            native.gaitType,
            native.footRaiseHeight,
            np.array([native.position[i] for i in range(3)]),
            native.bodyHeight,
            np.array([native.velocity[i] for i in range(3)]),
            native.yawSpeed,
            np.array([native.rangeObstacle[i] for i in range(8)]),
            [Cartesian.fromNative(native.footPosition2Body[i]) for i in range(4)],
            [Cartesian.fromNative(native.footSpeed2Body[i]) for i in range(4)],
            np.asarray(native.wirelessRemote, dtype=np.uint8),
            native.crc
        )

    def isValid(self):
        return (
            self.head[0] != 0 or
            self.head[1] != 0 or
            self.levelFlag != 0 or
            self.frameReserve != 0 or
            self.SN[0] != 0 or
            self.SN[1] != 0 or
            self.version[0] != 0 or
            self.version[1] != 0 or
            self.bandWidth != 0 or
            self.imu.isValid() or
            any([ms.isValid() for ms in self.motorState]) or
            self.bms.isValid() or
            np.any(self.footForce != 0) or
            np.any(self.footForceEst != 0) or
            self.mode != 0 or
            self.progress != 0 or
            self.gaitType != 0 or
            self.footRaiseHeight != 0 or
            np.any(self.position != 0) or
            self.bodyHeight != 0 or
            np.any(self.velocity != 0) or
            self.yawSpeed != 0 or
            np.any(self.rangeObstacle != 0) or
            any([cp.isValid() for cp in self.footPosition2Body]) or
            any([cp.isValid() for cp in self.footSpeed2Body]) or
            np.any(self.wirelessRemote != 0) or
            self.crc != 0
        )
    
class HighCmd(HasNativeRepr):
    def __init__(
        self,
        head: typing.Tuple[int, int],
        levelFlag: int,
        frameReserve: int,
        SN: typing.Tuple[int, int],
        version: typing.Tuple[int, int],
        bandWidth: int,
        mode: int, 
        # 0. idle, default stand
        # 1. force stand (controlled by dBodyHeight + ypr)
        # 2. target velocity walking (controlled by velocity + yawSpeed)
        # 3. target position walking (controlled by position + ypr[0])
        # 4. path mode walking (reserve for future release)
        # 5. position stand down. 
        # 6. position stand up 
        # 7. damping mode 
        # 8. recovery stand
        # 9. backflip
        # 10. jumpYaw
        # 11. straightHand
        # 12. dance1
        # 13. dance2
        gaitType:int, # 0.idle  1.trot  2.trot running  3.climb stair  4.trot obstacle
        speedLevel: int, # 0. default low speed. 1. medium speed 2. high speed. during walking, only respond MODE 3
        footRaiseHeight: float, # (unit: m, default: 0.08m), foot up height while walking, delta value
        bodyHeight: float, #(unit: m, default: 0.28m), delta value
        position: np.ndarray, #(unit: m), desired position in inertial frame
        euler: np.ndarray, #(unit: rad), roll pitch yaw in stand mode
        velocity: np.ndarray, #(unit: m/s), forwardSpeed, sideSpeed in body frame
        yawSpeed: float, #(unit: rad/s), rotateSpeed in body frame
        bms: BmsCmd,
        led: typing.List[LED],
        wirelessRemote: np.ndarray,
        crc: int        
    ):
        assert len(head) == 2
        assert len(SN) == 2
        assert len(version) == 2
        assert position.shape == (3,)
        assert euler.shape == (3,)
        assert velocity.shape == (2,)
        assert len(led) == 4
        assert wirelessRemote.shape == (40,)
        
        self.head = head
        self.levelFlag = levelFlag
        self.frameReserve = frameReserve
        self.SN = SN
        self.version = version
        self.bandWidth = bandWidth
        self.mode = mode
        self.gaitType = gaitType
        self.speedLevel = speedLevel
        self.footRaiseHeight = footRaiseHeight
        self.bodyHeight = bodyHeight
        self.position = position
        self.euler = euler
        self.velocity = velocity
        self.yawSpeed = yawSpeed
        self.bms = bms
        self.led = led
        self.wirelessRemote = wirelessRemote
        self.crc = crc
    
    def toNative(self):
        nat = robot_interface.HighCmd()
        return self.toNativeM(nat)
    
    def toNativeM(self, native):
        # assert len(self.head) == 2
        # assert len(self.SN) == 2
        # assert len(self.version) == 2
        # assert self.position.shape == (3,)
        # assert self.euler.shape == (3,)
        # assert self.velocity.shape == (2,)
        # assert len(self.led) == 4
        # assert self.wirelessRemote.shape == (40,)

        native.head[0] = self.head[0]
        native.head[1] = self.head[1]
        native.levelFlag = self.levelFlag
        native.frameReserve = self.frameReserve
        native.SN[0] = self.SN[0]
        native.SN[1] = self.SN[1]
        native.version[0] = self.version[0]
        native.version[1] = self.version[1]
        native.bandWidth = self.bandWidth
        native.mode = self.mode
        native.gaitType = self.gaitType
        native.speedLevel = self.speedLevel
        native.footRaiseHeight = self.footRaiseHeight
        native.bodyHeight = self.bodyHeight
        for i in range(3):
            native.position[i] = self.position[i]
        for i in range(3):
            native.euler[i] = self.euler[i]
        for i in range(3):
            native.velocity[i] = self.velocity[i]
        native.yawSpeed = self.yawSpeed
        #native.bms = self.bms.toNative()
        self.bms.toNativeM(native.bms)
        for i in range(4):
            #native.led[i] = self.led[i].toNative()
            self.led[i].toNativeM(native.led[i])
        for i in range(40):
            native.wirelessRemote[i] = self.wirelessRemote[i]
        native.crc = self.crc
    
        return native
    
    @staticmethod
    def fromNative(native):
        return __class__(
            (native.head[0], native.head[1]),
            native.levelFlag,
            native.frameReserve,
            (native.SN[0], native.SN[1]),
            (native.version[0], native.version[1]),
            native.bandWidth,
            native.mode,
            native.gaitType,
            native.speedLevel,
            native.footRaiseHeight,
            native.bodyHeight,
            np.array([native.position[i] for i in range(3)]),
            np.array([native.euler[i] for i in range(3)]),
            np.array([native.velocity[i] for i in range(3)]),
            native.yawSpeed,
            BmsCmd.fromNative(native.bms),
            [LED.fromNative(native.led[i]) for i in range(4)],
            np.asarray(native.wirelessRemote, dtype=np.uint8),
            native.crc
        )
unitree_go1_wrapper/comm.py
realquantumcookie/APRL
from enum import Enum
import math

"""
From quadruped.h
"""
class LegIndex(Enum):
    LEG_INDEX_FR = 0
    LEG_INDEX_FL = 1
    LEG_INDEX_RR = 2
    LEG_INDEX_RL = 3

class JointIndex(Enum):
    FR_0 = 0
    FR_1 = 1
    FR_2 = 2

    FL_0 = 3
    FL_1 = 4
    FL_2 = 5

    RR_0 = 6
    RR_1 = 7
    RR_2 = 8

    RL_0 = 9
    RL_1 = 10
    RL_2 = 11

"""
From comm.h
"""
class ControlLevel(Enum):
    HIGH = 0xee
    LOW = 0xff
    TRIGGER = 0xf0

"""
might be something else, see https://github.com/unitreerobotics/unitree_legged_sdk/blob/v3.8.0/example_py/example_position.py
also see https://github.com/unitreerobotics/unitree_legged_sdk/blob/v3.8.0/include/unitree_legged_sdk/comm.h
"""
POS_STOP_F = math.pow(10,9) #2.146e9 

VEL_STOP_F = 16000.0
unitree_go1_wrapper/common_constants.py
realquantumcookie/APRL
import numpy as np

"""
From go1_const.h
"""

GO1_HIP_MAX   = 1.047    # unit:radian ( = 60   degree)
GO1_HIP_MIN   = -1.047   # unit:radian ( = -60  degree)
GO1_THIGH_MAX = 2.966    # unit:radian ( = 170  degree)
GO1_THIGH_MIN = -0.663   # unit:radian ( = -38  degree)
GO1_CALF_MAX  = -0.837   # unit:radian ( = -48  degree)
GO1_CALF_MIN  = -2.721   # unit:radian ( = -156 degree)
GO1_HIP_INIT = 0.0 / 180 * np.pi
GO1_THIGH_INIT = 45.0 / 180 * np.pi
GO1_CALF_INIT = -90.0 / 180 * np.pi
unitree_go1_wrapper/go1_constants.py
realquantumcookie/APRL
# https://github.com/unitreerobotics/unitree_legged_sdk/blob/go1/include/unitree_legged_sdk/joystick.h
from .comm import HasNativeRepr, State
from .native_interface import robot_interface
import numpy as np
from typing import Tuple

class XKeySwitchUnionComponents(HasNativeRepr):
    def __init__(
        self,
        R1: int,
        L1: int,
        start: int,
        select: int,
        R2: int,
        L2: int, 
        F1: int,
        F2: int,
        A: int,
        B: int,
        X: int,
        Y: int,
        up: int,
        right: int,
        down: int,
        left: int
    ):
        self.R1 = R1
        self.L1 = L1
        self.start = start
        self.select = select
        self.R2 = R2
        self.L2 = L2
        self.F1 = F1
        self.F2 = F2
        self.A = A
        self.B = B
        self.X = X
        self.Y = Y
        self.up = up
        self.right = right
        self.down = down
        self.left = left
    
    def toNative(self):
        native = np.zeros(16, dtype=np.uint8)
        return self.toNativeM(native)

    def toNativeM(self, native : np.ndarray):
        assert native.shape == (16,) and native.dtype == np.uint8
        native[0] = self.R1
        native[1] = self.L1
        native[2] = self.start
        native[3] = self.select
        native[4] = self.R2
        native[5] = self.L2
        native[6] = self.F1
        native[7] = self.F2
        native[8] = self.A
        native[9] = self.B
        native[10] = self.X
        native[11] = self.Y
        native[12] = self.up
        native[13] = self.right
        native[14] = self.down
        native[15] = self.left
    
    @staticmethod
    def fromNative(native : np.ndarray):
        # 16 bytes uint8 array
        assert native.shape == (16,) and native.dtype == np.uint8
        return __class__(
            *native
        )




class XKeySwitchUnion(HasNativeRepr):
    def __init__(
        self,
        value : int,
        components : XKeySwitchUnionComponents
    ):
        # It's a union type in C++ but we can't do that in Python
        # So we'll just have to keep track of the value and the components
        self.value = value
        self.components = components
    
    def toNative(self):
        return self.components.toNative()
    
    def toNativeM(self, native):
        return self.components.toNativeM(native)

    @staticmethod
    def fromNative(native : np.ndarray):
        assert native.shape == (16,) and native.dtype == np.uint8
        return __class__(
            native[:2].view(np.uint16)[0],
            XKeySwitchUnionComponents.fromNative(native)
        )

class XRockerBtnDataStruct(HasNativeRepr):
    def __init__(
        self,
        head: Tuple[int, int],
        btn: XKeySwitchUnion,
        lx: float,
        ly: float,
        rx: float,
        ry: float,
        L2: float # not sure what this is for?
    ):
        assert len(head) == 2
        self.head = head
        self.btn = btn
        self.lx = lx
        self.ly = ly
        self.rx = rx
        self.ry = ry
        self.L2 = L2
    
    def toNative(self) -> np.ndarray:
        native = np.zeros(40, dtype=np.uint8)
        return self.toNativeM(native)

    def toNativeM(self, native : np.ndarray):
        assert native.shape == (40,) and native.dtype == np.uint8
        native[:2] = np.array(self.head, dtype=np.uint8)
        self.btn.toNativeM(native[2:18])
        native[4:4+4*5] = np.array([self.lx, self.rx, self.ry, self.L2, self.ly], dtype=np.float16).view(np.uint8)
        return native

    @staticmethod
    def fromNative(native : np.ndarray):
        assert native.shape == (40,) and native.dtype == np.uint8
        head = tuple(native[:2]) #byte aligned
        
        float_values = native[4:4 + 4*5].view(np.float32)
        lx, rx, ry, L2, ly = float_values

        btn = XKeySwitchUnion.fromNative(native[2:18]) # Doesn't work for now
        return __class__(
            head,
            btn,
            lx,
            ly,
            rx,
            ry,
            L2
        )

    @staticmethod
    def max_mag():
        return np.sqrt(2)
unitree_go1_wrapper/joystick.py
realquantumcookie/APRL
import sys
import os
import platform

current_dir = os.path.dirname(os.path.abspath(__file__))
#project_root_dir = os.path.dirname(current_dir)

machine_arch = platform.machine().lower()
if machine_arch == "x86_64":
    machine_arch = "amd64"

to_append = os.path.join(current_dir, "lib", "python", machine_arch)
#print("appending {} to sys.path".format(to_append))

sys.path.append(
    to_append
)

import robot_interface
unitree_go1_wrapper/native_interface.py
realquantumcookie/APRL
# https://github.com/unitreerobotics/unitree_legged_sdk/blob/go1/include/unitree_legged_sdk/safety.h
from .native_interface import robot_interface
from . import comm, common_constants
from .udp import UDP

class Safety:
    """
    Safety is a class that provides safety functions to call in each robot control loop
    When a command to be issued to the robot exceeds a certain safety limit, the function will halt the entire program
    """
    
    def __init__(self, udp: UDP):
        self._safety = robot_interface.Safety(robot_interface.LeggedType.Go1)
        self.udp = udp
    
    def positionLimit(self) -> None:
        if self.udp.control_level == common_constants.ControlLevel.LOW:
            self._safety.PositionLimit(self.udp._low_cmd_native)
    
    def powerProtect(self, factor : int) -> None:
        """
        Apply percentage power protection
        factor ranges in [1,10] meaning 10% to 100% repsectively
        """
        if self.udp.control_level == common_constants.ControlLevel.LOW:
            self._safety.PowerProtect(self.udp._low_cmd_native, self.udp._low_state_native, factor)
    
    def positionProtect(self, limit : float = 0.087) -> None:
        """
        limit is the maximum angle in radians, default is 5 deg
        """
        if self.udp.control_level == common_constants.ControlLevel.LOW:
            self._safety.PositionProtect(self.udp._low_cmd_native, self.udp._low_state_native, limit)
unitree_go1_wrapper/safety.py
realquantumcookie/APRL
# https://github.com/unitreerobotics/unitree_legged_sdk/blob/go1/include/unitree_legged_sdk/udp.h
from .native_interface import robot_interface
from . import common_constants, comm
import copy

UDP_CLIENT_PORT = 8080
UDP_SERVER_PORT_BASIC = 8007
UDP_SERVER_PORT_SPORT = 8082
UDP_SERVER_IP_BASIC = "192.168.123.10"
UDP_SERVER_IP_SPORT = "192.168.123.161"

class UDP:
    """
    Wrapper for robot_interface.UDP
    """

    def __init__(
            self,
            control_level: common_constants.ControlLevel, 
            #localPort : int = UDP_CLIENT_PORT, 
            #target_ip : str = UDP_SERVER_IP_BASIC, 
            #target_port : int = UDP_SERVER_PORT_BASIC
        ):
        #self._udp = robot_interface.UDP(control_level.value, localPort, target_ip, target_port)
        if control_level == common_constants.ControlLevel.LOW:
            self._udp = robot_interface.UDP(control_level.value, UDP_CLIENT_PORT, UDP_SERVER_IP_BASIC, UDP_SERVER_PORT_BASIC)
            self._low_state_native = robot_interface.LowState()
            self._low_cmd_native = robot_interface.LowCmd()
            self._udp.InitCmdData(self._low_cmd_native)
            self._low_cmd_default = comm.LowCmd.fromNative(self._low_cmd_native)
            self._low_state = None
        elif control_level == common_constants.ControlLevel.HIGH:
            self._udp = robot_interface.UDP(control_level.value, UDP_CLIENT_PORT, UDP_SERVER_IP_SPORT, UDP_SERVER_PORT_SPORT)
            self._high_state_native = robot_interface.HighState()
            self._high_cmd_native = robot_interface.HighCmd()
            self._udp.InitCmdData(self._high_cmd_native)
            self._high_cmd_default = comm.HighCmd.fromNative(self._high_cmd_native)
            self._high_state = None
        self.control_level = control_level
    
    def setRecvTimeout(self,timeInMs : int) -> None:
        """
        Sets the timeout for receiving data from the UDP server
        """
        self._udp.setRecvTimeout(timeInMs)
    
    def send(self) -> None:
        """
        Sends data waiting in the client buffer to the server
        """
        if self.control_level == common_constants.ControlLevel.LOW:
            self._udp.SetSend(self._low_cmd_native)
        else:
            self._udp.SetSend(self._high_cmd_native)
        self._udp.Send()

    def recv(self) -> None:
        """
        Receives data from UDP server and puts it in a temporary buffer inside the UDP client. 
        The data saved in the buffer can be retrieved by calling getRecvLow() or getRecvHigh()
        """
        self._udp.Recv()
        if self.control_level == common_constants.ControlLevel.LOW:
            self._udp.GetRecv(self._low_state_native)
            self._low_state = None # clear conversion cache
        else:
            self._udp.GetRecv(self._high_state_native)
            self._high_state = None # clear conversion cache
    
    def initLowCmdData(self) -> comm.LowCmd:
        """
        Returns an initialized LowCmd with default values
        """
        return comm.LowCmd(
            copy.copy(self._low_cmd_default.head),
            copy.copy(self._low_cmd_default.version),
            self._low_cmd_default.bandWidth,
            [copy.copy(self._low_cmd_default.motorCmd[i]) for i in range(len(self._low_cmd_default.motorCmd))],
            copy.copy(self._low_cmd_default.bms),
            self._low_cmd_default.wirelessRemote.copy(),
            self._low_cmd_default.crc
        )
    
    def initCommunicationLowCmdData(self) -> comm.LowCmd:
        init_cmd = self.initLowCmdData()
        for i in range(12):
            init_cmd.motorCmd[i].q = common_constants.POS_STOP_F
            init_cmd.motorCmd[i].dq = common_constants.VEL_STOP_F
            init_cmd.motorCmd[i].Kp = 0.0
            init_cmd.motorCmd[i].Kd = 0.0
            init_cmd.motorCmd[i].tau = 0.0
        return init_cmd

    
    def initHighCmdData(self) -> comm.HighCmd:
        """
        Takes in a pointer to a HighCmd struct and initializes it with default values
        """
        return comm.HighCmd(
            copy.copy(self._high_cmd_default.head),
            self._high_cmd_default.levelFlag,
            self._high_cmd_default.frameReserve,
            copy.copy(self._high_cmd_default.SN),
            copy.copy(self._high_cmd_default.version),
            self._high_cmd_default.bandWidth,
            self._high_cmd_default.mode,
            self._high_cmd_default.gaitType,
            self._high_cmd_default.speedLevel,
            self._high_cmd_default.footRaiseHeight,
            self._high_cmd_default.bodyHeight,
            self._high_cmd_default.position.copy(),
            self._high_cmd_default.euler.copy(),
            self._high_cmd_default.velocity.copy(),
            self._high_cmd_default.yawSpeed,
            copy.copy(self._high_cmd_default.bms),
            [copy.copy(self._high_cmd_default.led[i] for i in range(len(self._high_cmd_default.led)))],
            self._high_cmd_default.wirelessRemote.copy(),
            self._high_cmd_default.crc
        )
    
    def setToSendLow(self, lowCmd : comm.LowCmd) -> None:
        """
        Takes in a pointer to a LowCmd struct and sets it to be sent later
        To actually send this data after calling this function, call send()
        """
        self._low_cmd_native = lowCmd.toNativeM(self._low_cmd_native)

    def setToSendHigh(self, highCmd : comm.HighCmd) -> None:
        """
        Takes in a pointer to a HighCmd struct and sets it to be sent later
        To actually send this data after calling this function, call send()
        """
        self._high_cmd_native = highCmd.toNativeM(self._high_cmd_native)
    
    def getRecvLow(self) -> comm.LowState:
        """
        Constructs a LowState struct and fills it with the data received from the server
        """
        if self._low_state is None:
            self._low_state = comm.LowState.fromNative(self._low_state_native)
        return self._low_state
    
    def getRecvHigh(self) -> comm.HighState:
        """
        Constructs a HighState struct and fills it with the data received from the server
        """
        if self._high_state is None:
            self._high_state = comm.HighState.fromNative(self._high_state_native)
        return self._high_state
        


unitree_go1_wrapper/udp.py
realquantumcookie/APRL